Open abrown opened 1 year ago
The following is a sketch of what the higher-level API could look like:
struct Graph { ... }
impl Graph {
pub fn load(data: impl Iterator<Item = &[u8]>, encoding: GraphEncoding, target: ExecutionTarget) -> Result<Self>;
pub fn get_execution_context(&self) -> Result<ExecutionContext>;
pub fn get_input_types(&self) -> TensorTypes;
pub fn get_output_types(&self) -> TensorTypes
}
struct TensorTypes { ... }
impl TensorTypes {
pub fn len(&self) -> usize;
pub fn get(&self, index: u32) -> Result<TensorType>;
}
struct ExecutionContext { ... }
impl ExecutionContext {
pub fn set_input(&mut self, index: u32, tensor: Tensor) -> Result<()>;
pub fn compute(&mut self) -> Result<()>;
pub fn get_output(&self, index: u32, index: u32) -> Result<Tensor>;
}
Some considerations:
get_execution_context
must be a &mut self
, but it would be preferable if it were notExecutionContext
may need a lifetime parameter since a Tensor
may have an associated lifetimeget_output
consume self
since there may be multiple outputs that need to be retrievedunsafe
get_input_types_len
, get_input_type
, get_output_types_len
, get_output_type
)The functions defined in
wasi_nn::wasi_ephemeral_nn::*
are the raw Wasm symbols used to link with the host's implementation of wasi-nn. Thewasi-nn
Rust bindings should not expose these, since they should be providing a higher level of abstraction. These symbols are visible in the documentation here.
Once new bindings are created, the old low-level symbols should be removed from the public API.
We should also address https://github.com/bytecodealliance/wasi-nn/issues/13 as well. Being able to generate a compatible tensor server side seems like it will be the entry point to most workflows.
My initial thoughts are that wasi-nn shouldn't care about correctly implementing the mapping from images onto tensors, but we may want to provide a conversion api for bf16 if it is commonly used. I see no reason we can't provide helper functions for extracting tensor from the image, but they shouldn't be part of the spec.
We should also address https://github.com/bytecodealliance/wasi-nn/issues/13 as well.
@geekbeast, commented over in that issue.
The functions defined in
wasi_nn::wasi_ephemeral_nn::*
are the raw Wasm symbols used to link with the host's implementation of wasi-nn. Thewasi-nn
Rust bindings should not expose these, since they should be providing a higher level of abstraction. These symbols are visible in the documentation here.