|
32 | 32 | #include "qjl_block.h" |
33 | 33 |
|
34 | 34 | #include <math.h> |
35 | | -#include <stdatomic.h> |
36 | 35 | #include <stdint.h> |
37 | 36 | #include <stdlib.h> |
38 | 37 | #include <string.h> |
39 | 38 |
|
| 39 | +/* Portable atomic primitives for the lazy-init CAS below. |
| 40 | + * |
| 41 | + * MSVC's <stdatomic.h> emits `#error "C atomic support is not enabled"` |
| 42 | + * unless the project is built with `/experimental:c11atomics` (only on |
| 43 | + * very recent toolchains), so we cannot rely on `_Atomic` here. Use: |
| 44 | + * - MSVC: _InterlockedCompareExchange / _InterlockedExchange |
| 45 | + * (with MemoryBarrier() for acq/rel semantics). |
| 46 | + * - GCC/clang: __atomic_* builtins (available everywhere we target). |
| 47 | + * |
| 48 | + * Both wrappers operate on a `volatile long` cell so the MSVC intrinsics |
| 49 | + * see the correct width on 64-bit Windows (where `int` is 32-bit and |
| 50 | + * `long` is also 32-bit). */ |
| 51 | +#if defined(_MSC_VER) |
| 52 | +#include <windows.h> |
| 53 | +typedef volatile long qjl_atomic_int; |
| 54 | +static inline int qjl_atomic_load_acquire(qjl_atomic_int * p) { |
| 55 | + long v = *p; |
| 56 | + MemoryBarrier(); |
| 57 | + return (int) v; |
| 58 | +} |
| 59 | +static inline void qjl_atomic_store_release(qjl_atomic_int * p, int v) { |
| 60 | + MemoryBarrier(); |
| 61 | + _InterlockedExchange(p, (long) v); |
| 62 | +} |
| 63 | +static inline int qjl_atomic_cas_acq_rel(qjl_atomic_int * p, int expected, int desired) { |
| 64 | + long prev = _InterlockedCompareExchange(p, (long) desired, (long) expected); |
| 65 | + return prev == (long) expected; |
| 66 | +} |
| 67 | +#else |
| 68 | +typedef volatile int qjl_atomic_int; |
| 69 | +static inline int qjl_atomic_load_acquire(qjl_atomic_int * p) { |
| 70 | + return __atomic_load_n(p, __ATOMIC_ACQUIRE); |
| 71 | +} |
| 72 | +static inline void qjl_atomic_store_release(qjl_atomic_int * p, int v) { |
| 73 | + __atomic_store_n(p, v, __ATOMIC_RELEASE); |
| 74 | +} |
| 75 | +static inline int qjl_atomic_cas_acq_rel(qjl_atomic_int * p, int expected, int desired) { |
| 76 | + return __atomic_compare_exchange_n(p, &expected, desired, 0, |
| 77 | + __ATOMIC_ACQ_REL, __ATOMIC_ACQUIRE); |
| 78 | +} |
| 79 | +#endif |
| 80 | + |
40 | 81 | /* Confirm the two block layouts agree byte-for-byte. */ |
41 | 82 | _Static_assert(sizeof(qjl_block_qjl1_256) == sizeof(block_qjl1_256), |
42 | 83 | "qjl block layouts must agree (kernel lib vs ggml)"); |
@@ -68,30 +109,26 @@ _Static_assert(QJL_PROJECTION_DIM == QK_QJL, |
68 | 109 | #define QJL_INIT_RUNNING 1 |
69 | 110 | #define QJL_INIT_READY 2 |
70 | 111 |
|
71 | | -static _Atomic int g_qjl_prj_state = QJL_INIT_UNINIT; |
| 112 | +static qjl_atomic_int g_qjl_prj_state = QJL_INIT_UNINIT; |
72 | 113 | static float *g_qjl_prj = NULL; |
73 | 114 |
|
74 | 115 | static const float * qjl_default_projection(void) { |
75 | | - int state = atomic_load_explicit(&g_qjl_prj_state, memory_order_acquire); |
| 116 | + int state = qjl_atomic_load_acquire(&g_qjl_prj_state); |
76 | 117 | if (state == QJL_INIT_READY) { |
77 | 118 | return g_qjl_prj; |
78 | 119 | } |
79 | 120 |
|
80 | | - int expected = QJL_INIT_UNINIT; |
81 | | - if (atomic_compare_exchange_strong_explicit(&g_qjl_prj_state, &expected, |
82 | | - QJL_INIT_RUNNING, |
83 | | - memory_order_acq_rel, |
84 | | - memory_order_acquire)) { |
| 121 | + if (qjl_atomic_cas_acq_rel(&g_qjl_prj_state, QJL_INIT_UNINIT, QJL_INIT_RUNNING)) { |
85 | 122 | g_qjl_prj = (float *) malloc(sizeof(float) * QJL_DEFAULT_HEAD_DIM * QJL_DEFAULT_PROJ_DIM); |
86 | 123 | if (g_qjl_prj != NULL) { |
87 | 124 | qjl_make_projection_mt(g_qjl_prj, QJL_DEFAULT_HEAD_DIM, QJL_DEFAULT_PROJ_DIM, QJL_DEFAULT_SEED); |
88 | 125 | } |
89 | | - atomic_store_explicit(&g_qjl_prj_state, QJL_INIT_READY, memory_order_release); |
| 126 | + qjl_atomic_store_release(&g_qjl_prj_state, QJL_INIT_READY); |
90 | 127 | return g_qjl_prj; |
91 | 128 | } |
92 | 129 |
|
93 | 130 | /* Another thread is initializing — wait for it. */ |
94 | | - while (atomic_load_explicit(&g_qjl_prj_state, memory_order_acquire) != QJL_INIT_READY) { |
| 131 | + while (qjl_atomic_load_acquire(&g_qjl_prj_state) != QJL_INIT_READY) { |
95 | 132 | /* tiny pause; this loop runs at most once per process lifetime */ |
96 | 133 | } |
97 | 134 | return g_qjl_prj; |
|
0 commit comments