3131#define LLAMACPP_BACKEND_LIB "libwasi_nn_llamacpp" LIB_EXTENTION
3232
3333/* Global variables */
34+ static korp_mutex wasi_nn_lock ;
35+ /*
36+ * the "lookup" table is protected by wasi_nn_lock.
37+ *
38+ * an exception: during wasm_runtime_destroy, wasi_nn_destroy tears down
39+ * the table without acquiring the lock. it's ok because there should be
40+ * no other threads using the runtime at this point.
41+ */
3442struct backends_api_functions {
3543 void * backend_handle ;
3644 api_function functions ;
@@ -109,12 +117,18 @@ wasi_nn_initialize()
109117{
110118 NN_DBG_PRINTF ("[WASI NN General] Initializing wasi-nn" );
111119
120+ if (os_mutex_init (& wasi_nn_lock )) {
121+ NN_ERR_PRINTF ("Error while initializing global lock" );
122+ return false;
123+ }
124+
112125 // hashmap { instance: wasi_nn_ctx }
113126 hashmap = bh_hash_map_create (HASHMAP_INITIAL_SIZE , true, hash_func ,
114127 key_equal_func , key_destroy_func ,
115128 value_destroy_func );
116129 if (hashmap == NULL ) {
117130 NN_ERR_PRINTF ("Error while initializing hashmap" );
131+ os_mutex_destroy (& wasi_nn_lock );
118132 return false;
119133 }
120134
@@ -175,6 +189,8 @@ wasi_nn_destroy()
175189
176190 memset (& lookup [i ].functions , 0 , sizeof (api_function ));
177191 }
192+
193+ os_mutex_destroy (& wasi_nn_lock );
178194}
179195
180196/* Utils */
@@ -349,6 +365,8 @@ static bool
349365detect_and_load_backend (graph_encoding backend_hint ,
350366 graph_encoding * loaded_backend )
351367{
368+ bool ret ;
369+
352370 if (backend_hint > autodetect )
353371 return false;
354372
@@ -360,16 +378,23 @@ detect_and_load_backend(graph_encoding backend_hint,
360378
361379 * loaded_backend = backend_hint ;
362380
381+ os_mutex_lock (& wasi_nn_lock );
363382 /* if already loaded */
364- if (lookup [backend_hint ].backend_handle )
383+ if (lookup [backend_hint ].backend_handle ) {
384+ os_mutex_unlock (& wasi_nn_lock );
365385 return true;
386+ }
366387
367388 const char * backend_lib_name =
368389 graph_encoding_to_backend_lib_name (backend_hint );
369- if (!backend_lib_name )
390+ if (!backend_lib_name ) {
391+ os_mutex_unlock (& wasi_nn_lock );
370392 return false;
393+ }
371394
372- return prepare_backend (backend_lib_name , lookup + backend_hint );
395+ ret = prepare_backend (backend_lib_name , lookup + backend_hint );
396+ os_mutex_unlock (& wasi_nn_lock );
397+ return ret ;
373398}
374399
375400/* WASI-NN implementation */
0 commit comments