Skip to content

wasi-nn: add minimum serialization on WASINNContext #4387

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 20, 2025
Merged

wasi-nn: add minimum serialization on WASINNContext #4387

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
149 changes: 116 additions & 33 deletions core/iwasm/libraries/wasi-nn/src/wasi_nn.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,13 +98,16 @@ wasi_nn_ctx_destroy(WASINNContext *wasi_nn_ctx)
NN_DBG_PRINTF("-> is_model_loaded: %d", wasi_nn_ctx->is_model_loaded);
NN_DBG_PRINTF("-> current_encoding: %d", wasi_nn_ctx->backend);

bh_assert(!wasi_nn_ctx->busy);

/* deinit() the backend */
if (wasi_nn_ctx->is_backend_ctx_initialized) {
wasi_nn_error res;
call_wasi_nn_func(wasi_nn_ctx->backend, deinit, res,
wasi_nn_ctx->backend_ctx);
}

os_mutex_destroy(&wasi_nn_ctx->lock);
wasm_runtime_free(wasi_nn_ctx);
}

Expand Down Expand Up @@ -150,6 +153,11 @@ wasi_nn_initialize_context()
}

memset(wasi_nn_ctx, 0, sizeof(WASINNContext));
if (os_mutex_init(&wasi_nn_ctx->lock)) {
NN_ERR_PRINTF("Error when initializing a lock for WASI-NN context");
wasm_runtime_free(wasi_nn_ctx);
return NULL;
}
return wasi_nn_ctx;
}

Expand All @@ -176,6 +184,35 @@ wasm_runtime_get_wasi_nn_ctx(wasm_module_inst_t instance)
return wasi_nn_ctx;
}

static WASINNContext *
lock_ctx(wasm_module_inst_t instance)
{
WASINNContext *wasi_nn_ctx = wasm_runtime_get_wasi_nn_ctx(instance);
if (wasi_nn_ctx == NULL) {
return NULL;
}
os_mutex_lock(&wasi_nn_ctx->lock);
if (wasi_nn_ctx->busy) {
os_mutex_unlock(&wasi_nn_ctx->lock);
return NULL;
}
wasi_nn_ctx->busy = true;
os_mutex_unlock(&wasi_nn_ctx->lock);
return wasi_nn_ctx;
}

static void
unlock_ctx(WASINNContext *wasi_nn_ctx)
{
if (wasi_nn_ctx == NULL) {
return;
}
os_mutex_lock(&wasi_nn_ctx->lock);
bh_assert(wasi_nn_ctx->busy);
wasi_nn_ctx->busy = false;
os_mutex_unlock(&wasi_nn_ctx->lock);
}

