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++ Elixir backend has since been removed (commit bb94217); Nx.Vulkan.VulkanoBackend is the only backend. It was preferred 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

Test coverage (2026-08): 863 doctests, 361 tests, 0 failures on the fleet — GT 650M (Kepler, FreeBSD), GT 750M (mac-248), RTX 3060 Ti (Ampere, Linux). The spirit C++ backend and its test suite were dropped. Bench coverage committed to bench_results/.

Stage breakdown

Stages are sized to land in one focused session each.

Stage 1 — Elementwise binary (DONE)

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 (DONE)

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 (DONE — sum/reduce_max/reduce_min; non-trivial axis sets fall back to host)

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

Stage 4 — Shape / movement (PARTIAL DONE — reshape, squeeze, 2D transpose [1,0] on GPU; broadcast/slice/pad/concatenate/gather/take on host fallback per Tier 1 of SHAPE_C_PLAN.md)

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

Stage 5 — Linalg (PARTIAL — dot/matmul rank-2 f64 fast path on GPU; cholesky/solve/qr/svd via block/4 host fallback; native SPV impls TODO)

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 (PARTIAL — comparison, select, all, any on host fallback; Random TODO)

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

Stage 7 — Defn integration (DONE for Evaluator path — pin global default to VulkanoBackend at boot (Application.start), route Exmc.JIT.jit through Nx.Defn.Evaluator instead of Nx.Vulkan.jit (which would force the spirit backend). Custom Defn compiler TODO.)

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 (DONE — forward op coverage IS gradient coverage. Nx.Defn.grad is a graph transformation; once forward ops exist, gradients automatic. Validated end-to-end on Axon training step on the spirit backend; vulkano path inherits via the same Defn substrate.)

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

Stage 9 — Axon parity (DONE — Axon training loop ran end-to-end on VulkanoBackend; matches BinaryBackend reference to 8.6e-8 on the dense_0 kernel gradient sum.)

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 (DONE — Scholar LinearRegression smoke-test passed via the block/4 host-fallback path. Native SVD / cholesky impl TODO before declaring full parity.)

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

Stage 11 — Performance pass (IN FLIGHT — Tier 1 of SHAPE_C_PLAN.md landed: host-fallback ops skip the upload-back round trip and return BinaryBackend tensors. Consumer bench shows median ~1.25-1.3x speedup when result is read via to_flat_list. Persistent buffer pool, disk pipeline cache, native shaders for the bandwidth-bound four (broadcast, pad, concatenate, put_slice) all TODO — see SHAPE_C_PLAN.md Tier 2.)

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 (shipped) and f32 compute (shipped) — the hot ops (elementwise, matmul, conv, reduce, transpose) have native f32 and f64 shaders and dtype-dispatch on the tensor type. f64 is the default accumulator policy (correctness first; consumer GPUs are slower at f64), but f32 is native — no longer merely cast — and wins on bandwidth-bound ops.
• 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.

5. Multi-device on a single machine. mac-247 (FreeBSD 15 + 2013-era MacBook Pro) has the GT 650M Mac Edition AND an Intel HD Graphics 4000 (Ivy Bridge iGPU). pciconf -lv confirms both on the PCI bus:

   vgapci1: Intel HD Graphics 4000      (vendor 0x8086, dev 0x0166)
   vgapci0: NVIDIA GT 650M Mac Edition  (vendor 0x10de, dev 0x0fd5)

   Currently only NVIDIA is exposed to Vulkan. vulkaninfo shows llvmpipe as the second device (Mesa's software Vulkan, not the iGPU). To surface the Intel iGPU:

   • Load i915kms + drm-kmod
   • Confirm FreeBSD graphics/mesa-libs ships with the anv Intel Vulkan driver enabled for x86_64
   • Investigate Apple MUX state — early-2013 MBPs may hard-route the iGPU into low-power/standby when discrete is active

   Even with both surfaced, nx_vulkan's ctx() picks the first DiscreteGpu and ignores everything else. Multi-device routing would need:

   • Device selection in ctx() (env var, config, or runtime API)
   • Per-device pipeline cache + allocator
   • Either device-affinity tags on Nx.Vulkan.VulkanoBackend tensors, or a workload-router that chooses device per op

   Performance ceiling estimate: GT 650M = ~691 GFLOPS f32; HD 4000 = ~330 GFLOPS f32 (no f64). Theoretical +50% peak on mac-247; realistic +20-30% on well-partitioned workloads.

   Filed as long-term — not blocking M-II or W-stage work. Pick up if dual-device compute on the legacy MBP becomes interesting (e.g., for a "compute fabric from yesterday's hardware" demo).