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