Calling prebuilt stdlib modules from the interpreter

Experimental branch: exp-native-stdlib in the paugier/spy fork, based on upstream commit a39950ca.

Goal¶

The standard library of SPy is written in SPy. When the interpreter runs it, every operation of every loop is interpreted, which is slow. Idea: compile a few stdlib modules once, ahead of time, and make the interpreter call the compiled functions instead of interpreting their bodies.

The first target is random, whose interpreter tests are very slow:

Why it can work¶

The interpreter already owns a WASM instance, vm.ll, built from libspy (the C runtime of SPy). Heap objects allocated by interpreted code (gc_alloc) live in its linear memory, and the C backend allocates with the same malloc. The function libspy.get_LLMOD(extra_archives) can link more archives into this instance. So compiled stdlib code can be added to it, and both worlds then share one heap without any copy.

What a native call costs¶

Measured from Python, with a debug-built archive:

A native call costs about as much as a trivial interpreted call, so it never loses, but it gains only a factor 2 on a tiny function. The gain comes from functions which contain a loop: those are the ones to prebuild.

A second cost is paid once per process. Linking more archives into vm.ll produces a larger WASM module, and compiling it takes about 150 ms. The module is cached per process, like the plain libspy one, so only the first VM pays. The cache key of the archive also needs the version of zig, which is read once and memoized.

The manifest¶

The manifest is not written by hand: it is read off what the C backend emits for the redshifted modules. The generator runs the same steps as spy build (ImportAnalyzer, vm.redshift(), CBackend.cwrite()) on a stub program which imports the requested modules, then looks at every (fqn, w_obj) pair that the backend emitted.

User code:

from random import seed, random

def main() -> None:
    seed(42)
    print(random())

For the module random, the C backend emits these functions (lowered code):

void spy_random$seed(uint32_t value$0);
double spy_random$random(void);
double spy_random$uniform(double a$0, double b$0);

Which of them become native?

• A function is a candidate if it belongs to a requested module and the backend emits it, which is the case of every red function that is not @force_inline:

  if w_obj.color != "red" or w_obj.is_force_inline:
      continue

• It is native if its signature can be passed through a wasm call. Otherwise it stays interpreted, and the reason is recorded in the manifest:

  def _unsupported_reason(fqn: FQN, w_func: Any) -> Optional[str]:
      if not _is_plain(fqn):
          return "generic instantiation or closure"
      w_ft = w_func.w_functype
      for p in w_ft.params:
          if p.kind != "simple":
              return "variadic parameters"
          if p.w_T not in scalar_types():
              return f"parameter of type {type_name(p.w_T)}"
      if w_ft.w_restype not in result_types():
          return f"result of type {type_name(w_ft.w_restype)}"
      return None

  scalar_types() is i32, u32, f64, f32, bool and str; a result can also be None. A generic instantiation or a closure carries qualifiers or a suffix in its name (_is_plain is false), and we cannot be sure that the interpreter would give it the same name.

• A variable (a W_Cell) is native if its content is an i32 or a pointer, which are the two kinds that the C backend supports. All the variables of all the emitted modules are native, not only those of the requested modules, because the compiled code of any of them may read or write the C global.

The command python -m spy.build.native_stdlib random prints the result:

native functions:
  random::seed(builtins::u32) -> types::NoneType   [native_spy_random__seed]
  random::random() -> builtins::f64   [native_spy_random__random]
  random::uniform(builtins::f64, builtins::f64) -> builtins::f64   [native_spy_random__uniform]
native variables:
  _random::_ptr_global_random: unsafe::gc_ptr[_random::Random]   [native_spy__random___ptr_global_random]
kept interpreted:

It is saved as manifest.json next to the archive. An entry contains both the SPy and the C view of the function, taken from the backend (fqn.c_name and Context.w2c):

{
  "fqn": "random::seed",
  "export": "native_spy_random__seed",
  "c_name": "spy_random$seed",
  "header": "random.h",
  "params": ["builtins::u32"],
  "restype": "types::NoneType",
  "c_params": ["uint32_t"],
  "c_restype": "void"
}

The names of the C functions contain $ and are not visible from outside the archive, so the shim exports them under simple names. For random it is:

#include "spy.h"
#include "_random.h"
#include "random.h"

void WASM_EXPORT(native_spy_random__seed)(uint32_t a0) { spy_random$seed(a0); }
double WASM_EXPORT(native_spy_random__random)(void) { return spy_random$random(); }
double WASM_EXPORT(native_spy_random__uniform)(double a0, double a1) { return spy_random$uniform(a0, a1); }
void *WASM_EXPORT(native_spy__random___ptr_global_random)(void) { return &spy__random$_ptr_global_random; }

