gnh1201/wasm-micro-runtime
1
0
Fork 0
mirror of https://github.com/bytecodealliance/wasm-micro-runtime.git synced 2025-06-18 02:59:21 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

wasi-nn: protect the backend lookup table with a lock (#4319)

Browse Source 
this would avoid potential issues when multiple instances happen to
make an attempt to load a backend at the same time.

Fixes: https://github.com/bytecodealliance/wasm-micro-runtime/issues/4314
This commit is contained in:
YAMAMOTO Takashi 2025-06-05 10:19:46 +09:00 committed by GitHub
parent 93ef19b0ca
commit 85efe08431
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 28 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
core/iwasm/libraries/wasi-nn/src/wasi_nn.c
View File

@ -31,6 +31,14 @@
#define LLAMACPP_BACKEND_LIB "libwasi_nn_llamacpp" LIB_EXTENTION
/* Global variables */
static korp_mutex wasi_nn_lock;
/*
* the "lookup" table is protected by wasi_nn_lock.
*
* an exception: during wasm_runtime_destroy, wasi_nn_destroy tears down
* the table without acquiring the lock. it's ok because there should be
* no other threads using the runtime at this point.
*/
struct backends_api_functions {
void *backend_handle;
api_function functions;
@ -109,12 +117,18 @@ wasi_nn_initialize()
{
NN_DBG_PRINTF("[WASI NN General] Initializing wasi-nn");
if (os_mutex_init(&wasi_nn_lock)) {
NN_ERR_PRINTF("Error while initializing global lock");
return false;
}
// hashmap { instance: wasi_nn_ctx }
hashmap = bh_hash_map_create(HASHMAP_INITIAL_SIZE, true, hash_func,
key_equal_func, key_destroy_func,
value_destroy_func);
if (hashmap == NULL) {
NN_ERR_PRINTF("Error while initializing hashmap");
os_mutex_destroy(&wasi_nn_lock);
return false;
}
@ -175,6 +189,8 @@ wasi_nn_destroy()
memset(&lookup[i].functions, 0, sizeof(api_function));
}
os_mutex_destroy(&wasi_nn_lock);
}
/* Utils */
@ -349,6 +365,8 @@ static bool
detect_and_load_backend(graph_encoding backend_hint,
graph_encoding *loaded_backend)
{
bool ret;
if (backend_hint > autodetect)
return false;
@ -360,16 +378,23 @@ detect_and_load_backend(graph_encoding backend_hint,
*loaded_backend = backend_hint;
os_mutex_lock(&wasi_nn_lock);
/* if already loaded */
if (lookup[backend_hint].backend_handle)
if (lookup[backend_hint].backend_handle) {
os_mutex_unlock(&wasi_nn_lock);
return true;
}
const char *backend_lib_name =
graph_encoding_to_backend_lib_name(backend_hint);
if (!backend_lib_name)
if (!backend_lib_name) {
os_mutex_unlock(&wasi_nn_lock);
return false;
}
return prepare_backend(backend_lib_name, lookup + backend_hint);
ret = prepare_backend(backend_lib_name, lookup + backend_hint);
os_mutex_unlock(&wasi_nn_lock);
return ret;
}
/* WASI-NN implementation */