Nx.Vulkan.VulkanoBackend — Roadmap

Primary objective (2026-05-20 onward): make Nx.Vulkan.VulkanoBackend a viable Nx backend for the three target ecosystems — exmc (NUTS sampling on FreeBSD), Axon (neural networks, with autograd), Scholar (classical ML, with linalg).

Previously: Nx.Vulkan.Backend (C++ spirit) was the Vulkan backend. The C++ path is now legacy; new work targets the vulkano path because:

• Resource lifetimes are managed by Rust ownership (Arc<Buffer> + Subbuffer<u8>), eliminating the stale-handle bug class that bit the R4 cutover.
• vulkano builds + runs cleanly on FreeBSD 15.0 and Linux without vendor-specific shims.
• vulkano matches the C++ spirit path's dispatch latency within ~10% on the bench target (GT 650M).
• Per-op shaders (the existing SPV catalog under priv/shaders/) load and dispatch identically — no shader rewrite needed.

Where we are

Stage breakdown

Stages are sized to land in one focused session each.

Stage 1 — Elementwise binary

Ops: add, subtract, multiply, divide, pow, max, min.

NIF: apply_binary(out_ref, a_ref, b_ref, n, op_code, spv_path) — takes 3 buffer refs, dispatches elementwise_binary.spv (already in priv/shaders/) with the op selected via specialization constant. Push block: uint n. Workgroup 256, ceil(n/256) groups.

VulkanoBackend callbacks: 7 op handlers that allocate an output buffer and call apply_binary. Validation: head-to-head against Nx.BinaryBackend for each op on f32 tensors.

Stage 2 — Elementwise unary

Ops: exp, log, sqrt, abs, negate, sigmoid, tanh, relu (clamp to 0), ceil, floor, sign, reciprocal, square, erf, expm1.

NIF: apply_unary(out_ref, a_ref, n, op_code, spv_path). Same pattern as binary, one input. SPV: elementwise_unary.spv.

Stage 3 — Reductions

Ops: sum, reduce_max, reduce_min over all axes (full reduction to scalar). Then per-axis via reduce_axis.spv.

Stage 4 — Shape / movement

Ops: reshape (zero-copy ref rewrap), squeeze, broadcast (GPU-side broadcast shader for non-zero-stride cases), transpose, slice, gather.

Stage 5 — Linalg

Ops: dot/6 (matmul), cholesky, solve, qr, svd, determinant. Some of these need new shaders; matmul has multiple tilings already in priv/shaders/.

Stage 6 — Random + comparison + select

Ops: Nx.Random.* (Philox-backed), less/greater/equal/ not_equal, select.

Stage 7 — Defn integration

So defn blocks targeting Nx.Vulkan.VulkanoBackend work end-to- end. May require a custom Nx.Defn compiler or routing through the existing Vulkan-aware compiler with vulkano backend.

Stage 8 — Autograd primitives

For Axon: implement gradients of all stage-1–6 ops. Most are automatic via Nx.Defn.Grad once forward-pass ops exist; some need custom adjoint impls.

Stage 9 — Axon parity

Run a small Axon model (MLP, small CNN) end-to-end on Nx.Vulkan.VulkanoBackend. Compare loss + gradients against BinaryBackend reference.

Stage 10 — Scholar parity

Run k-means or PCA on Nx.Vulkan.VulkanoBackend. The linalg ops from stage 5 are the gate.

Stage 11 — Performance pass

Add persistent buffer pool, vulkano SubbufferAllocator integration, pipeline cache to disk (vulkano's PipelineCache::with_data). Compare to C++ spirit + EXLA on Axon training step / sec.

Performance target

For exmc on GT 650M: regime-model NUTS sample ≤500 ms (already met via the synthesised chain shader). For Axon on FreeBSD: at least half of EXLA's throughput on the same hardware where EXLA runs.

Non-goals

• f64 compute (vulkano supports it but most consumer GPUs are ~32× slower at f64). Storage f64 is fine; compute defaults to f32 with as_type cast.
• CUDA-specific features (tensor cores, mixed precision) — vulkano abstracts over them, but extracting them is out of scope until stages 1–10 are done.
• Multi-GPU. Single device per process for now.

Open architectural questions

1. Persistent buffer pool. Per-call alloc/free works but hits the allocator on every op. A SubbufferAllocator keyed by size class would amortise this. Defer until stage 11.

2. Pipeline cache. vulkano supports PipelineCache::with_data for disk-persisted compiled pipelines. Plumb through after stage 5.

3. Defn compiler. EXLA has its own; we'd need either a Nx.Defn.Compiler impl that knows how to dispatch through Nx.Vulkan.NativeV, or rely on Nx.Defn.Evaluator driving the backend op-by-op. Stage 7 decides.

4. Hex publish strategy. Once stages 1–6 land, publish a 0.1 nx_vulkan_vulkano package. Existing nx_vulkan keeps the C++ path until parity is comfortable.