wasmedge_nn/nn/
mod.rs

1pub mod backend;
2pub mod ctx;
3
4use std::fmt;
5
6pub type Tensor<'a> = wasi_nn::Tensor<'a>;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum GraphEncoding {
10    Openvino,
11}
12impl From<GraphEncoding> for wasi_nn::GraphEncoding {
13    fn from(encoding: GraphEncoding) -> Self {
14        match encoding {
15            GraphEncoding::Openvino => wasi_nn::GRAPH_ENCODING_OPENVINO,
16        }
17    }
18}
19impl fmt::Display for GraphEncoding {
20    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
21        match self {
22            GraphEncoding::Openvino => write!(f, "openvino"),
23        }
24    }
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum ExecutionTarget {
29    CPU,
30    GPU,
31}
32impl From<ExecutionTarget> for wasi_nn::ExecutionTarget {
33    fn from(target: ExecutionTarget) -> Self {
34        match target {
35            ExecutionTarget::CPU => wasi_nn::EXECUTION_TARGET_CPU,
36            ExecutionTarget::GPU => wasi_nn::EXECUTION_TARGET_GPU,
37        }
38    }
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum Dtype {
43    F32,
44    U8,
45}
46impl Dtype {
47    /// Returns the number of bytes occupied by the dtype.
48    pub fn bytes(&self) -> usize {
49        match self {
50            Dtype::U8 => 1,
51            Dtype::F32 => 4,
52        }
53    }
54}
55impl From<Dtype> for wasi_nn::TensorType {
56    fn from(dtype: Dtype) -> Self {
57        match dtype {
58            Dtype::F32 => wasi_nn::TENSOR_TYPE_F32,
59            Dtype::U8 => wasi_nn::TENSOR_TYPE_U8,
60        }
61    }
62}
63
64#[derive(Debug, Clone, Copy)]
65pub struct Dimension {
66    pub height: u32,
67    pub width: u32,
68    pub channels: u32,
69    pub dtype: Dtype,
70}
71impl Dimension {
72    pub fn new(height: u32, width: u32, channels: u32, dtype: Dtype) -> Self {
73        Self {
74            height,
75            width,
76            channels,
77            dtype,
78        }
79    }
80}
wasmedge_nn/nn/mod.rs

wasmedge_nn/nn/
mod.rs
