wasi-nn WASINNContext management seems broken

[yamt]

wasm_runtime_get_wasi_nn_ctx associates WASINNContext to wasm_module_inst_t
with a hash, using the pointer value of wasm_module_inst_t as a key.
the context and association is left as they are until runtime shutdown. (wasm_runtime_destroy)
it seems unsafe because a wasm_module_inst_t pointer does not identify a module instance.
once the pointer is freed by wasm_runtime_deinstantiate, the same pointer value can be used for an unrelated module instance.

[kylo5aby]

the root cause is the hash map approach it used isn't satisfied this scenario, and then it causes unsafe key and lifecycle mismatch.

iiuc, here we can replace the hash map approach with WAMR's native module instance context system, following the same pattern used by wasi itself.

in wasi_nn.c, use a safe context key:

static void *g_wasi_nn_context_key = NULL;

bool wasi_nn_initialize()
{
    // create context key with destructor callback
    g_wasi_nn_context_key = wasm_native_create_context_key(wasi_nn_context_dtor);
   ...
}

once module_inst is destroyed, callback wasi_nn_context_dtor will destroyed related wasi nn context.

rewrite context management:

static WASINNContext *
wasm_runtime_get_wasi_nn_ctx(wasm_module_inst_t instance)
{
    // use WAMR's native context system
    WASINNContext *wasi_nn_ctx = 
        (WASINNContext *)wasm_native_get_context(instance, g_wasi_nn_context_key);
    
    if (wasi_nn_ctx == NULL) {
         ...
        wasm_native_set_context(instance, g_wasi_nn_context_key, wasi_nn_ctx);
    }

    return wasi_nn_ctx;
}

do you have any suggestions?

[yamt]

iiuc, here we can replace the hash map approach with WAMR's native module instance context system, following the same pattern used by wasi itself.

i agree.
to take the approach, we should make wasi-nn host code thread safe first, because it would make a context shared among threads.
cf. #2430

[kylo5aby]

iiuc, here we can replace the hash map approach with WAMR's native module instance context system, following the same pattern used by wasi itself.

i agree. to take the approach, we should make wasi-nn host code thread safe first, because it would make a context shared among threads. cf. #2430

What are the usage scenarios of multi threads? for current implementation, wasi nn doesn't support multi threads, do you want to run several wasi nn app in wasm instances in several host threads, or any other usages?

[yamt]

iiuc, here we can replace the hash map approach with WAMR's native module instance context system, following the same pattern used by wasi itself.

i agree. to take the approach, we should make wasi-nn host code thread safe first, because it would make a context shared among threads. cf. #2430

What are the usage scenarios of multi threads? for current implementation, wasi nn doesn't support multi threads, do you want to run several wasi nn app in wasm instances in several host threads, or any other usages?

someone around me seemed to expect a wasi-nn model loaded in a thread could be used in another thread in the application.

having said that, my point was not about use cases.

if we use wasm_native_create_context_key here, we should either

• somehow prevent a context from being shared among threads.
• or, avoid thread-fragile implementation, probably by adding a few locks.