The last line exports the address of a C global variable, because a data symbol cannot be exported portably. It is explained in the section on module variables below.

The archive and the manifest are cached under a key made of the module sources, the libspy headers, the zig version and the sources of the SPy compiler, so editing any of them rebuilds the archive. Two directories with the same sources share the same archive.

The interpreter side¶

W_NativeFunc takes the place of the interpreted function with the same name:

class W_NativeFunc(W_Func):
    def __init__(self, w_orig: W_ASTFunc, export_name: str) -> None:
        self.w_functype = w_orig.w_functype
        self.fqn = w_orig.fqn
        self.def_loc = w_orig.def_loc
        self.w_origin = w_orig.w_origin
        self.export_name = export_name
        self.w_orig = w_orig

    def raw_call(self, vm: "SPyVM", args_w: Sequence[W_Object]) -> W_Object:
        args = [_to_wasm(w_arg, vm) for w_arg in args_w]
        res = vm.ll.call(self.export_name, *args)
        w_restype = self.w_functype.w_restype
        if w_restype is TYPES.w_NoneType:
            return vm.wrap(None)
        return _from_wasm(vm, w_restype, res)

• It copies w_functype (the signature) from the original function, so type checking and operator dispatch cannot tell the difference.

• raw_call converts the arguments, calls the export through vm.ll.call, and wraps the result back into an interpreter value.

The two conversion functions are the marshalling:

def _to_wasm(w_arg: W_Object, vm: "SPyVM") -> Any:
    if isinstance(w_arg, W_Str):
        return w_arg.ptr
    v = vm.unwrap(w_arg)
    if isinstance(w_arg, W_U32) and v >= 2**31:
        v -= 2**32
    elif isinstance(w_arg, W_Bool):
        v = int(v)
    return v


def _from_wasm(vm: "SPyVM", w_T: Any, v: Any) -> W_Object:
    if w_T is B.w_u32:
        return W_U32(v & 0xFFFFFFFF)
    elif w_T is B.w_i32:
        return W_I32(v)
    elif w_T is B.w_f64:
        return W_F64(v)
    elif w_T is B.w_f32:
        return W_F32(v)
    elif w_T is B.w_bool:
        return vm.wrap(bool(v))
    elif w_T is B.w_str:
        return W_Str.from_ptr(vm, v)
    raise NotImplementedError(f"cannot marshal result of type {w_T}")

• WebAssembly integers are signed, so a u32 above 2**31 is mapped to a negative number on the way in, and masked back on the way out. A bool travels as 0 or 1.

• A str travels as its address, as explained in the next section.

NativeRegistry is built from the manifest. Before replacing a function it compares the signature of the interpreted function with the one of the manifest, because calling compiled code with the wrong signature would silently produce garbage:

def maybe_replace_func(self, fqn: FQN, w_func: W_Object) -> W_Object:
    nf = self.funcs.get(str(fqn))
    if nf is None or not isinstance(w_func, W_ASTFunc):
        return w_func
    w_ft = w_func.w_functype
    params = [ns.type_name(p.w_T) for p in w_ft.params]
    restype = ns.type_name(w_ft.w_restype)
    if (params, restype) != (nf.params, nf.restype):
        raise RuntimeError(
            f"native stdlib out of sync: {fqn} is ({', '.join(params)}) -> "
            f"{restype} but the manifest says ({', '.join(nf.params)}) -> "
            f"{nf.restype}"
        )
    return W_NativeFunc(w_func, nf.export)

The replacement is done where the interpreter creates the function, before it is stored anywhere:

w_func = self.vm.maybe_native_func(fqn, w_func)
self.vm.add_global(fqn, w_func)

Replacing it later, inside add_global, would be too late: the frame keeps its own reference to the interpreted function, and from random import seed would still find that one.

Passing str¶

A str is already a pointer to a spy_StrObject in linear memory, and W_Str already stores that address in ptr. Passing a str to compiled code, or getting one back, is therefore a plain integer: no copy, and both sides see the same bytes. A str visible to SPy code is always sealed (its bytes are final, see str.h), so wrapping a returned address is safe.

User code:

def shout(s: str) -> str:
    return s.upper()

The stdlib function str.upper is _str::methods::upper. For _str, the C backend emits (lowered code):

