Skip to content

Commit 6bc7512

Browse files
committed
Move wasi_nn parameters parse logic into libc_wasi.c
1 parent 4b93110 commit 6bc7512

3 files changed

Lines changed: 134 additions & 103 deletions

File tree

core/iwasm/common/wasm_runtime_common.c

Lines changed: 58 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1696,65 +1696,6 @@ wasm_runtime_instantiation_args_destroy(struct InstantiationArgs2 *p)
16961696
wasm_runtime_free(p);
16971697
}
16981698

1699-
#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
1700-
typedef struct WASINNArguments WASINNArguments;
1701-
1702-
void
1703-
wasm_runtime_wasi_nn_graph_registry_args_set_defaults(WASINNArguments *args)
1704-
{
1705-
memset(args, 0, sizeof(*args));
1706-
}
1707-
1708-
bool
1709-
wasi_nn_graph_registry_set_args(WASINNArguments *registry, const char *encoding,
1710-
const char *target, uint32_t n_graphs,
1711-
const char **graph_paths)
1712-
{
1713-
if (!registry || !encoding || !target || !graph_paths) {
1714-
return false;
1715-
}
1716-
registry->encoding = strdup(encoding);
1717-
registry->target = strdup(target);
1718-
registry->n_graphs = n_graphs;
1719-
registry->graph_paths = (uint32_t **)malloc(sizeof(uint32_t *) * n_graphs);
1720-
memset(registry->graph_paths, 0, sizeof(uint32_t *) * n_graphs);
1721-
for (uint32_t i = 0; i < registry->n_graphs; i++)
1722-
registry->graph_paths[i] = strdup(graph_paths[i]);
1723-
1724-
return true;
1725-
}
1726-
1727-
int
1728-
wasi_nn_graph_registry_create(WASINNArguments **registryp)
1729-
{
1730-
WASINNArguments *args = wasm_runtime_malloc(sizeof(*args));
1731-
if (args == NULL) {
1732-
return -1;
1733-
}
1734-
wasm_runtime_wasi_nn_graph_registry_args_set_defaults(args);
1735-
*registryp = args;
1736-
return 0;
1737-
}
1738-
1739-
void
1740-
wasi_nn_graph_registry_destroy(WASINNArguments *registry)
1741-
{
1742-
if (registry) {
1743-
for (uint32_t i = 0; i < registry->n_graphs; i++)
1744-
if (registry->graph_paths[i]) {
1745-
// wasi_nn_graph_registry_unregister_graph(registry,
1746-
// registry->name[i]);
1747-
free(registry->graph_paths[i]);
1748-
}
1749-
if (registry->encoding)
1750-
free(registry->encoding);
1751-
if (registry->target)
1752-
free(registry->target);
1753-
free(registry);
1754-
}
1755-
}
1756-
#endif
1757-
17581699
void
17591700
wasm_runtime_instantiation_args_set_default_stack_size(
17601701
struct InstantiationArgs2 *p, uint32 v)
@@ -1853,7 +1794,65 @@ wasm_runtime_instantiation_args_set_wasi_ns_lookup_pool(
18531794
wasi_args->set_by_user = true;
18541795
}
18551796
#endif /* WASM_ENABLE_LIBC_WASI != 0 */
1797+
18561798
#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
1799+
typedef struct WASINNArguments WASINNArguments;
1800+
1801+
void
1802+
wasm_runtime_wasi_nn_graph_registry_args_set_defaults(WASINNArguments *args)
1803+
{
1804+
memset(args, 0, sizeof(*args));
1805+
}
1806+
1807+
bool
1808+
wasi_nn_graph_registry_set_args(WASINNArguments *registry, const char *encoding,
1809+
const char *target, uint32_t n_graphs,
1810+
const char **graph_paths)
1811+
{
1812+
if (!registry || !encoding || !target || !graph_paths) {
1813+
return false;
1814+
}
1815+
registry->encoding = strdup(encoding);
1816+
registry->target = strdup(target);
1817+
registry->n_graphs = n_graphs;
1818+
registry->graph_paths = (uint32_t **)malloc(sizeof(uint32_t *) * n_graphs);
1819+
memset(registry->graph_paths, 0, sizeof(uint32_t *) * n_graphs);
1820+
for (uint32_t i = 0; i < registry->n_graphs; i++)
1821+
registry->graph_paths[i] = strdup(graph_paths[i]);
1822+
1823+
return true;
1824+
}
1825+
1826+
int
1827+
wasi_nn_graph_registry_create(WASINNArguments **registryp)
1828+
{
1829+
WASINNArguments *args = wasm_runtime_malloc(sizeof(*args));
1830+
if (args == NULL) {
1831+
return -1;
1832+
}
1833+
wasm_runtime_wasi_nn_graph_registry_args_set_defaults(args);
1834+
*registryp = args;
1835+
return 0;
1836+
}
1837+
1838+
void
1839+
wasi_nn_graph_registry_destroy(WASINNArguments *registry)
1840+
{
1841+
if (registry) {
1842+
for (uint32_t i = 0; i < registry->n_graphs; i++)
1843+
if (registry->graph_paths[i]) {
1844+
// wasi_nn_graph_registry_unregister_graph(registry,
1845+
// registry->name[i]);
1846+
free(registry->graph_paths[i]);
1847+
}
1848+
if (registry->encoding)
1849+
free(registry->encoding);
1850+
if (registry->target)
1851+
free(registry->target);
1852+
free(registry);
1853+
}
1854+
}
1855+
18571856
void
18581857
wasm_runtime_instantiation_args_set_wasi_nn_graph_registry(
18591858
struct InstantiationArgs2 *p, WASINNArguments *registry)

