Module jit_dispatch 

Host-side implementations of tensor-wasm:jit/host::{dispatch, alloc, free}.

Registered on the wasmtime Linker when the auto-offload rewrite has swapped function bodies for dispatch trampolines (see tensor_wasm_jit::rewrite). The trio implements the v0.1.0 ABI documented in the tensor_wasm_jit::rewrite module-level docs.

dispatch

(func $dispatch
  (param i64 i64 i32 i32 i32) (result i32))
  ;;     fp_lo fp_hi sptr alen rlen   error

• Looks up the cached kernel by (fp_lo|fp_hi as u64, sm_version).
• On a cache hit under --features cuda: launches the real compiled kernel against the guest scratch region and returns 0 (DISPATCH_OK).
• On a cache hit WITHOUT CUDA: there is no kernel to run, so the dispatch deliberately deopts — it emits a tracing::warn! and returns DISPATCH_CACHE_MISS (the same code the genuine miss path uses) rather than echoing the guest’s own input back as if it were a real result. The rewrite trampoline turns any nonzero code into a wasm trap (after freeing the scratch slot), so a no-CUDA deployment fails loudly instead of silently returning wrong data.
• On a cache miss: emits a tracing::warn! and returns DISPATCH_CACHE_MISS; the trampoline traps the guest as above.

alloc / free

(func $alloc (param i32 (size)) (result i32 (ptr)))
(func $free  (param i32 (ptr)) (param i32 (size)))

Backed by a per-store bump arena reserved from the upper region of the guest’s first linear memory (memory 0). The arena starts SCRATCH_ARENA_BYTES bytes below the current memory length and grows downward; free returns slots to the arena in LIFO order. The first call grows the memory by one page if there isn’t already enough room. Multi-page reservations beyond SCRATCH_ARENA_BYTES return -1 and the trampoline’s subsequent dispatch call propagates the failure.

Structs

ArenaState

Per-instance state backing the alloc/free/dispatch imports.

Constants

DEFAULT_DISPATCH_SM_VERSION

Default sm_version the dispatcher looks up. Matches tensor_wasm_jit::rewrite::DEFAULT_SM_VERSION.

DISPATCH_BAD_SCRATCH

Return code: the dispatch failed because the caller’s scratch region pointed outside guest memory.

DISPATCH_CACHE_MISS

Return code: cache miss — no kernel was pre-populated for this fingerprint. The guest should treat this as a deopt signal.

DISPATCH_OK

Return code: cache hit (kernel would dispatch on the CUDA path).

SCRATCH_ARENA_BYTES

Bytes reserved for the scratch arena. Sized to comfortably hold the args+results of a multi-arg primitive function (each arg is at most 8 bytes; 64 KiB is enough for ~4000 i64 args). Operators tune via add_jit_dispatch_to_linker_with in test harnesses.

Traits

JitArenaProvider

Lets the JIT host imports reach the per-store ArenaState through caller.data_mut(). Implemented by crate::instance::InstanceState and by test-only payload types.

TenantContext

Trait for extracting the calling tenant from a wasmtime Store<T> payload.

Functions

add_jit_dispatch_to_linker

Register the tensor-wasm:jit/host::dispatch, alloc, and free imports on linker.

add_jit_dispatch_to_linker_with

Variant of add_jit_dispatch_to_linker that lets callers override the import module / field names and the target sm_version. Useful in tests and when the same linker hosts multiple offload generations side-by-side.