Nx.Vulkan.VulkanoBackend (nx_vulkan v0.3.0)

Copy Markdown View Source

Pure-Rust (vulkano) Nx.Backend implementation.

Tensors are represented by:

%Nx.Vulkan.VulkanoBackend{ref: ResourceArc<VulkanoTensor>,
                          shape: tuple, type: {kind, bits}}

The ref is a Rustler resource owning an Arc<Subbuffer<u8>> in vulkano. When the BEAM GCs the Elixir reference, vulkano's Drop runs vkDestroyBuffer + vkFreeMemory. Stale-handle bugs (where a freed VkBuf* is read back at the C++ layer) are structurally impossible: the Subbuffer cannot outlive its Buffer.

Dtypes

Native f32 and f64 shaders for the hot ops; the tensor's dtype picks the SPIR-V module at dispatch time. f64 is the default accumulator policy because correctness came first; f32 wins on bandwidth-bound work and is roughly 32× the rate for dot on consumer NVIDIA cards. Other dtypes (integers, u8, …) take the host fallback below.

Coverage

Storage callbacks: init/1, from_binary/3, to_binary/2, backend_copy/3, backend_transfer/3, backend_deallocate/1, inspect/2, constant/3, iota/3, eye/2.

Dispatched to the GPU as native shaders:

  • elementwise binary — add, multiply, subtract, divide, pow, max, min, plus a rank-≤4 broadcasting variant that keeps bias-add, scaling, and relu-via-max resident instead of host-falling-back
  • elementwise unary — exp, log, sqrt, abs, negate, sigmoid, tanh, floor, ceil, sign
  • dot, conv (im2col + GEMM), transpose
  • reductions — sum, reduce_max, reduce_min

Everything else Nx asks for is still implemented, via a host fallback that reads the tensor back, computes on Nx.BinaryBackend, and returns the result to the GPU — argsort, fft, triangular_solve, the window ops, the trig and bitwise families, and so on. Unsupported here means slower, not broken: every Nx callback returns a correct result on this backend.

Eager vs. fused

Used directly, this backend is eager — one dispatch per op, with an intermediate buffer between each. To fuse a chain of ops into a single shader and keep intermediates on-device, jit through Nx.Vulkan.Compiler instead.

Summary

Functions

Accumulator width for the f32 GPU GEMM path — governs both dot/matmul and conv's GEMM: :f64 (default, accuracy-safe) or :f32 (faster on f64-rate- limited GPUs, precision degrades ~√K). Set with put_f32_matmul_accumulator/1 or config :nx_vulkan, :f32_matmul_accumulator.

Set the f32 GEMM accumulator policy (:f64 | :f32). See f32_matmul_accumulator/0.

Functions

f32_matmul_accumulator()

Accumulator width for the f32 GPU GEMM path — governs both dot/matmul and conv's GEMM: :f64 (default, accuracy-safe) or :f32 (faster on f64-rate- limited GPUs, precision degrades ~√K). Set with put_f32_matmul_accumulator/1 or config :nx_vulkan, :f32_matmul_accumulator.

put_f32_matmul_accumulator(width)

Set the f32 GEMM accumulator policy (:f64 | :f32). See f32_matmul_accumulator/0.