Skip to content

Commit c724ce1

Browse files
FIXES #137: make calling wasm from python 7x faster
1 parent 4a52ebb commit c724ce1

File tree

1 file changed

+45
-33
lines changed

1 file changed

+45
-33
lines changed

wasmtime/_func.py

Lines changed: 45 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from typing import Callable, Optional, Generic, TypeVar, List, Union, Tuple, cast as cast_type, Sequence
88
from ._exportable import AsExtern
99
from ._store import Storelike
10-
10+
from ._bindings import wasmtime_val_raw_t
1111

1212
T = TypeVar('T')
1313
FUNCTIONS: "Slab[Tuple]"
@@ -27,11 +27,11 @@ def __init__(self, store: Storelike, ty: FuncType, func: Callable, access_caller
2727
set to `True` then the first argument given to `func` is an instance of
2828
type `Caller` below.
2929
"""
30-
3130
if not isinstance(store, Store):
3231
raise TypeError("expected a Store")
3332
if not isinstance(ty, FuncType):
3433
raise TypeError("expected a FuncType")
34+
self._func_call_init(ty)
3535
idx = FUNCTIONS.allocate((func, ty.results, access_caller))
3636
_func = ffi.wasmtime_func_t()
3737
ffi.wasmtime_func_new(
@@ -56,6 +56,33 @@ def type(self, store: Storelike) -> FuncType:
5656
ptr = ffi.wasmtime_func_type(store._context, byref(self._func))
5757
return FuncType._from_ptr(ptr, None)
5858

59+
def _func_call_init(self, ty):
60+
self._ty = ty
61+
ty_params = ty.params
62+
ty_results = ty.results
63+
self._params_str = (str(i) for i in ty_params)
64+
self._results_str = (str(i) for i in ty_results)
65+
params_n = len(ty_params)
66+
results_n = len(ty_results)
67+
self._params_n = params_n
68+
self._results_n = results_n
69+
n = max(params_n, results_n)
70+
self._vals_raw_type = wasmtime_val_raw_t*n
71+
72+
def _create_raw_vals(self, *params: IntoVal) -> ctypes.Array[wasmtime_val_raw_t]:
73+
raw = self._vals_raw_type()
74+
for i, param_str in enumerate(self._params_str):
75+
setattr(raw[i], param_str, params[i])
76+
return raw
77+
78+
def _extract_return(self, vals_raw: ctypes.Array[wasmtime_val_raw_t]) -> Union[IntoVal, Sequence[IntoVal], None]:
79+
if self._results_n==0:
80+
return None
81+
if self._results_n==1:
82+
return getattr(vals_raw[0], self._results_str[0])
83+
# we can use tuple construct, but I'm using list for compatability
84+
return [getattr(val_raw, ret_str) for val_raw, ret_str in zip(vals_raw, self._results_str)]
85+
5986
def __call__(self, store: Storelike, *params: IntoVal) -> Union[IntoVal, Sequence[IntoVal], None]:
6087
"""
6188
Calls this function with the given parameters
@@ -70,45 +97,30 @@ def __call__(self, store: Storelike, *params: IntoVal) -> Union[IntoVal, Sequenc
7097
Note that you can also use the `__call__` method and invoke a `Func` as
7198
if it were a function directly.
7299
"""
73-
74-
ty = self.type(store)
75-
param_tys = ty.params
76-
if len(params) > len(param_tys):
100+
params_n = len(params)
101+
if params_n > self._params_n:
77102
raise WasmtimeError("too many parameters provided: given %s, expected %s" %
78-
(len(params), len(param_tys)))
79-
if len(params) < len(param_tys):
103+
(params_n, self._params_n))
104+
if params_n < self._params_n:
80105
raise WasmtimeError("too few parameters provided: given %s, expected %s" %
81-
(len(params), len(param_tys)))
82-
83-
param_vals = [Val._convert(ty, params[i]) for i, ty in enumerate(param_tys)]
84-
params_ptr = (ffi.wasmtime_val_t * len(params))()
85-
for i, val in enumerate(param_vals):
86-
params_ptr[i] = val._unwrap_raw()
87-
88-
result_tys = ty.results
89-
results_ptr = (ffi.wasmtime_val_t * len(result_tys))()
90-
106+
(params_n, self._params_n))
107+
vals_raw = self._create_raw_vals(*params)
108+
vals_raw_ptr = ctypes.cast(vals_raw, ctypes.POINTER(wasmtime_val_raw_t))
109+
# according to https://docs.wasmtime.dev/c-api/func_8h.html#a3b54596199641a8647a7cd89f322966f
110+
# it's safe to call wasmtime_func_call_unchecked because
111+
# - we allocate enough space to hold all the parameters and all the results
112+
# - we set proper types
113+
# - but not sure about "Values such as externref and funcref are valid within the store being called"
91114
with enter_wasm(store) as trap:
92-
error = ffi.wasmtime_func_call(
115+
error = None
116+
ffi.wasmtime_func_call_unchecked(
93117
store._context,
94118
byref(self._func),
95-
params_ptr,
96-
len(params),
97-
results_ptr,
98-
len(result_tys),
119+
vals_raw_ptr,
99120
trap)
100121
if error:
101122
raise WasmtimeError._from_ptr(error)
102-
103-
results = []
104-
for i in range(0, len(result_tys)):
105-
results.append(Val(results_ptr[i]).value)
106-
if len(results) == 0:
107-
return None
108-
elif len(results) == 1:
109-
return results[0]
110-
else:
111-
return results
123+
return self._extract_return(vals_raw)
112124

113125
def _as_extern(self) -> ffi.wasmtime_extern_t:
114126
union = ffi.wasmtime_extern_union(func=self._func)

0 commit comments

Comments
 (0)