Make wasi_nn_load_by_name and wasi_nn_load_by_name_with_config share a common logic.

Author: zhanheng1

core/iwasm/libraries/wasi-nn/src/wasi_nn.c

@@ -535,38 +535,13 @@ copyin_and_nul_terminate(wasm_module_inst_t inst, char *name, uint32_t name_len,
     return success;
 }
 
-wasi_nn_error
-wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
-                     graph *g)
+static wasi_nn_error
+load_by_name_with_optional_config(WASINNContext *wasi_nn_ctx,
+                                  wasm_module_inst_t instance, bool use_config,
+                                  graph *g, const char *model_name,
+                                  const char *config, int32_t config_len)
 {
-    WASINNContext *wasi_nn_ctx = NULL;
-    char *nul_terminated_name = NULL;
-    wasi_nn_error res;
-
-    wasm_module_inst_t instance = wasm_runtime_get_module_inst(exec_env);
-    if (!instance) {
-        return runtime_error;
-    }
-
-    if (!wasm_runtime_validate_native_addr(instance, g,
-                                           (uint64)sizeof(graph))) {
-        NN_ERR_PRINTF("graph is invalid");
-        return invalid_argument;
-    }
-
-    res = copyin_and_nul_terminate(instance, name, name_len,
-                                   &nul_terminated_name);
-    if (res != success) {
-        goto fail;
-    }
-
-    NN_DBG_PRINTF("[WASI NN] LOAD_BY_NAME %s...", nul_terminated_name);
-
-    wasi_nn_ctx = lock_ctx(instance);
-    if (wasi_nn_ctx == NULL) {
-        res = busy;
-        goto fail;
-    }
+    wasi_nn_error res = success;
 
     WASINNRegistry *wasi_nn_registry =
         wasm_runtime_get_wasi_nn_registry(instance);
@@ -580,9 +555,9 @@ wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
     uint32 model_idx = 0;
     uint32_t global_n_graphs = wasi_nn_registry->n_graphs;
     for (model_idx = 0; model_idx < global_n_graphs; model_idx++) {
-        char *model_name = wasi_nn_registry->model_names[model_idx];
+        char *model_name_i = wasi_nn_registry->model_names[model_idx];
 
-        if (strcmp(nul_terminated_name, model_name) != 0) {
+        if (strcmp(model_name, model_name_i) != 0) {
             continue;
         }
 
@@ -594,7 +569,6 @@ wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
         execution_target target =
             (execution_target)(wasi_nn_registry->target[model_idx]);
 
-        // res = ensure_backend(instance, autodetect, wasi_nn_ctx);
         res = ensure_backend(instance, encoding, wasi_nn_ctx);
         if (res != success)
             goto fail;
@@ -603,9 +577,17 @@ wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
             && (model_idx < global_n_graphs)) {
             NN_DBG_PRINTF(
                 "Model is not yet loaded, will add to global context");
-            call_wasi_nn_func(wasi_nn_ctx->backend, load_by_name, res,
-                              wasi_nn_ctx->backend_ctx, global_model_path_i,
-                              strlen(global_model_path_i), target, g);
+            if (use_config && config && config_len > 0) {
+                call_wasi_nn_func(
+                    wasi_nn_ctx->backend, load_by_name_with_config, res,
+                    wasi_nn_ctx->backend_ctx, global_model_path_i,
+                    strlen(global_model_path_i), config, config_len, g);
+            }
+            else {
+                call_wasi_nn_func(wasi_nn_ctx->backend, load_by_name, res,
+                                  wasi_nn_ctx->backend_ctx, global_model_path_i,
+                                  strlen(global_model_path_i), target, g);
+            }
             if (res != success)
                 goto fail;
 
@@ -627,9 +609,51 @@ wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
     else if (model_idx >= global_n_graphs) {
         NN_ERR_PRINTF("Model %s is not loaded, you should pass its path "
                       "through --wasi-nn-graph",
-                      nul_terminated_name);
+                      model_name);
         res = not_found;
     }
+
+fail:
+
+    return res;
+}
+
+wasi_nn_error
+wasi_nn_load_by_name(wasm_exec_env_t exec_env, char *name, uint32_t name_len,
+                     graph *g)
+{
+    WASINNContext *wasi_nn_ctx = NULL;
+    char *nul_terminated_name = NULL;
+    wasi_nn_error res;
+
+    wasm_module_inst_t instance = wasm_runtime_get_module_inst(exec_env);
+    if (!instance) {
+        return runtime_error;
+    }
+
+    if (!wasm_runtime_validate_native_addr(instance, g,
+                                           (uint64)sizeof(graph))) {
+        NN_ERR_PRINTF("graph is invalid");
+        return invalid_argument;
+    }
+
+    res = copyin_and_nul_terminate(instance, name, name_len,
+                                   &nul_terminated_name);
+    if (res != success) {
+        goto fail;
+    }
+
+    NN_DBG_PRINTF("[WASI NN] LOAD_BY_NAME %s...", nul_terminated_name);
+
+    wasi_nn_ctx = lock_ctx(instance);
+    if (wasi_nn_ctx == NULL) {
+        res = busy;
+        goto fail;
+    }
+
+    res = load_by_name_with_optional_config(wasi_nn_ctx, instance, false, g,
+                                            nul_terminated_name, NULL, 0);
+
 fail:
     if (nul_terminated_name != NULL) {
         wasm_runtime_free(nul_terminated_name);
@@ -686,9 +710,9 @@ wasi_nn_load_by_name_with_config(wasm_exec_env_t exec_env, char *name,
         goto fail;
     ;
 
-    call_wasi_nn_func(wasi_nn_ctx->backend, load_by_name_with_config, res,
-                      wasi_nn_ctx->backend_ctx, nul_terminated_name, name_len,
-                      nul_terminated_config, config_len, g);
+    res = load_by_name_with_optional_config(wasi_nn_ctx, instance, true, g,
+                                            nul_terminated_name,
+                                            nul_terminated_config, config_len);
     if (res != success)
         goto fail;