add a set of apis to configure wasi-nn via InstantiationArgs2

core/iwasm/common/wasm_native.c

@@ -25,6 +25,10 @@ static NativeSymbolsList g_native_symbols_list = NULL;
 static void *g_wasi_context_key;
 #endif /* WASM_ENABLE_LIBC_WASI */
 
+#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
+static void *g_wasi_nn_registry_key;
+#endif
+
 uint32
 get_libc_builtin_export_apis(NativeSymbol **p_libc_builtin_apis);
 
@@ -473,6 +477,32 @@ wasi_context_dtor(WASMModuleInstanceCommon *inst, void *ctx)
 }
 #endif /* end of WASM_ENABLE_LIBC_WASI */
 
+#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
+WASINNRegistry *
+wasm_runtime_get_wasi_nn_registry(WASMModuleInstanceCommon *module_inst_comm)
+{
+    return wasm_native_get_context(module_inst_comm, g_wasi_nn_registry_key);
+}
+
+void
+wasm_runtime_set_wasi_nn_registry(WASMModuleInstanceCommon *module_inst_comm,
+                                  WASINNRegistry *wasi_nn_ctx)
+{
+    wasm_native_set_context(module_inst_comm, g_wasi_nn_registry_key,
+                            wasi_nn_ctx);
+}
+
+static void
+wasi_nn_context_dtor(WASMModuleInstanceCommon *inst, void *ctx)
+{
+    if (ctx == NULL) {
+        return;
+    }
+
+    wasm_runtime_wasi_nn_registry_destroy(ctx);
+}
+#endif
+
 #if WASM_ENABLE_QUICK_AOT_ENTRY != 0
 static bool
 quick_aot_entry_init(void);
@@ -582,6 +612,12 @@ wasm_native_init()
 #endif /* WASM_ENABLE_LIB_RATS */
 
 #if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
+    g_wasi_nn_registry_key =
+        wasm_native_create_context_key(wasi_nn_context_dtor);
+    if (g_wasi_nn_registry_key == NULL) {
+        goto fail;
+    }
+
     if (!wasi_nn_initialize())
         goto fail;
 
@@ -648,6 +684,10 @@ wasm_native_destroy()
 #endif
 
 #if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
+    if (g_wasi_nn_registry_key != NULL) {
+        wasm_native_destroy_context_key(g_wasi_nn_registry_key);
+        g_wasi_nn_registry_key = NULL;
+    }
     wasi_nn_destroy();
 #endif
 

core/iwasm/common/wasm_runtime_common.c

@@ -1795,6 +1795,109 @@ wasm_runtime_instantiation_args_set_wasi_ns_lookup_pool(
 }
 #endif /* WASM_ENABLE_LIBC_WASI != 0 */
 
