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 →
11.3 KiB284 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
8#define GGML_COMMON_DECL_C
9#include "ggml-common.h"
10#include "hex-common.h"
11#include "hex-profile.h"
12#include "htp-ctx.h"
13#include "htp-ops.h"
14#include "htp-tensor.h"
15#include "hvx-types.h"
16#include "hvx-utils.h"
17#include "hex-dma.h"
18
19#define htp_cumsum_tensors_preamble                         \
20    const struct htp_tensor * restrict src0 = octx->src[0]; \
21    const struct htp_tensor * restrict dst  = octx->dst;    \
22                                                            \
23    const uint32_t ne00 = src0->ne[0];                      \
24    const uint32_t ne01 = src0->ne[1];                      \
25    const uint32_t ne02 = src0->ne[2];                      \
26    const uint32_t ne03 = src0->ne[3];                      \
27                                                            \
28    const uint32_t ne0 = dst->ne[0];                        \
29    const uint32_t ne1 = dst->ne[1];                        \
30    const uint32_t ne2 = dst->ne[2];                        \
31    const uint32_t ne3 = dst->ne[3];                        \
32                                                            \
33    const uint32_t nb00 = src0->nb[0];                      \
34    const uint32_t nb01 = src0->nb[1];                      \
35    const uint32_t nb02 = src0->nb[2];                      \
36    const uint32_t nb03 = src0->nb[3];                      \
37                                                            \
38    const uint32_t nb0 = dst->nb[0];                        \
39    const uint32_t nb1 = dst->nb[1];                        \
40    const uint32_t nb2 = dst->nb[2];                        \
41    const uint32_t nb3 = dst->nb[3];
42
43struct htp_cumsum_context {
44    struct htp_ops_context * octx;
45    size_t          src_row_size;
46    size_t          dst_row_size;
47    size_t          src_row_size_aligned;
48    size_t          dst_row_size_aligned;
49    uint32_t        rows_per_thread;
50    uint32_t        total_rows;
51    uint32_t        row_start;
52};
53
54#define htp_cumsum_preamble                                                \
55    struct htp_cumsum_context * cctx = (struct htp_cumsum_context *) data; \
56    struct htp_ops_context *    octx = cctx->octx;                         \
57    htp_cumsum_tensors_preamble;                                           \
58    dma_queue * dma_queue = octx->ctx->dma[ith];
59
60// ---------------------------------------------------------------------------
61// HVX prefix scan helpers
62// ---------------------------------------------------------------------------
63
64#if __HVX_ARCH__ > 75
65static inline HVX_Vector hvx_cumsum_vadd(HVX_Vector a, HVX_Vector b) {
66    return Q6_Vsf_vadd_VsfVsf(a, b);
67}
68#else
69static inline HVX_Vector hvx_cumsum_vadd(HVX_Vector a, HVX_Vector b) {
70    return Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b));
71}
72#endif  // __HVX_ARCH__ > 75
73
74static inline HVX_Vector hvx_prefix_scan_f32(HVX_Vector v, HVX_Vector carry_in) {
75    const HVX_Vector zero = Q6_V_vsplat_R(0);
76
77    v = hvx_cumsum_vadd(v, Q6_V_vlalign_VVR(v, zero,  4));
78    v = hvx_cumsum_vadd(v, Q6_V_vlalign_VVR(v, zero,  8));
79    v = hvx_cumsum_vadd(v, Q6_V_vlalign_VVR(v, zero, 16));
80    v = hvx_cumsum_vadd(v, Q6_V_vlalign_VVR(v, zero, 32));
81    v = hvx_cumsum_vadd(v, Q6_V_vlalign_VVR(v, zero, 64));
82    v = hvx_cumsum_vadd(v, carry_in);
83
84    return v;
85}
86
87static inline HVX_Vector hvx_splat_last_f32(HVX_Vector v) {
88    return hvx_vec_repl4(Q6_V_vror_VR(v, 124));
89}
90
91static inline void hvx_cumsum_row_f32(const float * restrict src, float * restrict dst, uint32_t n) {
92    const uint32_t nvec = n / VLEN_FP32;
93    const uint32_t nloe = n % VLEN_FP32;
94
95    HVX_Vector carry = Q6_V_vsplat_R(0);
96
97    for (uint32_t i = 0; i < nvec; i++) {
98        HVX_Vector v = *((const HVX_UVector *) (src + i * VLEN_FP32));
99        v = hvx_prefix_scan_f32(v, carry);
100        hvx_vec_store_u(dst + i * VLEN_FP32, VLEN, v);
101        carry = hvx_splat_last_f32(v);
102    }
103
104    if (nloe) {
105        float acc = hvx_vec_get_f32(carry);
106        const float * src_tail = src + nvec * VLEN_FP32;
107        float       * dst_tail = dst + nvec * VLEN_FP32;
108        for (uint32_t i = 0; i < nloe; i++) {
109            acc        += src_tail[i];
110            dst_tail[i] = acc;
111        }
112    }
113}
114
115// ---------------------------------------------------------------------------
116// Per thread worker: Double-buffered DMA
117// ---------------------------------------------------------------------------
118
119static void cumsum_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) {
120    htp_cumsum_preamble;
121
122    const uint32_t ir0 = cctx->row_start + cctx->rows_per_thread * ith;
123    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->row_start + cctx->total_rows);
124
125    if (ir0 >= ir1) {
126        return;
127    }
128
129    const size_t src_row_size         = cctx->src_row_size;
130    const size_t dst_row_size         = cctx->dst_row_size;
131    const size_t src_row_size_aligned = cctx->src_row_size_aligned;
132    const size_t dst_row_size_aligned = cctx->dst_row_size_aligned;
133
134    const uint8_t * src_data = (const uint8_t *) src0->data;
135    uint8_t *       dst_data = (uint8_t *) dst->data;
136
137    uint8_t * src_spad = octx->src0_spad.data + (ith * src_row_size_aligned * 2);
138    uint8_t * dst_spad = octx->dst_spad.data  + (ith * dst_row_size_aligned * 2);
139
140    for (uint32_t ir = ir0, spad_idx = 0; ir < ir1 && spad_idx < 2; ir++, spad_idx++) {
141        // Dummy dst writeback to establish queue ordering
142        dma_queue_push_vtcm_to_ddr(dma_queue,
143                                   dma_make_ptr(dst_data, dst_spad + (spad_idx * dst_row_size_aligned)),
144                                   dst_row_size, dst_row_size_aligned, 0);
145
146        dma_queue_push_ddr_to_vtcm(dma_queue,
147                                   dma_make_ptr(src_spad + (spad_idx * src_row_size_aligned),
148                                                src_data + (ir * src_row_size)),
149                                   src_row_size_aligned, src_row_size, 1);
150    }
151
152    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
153
154    for (uint32_t ir = ir0; ir < ir1; ir++) {
155        float * dst_spad_row = (float *) dma_queue_pop(dma_queue).src;
156        float * src_spad_row = (float *) dma_queue_pop(dma_queue).dst;
157
158        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
159        hvx_cumsum_row_f32(src_spad_row, dst_spad_row, ne00);
160        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
161
162        dma_queue_push_vtcm_to_ddr(dma_queue,
163                                   dma_make_ptr(dst_data + (ir * dst_row_size), (uint8_t *) dst_spad_row),
164                                   dst_row_size, dst_row_size_aligned, 1);
165
166        const uint32_t next_row = ir + 2;
167        if (next_row < ir1) {
168            dma_queue_push_ddr_to_vtcm(dma_queue,
169                                       dma_make_ptr((uint8_t *) src_spad_row, src_data + (next_row * src_row_size)),
170                                       src_row_size_aligned, src_row_size, 1);
171        }
172    }
173
174    dma_queue_flush(dma_queue);
175
176    FARF(HIGH, "cumsum-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
177         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1,
178         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
179}
180
181// ---------------------------------------------------------------------------
182// Per thread worker: Direct HVX (no DMA)
183// ---------------------------------------------------------------------------
184
185static void cumsum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
186    htp_cumsum_preamble;
187
188    const uint8_t * src_data = (const uint8_t *) src0->data;
189    uint8_t *       dst_data = (uint8_t *) dst->data;
190
191    const uint32_t ir0 = cctx->row_start + cctx->rows_per_thread * ith;
192    const uint32_t ir1 = MIN(ir0 + cctx->rows_per_thread, cctx->row_start + cctx->total_rows);
193
194    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
195    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
196
197    for (uint32_t ir = ir0; ir < ir1; ir++) {
198        const float * restrict src_row = (const float *) (src_data + ir * cctx->src_row_size);
199        float       * restrict dst_row = (float *)       (dst_data + ir * cctx->dst_row_size);
200        hvx_cumsum_row_f32(src_row, dst_row, ne00);
201    }
202
203    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
204
205    FARF(HIGH, "cumsum-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
206         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ir0, ir1,
207         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
208}
209
210int op_cumsum_f32(struct htp_ops_context * octx) {
211    const struct htp_tensor * src0 = octx->src[0];
212    const struct htp_tensor * dst  = octx->dst;
213
214    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
215        return HTP_STATUS_OK;
216    }
217
218    const uint32_t total_rows      = src0->ne[1] * src0->ne[2] * src0->ne[3];
219    const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
220
221    uint32_t row_start = 0;
222    uint32_t nrows     = total_rows;
223
224    if (octx->ctx->mdev.count > 1) {
225        uint32_t rows_per_chunk = 0;
226        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_data_row_size, &rows_per_chunk);
227        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
228        row_start = range.start;
229        nrows     = range.count;
230    }
231
232    if (nrows == 0) {
233        return HTP_STATUS_OK;
234    }
235
236    const uint32_t n_threads = octx->n_threads;
237
238    const size_t src_row_size         = src0->nb[1];
239    const size_t dst_row_size         = dst->nb[1];
240    const size_t src_row_size_aligned = hex_round_up(src_row_size, VLEN);
241    const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN);
242
243    // 2 ping-pong buffers per thread for src and dst
244    const size_t spad_per_thread = 2 * (src_row_size_aligned + dst_row_size_aligned);
245
246    octx->src0_spad.size_per_thread = src_row_size_aligned * 2;
247    octx->dst_spad.size_per_thread  = dst_row_size_aligned * 2;
248
249    octx->src0_spad.size  = n_threads * octx->src0_spad.size_per_thread;
250    octx->dst_spad.size   = n_threads * octx->dst_spad.size_per_thread;
251
252    octx->src0_spad.data  = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
253    octx->dst_spad.data   = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
254
255    struct htp_cumsum_context cctx = {
256        .octx                 = octx,
257        .src_row_size         = src_row_size,
258        .dst_row_size         = dst_row_size,
259        .src_row_size_aligned = src_row_size_aligned,
260        .dst_row_size_aligned = dst_row_size_aligned,
261        .rows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
262        .total_rows           = nrows,
263        .row_start            = row_start,
264    };
265
266    if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
267        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32, &cctx, n_threads);
268    } else {
269        work_queue_run(octx->ctx->work_queue, cumsum_thread_f32_dma, &cctx, n_threads);
270    }
271
272    return HTP_STATUS_OK;
273}
274
275int op_cumsum(struct htp_ops_context * octx) {
276    const struct htp_tensor * dst = octx->dst;
277
278    switch (dst->type) {
279        case HTP_TYPE_F32:
280            return op_cumsum_f32(octx);
281        default:
282            return HTP_STATUS_NO_SUPPORT;
283    }
284}