sync up with latest wasi-nn spec (#3530)

This commit is contained in:
liang.he
2024-06-17 14:58:09 +08:00
committed by GitHub
parent d3e89895be
commit db025e457a
9 changed files with 209 additions and 101 deletions

View File

@ -29,6 +29,10 @@ load(graph_builder_array *builder, graph_encoding encoding,
execution_target target, graph *g)
__attribute__((import_module("wasi_nn")));
wasi_nn_error
load_by_name(const char *name, graph *g)
__attribute__((import_module("wasi_nn")));
/**
* INFERENCE
*

View File

@ -18,7 +18,9 @@ extern "C" {
*
*/
// Error codes returned by functions in this API.
// sync up with
// https://github.com/WebAssembly/wasi-nn/blob/main/wit/wasi-nn.wit#L136 Error
// codes returned by functions in this API.
typedef enum {
// No error occurred.
success = 0,
@ -26,12 +28,21 @@ typedef enum {
invalid_argument,
// Invalid encoding.
invalid_encoding,
// Caller module is missing a memory export.
missing_memory,
// Device or resource busy.
busy,
// The operation timed out.
timeout,
// Runtime Error.
runtime_error,
// Unsupported operation.
unsupported_operation,
// Graph is too large.
too_large,
// Graph not found.
not_found,
// The operation is insecure or has insufficient privilege to be performed.
// e.g., cannot access a hardware feature requested
security,
// The operation failed for an unspecified reason.
unknown,
} wasi_nn_error;
/**
@ -48,8 +59,14 @@ typedef struct {
uint32_t size;
} tensor_dimensions;
#if WASM_ENABLE_WASI_EPHEMERAL_NN != 0
// sync up with
// https://github.com/WebAssembly/wasi-nn/blob/main/wit/wasi-nn.wit#L27
// The type of the elements in a tensor.
typedef enum { fp16 = 0, fp32, fp64, bf16, u8, i32, i64 } tensor_type;
#else
typedef enum { fp16 = 0, fp32, up8, ip32 } tensor_type;
#endif /* WASM_ENABLE_WASI_EPHEMERAL_NN != 0 */
// The tensor data.
//
@ -96,6 +113,8 @@ typedef struct {
// An execution graph for performing inference (i.e., a model).
typedef uint32_t graph;
// sync up with
// https://github.com/WebAssembly/wasi-nn/blob/main/wit/wasi-nn.wit#L75
// Describes the encoding of the graph. This allows the API to be implemented by
// various backends that encode (i.e., serialize) their graph IR with different
// formats.
@ -105,7 +124,8 @@ typedef enum {
tensorflow,
pytorch,
tensorflowlite,
backend_amount
ggml,
autodetect,
} graph_encoding;
// Define where the graph should be executed.
@ -118,6 +138,7 @@ typedef uint32_t graph_execution_context;
typedef wasi_nn_error (*LOAD)(void *, graph_builder_array *, graph_encoding,
execution_target, graph *);
typedef wasi_nn_error (*LOAD_BY_NAME)(void *, const char *, uint32_t, graph *);
typedef wasi_nn_error (*INIT_EXECUTION_CONTEXT)(void *, graph,
graph_execution_context *);
typedef wasi_nn_error (*SET_INPUT)(void *, graph_execution_context, uint32_t,
@ -126,11 +147,12 @@ typedef wasi_nn_error (*COMPUTE)(void *, graph_execution_context);
typedef wasi_nn_error (*GET_OUTPUT)(void *, graph_execution_context, uint32_t,
tensor_data, uint32_t *);
/* wasi-nn general APIs */
typedef void (*BACKEND_INITIALIZE)(void **);
typedef void (*BACKEND_DEINITIALIZE)(void *);
typedef wasi_nn_error (*BACKEND_INITIALIZE)(void **);
typedef wasi_nn_error (*BACKEND_DEINITIALIZE)(void *);
typedef struct {
LOAD load;
LOAD_BY_NAME load_by_name;
INIT_EXECUTION_CONTEXT init_execution_context;
SET_INPUT set_input;
COMPUTE compute;
@ -140,7 +162,7 @@ typedef struct {
} api_function;
bool
wasi_nn_register_backend(graph_encoding backend_code, api_function apis);
wasi_nn_register_backend(api_function apis);
#ifdef __cplusplus
}