1 use super::PREOPENED_DIR_NAME;
2 use crate::check::artifacts_dir;
3 use std::path::Path;
4 use wasmtime::Result;
5 use wasmtime::{Config, Engine, Linker, Module, Store};
6 use wasmtime_wasi::p1::WasiP1Ctx;
7 use wasmtime_wasi::{DirPerms, FilePerms, WasiCtxBuilder};
8 use wasmtime_wasi_nn::{Backend, InMemoryRegistry, witx::WasiNnCtx};
9
10 /// Run a wasi-nn test program. This is modeled after
11 /// `crates/wasi/tests/all/main.rs` but still uses the older p1 API
12 /// for file reads.
run(path: &str, backend: Backend, preload_model: bool) -> Result<()>13 pub fn run(path: &str, backend: Backend, preload_model: bool) -> Result<()> {
14 let path = Path::new(path);
15 let engine = Engine::new(&Config::new())?;
16 let mut linker = Linker::new(&engine);
17 wasmtime_wasi_nn::witx::add_to_linker(&mut linker, |s: &mut Ctx| &mut s.wasi_nn)?;
18 wasmtime_wasi::p1::add_to_linker_sync(&mut linker, |s: &mut Ctx| &mut s.wasi)?;
19 let module = Module::from_file(&engine, path)?;
20 let mut store = Store::new(&engine, Ctx::new(&artifacts_dir(), preload_model, backend)?);
21 let instance = linker.instantiate(&mut store, &module)?;
22 let start = instance.get_typed_func::<(), ()>(&mut store, "_start")?;
23 start.call(&mut store, ())?;
24 Ok(())
25 }
26
27 /// The host state for running wasi-nn tests.
28 struct Ctx {
29 wasi: WasiP1Ctx,
30 wasi_nn: WasiNnCtx,
31 }
32
33 impl Ctx {
new(preopen_dir: &Path, preload_model: bool, mut backend: Backend) -> Result<Self>34 fn new(preopen_dir: &Path, preload_model: bool, mut backend: Backend) -> Result<Self> {
35 let mut builder = WasiCtxBuilder::new();
36 builder.inherit_stdio().preopened_dir(
37 preopen_dir,
38 PREOPENED_DIR_NAME,
39 DirPerms::READ,
40 FilePerms::READ,
41 )?;
42 let wasi = builder.build_p1();
43
44 let mut registry = InMemoryRegistry::new();
45 let mobilenet_dir = artifacts_dir();
46 if preload_model {
47 registry.load((backend).as_dir_loadable().unwrap(), &mobilenet_dir)?;
48 }
49 let wasi_nn = WasiNnCtx::new([backend], registry.into());
50
51 Ok(Self { wasi, wasi_nn })
52 }
53 }
54