FIXES #137: make calling wasm from python 7x faster

examples/gcd.py

@@ -1,10 +1,13 @@
 # Example of instantiating a wasm module and calling an export on it
 
 from wasmtime import Store, Module, Instance
-
+from functools import partial
 store = Store()
 module = Module.from_file(store.engine, './examples/gcd.wat')
 instance = Instance(store, module, [])
 gcd = instance.exports(store)["gcd"]
-
+gcd_func = partial(gcd, store)
+gcd_func_val = partial(gcd._call_val, store)
 print("gcd(6, 27) = %d" % gcd(store, 6, 27))
+print("gcd(6, 27) = %d" % gcd_func(6, 27))
+print("gcd(6, 27) = %d" % gcd_func_val(6, 27))

examples/gcd_perf.py

@@ -0,0 +1,25 @@
+import time
+from math import gcd as math_gcd
+from gcd import gcd_func as wasm_gcd, gcd_func_val as wasm_gcd_old
+
+
+def python_gcd(x, y):
+    while y:
+        x, y = y, x % y
+    return abs(x)
+
+
+a = 16516842
+b = 154654684
+
+print(math_gcd(a, b), python_gcd(a, b), wasm_gcd(a, b), wasm_gcd_old(a, b))
+
+N = 1_000
+by_name = locals()
+for name in "math_gcd", "python_gcd", "wasm_gcd", "wasm_gcd_old":
+    gcdf = by_name[name]
+    start_time = time.perf_counter()
+    for _ in range(N):
+        g = gcdf(a, b)
+    total_time = time.perf_counter() - start_time
+    print(total_time, "\t\t", name)

examples/simd_i8x16.py

@@ -0,0 +1,36 @@
+"""
+how to call v128 SIMD operations
+for more details see https://github.com/WebAssembly/simd/blob/main/proposals/simd/SIMD.md#integer-addition
+"""
+import ctypes
+
+from functools import partial
+from wasmtime import Store, Module, Instance, Func
+
+
+store = Store()
+module = Module(
+    store.engine,
+    """
+(module
+  (func $add_v128 (param $a v128) (param $b v128) (result v128)
+    local.get $a
+    local.get $b
+    i8x16.add
+  )
+  (export "add_v128" (func $add_v128))
+)
+""",
+)
+
+instance = Instance(store, module, [])
+vector_type = ctypes.c_uint8 * 16
+add_v128_f = instance.exports(store)["add_v128"]
+if not isinstance(add_v128_f, Func):
+    raise TypeError("expecting Func")
+add_v128 = partial(add_v128_f, store)
+a = vector_type(*(i for i in range(16)))
+b = vector_type(*(40 + i for i in range(16)))
+c: list[int] = add_v128(a, b)  # type: ignore
+print([v for v in c])
+print([v for v in c] == [i + j for i, j in zip(a, b)])

pytest.ini

@@ -1,2 +1,3 @@
 [pytest]
 addopts = --doctest-modules --flake8 --mypy
+

tests/test_func.py

@@ -1,5 +1,7 @@
 import unittest
+import ctypes
 
+from functools import partial
 from wasmtime import *
 
 
@@ -17,6 +19,34 @@ def test_add(self):
         func = Func(store, ty, lambda a, b: a + b)
         self.assertEqual(func(store, 1, 2), 3)
 
+    def test_simd_i8x16_add(self):
+        # i8x16.add is SIMD 128-bit vector of i8 items of size 16
+        store = Store()
+        module = Module(
+            store.engine,
+            """
+        (module
+        (func $add_v128 (param $a v128) (param $b v128) (result v128)
+            local.get $a
+            local.get $b
+            i8x16.add
+        )
+        (export "add_v128" (func $add_v128))
+        )
+        """,
+        )
+
+        instance = Instance(store, module, [])
+        vector_type = ctypes.c_uint8 * 16
+        add_v128_f = instance.exports(store)["add_v128"]
+        if not isinstance(add_v128_f, Func):
+            raise TypeError("expecting Func")
+        add_v128 = partial(add_v128_f, store)
+        a = vector_type(*(i for i in range(16)))
+        b = vector_type(*(40 + i for i in range(16)))
+        c: list[int] = add_v128(a, b)  # type: ignore
+        self.assertEqual([v for v in c], [i + j for i, j in zip(a, b)])
+
     def test_calls(self):
         store = Store()
         ty = FuncType([ValType.i32()], [])

wasmtime/_func.py

@@ -7,7 +7,8 @@
 from typing import Callable, Optional, Generic, TypeVar, List, Union, Tuple, cast as cast_type, Sequence
 from ._exportable import AsExtern
 from ._store import Storelike
-
+from ._bindings import wasmtime_val_raw_t
+from ._value import get_valtype_attr, val_getter, val_setter
 
 T = TypeVar('T')
 FUNCTIONS: "Slab[Tuple]"
