wasm-micro-runtime/core/iwasm/libraries/wasi-nn/include/wasi_nn.h at 9d8e31a494c6b84d5c8a946ab176f96fef8f164f · bytecodealliance/wasm-micro-runtime · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
/*
 * Copyright (C) 2019 Intel Corporation.  All rights reserved.
 * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 */

/**
 * Following definition from:
 * [Oct 25th, 2022]
 * https://github.com/WebAssembly/wasi-nn/blob/0f77c48ec195748990ff67928a4b3eef5f16c2de/wasi-nn.wit.md
 */

#ifndef WASI_NN_H
#define WASI_NN_H

#include <stdint.h>
#include "wasi_nn_types.h"

#if WASM_ENABLE_WASI_EPHEMERAL_NN != 0
#define WASI_NN_IMPORT(name) \
    __attribute__((import_module("wasi_ephemeral_nn"), import_name(name)))
#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 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

/**
 * @brief Load an opaque sequence of bytes to use for inference.
 *
 * @param builder   Model builder.
 * @param builder_len The size of model builder.
 * @param encoding  Model encoding.
 * @param target    Execution target.
 * @param g         Graph.
 * @return wasi_nn_error    Execution status.
 */
#if WASM_ENABLE_WASI_EPHEMERAL_NN != 0
WASI_NN_ERROR_TYPE
WASI_NN_NAME(load)
(WASI_NN_NAME(graph_builder) * builder, uint32_t builder_len,
 WASI_NN_NAME(graph_encoding) encoding, WASI_NN_NAME(execution_target) target,
 WASI_NN_NAME(graph) * g) WASI_NN_IMPORT("load");
#else
WASI_NN_ERROR_TYPE
WASI_NN_NAME(load)
(WASI_NN_NAME(graph_builder_array) * builder,
 WASI_NN_NAME(graph_encoding) encoding, WASI_NN_NAME(execution_target) target,
 WASI_NN_NAME(graph) * g) WASI_NN_IMPORT("load");
#endif

WASI_NN_ERROR_TYPE
WASI_NN_NAME(load_by_name)
(const char *name, uint32_t name_len, WASI_NN_NAME(graph) * g)
    WASI_NN_IMPORT("load_by_name");

/**
 * INFERENCE
 *
 */

/**
 * @brief Create an execution instance of a loaded graph.
 *
 * @param g         Graph.
 * @param ctx       Execution context.
 * @return wasi_nn_error    Execution status.
 */
WASI_NN_ERROR_TYPE
WASI_NN_NAME(init_execution_context)
(WASI_NN_NAME(graph) g, WASI_NN_NAME(graph_execution_context) * ctx)
    WASI_NN_IMPORT("init_execution_context");

/**
 * @brief Define the inputs to use for inference.
 *
 * @param ctx       Execution context.
 * @param index     Input tensor index.
 * @param tensor    Input tensor.
 * @return wasi_nn_error    Execution status.
 */
WASI_NN_ERROR_TYPE
WASI_NN_NAME(set_input)
(WASI_NN_NAME(graph_execution_context) ctx, uint32_t index,
 WASI_NN_NAME(tensor) * tensor) WASI_NN_IMPORT("set_input");

/**
 * @brief Compute the inference on the given inputs.
 *
 * @param ctx       Execution context.
 * @return wasi_nn_error    Execution status.
 */
WASI_NN_ERROR_TYPE
WASI_NN_NAME(compute)
(WASI_NN_NAME(graph_execution_context) ctx) WASI_NN_IMPORT("compute");

/**
 * @brief Extract the outputs after inference.
 *
 * @param ctx                   Execution context.
 * @param index                 Output tensor index.
 * @param output_tensor         Buffer where output tensor with index `index` is
 * copied.
 * @param output_tensor_size    Pointer to `output_tensor` maximum size.
 *                              After the function call it is updated with the
 * copied number of bytes.
 * @return wasi_nn_error                Execution status.
 */
#if WASM_ENABLE_WASI_EPHEMERAL_NN != 0
WASI_NN_ERROR_TYPE
WASI_NN_NAME(get_output)
(WASI_NN_NAME(graph_execution_context) ctx, uint32_t index,
 uint8_t *output_tensor, uint32_t output_tensor_max_size,
 uint32_t *output_tensor_size) WASI_NN_IMPORT("get_output");
#else
WASI_NN_ERROR_TYPE
WASI_NN_NAME(get_output)
(graph_execution_context ctx, uint32_t index, uint8_t *output_tensor,
 uint32_t *output_tensor_size) WASI_NN_IMPORT("get_output");
#endif

#endif