spy_StrObject * spy__str$methods$upper(spy_StrObject * self$0);

and the shim exports it as:

spy_StrObject * WASM_EXPORT(native_spy__str__methods__upper)(spy_StrObject * a0)
{ return spy__str$methods$upper(a0); }

In WebAssembly a spy_StrObject * is an i32, so raw_call just passes w_s.ptr. With str allowed, 25 functions of _str are native: __len__, __add__, __mul__, the comparisons, upper, replace, isspace, __contains__, __repr__, and the helpers behind find, count and strip. The compile-time dispatch of find, count, strip and split stays interpreted and picks a native helper such as _find_1.

Time of one call (ms), interpreted → native, on strings of 10, 100 and 1000 characters:

• The interpreted time grows linearly, 0.5 to 3.5 ms per character. The native time is flat: on 100000 characters, upper takes 4.7 ms, find 1.7 ms, count 2.7 ms and replace 2.8 ms (debug build, not optimized).

• The floor of about 0.1 ms is mostly the interpreted caller, not the native call.

• split does not change: it returns a list[str], which is not marshalled yet (see the section on aggregates).

On test_str.py -m interp, which uses short strings, the native mode is at parity with the interpreter (9.2 s against 8.6 s): the gain only shows on long strings.

Which other stdlib modules¶

Result of the manifest generator on each module, with str allowed:

_tuple, _dict, array and simd give nothing: they only contain generics, which emit no code until something instantiates them.

Over all modules, the functions kept interpreted are blocked by: generic instantiations or closures (47), stack structs by value (46, mostly iterators, StrBuilder methods and the __eq__ and __ne__ generated by @struct), range and slice (23), bytes (14), u8, i8, i64 and u64 (2).

Two blockers are cheap:

• bytes has the same representation as str: W_Bytes stores the address of a spy_BytesObject, which mirrors spy_StrObject.

• u8 and i8 are WASM i32, i64 and u64 are WASM i64: they only need to be added to scalar_types().

Aggregates: structs, tuples, lists¶

A heap struct is bytes in linear memory with the C layout, so it crosses as a pointer. A stack struct is a Python dictionary in the interpreter (W_Struct.values_w), so it must be copied. The typed containers are stack structs wrapping heap data:

@struct
class _ListImpl:
    __ll__: gc_ptr[ListData]       # the only field: a pointer to the heap data

A list[T] value is therefore a one-field struct holding a pointer, and a tuple is a struct of its items. The items of a list are already shared; only the small struct is copied.

The C ABI cannot be relied on. Compiling a struct by value for wasm32 shows that the convention changes with the build type. The probe is:

#ifdef SPY_DEBUG
typedef struct { uint32_t *p; ptrdiff_t length; } gcp;   // gc_ptr, debug build
#else
typedef struct { uint32_t *p; } gcp;                     // gc_ptr, release build
#endif
typedef struct { gcp ll; } List;                         // like _ListImpl
typedef struct { int32_t a; int32_t b; } Tup2;           // like tuple[i32, i32]

int32_t f_list(List l);     List r_list(void);
int32_t f_tup2(Tup2 t);     Tup2 r_tup2(void);

WASM signature of each function:

A struct with one scalar field is passed directly and returned directly. A larger struct is passed as the address of a copy and returned through a hidden extra parameter, an output address. List changes convention when SPY_DEBUG adds the field length to the pointer, and in both cases the signature shows an i32, so it cannot be told apart from a plain integer. The generated shim must therefore not depend on it: it takes the address of every aggregate and writes every aggregate result through an output address. For _split_sep(self: str, sep: str) -> list[str], whose result is a _ListImpl:

void WASM_EXPORT(native_spy__str___split_sep)(
    spy__list$_ListImpl *ret, spy_StrObject *a0, spy_StrObject *a1)
{
    *ret = spy__str$_split_sep(a0, a1);
}

On the interpreter side, raw_call:

1. reserves scratch memory for each aggregate argument and for the result;

2. writes the fields of each argument at the offsets given by calc_layout, which are the offsets C uses;

3. calls the export with the addresses;

4. reads the result back into a W_Struct.

This costs a few word accesses, about 10 µs for a list. A method which modifies a list, like append, changes the heap data ListData, so nothing has to be copied back.

Values which have no representation in linear memory never cross: type, object, functions, and interp_list, interp_tuple and interp_dict. Functions using them stay interpreted.

Generic instantiations¶

