char-fork/llama.cpp

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

Max Krasnyanskyhexagon: L2 cache handling rework (dirty bit tracking with lazy flushing) and more MUL_MAT updates (#25762)b2dd28a3b

char/aliceai-t5-moe
7.8 KiB244 linesraw
1#include "work-queue.h"
2#include "hex-utils.h"
3
4#include <qurt.h>
5#include <qurt_hvx.h>
6
7#include <stdatomic.h>
8#include <stdint.h>
9#include <stdio.h>
10#include <stdlib.h>
11#include <string.h>
12
13#include "HAP_farf.h"
14
15#define LOWEST_USABLE_QURT_PRIO (254)
16
17// internal structure kept in thread-local storage per instance of work queue
18typedef struct {
19    work_queue_t  queue;
20    unsigned int  id;
21} worker_context_t;
22
23struct work_queue_task_s {
24    work_queue_func_t func;
25    void *            data;
26    unsigned int      n_threads;
27    atomic_uint       barrier;
28};
29
30// internal structure kept in thread-local storage per instance of work queue
31struct work_queue_s {
32    atomic_uint        seqn;      // seqno used to detect new jobs
33    atomic_uint        idx_read;  // Updated by producer (pop/reclaim)
34    unsigned int       idx_write; // Updated by producer (push)
35    uint32_t           idx_mask;
36    uint32_t           capacity;
37
38    qurt_thread_t      thread[WORK_QUEUE_MAX_N_THREADS];   // thread ID's of the workers
39    worker_context_t   context[WORK_QUEUE_MAX_N_THREADS];  // worker contexts
40    void *             stack[WORK_QUEUE_MAX_N_THREADS];    // thread stack pointers
41    unsigned int       n_threads;                          // total threads (workers + main)
42    unsigned int       n_workers;                          // number of active threads (just workers)
43
44    atomic_bool        active;                             // workers are polling/active
45    atomic_bool        killed;                             // threads need to exit
46    bool               external_mem;                       // memory owned externally
47
48    struct work_queue_task_s queue[] __attribute__((aligned(HEX_L2_LINE_SIZE)));
49};
50
51static void work_queue_thread(void * context) {
52    worker_context_t * me = (worker_context_t *) context;
53    work_queue_t       q  = me->queue;
54
55    FARF(HIGH, "work-queue: thread %u started", me->id);
56
57    unsigned int prev_seqn = 0;
58
59    while (!atomic_load_explicit(&q->killed, memory_order_relaxed)) {
60        unsigned int seqn = atomic_load_explicit(&q->seqn, memory_order_acquire);
61        if (seqn == prev_seqn) {
62            if (atomic_load_explicit(&q->active, memory_order_relaxed)) {
63                hex_pause();
64            } else {
65                qurt_futex_wait(&q->seqn, prev_seqn);
66            }
67            continue;
68        }
69
70        prev_seqn = seqn;
71
72        // Process all active tasks in the queue
73        unsigned int ir = atomic_load_explicit(&q->idx_read, memory_order_relaxed);
74        unsigned int iw = q->idx_write;
75
76        while (ir != iw) {
77            struct work_queue_task_s * task = &q->queue[ir];
78
79            unsigned int n = task->n_threads;
80            unsigned int i = me->id;
81            if (i < n) {
82                task->func(n, i, task->data);
83
84                atomic_fetch_sub_explicit(&task->barrier, 1, memory_order_release);
85            } else {
86                while (atomic_load_explicit(&task->barrier, memory_order_relaxed) > 0) {
87                    hex_pause();
88                }
89            }
90
91            ir = (ir + 1) & q->idx_mask;
92        }
93    }
94
95    FARF(HIGH, "work-queue: thread %u stopped", me->id);
96}
97
98bool work_queue_run_async(work_queue_t q, work_queue_func_t func, void * data, unsigned int n) {
99    if (n > q->n_threads) {
100        FARF(ERROR, "work-queue: invalid number of jobs %u for n-threads %u", n, q->n_threads);
101        return false;
102    }
103
104    unsigned int ir = atomic_load_explicit(&q->idx_read, memory_order_relaxed);
105    unsigned int iw = q->idx_write;
106
107    if (((iw + 1) & q->idx_mask) == ir) {
108        FARF(ERROR, "work-queue-push: queue is full\n");
109        return false;
110    }
111
112    struct work_queue_task_s * task = &q->queue[iw];
113    task->func      = func;
114    task->data      = data;
115    task->n_threads = n;
116    atomic_store_explicit(&task->barrier, n, memory_order_relaxed);
117
118    q->idx_write = (iw + 1) & q->idx_mask;
119
120    // publish job to workers (already awake and polling)
121    atomic_fetch_add_explicit(&q->seqn, 1, memory_order_release);
122
123    // main thread runs job #0
124    func(n, 0, data);
125
126    atomic_fetch_sub_explicit(&task->barrier, 1, memory_order_release);
127
128    while (atomic_load_explicit(&task->barrier, memory_order_relaxed) > 0) {
129        hex_pause();
130    }
131
132    atomic_thread_fence(memory_order_acquire);
133
134    atomic_store_explicit(&q->idx_read, (ir + 1) & q->idx_mask, memory_order_relaxed);
135
136    return true;
137}
138
139size_t work_queue_sizeof(uint32_t n_threads, uint32_t capacity, uint32_t stack_size) {
140    capacity = hex_ceil_pow2(capacity);
141    uint32_t n_workers = n_threads > 1 ? n_threads - 1 : 0;
142    size_t size_stacks = stack_size * n_workers;
143    size_t size_q = hex_align_up(sizeof(struct work_queue_s) + capacity * sizeof(struct work_queue_task_s), HEX_L2_LINE_SIZE);
144    return size_stacks + size_q;
145}
146
147size_t work_queue_alignof(void) {
148    return 4096;
149}
150
151work_queue_t work_queue_init(void * ptr, uint32_t n_threads, uint32_t capacity, uint32_t stack_size) {
152    capacity = hex_ceil_pow2(capacity);
153    uint32_t n_workers = n_threads > 1 ? n_threads - 1 : 0;
154    unsigned char * mem_blob = (unsigned char *) ptr;
155
156    work_queue_t q = (work_queue_t) (mem_blob + stack_size * n_workers);
157    memset(q, 0, sizeof(struct work_queue_s) + capacity * sizeof(struct work_queue_task_s));
158
159    q->n_threads = n_threads;
160    q->n_workers = n_workers;
161    q->external_mem = true;
162    q->capacity = capacity;
163
164    for (unsigned int i = 0; i < n_workers; i++) {
165        q->stack[i]  = mem_blob; mem_blob += stack_size;
166        q->thread[i] = 0;
167        q->context[i].id    = i + 1;
168        q->context[i].queue = q;
169    }
170
171    atomic_init(&q->idx_read, 0);
172    atomic_init(&q->seqn,     0);
173    atomic_init(&q->active,   false);
174    q->idx_write = 0;
175    q->idx_mask  = capacity - 1;
176    q->killed    = 0;
177    for (int i = 0; i < (int) capacity; i++) {
178        atomic_init(&q->queue[i].barrier, 0);
179        q->queue[i].func      = NULL;
180        q->queue[i].data      = NULL;
181        q->queue[i].n_threads = 0;
182    }
183
184    // launch the workers
185    qurt_thread_attr_t attr;
186    qurt_thread_attr_init(&attr);
187
188    for (unsigned int i = 0; i < n_workers; i++) {
189        qurt_thread_attr_set_stack_addr(&attr, q->stack[i]);
190        qurt_thread_attr_set_stack_size(&attr, stack_size);
191
192        char thread_name[32];
193        snprintf(thread_name, sizeof(thread_name), "work-queue:%u", i);
194        qurt_thread_attr_set_name(&attr, thread_name);
195
196        // set up priority - by default, match the creating thread's prio
197        int prio = qurt_thread_get_priority(qurt_thread_get_id());
198        if (prio < 1) {
199            prio = 1;
200        }
201        if (prio > LOWEST_USABLE_QURT_PRIO) {
202            prio = LOWEST_USABLE_QURT_PRIO;
203        }
204
205        qurt_thread_attr_set_priority(&attr, prio);
206
207        int err = qurt_thread_create(&q->thread[i], &attr, work_queue_thread, (void *) &q->context[i]);
208        if (err) {
209            FARF(ERROR, "Could not launch worker threads!");
210            work_queue_free(q);
211            return NULL;
212        }
213    }
214
215    return q;
216}
217
218void work_queue_free(work_queue_t q) {
219    if (!q) { return; }
220
221    atomic_store_explicit(&q->killed,   1, memory_order_relaxed);
222    atomic_fetch_add_explicit(&q->seqn, 1, memory_order_release);
223    qurt_futex_wake(&q->seqn, q->n_workers);
224
225    for (unsigned int i = 0; i < q->n_workers; i++) {
226        if (q->thread[i]) {
227            int status;
228            (void) qurt_thread_join(q->thread[i], &status);
229        }
230    }
231}
232
233void work_queue_wakeup(work_queue_t q) {
234    if (!atomic_load_explicit(&q->active, memory_order_relaxed)) {
235        atomic_store_explicit(&q->active, true, memory_order_release);
236        // Increment seqn and wake workers to transition them out of sleep
237        atomic_fetch_add_explicit(&q->seqn, 1, memory_order_release);
238        qurt_futex_wake(&q->seqn, q->n_workers);
239    }
240}
241
242void work_queue_suspend(work_queue_t q) {
243    atomic_store_explicit(&q->active, false, memory_order_release);
244}