product-mini/platforms/common/libc_wasi.c

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,13 @@ typedef struct {
2121
uint32 ns_lookup_pool_size;
2222
} libc_wasi_parse_context_t;
2323

24+
typedef struct {
25+
const char *encoding;
26+
const char *target;
27+
const char *graph_paths[10];
28+
uint32 n_graphs;
29+
} wasi_nn_parse_context_t;
30+
2431
typedef enum {
2532
LIBC_WASI_PARSE_RESULT_OK = 0,
2633
LIBC_WASI_PARSE_RESULT_NEED_HELP,
@@ -177,3 +184,58 @@ libc_wasi_set_init_args(struct InstantiationArgs2 *args, int argc, char **argv,
177184
wasm_runtime_instantiation_args_set_wasi_ns_lookup_pool(
178185
args, ctx->ns_lookup_pool, ctx->ns_lookup_pool_size);
179186
}
187+
188+
#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
189+
libc_wasi_parse_result_t
190+
wasi_nn_parse(char **argv, wasi_nn_parse_context_t *ctx)
191+
{
192+
char *token;
193+
char *saveptr = NULL;
194+
int token_count = 0, ret = 0;
195+
char *tokens[12] = { 0 };
196+
197+
// encoding:tensorflowlite|openvino|llama target:cpu|gpu|tpu
198+
// --wasi-nn-graph=encoding:target:model_file_path1:model_file_path2:model_file_path3:......
199+
token = strtok_r(argv[0] + 16, ":", &saveptr);
200+
while (token) {
201+
tokens[token_count] = token;
202+
token_count++;
203+
token = strtok_r(NULL, ":", &saveptr);
204+
}
205+
206+
if (token_count < 2) {
207+
ret = LIBC_WASI_PARSE_RESULT_NEED_HELP;
208+
goto fail;
209+
}
210+
211+
ctx->n_graphs = token_count - 2;
212+
ctx->encoding = strdup(tokens[0]);
213+
ctx->target = strdup(tokens[1]);
214+
for (int i = 0; i < ctx->n_graphs; i++) {
215+
ctx->graph_paths[i] = strdup(tokens[i + 2]);
216+
}
217+
218+
fail:
219+
if (token)
220+
free(token);
221+
222+
223+
return ret;
224+
}
225+
226+
static void
227+
wasi_nn_set_init_args(struct InstantiationArgs2 *args, struct WASINNArguments *nn_registry, wasi_nn_parse_context_t *ctx)
228+
{
229+
wasi_nn_graph_registry_create(&nn_registry);
230+
wasi_nn_graph_registry_set_args(nn_registry, ctx->encoding, ctx->target, ctx->n_graphs,
231+
ctx->graph_paths);
232+
wasm_runtime_instantiation_args_set_wasi_nn_graph_registry(args,
233+
nn_registry);
234+
235+
for (uint32_t i = 0; i < ctx->n_graphs; i++)
236+
if (ctx->graph_paths[i])
237+
free(ctx->graph_paths[i]);
238+
free(ctx->encoding);
239+
free(ctx->target);
240+
}
241+
#endif

product-mini/platforms/posix/main.c

Lines changed: 14 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -648,10 +648,10 @@ main(int argc, char *argv[])
648648
#endif
649649

650650
#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
651+
wasi_nn_parse_context_t wasi_nn_parse_ctx;
651652
struct WASINNArguments *nn_registry;
652-
char *encoding, *target;
653-
uint32_t n_models = 0;
654-
char **model_paths;
653+
654+
memset(&wasi_nn_parse_ctx, 0, sizeof(wasi_nn_parse_ctx));
655655
#endif
656656

657657
#if WASM_ENABLE_LIBC_WASI != 0
@@ -844,35 +844,18 @@ main(int argc, char *argv[])
844844
wasm_proposal_print_status();
845845
return 0;
846846
}
847-
#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
847+
#if WASM_ENABLE_LIBC_WASI != 0 && (WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0)
848848
else if (!strncmp(argv[0], "--wasi-nn-graph=", 16)) {
849-
char *token;
850-
char *saveptr = NULL;
851-
int token_count = 0;
852-
char *tokens[12] = { 0 };
853-
854-
// encoding:tensorflowlite|openvino|llama target:cpu|gpu|tpu
855-
// --wasi-nn-graph=encoding:target:model_file_path1:model_file_path2:model_file_path3:......
856-
token = strtok_r(argv[0] + 16, ":", &saveptr);
857-
while (token) {
858-
tokens[token_count] = token;
859-
token_count++;
860-
token = strtok_r(NULL, ":", &saveptr);
861-
}
862-
863-
if (token_count < 2) {
864-
return print_help();
865-
}
866-
867-
n_models = token_count - 2;
868-
encoding = strdup(tokens[0]);
869-
target = strdup(tokens[1]);
870-
model_paths = malloc(n_models * sizeof(void *));
871-
for (int i = 0; i < n_models; i++) {
872-
model_paths[i] = strdup(tokens[i + 2]);
849+
libc_wasi_parse_result_t result =
850+
wasi_nn_parse(argv, &wasi_nn_parse_ctx);
851+
switch (result) {
852+
case LIBC_WASI_PARSE_RESULT_OK:
853+
continue;
854+
case LIBC_WASI_PARSE_RESULT_NEED_HELP:
855+
return print_help();
856+
case LIBC_WASI_PARSE_RESULT_BAD_PARAM:
857+
return 1;
873858
}
874-
if (token)
875-
free(token);
876859
}
877860
#endif
878861
else {
@@ -1025,11 +1008,7 @@ main(int argc, char *argv[])
10251008
#endif
10261009

10271010
#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
1028-
wasi_nn_graph_registry_create(&nn_registry);
1029-
wasi_nn_graph_registry_set_args(nn_registry, encoding, target, n_models,
1030-
model_paths);
1031-
wasm_runtime_instantiation_args_set_wasi_nn_graph_registry(inst_args,
1032-
nn_registry);
1011+
wasi_nn_set_init_args(inst_args, nn_registry, &wasi_nn_parse_ctx);
10331012
#endif
10341013
/* instantiate the module */
10351014
wasm_module_inst = wasm_runtime_instantiate_ex2(
@@ -1149,15 +1128,6 @@ main(int argc, char *argv[])
11491128
#endif
11501129
#if WASM_ENABLE_DEBUG_INTERP != 0
11511130
fail4:
1152-
#endif
1153-
#if WASM_ENABLE_WASI_NN != 0 || WASM_ENABLE_WASI_EPHEMERAL_NN != 0
1154-
wasi_nn_graph_registry_destroy(nn_registry);
1155-
for (uint32_t i = 0; i < n_models; i++)
1156-
if (model_paths[i])
1157-
free(model_paths[i]);
1158-
free(model_paths);
1159-
free(encoding);
1160-
free(target);
11611131
#endif
11621132
/* destroy the module instance */
11631133
wasm_runtime_deinstantiate(wasm_module_inst);

0 commit comments

Comments
 (0)