diff --git a/core/iwasm/libraries/wasi-nn/src/wasi_nn.c b/core/iwasm/libraries/wasi-nn/src/wasi_nn.c index 15f5fb5509..724485449f 100644 --- a/core/iwasm/libraries/wasi-nn/src/wasi_nn.c +++ b/core/iwasm/libraries/wasi-nn/src/wasi_nn.c @@ -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 */