void
wasi_nn_destroy()
{
Expand Down Expand Up @@ -401,7 +438,7 @@ detect_and_load_backend(graph_encoding backend_hint,

static wasi_nn_error
ensure_backend(wasm_module_inst_t instance, graph_encoding encoding,
WASINNContext **wasi_nn_ctx_ptr)
WASINNContext *wasi_nn_ctx)
{
wasi_nn_error res;

Expand All @@ -412,7 +449,6 @@ ensure_backend(wasm_module_inst_t instance, graph_encoding encoding,
goto fail;
}

WASINNContext *wasi_nn_ctx = wasm_runtime_get_wasi_nn_ctx(instance);
if (wasi_nn_ctx->is_backend_ctx_initialized) {
if (wasi_nn_ctx->backend != loaded_backend) {
res = unsupported_operation;
Expand All @@ -430,7 +466,6 @@ ensure_backend(wasm_module_inst_t instance, graph_encoding encoding,

wasi_nn_ctx->is_backend_ctx_initialized = true;
}
*wasi_nn_ctx_ptr = wasi_nn_ctx;
return success;
fail:
return res;
Expand Down Expand Up @@ -458,17 +493,23 @@ wasi_nn_load(wasm_exec_env_t exec_env, graph_builder_array_wasm *builder,
if (!instance)
return runtime_error;

WASINNContext *wasi_nn_ctx = lock_ctx(instance);
if (wasi_nn_ctx == NULL) {
res = busy;
goto fail;
}

graph_builder_array builder_native = { 0 };
#if WASM_ENABLE_WASI_EPHEMERAL_NN != 0
if (success
!= (res = graph_builder_array_app_native(
instance, builder, builder_wasm_size, &builder_native)))
return res;
goto fail;
#else /* WASM_ENABLE_WASI_EPHEMERAL_NN == 0 */
if (success
!= (res = graph_builder_array_app_native(instance, builder,
&builder_native)))
return res;
goto fail;
#endif /* WASM_ENABLE_WASI_EPHEMERAL_NN != 0 */

if (!wasm_runtime_validate_native_addr(instance, g,
Expand All @@ -478,8 +519,7 @@ wasi_nn_load(wasm_exec_env_t exec_env, graph_builder_array_wasm *builder,
goto fail;
}

WASINNContext *wasi_nn_ctx;
res = ensure_backend(instance, encoding, &wasi_nn_ctx);
res = ensure_backend(instance, encoding, wasi_nn_ctx);
if (res != success)
goto fail;

Expand All @@ -494,6 +534,7 @@ wasi_nn_load(wasm_exec_env_t exec_env, graph_builder_array_wasm *builder,
// XXX: Free intermediate structure pointers
if (builder_native.buf)
wasm_runtime_free(builder_native.buf);
unlock_ctx(wasi_nn_ctx);

return res;
}
Expand Down Expand Up @@ -527,18 +568,26 @@ wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,

NN_DBG_PRINTF("[WASI NN] LOAD_BY_NAME %s...", name);

WASINNContext *wasi_nn_ctx;
res = ensure_backend(instance, autodetect, &wasi_nn_ctx);
WASINNContext *wasi_nn_ctx = lock_ctx(instance);
if (wasi_nn_ctx == NULL) {
res = busy;
goto fail;
}

res = ensure_backend(instance, autodetect, wasi_nn_ctx);
if (res != success)
return res;
goto fail;

call_wasi_nn_func(wasi_nn_ctx->backend, load_by_name, res,
wasi_nn_ctx->backend_ctx, name, name_len, g);
if (res != success)
return res;
goto fail;

wasi_nn_ctx->is_model_loaded = true;
return success;
res = success;
fail:
unlock_ctx(wasi_nn_ctx);
return res;
}

wasi_nn_error
Expand Down Expand Up @@ -576,19 +625,28 @@ wasi_nn_load_by_name_with_config(wasm_exec_env_t exec_env, char *name,

NN_DBG_PRINTF("[WASI NN] LOAD_BY_NAME_WITH_CONFIG %s %s...", name, config);

WASINNContext *wasi_nn_ctx;
res = ensure_backend(instance, autodetect, &wasi_nn_ctx);
WASINNContext *wasi_nn_ctx = lock_ctx(instance);
if (wasi_nn_ctx == NULL) {
res = busy;
goto fail;
}

res = ensure_backend(instance, autodetect, wasi_nn_ctx);
if (res != success)
return res;
goto fail;
;

call_wasi_nn_func(wasi_nn_ctx->backend, load_by_name_with_config, res,
wasi_nn_ctx->backend_ctx, name, name_len, config,
config_len, g);
if (res != success)
return res;
goto fail;

wasi_nn_ctx->is_model_loaded = true;
return success;
res = success;
fail:
unlock_ctx(wasi_nn_ctx);
return res;
}

wasi_nn_error
Expand All @@ -602,20 +660,27 @@ wasi_nn_init_execution_context(wasm_exec_env_t exec_env, graph g,
return runtime_error;
}

WASINNContext *wasi_nn_ctx = wasm_runtime_get_wasi_nn_ctx(instance);

wasi_nn_error res;
WASINNContext *wasi_nn_ctx = lock_ctx(instance);
if (wasi_nn_ctx == NULL) {
res = busy;
goto fail;
}

if (success != (res = is_model_initialized(wasi_nn_ctx)))
return res;
goto fail;

if (!wasm_runtime_validate_native_addr(
instance, ctx, (uint64)sizeof(graph_execution_context))) {
NN_ERR_PRINTF("ctx is invalid");
return invalid_argument;
res = invalid_argument;
goto fail;
}

call_wasi_nn_func(wasi_nn_ctx->backend, init_execution_context, res,
wasi_nn_ctx->backend_ctx, g, ctx);
fail:
unlock_ctx(wasi_nn_ctx);
return res;
}

Expand All @@ -630,25 +695,30 @@ wasi_nn_set_input(wasm_exec_env_t exec_env, graph_execution_context ctx,
return runtime_error;
}

WASINNContext *wasi_nn_ctx = wasm_runtime_get_wasi_nn_ctx(instance);

wasi_nn_error res;
WASINNContext *wasi_nn_ctx = lock_ctx(instance);
if (wasi_nn_ctx == NULL) {
res = busy;
goto fail;
}

if (success != (res = is_model_initialized(wasi_nn_ctx)))
return res;
goto fail;

tensor input_tensor_native = { 0 };
if (success
!= (res = tensor_app_native(instance, input_tensor,
&input_tensor_native)))
return res;
goto fail;

call_wasi_nn_func(wasi_nn_ctx->backend, set_input, res,
wasi_nn_ctx->backend_ctx, ctx, index,
&input_tensor_native);
// XXX: Free intermediate structure pointers
if (input_tensor_native.dimensions)
wasm_runtime_free(input_tensor_native.dimensions);

fail:
unlock_ctx(wasi_nn_ctx);
return res;
}

Expand All @@ -662,14 +732,20 @@ wasi_nn_compute(wasm_exec_env_t exec_env, graph_execution_context ctx)
return runtime_error;
}

WASINNContext *wasi_nn_ctx = wasm_runtime_get_wasi_nn_ctx(instance);

wasi_nn_error res;
WASINNContext *wasi_nn_ctx = lock_ctx(instance);
if (wasi_nn_ctx == NULL) {
res = busy;
goto fail;
}

if (success != (res = is_model_initialized(wasi_nn_ctx)))
return res;
goto fail;

call_wasi_nn_func(wasi_nn_ctx->backend, compute, res,
wasi_nn_ctx->backend_ctx, ctx);
fail:
unlock_ctx(wasi_nn_ctx);
return res;
}

Expand All @@ -692,16 +768,21 @@ wasi_nn_get_output(wasm_exec_env_t exec_env, graph_execution_context ctx,
return runtime_error;
}

WASINNContext *wasi_nn_ctx = wasm_runtime_get_wasi_nn_ctx(instance);

wasi_nn_error res;
WASINNContext *wasi_nn_ctx = lock_ctx(instance);
if (wasi_nn_ctx == NULL) {
res = busy;
goto fail;
}

if (success != (res = is_model_initialized(wasi_nn_ctx)))
return res;
goto fail;

if (!wasm_runtime_validate_native_addr(instance, output_tensor_size,
(uint64)sizeof(uint32_t))) {
NN_ERR_PRINTF("output_tensor_size is invalid");
return invalid_argument;
res = invalid_argument;
goto fail;
}

#if WASM_ENABLE_WASI_EPHEMERAL_NN != 0
Expand All @@ -714,6 +795,8 @@ wasi_nn_get_output(wasm_exec_env_t exec_env, graph_execution_context ctx,
wasi_nn_ctx->backend_ctx, ctx, index, output_tensor,
output_tensor_size);
#endif /* WASM_ENABLE_WASI_EPHEMERAL_NN != 0 */
fail:
unlock_ctx(wasi_nn_ctx);
return res;
}

Expand Down
4 changes: 4 additions & 0 deletions core/iwasm/libraries/wasi-nn/src/wasi_nn_private.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,11 @@
#include "wasi_nn_types.h"
#include "wasm_export.h"

#include "bh_platform.h"

typedef struct {
korp_mutex lock;
bool busy;
bool is_backend_ctx_initialized;
bool is_model_loaded;
graph_encoding backend;
Expand Down
Loading