tensor-wasm-jit

Just-in-time compilation pipeline that opportunistically offloads hot Wasm code to the GPU. Uses a Cranelift CLIF block analyser to detect offload candidates, lowers them through TensorWasmIR (a normalised intermediate representation), emits PTX text, and caches compiled modules in a DashMap-backed cache with a parking_lot::Mutex<LruCache> eviction queue. A DeoptGuard provides a safe CPU fallback path whenever GPU offload fails or guards trip at runtime.

v0.1.0 ABI: dispatch trampoline contract

When rewrite_wasm swaps a hot function body for an offload trampoline, the trampoline marshals arguments through a scratch region in the guest's first linear memory:

1. call __tensor_wasm_jit_alloc(scratch_size: i32) -> i32 — host-side arena returns a non-overlapping i32 pointer.
2. For each parameter, store its raw bytes at scratch_ptr + arg_off using the appropriate <type>.store. Offsets are byte-packed in declared order with natural alignment (i64/f64 = 8B, i32/f32 = 4B). No inter-arg padding.
3. call __tensor_wasm_jit_dispatch(fp_lo: i64, fp_hi: i64, scratch_ptr: i32, args_byte_len: i32, results_byte_len: i32) -> i32 — runs the cached PTX kernel (CUDA path) or the host reference (no-CUDA path); writes result bytes to scratch_ptr + args_byte_len. Returns 0 on success, nonzero on error.
4. For each result, load its bytes from scratch_ptr + args_byte_len + result_off and push to the wasm stack.
5. call __tensor_wasm_jit_free(scratch_ptr: i32, scratch_size: i32) — returns the slot to the arena.

The fingerprint is the 64-bit BLAKE3-truncated blueprint hash. We split it into two i64 halves so a Wasm 1.0 module without the multi-value proposal can still carry the full 64-bit value.

Supported types (v0.1.0)

Only primitive scalars: i32, i64, f32, f64. v128 and reftypes in either parameters or results disqualify a candidate; the rewriter falls back to leaving the original body intact and Wasmtime executes it on the CPU. Lifting this restriction is a v0.2 task — it needs richer marshalling and an SoA layout decision on the GPU side.

Host import signatures

(func $tensor-wasm:jit/host::alloc    (param i32)                            (result i32))
(func $tensor-wasm:jit/host::free     (param i32 i32))
(func $tensor-wasm:jit/host::dispatch (param i64 i64 i32 i32 i32)            (result i32))

The host implementation lives in tensor_wasm_exec::jit_dispatch.

Feature flags

Unstable pliron_* surface

The pliron_dialect, pliron_lowering, and pliron_ptx modules are gated on cuda-oxide-backend and additionally marked #[doc(hidden)]. They host the in-progress W3.x/W4.x scaffolding (today: ~700 lines of NotYetWired stubs) that the v0.4 author will fill in. Names, signatures, and behaviour inside these modules are NOT covered by semver compatibility until tensor-wasm-jit reaches 1.0 — external consumers should not import from them.

See docs/BUILD.md for the project-wide flag taxonomy.

Dependencies

External crates this crate depends on (pinned at workspace root):

• blake3 — stable cross-platform fingerprints for the kernel cache key.
• cranelift-codegen — CLIF block analysis driving offload detection.
• dashmap — lock-free storage for the kernel cache.
• lru — eviction policy queue.
• parking_lot — non-poisoning mutex backing the eviction queue.
• thiserror — derive macro for the JIT pipeline's error enum.
• tracing — structured spans/events for compilation and deopt.
• wasm-encoder / wasmparser — the rewrite pipeline's read/write Wasm machinery.

Internal: tensor-wasm-core (only for the optional TensorWasmMetrics handle in DeoptGuard::with_metrics).