1 //! This module attempts to paper over the differences between the two 2 //! implementations of wasi-nn: the legacy WITX-based version (`mod witx`) and 3 //! the up-to-date WIT version (`mod wit`). Since the tests are mainly a simple 4 //! classifier, this exposes a high-level `classify` function to go along with 5 //! `load`, etc. 6 //! 7 //! This module exists solely for convenience--e.g., reduces test duplication. 8 //! In the future can be safely disposed of or altered as more tests are added. 9 10 /// Call `wasi-nn` functions from WebAssembly using the canonical ABI of the 11 /// component model via WIT-based tooling. Used by `bin/nn_wit_*.rs` tests. 12 pub mod wit { 13 use anyhow::{anyhow, Result}; 14 use std::time::Instant; 15 16 // Generate the wasi-nn bindings based on the `*.wit` files. 17 wit_bindgen::generate!({ 18 path: "../wasi-nn/wit", 19 world: "ml", 20 default_bindings_module: "test_programs::ml" 21 }); 22 use self::wasi::nn::errors; 23 use self::wasi::nn::graph::{self, Graph}; 24 pub use self::wasi::nn::graph::{ExecutionTarget, GraphEncoding}; // Used by tests. 25 use self::wasi::nn::tensor::{Tensor, TensorType}; 26 27 /// Load a wasi-nn graph from a set of bytes. 28 pub fn load( 29 bytes: &[Vec<u8>], 30 encoding: GraphEncoding, 31 target: ExecutionTarget, 32 ) -> Result<Graph> { 33 graph::load(bytes, encoding, target).map_err(err_as_anyhow) 34 } 35 36 /// Load a wasi-nn graph by name. 37 pub fn load_by_name(name: &str) -> Result<Graph> { 38 graph::load_by_name(name).map_err(err_as_anyhow) 39 } 40 41 /// Run a wasi-nn inference using a simple classifier model (single input, 42 /// single output). 43 pub fn classify(graph: Graph, input: (&str, Vec<u8>), output: &str) -> Result<Vec<f32>> { 44 let context = graph.init_execution_context().map_err(err_as_anyhow)?; 45 println!( 46 "[nn] created wasi-nn execution context with ID: {:?}", 47 context 48 ); 49 50 // Many classifiers have a single input; currently, this test suite also 51 // uses tensors of the same shape, though this is not usually the case. 52 let tensor = Tensor::new(&vec![1, 3, 224, 224], TensorType::Fp32, &input.1); 53 context.set_input(input.0, tensor).map_err(err_as_anyhow)?; 54 println!("[nn] set input tensor: {} bytes", input.1.len()); 55 56 let before = Instant::now(); 57 context.compute().map_err(err_as_anyhow)?; 58 println!( 59 "[nn] executed graph inference in {} ms", 60 before.elapsed().as_millis() 61 ); 62 63 // Many classifiers emit probabilities as floating point values; here we 64 // convert the raw bytes to `f32` knowing all models used here use that 65 // type. 66 let output = context.get_output(output).map_err(err_as_anyhow)?; 67 println!( 68 "[nn] retrieved output tensor: {} bytes", 69 output.data().len() 70 ); 71 let output: Vec<f32> = output 72 .data() 73 .chunks(4) 74 .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) 75 .collect(); 76 Ok(output) 77 } 78 79 fn err_as_anyhow(e: errors::Error) -> anyhow::Error { 80 anyhow!("error: {e:?}") 81 } 82 } 83 84 /// Call `wasi-nn` functions from WebAssembly using the legacy WITX-based 85 /// tooling. This older API has been deprecated for the newer WIT-based API but 86 /// retained for backwards compatibility testing--i.e., `bin/nn_witx_*.rs` 87 /// tests. 88 pub mod witx { 89 use anyhow::Result; 90 use std::time::Instant; 91 pub use wasi_nn::{ExecutionTarget, GraphEncoding}; 92 use wasi_nn::{Graph, GraphBuilder, TensorType}; 93 94 /// Load a wasi-nn graph from a set of bytes. 95 pub fn load( 96 bytes: &[&[u8]], 97 encoding: GraphEncoding, 98 target: ExecutionTarget, 99 ) -> Result<Graph> { 100 Ok(GraphBuilder::new(encoding, target).build_from_bytes(bytes)?) 101 } 102 103 /// Load a wasi-nn graph by name. 104 pub fn load_by_name( 105 name: &str, 106 encoding: GraphEncoding, 107 target: ExecutionTarget, 108 ) -> Result<Graph> { 109 Ok(GraphBuilder::new(encoding, target).build_from_cache(name)?) 110 } 111 112 /// Run a wasi-nn inference using a simple classifier model (single input, 113 /// single output). 114 pub fn classify(graph: Graph, tensor: Vec<u8>) -> Result<Vec<f32>> { 115 let mut context = graph.init_execution_context()?; 116 println!( 117 "[nn] created wasi-nn execution context with ID: {}", 118 context 119 ); 120 121 // Many classifiers have a single input; currently, this test suite also 122 // uses tensors of the same shape, though this is not usually the case. 123 context.set_input(0, TensorType::F32, &[1, 3, 224, 224], &tensor)?; 124 println!("[nn] set input tensor: {} bytes", tensor.len()); 125 126 let before = Instant::now(); 127 context.compute()?; 128 println!( 129 "[nn] executed graph inference in {} ms", 130 before.elapsed().as_millis() 131 ); 132 133 // Many classifiers emit probabilities as floating point values; here we 134 // convert the raw bytes to `f32` knowing all models used here use that 135 // type. 136 let mut output_buffer = vec![0u8; 1001 * std::mem::size_of::<f32>()]; 137 let num_bytes = context.get_output(0, &mut output_buffer)?; 138 println!("[nn] retrieved output tensor: {} bytes", num_bytes); 139 let output: Vec<f32> = output_buffer[..num_bytes] 140 .chunks(4) 141 .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) 142 .collect(); 143 Ok(output) 144 } 145 } 146 147 /// Sort some classification probabilities. 148 /// 149 /// Many classification models output a buffer of probabilities for each class, 150 /// placing the match probability for each class at the index for that class 151 /// (the probability of class `N` is stored at `probabilities[N]`). 152 pub fn sort_results(probabilities: &[f32]) -> Vec<InferenceResult> { 153 let mut results: Vec<InferenceResult> = probabilities 154 .iter() 155 .enumerate() 156 .map(|(c, p)| InferenceResult(c, *p)) 157 .collect(); 158 results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); 159 results 160 } 161 162 // A wrapper for class ID and match probabilities. 163 #[derive(Debug, PartialEq)] 164 pub struct InferenceResult(usize, f32); 165 impl InferenceResult { 166 pub fn class_id(&self) -> usize { 167 self.0 168 } 169 pub fn probability(&self) -> f32 { 170 self.1 171 } 172 } 173