1919
2020use anyhow:: Result ;
2121use wasmtime:: * ;
22- use wasmtime_wasi:: { Wasi , WasiCtx } ;
22+ use wasmtime_wasi:: { WasiCtx , WasiCtxBuilder } ;
2323
2424use super :: Tensor ;
2525
2626pub struct GraphExecutor {
2727 pub ( crate ) wasm_addr : i32 ,
2828 pub ( crate ) input_size : i32 ,
2929 pub ( crate ) output_size : i32 ,
30+ pub ( crate ) store : Option < Store < WasiCtx > > ,
31+ // None-WASI version:
32+ // pub(crate) store: Option<Store<()>>,
3033 pub ( crate ) instance : Option < Instance > ,
3134}
3235
@@ -37,25 +40,44 @@ impl GraphExecutor {
3740 wasm_addr : 0 ,
3841 input_size : 0 ,
3942 output_size : 0 ,
43+ store : None ,
4044 instance : None ,
4145 }
4246 }
4347
4448 pub fn instantiate ( & mut self , wasm_graph_file : String ) -> Result < ( ) > {
45- let engine = Engine :: new ( Config :: new ( ) . wasm_simd ( true ) ) ;
46- let store = Store :: new ( & engine) ;
49+ // It seems WASI in this example is not necessary
4750
51+ // None WASI version: works with no SIMD
52+ // let engine = Engine::new(Config::new().wasm_simd(true)).unwrap();
53+ // let mut store = Store::new(&engine, ());
54+ // let module = Module::from_file(store.engine(), &wasm_graph_file)?;
55+
56+ // let instance = Instance::new(&mut store, &module, &[])?;
57+
58+ // self.instance = Some(instance);
59+ // self.store = Some(store);
60+
61+ // Ok(())
62+
63+ // WASI version:
64+ let engine = Engine :: new ( Config :: new ( ) . wasm_simd ( true ) ) . unwrap ( ) ;
4865 // First set up our linker which is going to be linking modules together. We
4966 // want our linker to have wasi available, so we set that up here as well.
50- let mut linker = Linker :: new ( & store) ;
67+ let mut linker = Linker :: new ( & engine) ;
68+ wasmtime_wasi:: add_to_linker ( & mut linker, |s| s) ?;
5169 // Create an instance of `Wasi` which contains a `WasiCtx`. Note that
5270 // `WasiCtx` provides a number of ways to configure what the target program
5371 // will have access to.
54- let wasi = Wasi :: new ( & store, WasiCtx :: new ( std:: env:: args ( ) ) ?) ;
55- wasi. add_to_linker ( & mut linker) ?;
72+ let wasi = WasiCtxBuilder :: new ( )
73+ . inherit_stdio ( )
74+ . inherit_args ( ) ?
75+ . build ( ) ;
76+ let mut store = Store :: new ( & engine, wasi) ;
5677
57- let module = Module :: from_file ( & store, & wasm_graph_file) ?;
58- self . instance = Some ( linker. instantiate ( & module) ?) ;
78+ let module = Module :: from_file ( & engine, & wasm_graph_file) ?;
79+ self . instance = Some ( linker. instantiate ( & mut store, & module) ?) ;
80+ self . store = Some ( store) ;
5981
6082 Ok ( ( ) )
6183 }
@@ -65,26 +87,24 @@ impl GraphExecutor {
6587 . instance
6688 . as_ref ( )
6789 . unwrap ( )
68- . get_memory ( "memory" )
90+ . get_memory ( self . store . as_mut ( ) . unwrap ( ) , "memory" )
6991 . ok_or_else ( || anyhow:: format_err!( "failed to find `memory` export" ) ) ?;
7092
7193 // Specify the wasm address to access the wasm memory.
72- let wasm_addr = memory. data_size ( ) ;
94+ let wasm_addr = memory. data_size ( self . store . as_mut ( ) . unwrap ( ) ) ;
95+
7396 // Serialize the data into a JSON string.
7497 let in_data = serde_json:: to_vec ( & input_data) ?;
7598 let in_size = in_data. len ( ) ;
99+
76100 // Grow up memory size according to in_size to avoid memory leak.
77- memory. grow ( ( in_size >> 16 ) as u32 + 1 ) ?;
101+ memory. grow ( self . store . as_mut ( ) . unwrap ( ) , ( in_size >> 16 ) as u32 + 1 ) ?;
78102
79- // Insert the input data into wasm memory.
80- for i in 0 ..in_size {
81- unsafe {
82- memory. data_unchecked_mut ( ) [ wasm_addr + i] = * in_data. get ( i) . unwrap ( ) ;
83- }
84- }
103+ memory. write ( self . store . as_mut ( ) . unwrap ( ) , wasm_addr, & in_data) ?;
85104
86105 self . wasm_addr = wasm_addr as i32 ;
87106 self . input_size = in_size as i32 ;
107+
88108 Ok ( ( ) )
89109 }
90110
@@ -94,11 +114,12 @@ impl GraphExecutor {
94114 . instance
95115 . as_ref ( )
96116 . unwrap ( )
97- . get_func ( "run" )
98- . ok_or_else ( || anyhow:: format_err!( "failed to find `run` function export!" ) ) ?
99- . get2 :: < i32 , i32 , i32 > ( ) ?;
117+ . get_func ( self . store . as_mut ( ) . unwrap ( ) , "run" )
118+ . ok_or_else ( || anyhow:: format_err!( "failed to find `run` function export!" ) ) ?;
100119
101- let out_size = run ( self . wasm_addr , self . input_size ) ?;
120+ let params = [ Val :: I32 ( self . wasm_addr ) , Val :: I32 ( self . input_size ) ] ;
121+ let out_size = run. call ( self . store . as_mut ( ) . unwrap ( ) , & params[ ..] ) ?;
122+ let out_size = ( * out_size) [ 0 ] . unwrap_i32 ( ) ;
102123 if out_size == 0 {
103124 panic ! ( "graph run failed!" ) ;
104125 }
@@ -107,18 +128,22 @@ impl GraphExecutor {
107128 Ok ( ( ) )
108129 }
109130
110- pub fn get_output ( & self ) -> Result < Tensor > {
131+ pub fn get_output ( & mut self ) -> Result < Tensor > {
111132 let memory = self
112133 . instance
113134 . as_ref ( )
114135 . unwrap ( )
115- . get_memory ( "memory" )
136+ . get_memory ( self . store . as_mut ( ) . unwrap ( ) , "memory" )
116137 . ok_or_else ( || anyhow:: format_err!( "failed to find `memory` export" ) ) ?;
117138
118- let out_data = unsafe {
119- & memory. data_unchecked ( ) [ self . wasm_addr as usize ..] [ ..self . output_size as usize ]
120- } ;
121- let out_vec: Tensor = serde_json:: from_slice ( out_data) . unwrap ( ) ;
139+ let mut out_data = vec ! [ 0 as u8 ; self . output_size as _] ;
140+ memory. read (
141+ self . store . as_mut ( ) . unwrap ( ) ,
142+ self . wasm_addr as _ ,
143+ & mut out_data,
144+ ) ?;
145+
146+ let out_vec: Tensor = serde_json:: from_slice ( & out_data) . unwrap ( ) ;
122147 Ok ( out_vec)
123148 }
124149}
0 commit comments