+#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
+void
+wasm_runtime_wasi_nn_graph_registry_args_set_defaults(WASINNRegistry *args)
+{
+    memset(args, 0, sizeof(*args));
+}
+
+bool
+wasm_runtime_wasi_nn_registry_set_args(WASINNRegistry *registry,
+                                       const char **model_names,
+                                       const uint32_t **encoding,
+                                       const uint32_t **target,
+                                       uint32_t n_graphs,
+                                       const char **graph_paths)
+{
+    if (!registry || !model_names || !encoding || !target || !graph_paths) {
+        return false;
+    }
+
+    if ((sizeof(uint32_t *) * n_graphs) >= UINT32_MAX) {
+        LOG_ERROR("Invalid size for wasm_runtime_malloc.");
+        return NULL;
+    }
+
+    registry->n_graphs = n_graphs;
+    registry->target =
+        (uint32_t **)wasm_runtime_malloc(sizeof(uint32_t *) * n_graphs);
+    registry->encoding =
+        (uint32_t **)wasm_runtime_malloc(sizeof(uint32_t *) * n_graphs);
+    registry->loaded =
+        (uint32_t **)wasm_runtime_malloc(sizeof(uint32_t *) * n_graphs);
+    registry->model_names =
+        (uint32_t **)wasm_runtime_malloc(sizeof(uint32_t *) * n_graphs);
+    registry->graph_paths =
+        (uint32_t **)wasm_runtime_malloc(sizeof(uint32_t *) * n_graphs);
+    memset(registry->target, 0, sizeof(uint32_t *) * n_graphs);
+    memset(registry->encoding, 0, sizeof(uint32_t *) * n_graphs);
+    memset(registry->loaded, 0, sizeof(uint32_t *) * n_graphs);
+    memset(registry->model_names, 0, sizeof(uint32_t *) * n_graphs);
+    memset(registry->graph_paths, 0, sizeof(uint32_t *) * n_graphs);
+
+    for (uint32_t i = 0; i < registry->n_graphs; i++) {
+        registry->graph_paths[i] = bh_strdup(graph_paths[i]);
+        registry->model_names[i] = bh_strdup(model_names[i]);
+        registry->encoding[i] = encoding[i];
+        registry->target[i] = target[i];
+    }
+
+    return true;
+}
+
+int
+wasm_runtime_wasi_nn_registry_create(WASINNRegistry **registryp)
+{
+    WASINNRegistry *args = wasm_runtime_malloc(sizeof(*args));
+    if (args == NULL) {
+        return -1;
+    }
+    wasm_runtime_wasi_nn_graph_registry_args_set_defaults(args);
+    *registryp = args;
+    return 0;
+}
+
+void
+wasm_runtime_wasi_nn_registry_destroy(WASINNRegistry *registry)
+{
+    if (registry) {
+        for (uint32_t i = 0; i < registry->n_graphs; i++)
+            if (registry->graph_paths[i]) {
+                wasm_runtime_free(registry->graph_paths[i]);
+                if (registry->model_names[i])
+                    wasm_runtime_free(registry->model_names[i]);
+            }
+        if (registry->encoding)
+            wasm_runtime_free(registry->encoding);
+        if (registry->target)
+            wasm_runtime_free(registry->target);
+        if (registry->loaded)
+            wasm_runtime_free(registry->loaded);
+        wasm_runtime_free(registry);
+    }
+}
+
+void
+wasm_runtime_instantiation_args_set_wasi_nn_registry(
+    struct InstantiationArgs2 *p, WASINNRegistry *registry)
+{
+    if (!registry)
+        return;
+    WASINNRegistry *wasi_nn_registry = &p->nn_registry;
+
+    wasi_nn_registry->n_graphs = registry->n_graphs;
+
+    for (uint32_t i = 0; i < registry->n_graphs; i++) {
+        registry->graph_paths[i] = bh_strdup(registry->graph_paths[i]);
+        registry->model_names[i] = bh_strdup(registry->model_names[i]);
+        registry->encoding[i] = registry->encoding[i];
+        registry->target[i] = registry->target[i];
+        wasi_nn_registry->loaded = registry->loaded;
+    }
+}
+#endif
+
 WASMModuleInstanceCommon *
 wasm_runtime_instantiate_ex2(WASMModuleCommon *module,
                              const struct InstantiationArgs2 *args,

core/iwasm/common/wasm_runtime_common.h

@@ -545,6 +545,18 @@ typedef struct WASMModuleInstMemConsumption {
     uint32 exports_size;
 } WASMModuleInstMemConsumption;
 
+#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
+typedef struct WASINNRegistry {
+    char **model_names;
+    uint32_t **encoding;
+    uint32_t **target;
+
+    uint32_t n_graphs;
+    uint32_t **loaded;
+    char **graph_paths;
+} WASINNRegistry;
+#endif
+
 #if WASM_ENABLE_LIBC_WASI != 0
 #if WASM_ENABLE_UVWASI == 0
 typedef struct WASIContext {
@@ -612,11 +624,22 @@ WASMExecEnv *
 wasm_runtime_get_exec_env_tls(void);
 #endif
 
+#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
+WASM_RUNTIME_API_EXTERN int
+wasm_runtime_wasi_nn_registry_create(WASINNRegistry **registryp);
+
+WASM_RUNTIME_API_EXTERN void
+wasm_runtime_wasi_nn_registry_destroy(WASINNRegistry *registry);
+#endif
+
 struct InstantiationArgs2 {
     InstantiationArgs v1;
 #if WASM_ENABLE_LIBC_WASI != 0
     WASIArguments wasi;
 #endif
+#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
+    WASINNRegistry nn_registry;
+#endif
 };
 
 void
@@ -775,6 +798,20 @@ wasm_runtime_instantiation_args_set_wasi_ns_lookup_pool(
     struct InstantiationArgs2 *p, const char *ns_lookup_pool[],
     uint32 ns_lookup_pool_size);
 
+#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
+WASM_RUNTIME_API_EXTERN void
+wasm_runtime_instantiation_args_set_wasi_nn_registry(
+    struct InstantiationArgs2 *p, WASINNRegistry *registry);
+
+WASM_RUNTIME_API_EXTERN bool
+wasm_runtime_wasi_nn_registry_set_args(WASINNRegistry *registry,
+                                       const char **model_names,
+                                       const uint32_t **encoding,
+                                       const uint32_t **target,
+                                       uint32_t n_graphs,
+                                       const char **graph_paths);
+#endif
+
 /* See wasm_export.h for description */
 WASM_RUNTIME_API_EXTERN WASMModuleInstanceCommon *
 wasm_runtime_instantiate_ex2(WASMModuleCommon *module,
@@ -1427,6 +1464,19 @@ wasm_runtime_check_and_update_last_used_shared_heap(
     uint8 **shared_heap_base_addr_adj_p, bool is_memory64);
 #endif
 
+#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
+WASM_RUNTIME_API_EXTERN void
+wasm_runtime_set_wasi_nn_registry(WASMModuleInstanceCommon *module_inst,
+                                  WASINNRegistry *wasi_ctx);
+
+WASM_RUNTIME_API_EXTERN WASINNRegistry *
+wasm_runtime_get_wasi_nn_registry(WASMModuleInstanceCommon *module_inst_comm);
+
+WASM_RUNTIME_API_EXTERN void
+wasm_runtime_set_wasi_nn_registry(WASMModuleInstanceCommon *module_inst_comm,
+                                  WASINNRegistry *wasi_nn_ctx);
+#endif
+
 #ifdef __cplusplus
 }
 #endif

core/iwasm/include/wasm_export.h

@@ -290,6 +290,7 @@ typedef struct InstantiationArgs {
 #endif /* INSTANTIATION_ARGS_OPTION_DEFINED */
 
 struct InstantiationArgs2;
+typedef struct WASINNRegistry WASINNRegistry;
 
 #ifndef WASM_VALKIND_T_DEFINED
 #define WASM_VALKIND_T_DEFINED
@@ -796,6 +797,19 @@ wasm_runtime_instantiation_args_set_wasi_ns_lookup_pool(
     struct InstantiationArgs2 *p, const char *ns_lookup_pool[],
     uint32_t ns_lookup_pool_size);
 
+WASM_RUNTIME_API_EXTERN struct WASINNRegistry *
+wasm_runtime_get_wasi_nn_registry(const wasm_module_inst_t module_inst);
+
+WASM_RUNTIME_API_EXTERN void
+wasm_runtime_set_wasi_nn_registry(wasm_module_inst_t module_inst,
+                                  struct WASINNRegistry *wasi_ctx);
+
+WASM_RUNTIME_API_EXTERN int
+wasm_runtime_wasi_nn_registry_create(struct WASINNRegistry **registryp);
+
+WASM_RUNTIME_API_EXTERN void
+wasm_runtime_wasi_nn_registry_destroy(struct WASINNRegistry *registry);
+
 /**
  * Instantiate a WASM module, with specified instantiation arguments
  *

core/iwasm/libraries/wasi-nn/include/wasi_ephemeral_nn.h

@@ -7,6 +7,3 @@
 #define WASI_NN_NAME(name) wasi_ephemeral_nn_##name
 
 #include "wasi_nn.h"
-
-#undef WASM_ENABLE_WASI_EPHEMERAL_NN
-#undef WASI_NN_NAME

core/iwasm/libraries/wasi-nn/include/wasi_nn.h

@@ -21,7 +21,7 @@
 #else
 #define WASI_NN_IMPORT(name) \
     __attribute__((import_module("wasi_nn"), import_name(name)))
-#warning You are using "wasi_nn", which is a legacy WAMR-specific ABI. It's deperecated and will likely be removed in future versions of WAMR. Please use "wasi_ephemeral_nn" instead. (For a WASM module, use the wasi_ephemeral_nn.h header instead. For the runtime configurations, enable WASM_ENABLE_WASI_EPHEMERAL_NN/WAMR_BUILD_WASI_EPHEMERAL_NN.)
+#warning You are using "wasi_nn", which is a legacy WAMR-specific ABI. It is deperecated and will likely be removed in future versions of WAMR. Please use "wasi_ephemeral_nn" instead. (For a WASM module, use the wasi_ephemeral_nn.h header instead. For the runtime configurations, enable WASM_ENABLE_WASI_EPHEMERAL_NN/WAMR_BUILD_WASI_EPHEMERAL_NN.)
 #endif
 
 /**

core/iwasm/libraries/wasi-nn/include/wasi_nn_types.h

@@ -27,7 +27,7 @@ extern "C" {
 #define WASI_NN_TYPE_NAME(name) WASI_NN_NAME(type_##name)
 #define WASI_NN_ENCODING_NAME(name) WASI_NN_NAME(encoding_##name)
 #define WASI_NN_TARGET_NAME(name) WASI_NN_NAME(target_##name)
-#define WASI_NN_ERROR_TYPE WASI_NN_NAME(error);
+#define WASI_NN_ERROR_TYPE WASI_NN_NAME(error)
 #endif
 
 /**
@@ -169,6 +169,7 @@ typedef enum WASI_NN_NAME(execution_target) {
     WASI_NN_TARGET_NAME(cpu) = 0,
     WASI_NN_TARGET_NAME(gpu),
     WASI_NN_TARGET_NAME(tpu),
+    WASI_NN_TARGET_NAME(unsupported_target),
 } WASI_NN_NAME(execution_target);
 
 // Bind a `graph` to the input and output tensors for an inference.