char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
b7feacf7f
char/aliceai-t5-moe
1#include "virtgpu-utils.h" 2 3#include <malloc.h> 4#include <stdlib.h> 5 6#include <cstring> 7 8#define NODE_ALLOC_ALIGN 64 9#define NODE_PTR_MASK (~((uintptr_t) NODE_ALLOC_ALIGN - 1)) 10#define NODE_LEVEL_MASK ((uintptr_t) NODE_ALLOC_ALIGN - 1) 11#define NULL_NODE 0 12 13#define os_malloc_aligned (_size ,_align ) _aligned_malloc(_size, _align) 14#define os_free_aligned (_ptr ) free(_ptr) 15#define p_atomic_cmpxchg (v ,old ,_new ) __sync_val_compare_and_swap((v), (old), (_new)) 16 17static inline uint64_t util_logbase2_64 (uint64_t n ) { 18#if defined(HAVE___BUILTIN_CLZLL ) 19return ((sizeof (uint64_t )* 8 - 1 )- __builtin_clzll (n |1 )); 20#else 21uint64_t pos = 0ull ; 22if (n >=1ull <<32 ) { 23n >>=32 ; 24pos += 32 ; 25 } 26if (n >=1ull <<16 ) { 27n >>=16 ; 28pos += 16 ; 29 } 30if (n >=1ull <<8 ) { 31n >>=8 ; 32pos += 8 ; 33 } 34if (n >=1ull <<4 ) { 35n >>=4 ; 36pos += 4 ; 37 } 38if (n >=1ull <<2 ) { 39n >>=2 ; 40pos += 2 ; 41 } 42if (n >=1ull <<1 ) { 43pos += 1 ; 44 } 45return pos ; 46#endif 47} 48 49void util_sparse_array_init (util_sparse_array * arr ,size_t elem_size ,size_t node_size ) { 50memset (arr ,0 ,sizeof (* arr )); 51arr -> elem_size = elem_size ; 52arr -> node_size_log2 = util_logbase2_64 (node_size ); 53assert (node_size >=2 && node_size == (1ull <<arr -> node_size_log2 )); 54} 55 56static inline void * os_malloc_aligned (size_t size ,size_t alignment ) { 57void * ptr ; 58alignment = (alignment + sizeof (void * )- 1 )& ~(sizeof (void * )- 1 ); 59if (posix_memalign (& ptr ,alignment ,size )!= 0 ) { 60return NULL ; 61 } 62return ptr ; 63} 64 65static inline void * _util_sparse_array_node_data (uintptr_t handle ) { 66return (void * ) (handle & NODE_PTR_MASK ); 67} 68 69static inline unsigned _util_sparse_array_node_level (uintptr_t handle ) { 70return handle & NODE_LEVEL_MASK ; 71} 72 73static inline void _util_sparse_array_node_finish (util_sparse_array * arr ,uintptr_t node ) { 74if (_util_sparse_array_node_level (node )> 0 ) { 75uintptr_t * children = (uintptr_t * )_util_sparse_array_node_data (node ); 76size_t node_size = 1ull <<arr -> node_size_log2 ; 77for (size_t i = 0 ;i < node_size ;i ++ ) { 78if (children [i ]) { 79_util_sparse_array_node_finish (arr ,children [i ]); 80 } 81 } 82 } 83 84os_free_aligned (_util_sparse_array_node_data (node )); 85} 86 87static inline uintptr_t _util_sparse_array_node (void * data ,unsigned level ) { 88assert (data != NULL ); 89assert (((uintptr_t )data & NODE_LEVEL_MASK )== 0 ); 90assert ((level & NODE_PTR_MASK )== 0 ); 91return (uintptr_t )data |level ; 92} 93 94inline uintptr_t _util_sparse_array_node_alloc (util_sparse_array * arr ,unsigned level ) { 95size_t size ; 96if (level == 0 ) { 97size = arr -> elem_size <<arr -> node_size_log2 ; 98 }else { 99size = sizeof (uintptr_t ) <<arr -> node_size_log2 ; 100 } 101 102void * data = os_malloc_aligned (size ,NODE_ALLOC_ALIGN ); 103memset (data ,0 ,size ); 104 105return _util_sparse_array_node (data ,level ); 106} 107 108static inline uintptr_t _util_sparse_array_set_or_free_node (uintptr_t * node_ptr ,uintptr_t cmp_node ,uintptr_t node ) { 109uintptr_t prev_node = p_atomic_cmpxchg (node_ptr ,cmp_node ,node ); 110 111if (prev_node != cmp_node ) { 112/* We lost the race. Free this one and return the one that was already 113* allocated. 114*/ 115os_free_aligned (_util_sparse_array_node_data (node )); 116return prev_node ; 117 }else { 118return node ; 119 } 120} 121 122void * util_sparse_array_get (util_sparse_array * arr ,uint64_t idx ) { 123const unsigned node_size_log2 = arr -> node_size_log2 ; 124uintptr_t root = p_atomic_read (& arr -> root ); 125if (unlikely (!root )) { 126unsigned root_level = 0 ; 127uint64_t idx_iter = idx >>node_size_log2 ; 128while (idx_iter ) { 129idx_iter >>=node_size_log2 ; 130root_level ++ ; 131 } 132uintptr_t new_root = _util_sparse_array_node_alloc (arr ,root_level ); 133root = _util_sparse_array_set_or_free_node (& arr -> root ,NULL_NODE ,new_root ); 134 } 135 136while (1 ) { 137unsigned root_level = _util_sparse_array_node_level (root ); 138uint64_t root_idx = idx >> (root_level * node_size_log2 ); 139if (likely (root_idx < (1ull <<node_size_log2 ))) { 140break ; 141 } 142 143/* In this case, we have a root but its level is low enough that the 144* requested index is out-of-bounds. 145*/ 146uintptr_t new_root = _util_sparse_array_node_alloc (arr ,root_level + 1 ); 147 148uintptr_t * new_root_children = (uintptr_t * )_util_sparse_array_node_data (new_root ); 149new_root_children [0 ]= root ; 150 151/* We only add one at a time instead of the whole tree because it's 152* easier to ensure correctness of both the tree building and the 153* clean-up path. Because we're only adding one node we never have to 154* worry about trying to free multiple things without freeing the old 155* things. 156*/ 157root = _util_sparse_array_set_or_free_node (& arr -> root ,root ,new_root ); 158 } 159 160void * node_data = _util_sparse_array_node_data (root ); 161unsigned node_level = _util_sparse_array_node_level (root ); 162while (node_level > 0 ) { 163uint64_t child_idx = (idx >> (node_level * node_size_log2 ))& ((1ull <<node_size_log2 )- 1 ); 164 165uintptr_t * children = (uintptr_t * )node_data ; 166uintptr_t child = p_atomic_read (& children [child_idx ]); 167 168if (unlikely (!child )) { 169child = _util_sparse_array_node_alloc (arr ,node_level - 1 ); 170child = _util_sparse_array_set_or_free_node (& children [child_idx ],NULL_NODE ,child ); 171 } 172 173node_data = _util_sparse_array_node_data (child ); 174node_level = _util_sparse_array_node_level (child ); 175 } 176 177uint64_t elem_idx = idx & ((1ull <<node_size_log2 )- 1 ); 178return (void * ) ((char * )node_data + (elem_idx * arr -> elem_size )); 179}