The functions of list[str] have a name with a qualifier, for instance _list::list[str]::_ListImpl::append. The compiled archive can only contain the instantiations which the stub program creates, so it holds a fixed menu (for instance list[str], list[i32] and list[f64]), never list[Point] for a type defined by the user. Any function not in the manifest stays interpreted, and the signature check above rejects a name which would match with another meaning.

A closure has a numeric suffix (_getitem_int#1, #2) which depends on the order of creation, so it cannot be matched until its name is made deterministic.

Iteration is not a good target: a for loop over a list calls three methods per iteration on a stack struct, and each call would cost 21 µs plus the copy of the struct. The functions worth prebuilding are the ones which contain a whole loop: join, split, sorting and the array kernels.

Module variables¶

_random.spy keeps the generator in a module variable:

var _ptr_global_random = gc_ptr[Random].NULL

def _get_global_random(value: u32 = 0, reset: bool = False) -> Random:
    global _ptr_global_random
    if _ptr_global_random == gc_ptr[Random].NULL:
        reset = True
        _ptr_global_random = gc_alloc[Random](1)
        ll = gc_alloc[RandomLowLevel](1)
        _ptr_global_random[0].__ll__ = ll
    if reset:
        ll = _ptr_global_random[0].__ll__
        ll.index = N
        ll.state = Random.init_genrand_state(value)
    return _ptr_global_random[0]

In the interpreter, a module variable is a Python object (W_Cell). In the compiled code, it is a C global variable in linear memory:

spy_unsafe$gc_ptr___random$Random spy__random$_ptr_global_random = {0};

If both existed independently, native seed and interpreted _get_global_random would work on two different generators. So the variables of the archive become W_NativeCell, which read and write the C global through the exported address.

class W_NativeCell(W_Cell):
    def __init__(
        self, vm: "SPyVM", fqn: FQN, w_initial: W_Object, nc: ns.NativeCell
    ) -> None:
        self.fqn = fqn
        self._vm = vm
        self._kind = nc.kind
        self._w_T = vm.dynamic_type(w_initial)
        if nc.kind == "ptr":
            assert isinstance(self._w_T, W_PtrType)
        elif nc.kind == "i32":
            assert self._w_T is B.w_i32
        else:
            raise NotImplementedError(f"unsupported kind of native cell: {nc.kind}")
        self._addr = vm.ll.call(nc.export)

    def get(self) -> W_Object:
        mem = self._vm.ll.mem
        if self._kind == "i32":
            return W_I32(mem.read_i32(self._addr))
        addr, length = mem.read_ptr(self._addr)
        return W_Ptr(self._w_T, addr, length)

    def set(self, w_val: W_Object) -> None:
        mem = self._vm.ll.mem
        if self._kind == "i32":
            assert isinstance(w_val, W_I32)
            mem.write_i32(self._addr, int(w_val.value))
        else:
            assert isinstance(w_val, W_Ptr)
            mem.write_i32(self._addr, int(w_val.addr))
            mem.write_i32(self._addr + 4, int(w_val.length))

The initial value computed by the interpreter is ignored: the C global is already initialized by the compiled module. A pointer is stored as the two 32-bit words {addr, length}.

Both sides must compute the same memory layout¶

Compiled code builds structs in memory and the interpreter reads them back, so both must place every field at the same offset. Example, from array.spy:

@struct
class ArrayData:
    l: i32
    items: gc_ptr[DTYPE]

Its lowered C struct is:

struct spy_array$_array1__builtins$u32$ArrayData {
    int32_t l;
    spy_unsafe$gc_ptr__builtins$u32 items;
};

typedef struct spy_unsafe$gc_ptr__builtins$u32 {
    uint32_t *p;
#ifdef SPY_DEBUG
    ptrdiff_t length;
#endif
} spy_unsafe$gc_ptr__builtins$u32;

In debug builds a pointer is 8 bytes (p and length), but on wasm32 it only needs 4-byte alignment. C therefore places items at offset 4 and the struct is 12 bytes. calc_layout follows the same rules: each field is aligned on its own alignment, and the size is padded to the alignment of the struct. Aligning on the size of the field instead (8 for a pointer) would put items at offset 8 in a struct of 16 bytes, and the interpreter would read garbage.

def calc_layout(fields_w):
    from spy.vm.modules.unsafe.misc import alignof, sizeof

    offset = 0
    struct_align = 1
    struct_fields_w = []
    for name, w_field in fields_w.items():
        field_size = sizeof(w_field.w_T)
        field_align = alignof(w_field.w_T)
        struct_align = max(struct_align, field_align)
        offset = (offset + (field_align - 1)) & ~(field_align - 1)
        struct_fields_w.append(W_StructField(name, w_field.w_T, offset, w_field.loc))
        offset += field_size
    size = (offset + (struct_align - 1)) & ~(struct_align - 1)
    return struct_fields_w, size

Because length exists only under SPY_DEBUG, the archive must be built in debug mode, like the libspy used by the interpreter.

The interpreter keeps stack structs as dictionaries. Putting them in linear memory would need a shadow stack in every frame and rules for structs which are returned or captured, and it would also change a rule of the language: today a stack struct is immutable and has no address, which is a decision about SPy itself. It is not needed here, since copying at the boundary covers lists and tuples. It would not be expensive either: reading a word costs 2.5 µs, against 52 µs or more for an interpreted call. What protects the two sides from drifting apart is a check, described in the plan.

Errors¶

A C panic (for instance uniform(2, 1)) first stores its message through a host import, then traps. LLSPyInstance.call catches the trap and raises the corresponding SPyError, here a W_ValueError. The interpreter sees a normal exception and stays usable.

State of the prototype¶

Works and is tested (spy/tests/test_native_stdlib.py, spy/tests/test_native_manifest.py and spy/tests/test_native_str.py):

• the manifest of random, and of a test module with i32, u32, f64, f32, bool, str and None signatures, an i32 variable, and a function which must stay interpreted (a struct parameter);

• results identical to the interpreter for random and for the 25 functions of _str, u32 above 2**31, errors, and state shared in both directions (native seed then interpreted Random.random, the reverse, and an i32 variable);

• the signature check, and the cache key (same sources: same archive).

Limits:

• the archive contains every C file of the program, not only the closure of the requested modules, which makes it about 700 KiB;

• not marshalled, so kept interpreted: bytes, u8, i8, i64, u64, stack structs (hence list, tuple, split and join), generic instantiations, and variables that are neither i32 nor pointers;

• the timings are measured on a debug build;

• WebAssembly in the browser: emscripten bundling does not exist yet.

Plan¶

Each step is a separate, testable change.

☐ 1. Land the layout fix. The commit Fix calc_layout, consistently with C rules with its test, as a PR on its own; it is useful even without this work.

☐ 2. Startup. Cache the linked module per process in get_LLMOD, and memoize the zig version. Without it, every VM of a test suite pays 150 ms.

☐ 3. str, bytes and the other scalars. str is done in the prototype; add bytes, u8, i8, i64 and u64.

☐ 4. Layout guard. The C backend generates a table of sizeof and offsetof for every struct of the archive. A test compares it with the interpreter’s layout, so that a mismatch fails at test time and not as garbage at run time.

☐ 5. One archive, minimal closure. Build a single archive for all prebuilt modules in one C backend run, so that shared instantiations such as array[u32] exist once (two archives defining the same symbol would not link), and keep only what the exported functions reach.

☐ 6. Aggregates by copy. Make the shim take addresses and return through an output address, and marshal stack structs with calc_layout. Start with list[str], which unblocks split and join.

☐ 7. Menu of generic instantiations. Prebuild a fixed set (list[str], list[i32], list[f64], …) and give closures deterministic names.

☐ 8. Distribution. Ship the archive the same way as libspy. For the browser, link it into libspy at build time.

☐ 9. Test policy. Keep the feature off by default. CI runs the stdlib tests in both modes, because in native mode the interp tests of a prebuilt module exercise the compiled code, not the interpreter.

☐ 10. Decide the default. Turn it on once steps 2 to 5 and 8 are done and the startup cost is measured on the whole test suite.

Code¶

The branch adds, on top of upstream:

• spy/vm/struct.py: calc_layout, with its test in spy/tests/compiler/test_struct.py;

• spy/build/native_stdlib.py: manifest generation, shim, archive build and cache;

• spy/vm/native.py: W_NativeFunc, W_NativeCell, NativeRegistry, and the marshalling of scalars and str;

• spy/libspy/__init__.py and spy/build/wasm_bundle.py: the per-process caches;

• the hooks in vm.py, astframe.py and modframe.py.

pytest spy/tests/test_native_manifest.py spy/tests/test_native_stdlib.py spy/tests/test_native_str.py
SPY_NATIVE_STDLIB=1 pytest spy/tests/stdlib/test_random.py -m interp --slow-tests
python -m spy.build.native_stdlib random
python -m spy.build.native_stdlib _str