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.8 KiB351 linesraw
1#include "htp-tensor.h"
2
3#include <qurt.h>
4#include <qurt_memory.h>
5#include <HAP_farf.h>
6
7#include "hex-common.h"
8#include "hex-utils.h"
9#include "hex-fastdiv.h"
10#include "hex-profile.h"
11#include "htp-ctx.h"
12#include "work-queue.h"
13
14struct l2flush_range {
15    uint32_t start;       // line-aligned start address
16    uint32_t end;         // line-aligned end address
17    uint32_t block_first; // global block index of this range's first block
18    uint32_t n_blocks;    // number of HEX_L2_BLOCK_SIZE chunks (last may be partial)
19};
20
21struct l2flush_multi_task {
22    struct htp_thread_trace * trace;
23    struct l2flush_range      ranges[HTP_MAX_DIRTY_RANGES];
24    uint32_t                  n_ranges;
25    uint32_t                  total_blocks;
26    uint32_t                  blocks_per_thread;
27};
28
29static void flush_all_dcache(struct htp_context * ctx) {
30    struct htp_thread_trace * tr = &ctx->trace[0];
31    htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
32    qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE);
33    hex_l2fetch_block(ctx, ctx->footprint);
34    htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
35    memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
36}
37
38static void l2flush_multi_worker(unsigned int n, unsigned int i, void * data) {
39    struct l2flush_multi_task * task = (struct l2flush_multi_task *) data;
40    (void) n;
41
42    const uint32_t gb_first = i * task->blocks_per_thread;
43    uint32_t       gb_last  = gb_first + task->blocks_per_thread;
44    if (gb_last > task->total_blocks) {
45        gb_last = task->total_blocks;
46    }
47    if (gb_first >= gb_last) {
48        return;
49    }
50
51    struct htp_thread_trace * tr = &task->trace[i];
52    htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
53
54    for (uint32_t r = 0; r < task->n_ranges; r++) {
55        const struct l2flush_range * rg = &task->ranges[r];
56        const uint32_t rb_first = rg->block_first;
57        const uint32_t rb_last  = rg->block_first + rg->n_blocks;
58
59        const uint32_t lo = gb_first > rb_first ? gb_first : rb_first;
60        const uint32_t hi = gb_last  < rb_last  ? gb_last  : rb_last;
61        if (lo >= hi) {
62            continue;
63        }
64
65        const uint32_t s = rg->start + (lo - rb_first) * HEX_L2_BLOCK_SIZE;
66        uint32_t       e = rg->start + (hi - rb_first) * HEX_L2_BLOCK_SIZE;
67        if (e > rg->end) {
68            e = rg->end;
69        }
70        hex_l2flush((void *) (uintptr_t) s, e - s);
71    }
72
73    htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, gb_first);
74}
75
76static void merge_dirty_ranges(struct htp_context * ctx) {
77    for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
78        struct htp_dirty_range * r = &ctx->dirty_ranges[i];
79        if (!r->start) continue;
80
81        for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES;) {
82            struct htp_dirty_range * s = &ctx->dirty_ranges[j];
83            if (i == j || !s->start || r->end < s->start || s->end < r->start) {
84                j++;
85                continue;
86            }
87
88            r->start = MIN(r->start, s->start);
89            r->end   = MAX(r->end, s->end);
90            s->start = 0;
91            s->end   = 0;
92            j = 0;
93        }
94    }
95}
96
97void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
98    const struct htp_tensor * pending[HTP_OP_MAX_OUTPUTS];
99    uint32_t n_pending = 0;
100
101    for (uint32_t i = 0; i < n; i++) {
102        const struct htp_tensor * t = tensors[i];
103        if (!t || (t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE))) {
104            continue;
105        }
106
107        uint32_t t_start = t->data;
108        uint32_t t_end   = t_start + t->size;
109
110        bool merged = false;
111        for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES; j++) {
112            struct htp_dirty_range * r = &ctx->dirty_ranges[j];
113            if (!r->start) continue;
114
115            if (r->start <= t_end && t_start <= r->end) {
116                uint32_t new_start = (t_start < r->start) ? t_start : r->start;
117                uint32_t new_end   = (t_end > r->end)     ? t_end   : r->end;
118                r->start = new_start;
119                r->end   = new_end;
120                merged   = true;
121            }
122        }
123
124        if (!merged) {
125            pending[n_pending++] = t;
126        }
127    }
128
129    merge_dirty_ranges(ctx);
130
131    if (n_pending == 0) {
132        return;
133    }
134
135    uint32_t empty_indices[HTP_MAX_DIRTY_RANGES];
136    uint32_t active_indices[HTP_MAX_DIRTY_RANGES];
137    uint32_t n_active = 0;
138    uint32_t n_empty  = 0;
139    for (uint32_t j = 0; j < HTP_MAX_DIRTY_RANGES; j++) {
140        if (ctx->dirty_ranges[j].start) {
141            active_indices[n_active++] = j;
142        } else {
143            empty_indices[n_empty++]   = j;
144        }
145    }
146
147    if (n_pending <= n_empty) {
148        for (uint32_t i = 0; i < n_pending; i++) {
149            uint32_t idx = empty_indices[i];
150            struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
151            r->start = pending[i]->data;
152            r->end   = pending[i]->data + pending[i]->size;
153        }
154        merge_dirty_ranges(ctx);
155        return;
156    }
157
158    uint32_t n_evict = n_pending - n_empty;
159    uint32_t total_evict_size = 0;
160    for (uint32_t i = 0; i < n_evict; i++) {
161        uint32_t idx = active_indices[i];
162        struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
163        total_evict_size += r->end - r->start;
164    }
165
166    if (total_evict_size > HEX_L2_FLUSH_ALL_THRESHOLD) {
167        flush_all_dcache(ctx);
168        for (uint32_t i = 0; i < n_pending; i++) {
169            struct htp_dirty_range * r = &ctx->dirty_ranges[i];
170            r->start = pending[i]->data;
171            r->end   = pending[i]->data + pending[i]->size;
172        }
173        merge_dirty_ranges(ctx);
174        return;
175    }
176
177    if (total_evict_size > HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1 && n_evict <= HTP_MAX_DIRTY_RANGES) {
178        struct l2flush_multi_task task;
179        task.trace    = ctx->trace;
180        task.n_ranges = n_evict;
181
182        uint32_t block_acc = 0;
183        for (uint32_t i = 0; i < n_evict; i++) {
184            uint32_t idx = active_indices[i];
185            struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
186
187            struct l2flush_range * rg = &task.ranges[i];
188            rg->start = hex_align_down((size_t) r->start, HEX_L2_LINE_SIZE);
189            rg->end   = hex_align_up((size_t) r->end, HEX_L2_LINE_SIZE);
190            rg->block_first = block_acc;
191            rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
192            block_acc += rg->n_blocks;
193        }
194
195        task.total_blocks      = block_acc;
196        task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
197
198        work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
199    } else {
200        struct htp_thread_trace * tr = &ctx->trace[0];
201        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
202        for (uint32_t i = 0; i < n_evict; i++) {
203            uint32_t idx = active_indices[i];
204            struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
205            uint32_t size = r->end - r->start;
206            hex_l2flush((void *) (uintptr_t) r->start, size);
207        }
208        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
209    }
210
211    for (uint32_t i = 0; i < n_evict; i++) {
212        uint32_t idx = active_indices[i];
213        struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
214        r->start = pending[i]->data;
215        r->end   = pending[i]->data + pending[i]->size;
216    }
217
218    for (uint32_t i = 0; i < n_empty; i++) {
219        uint32_t idx = empty_indices[i];
220        struct htp_dirty_range * r = &ctx->dirty_ranges[idx];
221        r->start = pending[n_evict + i]->data;
222        r->end   = pending[n_evict + i]->data + pending[n_evict + i]->size;
223    }
224
225    merge_dirty_ranges(ctx);
226}
227
228static void make_tensor_clean(struct htp_context * ctx, const struct htp_tensor * t) {
229    uint32_t t_start = t->data;
230    uint32_t t_end   = t_start + t->size;
231
232    for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
233        struct htp_dirty_range * r = &ctx->dirty_ranges[i];
234        if (!r->start) continue;
235
236        if (r->start < t_end && t_start < r->end) {
237            if (t_start <= r->start && r->end <= t_end) {
238                r->start = 0;
239            } else if (t_start <= r->start) {
240                r->start = t_end;
241            } else if (r->end <= t_end) {
242                r->end = t_start;
243            }
244        }
245    }
246}
247
248static inline bool is_tensor_dirty(struct htp_context * ctx, const struct htp_tensor * t) {
249    uint32_t t_start = t->data;
250    uint32_t t_end   = t_start + t->size;
251
252    for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
253        struct htp_dirty_range * r = &ctx->dirty_ranges[i];
254        if (!r->start) continue;
255
256        if (r->start < t_end && t_start < r->end) {
257            return true;
258        }
259    }
260    return false;
261}
262
263static void flush_dirty_ranges(struct htp_context * ctx, const struct htp_dirty_range * ranges, uint32_t n_ranges, uint64_t total_dirty) {
264    if (total_dirty >= HEX_L2_FLUSH_WQ_THRESHOLD && ctx->n_threads > 1) {
265        struct l2flush_multi_task task;
266        task.trace    = ctx->trace;
267        task.n_ranges = n_ranges;
268
269        uint32_t block_acc = 0;
270        for (uint32_t i = 0; i < n_ranges; i++) {
271            const struct htp_dirty_range * r = &ranges[i];
272            struct l2flush_range * rg = &task.ranges[i];
273            rg->start = hex_align_down((size_t) r->start, HEX_L2_LINE_SIZE);
274            rg->end   = hex_align_up((size_t) r->end, HEX_L2_LINE_SIZE);
275            rg->block_first = block_acc;
276            rg->n_blocks = (rg->end - rg->start + HEX_L2_BLOCK_SIZE - 1) / HEX_L2_BLOCK_SIZE;
277            block_acc += rg->n_blocks;
278        }
279
280        task.total_blocks      = block_acc;
281        task.blocks_per_thread = fastdiv(block_acc + ctx->n_threads - 1, &ctx->n_threads_div);
282
283        work_queue_run(ctx->work_queue, l2flush_multi_worker, &task, ctx->n_threads);
284    } else {
285        struct htp_thread_trace * tr = &ctx->trace[0];
286        htp_trace_event_start(tr, HTP_TRACE_EVT_L2FLUSH, 0);
287        for (uint32_t i = 0; i < n_ranges; i++) {
288            const struct htp_dirty_range * r = &ranges[i];
289            hex_l2flush((void *) (uintptr_t) r->start, r->end - r->start);
290        }
291        htp_trace_event_stop(tr, HTP_TRACE_EVT_L2FLUSH, 0);
292    }
293}
294
295void htp_flush_dirty_ranges(struct htp_context * ctx) {
296    struct htp_dirty_range ranges[HTP_MAX_DIRTY_RANGES];
297    uint32_t n_ranges = 0;
298    uint64_t total_dirty = 0;
299
300    for (uint32_t i = 0; i < HTP_MAX_DIRTY_RANGES; i++) {
301        const struct htp_dirty_range * r = &ctx->dirty_ranges[i];
302        if (!r->start) {
303            continue;
304        }
305        ranges[n_ranges++] = *r;
306        total_dirty += r->end - r->start;
307    }
308
309    if (total_dirty == 0) {
310        return;
311    }
312
313    if (total_dirty > HEX_L2_FLUSH_ALL_THRESHOLD) {
314        flush_all_dcache(ctx);
315        return;
316    }
317
318    flush_dirty_ranges(ctx, ranges, n_ranges, total_dirty);
319    memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges));
320}
321
322void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) {
323    const struct htp_tensor * dirty_tensors[HTP_OP_MAX_INPUTS];
324    struct htp_dirty_range ranges[HTP_OP_MAX_INPUTS];
325    uint32_t n_dirty = 0;
326    uint64_t total_dirty = 0;
327
328    for (uint32_t i = 0; i < n; i++) {
329        const struct htp_tensor * t = tensors[i];
330        if (t && is_tensor_dirty(ctx, t)) {
331            dirty_tensors[n_dirty++] = t;
332            ranges[n_dirty - 1].start = t->data;
333            ranges[n_dirty - 1].end   = t->data + t->size;
334            total_dirty += t->size;
335        }
336    }
337
338    if (total_dirty == 0) {
339        return;
340    }
341
342    if (total_dirty > HEX_L2_FLUSH_ALL_THRESHOLD) {
343        flush_all_dcache(ctx);
344        return;
345    }
346
347    flush_dirty_ranges(ctx, ranges, n_dirty, total_dirty);
348    for (uint32_t i = 0; i < n_dirty; i++) {
349        make_tensor_clean(ctx, dirty_tensors[i]);
350    }
351}