@@ -16,6 +17,13 @@
 
 class Func:
     _func: ffi.wasmtime_func_t
+    _ty: FuncType
+    _params_n: int
+    _results_n: int
+    _params_str: list[str]
+    _results_str: list[str]
+    _results_str0: str
+    _vals_raw_type: type[ctypes.Array[wasmtime_val_raw_t]]
 
     def __init__(self, store: Storelike, ty: FuncType, func: Callable, access_caller: bool = False):
         """
@@ -27,11 +35,11 @@ def __init__(self, store: Storelike, ty: FuncType, func: Callable, access_caller
         set to `True` then the first argument given to `func` is an instance of
         type `Caller` below.
         """
-
         if not isinstance(store, Store):
             raise TypeError("expected a Store")
         if not isinstance(ty, FuncType):
             raise TypeError("expected a FuncType")
+        self._init_call(ty)
         idx = FUNCTIONS.allocate((func, ty.results, access_caller))
         _func = ffi.wasmtime_func_t()
         ffi.wasmtime_func_new(
@@ -56,21 +64,39 @@ def type(self, store: Storelike) -> FuncType:
         ptr = ffi.wasmtime_func_type(store._context, byref(self._func))
         return FuncType._from_ptr(ptr, None)
 
-    def __call__(self, store: Storelike, *params: IntoVal) -> Union[IntoVal, Sequence[IntoVal], None]:
-        """
-        Calls this function with the given parameters
-
-        Parameters can either be a `Val` or a native python value which can be
-        converted to a `Val` of the corresponding correct type
-
-        Returns `None` if this func has 0 return types
-        Returns a single value if the func has 1 return type
-        Returns a list if the func has more than 1 return type
+    def _create_raw_vals(self, *params: IntoVal) -> ctypes.Array[wasmtime_val_raw_t]:
+        raw = self._vals_raw_type()
+        for i, param_str in enumerate(self._params_str):
+            val_setter(self._func.store_id, raw[i], param_str, params[i])
+        return raw
 
-        Note that you can also use the `__call__` method and invoke a `Func` as
-        if it were a function directly.
+    def _extract_return(self, vals_raw: ctypes.Array[wasmtime_val_raw_t]) -> Union[IntoVal, Sequence[IntoVal], None]:
+        if self._results_n == 0:
+            return None
+        if self._results_n == 1:
+            return val_getter(self._func.store_id, vals_raw[0], self._results_str0)
+        # we can use tuple construct, but I'm using list for compatability
+        return [val_getter(self._func.store_id, val_raw, ret_str) for val_raw, ret_str in zip(vals_raw, self._results_str)]
+
+    def _init_call(self, ty: FuncType) -> None:
+        """init signature properties used by call"""
+        self._ty = ty
+        ty_params = ty.params
+        ty_results = ty.results
+        params_n = len(ty_params)
+        results_n = len(ty_results)
+        self._params_str = [get_valtype_attr(i) for i in ty_params]
+        self._results_str = [get_valtype_attr(i) for i in ty_results]
+        self._results_str0 = get_valtype_attr(ty_results[0]) if results_n else 'i32'
+        self._params_n = params_n
+        self._results_n = results_n
+        n = max(params_n, results_n)
+        self._vals_raw_type = wasmtime_val_raw_t * n
+
+    def _call_val(self, store: Storelike, *params: IntoVal) -> Union[IntoVal, Sequence[IntoVal], None]:
+        """
+        internal implementation of calling a function that uses `wasmtime_func_call`
         """
-
         ty = self.type(store)
         param_tys = ty.params
         if len(params) > len(param_tys):
@@ -110,6 +136,52 @@ def __call__(self, store: Storelike, *params: IntoVal) -> Union[IntoVal, Sequenc
         else:
             return results
 
+    def _call_raw(self, store: Storelike, *params: IntoVal) -> Union[IntoVal, Sequence[IntoVal], None]:
+        """
+        internal implementation of calling a function that uses `wasmtime_func_call_unchecked`
+        """
+        if getattr(self, "_ty", None) is None:
+            self._init_call(self.type(store))
+        params_n = len(params)
+        if params_n > self._params_n:
+            raise WasmtimeError("too many parameters provided: given %s, expected %s" %
+                                (params_n, self._params_n))
+        if params_n < self._params_n:
+            raise WasmtimeError("too few parameters provided: given %s, expected %s" %
+                                (params_n, self._params_n))
+        vals_raw = self._create_raw_vals(*params)
+        vals_raw_ptr = ctypes.cast(vals_raw, ctypes.POINTER(wasmtime_val_raw_t))
+        # according to https://docs.wasmtime.dev/c-api/func_8h.html#a3b54596199641a8647a7cd89f322966f
+        # it's safe to call wasmtime_func_call_unchecked because
+        # - we allocate enough space to hold all the parameters and all the results
+        # - we set proper types by reading types from ty
+        # - externref and funcref are valid within the store being called
+        with enter_wasm(store) as trap:
+            error = ffi.wasmtime_func_call_unchecked(
+                store._context,
+                byref(self._func),
+                vals_raw_ptr,
+                trap)
+            if error:
+                raise WasmtimeError._from_ptr(error)
+        return self._extract_return(vals_raw)
+
+    def __call__(self, store: Storelike, *params: IntoVal) -> Union[IntoVal, Sequence[IntoVal], None]:
+        """
+        Calls this function with the given parameters
+
+        Parameters can either be a `Val` or a native python value which can be
+        converted to a `Val` of the corresponding correct type
+
+        Returns `None` if this func has 0 return types
+        Returns a single value if the func has 1 return type
+        Returns a list if the func has more than 1 return type
+
+        Note that you can also use the `__call__` method and invoke a `Func` as
+        if it were a function directly.
+        """
+        return self._call_raw(store, *params)
+
     def _as_extern(self) -> ffi.wasmtime_extern_t:
         union = ffi.wasmtime_extern_union(func=self._func)
         return ffi.wasmtime_extern_t(ffi.WASMTIME_EXTERN_FUNC, union)

wasmtime/_types.py

@@ -81,7 +81,7 @@ def __str__(self) -> str:
             return 'anyref'
         if kind == ffi.WASM_FUNCREF.value:
             return 'funcref'
-        return 'ValType(%d)' % kind.value
+        return 'ValType(%d)' % kind  # type: ignore
 
     def __del__(self) -> None:
         if not hasattr(self, '_owner') or not hasattr(self, '_ptr'):

wasmtime/_value.py

@@ -4,6 +4,18 @@
 import ctypes
 import typing
 
+val_id2attr = {
+    WASMTIME_I32.value: 'i32',
+    WASMTIME_I64.value: 'i64',
+    WASMTIME_F32.value: 'f32',
+    WASMTIME_F64.value: 'f64',
+    WASMTIME_V128.value: 'v128',
+    WASMTIME_FUNCREF.value: 'funcref',
+    WASMTIME_EXTERNREF.value: 'externref',
+    WASM_FUNCREF.value: 'funcref',
+    WASM_ANYREF.value: 'externref',
+}
+
 
 @ctypes.CFUNCTYPE(None, c_void_p)
 def _externref_finalizer(extern_id: int) -> None:
@@ -25,6 +37,63 @@ def _unintern(val: int) -> typing.Any:
     return Val._id_to_extern.get(val)
 
 
+def get_valtype_attr(ty: ValType) -> str:
+    return val_id2attr[wasm_valtype_kind(ty._ptr)]  # type: ignore
+
+
+def val_getter(store_id: int, val_raw: wasmtime_val_raw_t, attr: str) -> typing.Union[int, float, "wasmtime.Func", typing.Any]:
+    val = getattr(val_raw, attr)
+
+    if attr == 'externref':
+        ptr = ctypes.POINTER(wasmtime_externref_t)
+        if not val:
+            return None
+        ffi = ptr.from_address(val)
+        if not ffi:
+            return None
+        extern_id = wasmtime_externref_data(ffi)
+        return _unintern(extern_id)
+    elif attr == 'funcref':
+        if val == 0:
+            return None
+        f = wasmtime_func_t()
+        f.store_id = store_id
+        f.index = val
+        return wasmtime.Func._from_raw(f)
+    return val
+
+
+def val_setter(store_id: int, dst: wasmtime_val_raw_t, attr: str, val: "IntoVal") -> None:
+    casted: typing.Union[Any, int, float, None, wasmtime.Func]
+    if attr == 'externref':
+        if isinstance(val, Val) and val._raw and val._raw.kind == WASMTIME_EXTERNREF.value:
+            casted = ctypes.addressof(val._raw.of.externref)
+        else:
+            ex = Val.externref(val)._raw
+            if ex:
+                casted = ctypes.addressof(ex.of.externref)
+            else:
+                casted = 0
+    elif attr == 'funcref':
+        if isinstance(val, Val) and val._raw and val._raw.kind == WASMTIME_FUNCREF.value:
+            casted = val._raw.of.funcref.index
+        elif isinstance(val, wasmtime.Func):
+            if val._func.store_id != store_id:
+                raise WasmtimeError("passed funcref does not belong to same store")
+            casted = val._func.index
+        else:
+            raise WasmtimeError("expecting param of type funcref got " + type(val).__name__)
+    else:
+        if isinstance(val, Val):
+            if val._raw:
+                casted = getattr(val._raw.of, attr)
+            else:
+                casted = 0
+        else:
+            casted = val
+    setattr(dst, attr, casted)
+
+
 class Val:
     # We can't let the extern values we wrap `externref`s around be GC'd, so we
     # pin them in `_id_to_extern`. Additionally, we might make multiple