Nx.Vulkan.NativeV (nx_vulkan v0.2.0)

Copy Markdown View Source

Rustler NIF for the vulkano compute backend.

Resource lifetimes are managed by Rust ownership (Arc<Buffer>). When the BEAM GCs the Elixir reference, vulkano's Drop runs and the GPU memory is freed.

Builds on Linux + FreeBSD 15.0 with vulkano 0.34.

Summary

Functions

Elementwise binary op. op_code selects which operation the shader executes via a specialisation constant

Broadcasting elementwise binary op (rank <= 4). Bindings a/b/out/params; params is [rank, out[4], a[4], b[4]] int32; n = output element count. Keeps bias-add / scaling / relu-via-max on the GPU. op codes: 0=add 1=mul 2=sub 3=div 5=max 6=min (pow excluded — fp64 has no pow).

Broadcasting comparison -> u8 (packed as u32). Bindings a/b/out/params; op codes 0=eq 1=ne 2=lt 3=le 4=gt 5=ge. out buffer must be padded to a 4-byte multiple (ceil(n/4) u32 words). Keeps mask-producing ops on the GPU.

Gather (leading-prefix axes). Bindings in/out/indices/params; params is [K, ews, idx_words, count, stride[4]] int32 (ews = value_bytes/4, idx_words = index_bytes/4, count = product of trailing non-indexed dims); n = output element count, k = number of indexed leading axes. Keeps the common gather (default all-axes / leading-prefix axes) on the GPU.

Pad (type-generic copy). Bindings in/out/params/pad-value; params is [rank, ews, S[4], O[4], low[4], interior[4]] int32 (ews = element_bytes/4); pad_value buffer is one element (ews u32 words); n = output element count. Elements landing in an edge pad, interior gap, or outside the source get the pad value. Keeps pad on the GPU.

Broadcasting select: out = pred ? t : f. Bindings pred/t/f/out/params; params is [rank, out[4], pred[4], t[4], f[4]] int32; n = output element count. pred is u8 (read as u32 words in the shader). Keeps masking / where / relu-grad on the GPU.

Strided slice (type-generic u32-word copy). Bindings in/out/params; params is [rank, ews, S[4], O[4], start[4], stride[4]] int32 (ews = element_bytes/4); n = output element count. Keeps static-start slices on the GPU.

Elementwise unary op. op_code selects

Explicit broadcast, rank <= 4. Semantics match Nx.broadcast(t, shape, axes: axes) — input axis i lands on output axis axes[i], and an input axis of size 1 repeats.

Allocate a zero-initialised device buffer of n_bytes. Returns {:ok, ref}.

Buffer size in bytes (returns integer, never crashes on a valid resource).

Read a device buffer back to a host binary. Returns {:ok, binary}.

Allocate a device buffer + upload data to it. Returns {:ok, ref}. The ref is a Rustler resource that owns the underlying Arc<Buffer> — when the BEAM GCs it, vulkano's Drop runs and the GPU memory is freed.

Overwrite an existing device buffer with new host data. Returns :ok or {:error, :size_mismatch} when sizes disagree.

Elementwise dtype cast (in binding 0 -> out binding 1). The shader defines the source/dest types; n = element count. Used for f32<->f64 on the GPU.

Concatenate N device buffers into a single fresh buffer via vkCmdCopyBuffer (no shader). Inputs are copied in list order into the destination starting at offset 0; total output size = Σ inputs[i].n_bytes. Returns {:ok, output_ref}.

Conv GEMM: out{N,Cout,O_total} = im2col(A){M,K} · kernel{Cout,K}, written in canonical output layout. f64.

im2col unfold for conv (spatial rank <= 3, groups == 1), f64. Fills the column matrix col (M×K = NO_total × CinK_total) from in. params is a 21-int buffer of per-dim [D,O,K,stride,pad_lo,input_dil,kernel_dil].

Physical Vulkan device name + type, e.g. {:ok, "NVIDIA GeForce GT 650M", "DiscreteGpu"}. Used to label benchmark/parity reports per host.

Whether the physical device advertises shaderFloat64{:ok, boolean}. The _f64.spv shaders and any generated f64 kernel require it; on a device without it, pipeline creation for those fails at dispatch time, so callers must gate on this and take a host fallback instead.

Generic JIT-shader dispatch (thrust 3 fusion compiler). Runs a runtime- generated shader: inputs at bindings 0..k-1 (in the order of in_refs), output at binding k, push {n = element count}. One dispatch replaces a whole fused elementwise chain generated by Nx.Vulkan.Codegen.

Generic JIT fused-reduce dispatch (thrust 3). Runtime-generated shader that fuses an elementwise chain into a reduction: inputs at bindings 0..k-1, output at k, push {outer, reduce_size, inner}. One invocation per output slot; the reduce op (sum f64-acc / max / min) is baked into the generated shader.

Radix-2 Cooley-Tukey FFT (power-of-two, last-axis, batched) in f64.

Boundary-cast f64 leapfrog chain dispatch.

Register-blocked matmul dispatch (32-wide output tiles) for the *_rb32 shaders. Benchmark-only — the register-blocked kernels regressed on Kepler; used by examples/matmul_rb_race.exs to evaluate them on other GPUs. See F32_PLAN.md.

2D matmul. C = A · B where A is M×K row-major, B is K×N row-major, C is M×N row-major. All f32. Buffers: a (mk4), b (kn4), out (mn4).

Per-axis reduction. op_code: 0=sum, 1=max, 2=min.

Reverse along a set of axes, rank <= 4. Semantics match Nx.reverse(t, axes: axes).

2D transpose. Input A is M×N row-major; output is N×M row-major. Buffers: a (m*n*4 bytes), out (m*n*4 bytes).

Generic permuted transpose for rank <= 4. Semantics match Nx.transpose(t, axes: perm): output axis d is input axis perm[d].

Functions

apply_binary(out, a, b, n, op_code, spv_path)

Elementwise binary op. op_code selects which operation the shader executes via a specialisation constant:

0=add  1=mul  2=sub  3=div  4=pow  5=max  6=min

Buffers must all be the same byte size. Returns :ok or {:error, :size_mismatch} / {:error, :dispatch_failed, msg}.

apply_binary_broadcast(out, a, b, params, n, rank, op_code, spv_path)

Broadcasting elementwise binary op (rank <= 4). Bindings a/b/out/params; params is [rank, out[4], a[4], b[4]] int32; n = output element count. Keeps bias-add / scaling / relu-via-max on the GPU. op codes: 0=add 1=mul 2=sub 3=div 5=max 6=min (pow excluded — fp64 has no pow).

apply_compare(out, a, b, params, n, rank, op_code, spv_path)

Broadcasting comparison -> u8 (packed as u32). Bindings a/b/out/params; op codes 0=eq 1=ne 2=lt 3=le 4=gt 5=ge. out buffer must be padded to a 4-byte multiple (ceil(n/4) u32 words). Keeps mask-producing ops on the GPU.

apply_gather(out, in, idx, params, n, k, spv_path)

Gather (leading-prefix axes). Bindings in/out/indices/params; params is [K, ews, idx_words, count, stride[4]] int32 (ews = value_bytes/4, idx_words = index_bytes/4, count = product of trailing non-indexed dims); n = output element count, k = number of indexed leading axes. Keeps the common gather (default all-axes / leading-prefix axes) on the GPU.

apply_pad(out, in, params, padval, n, rank, spv_path)

Pad (type-generic copy). Bindings in/out/params/pad-value; params is [rank, ews, S[4], O[4], low[4], interior[4]] int32 (ews = element_bytes/4); pad_value buffer is one element (ews u32 words); n = output element count. Elements landing in an edge pad, interior gap, or outside the source get the pad value. Keeps pad on the GPU.

apply_select(out, pred, t, f, params, n, rank, spv_path)

Broadcasting select: out = pred ? t : f. Bindings pred/t/f/out/params; params is [rank, out[4], pred[4], t[4], f[4]] int32; n = output element count. pred is u8 (read as u32 words in the shader). Keeps masking / where / relu-grad on the GPU.

apply_slice(out, in, params, n, rank, spv_path)

Strided slice (type-generic u32-word copy). Bindings in/out/params; params is [rank, ews, S[4], O[4], start[4], stride[4]] int32 (ews = element_bytes/4); n = output element count. Keeps static-start slices on the GPU.

apply_unary(out, a, n, op_code, spv_path)

Elementwise unary op. op_code selects:

0=exp  1=log  2=sqrt  3=abs  4=neg  5=sigmoid  6=tanh  7=relu
8=ceil  9=floor  10=sign  11=reciprocal  12=square

Buffers must be the same byte size.

broadcast_nd(out, a, params, n, spv_path)

Explicit broadcast, rank <= 4. Semantics match Nx.broadcast(t, shape, axes: axes) — input axis i lands on output axis axes[i], and an input axis of size 1 repeats.

Buffers: a, out, params — an int32 buffer laid out as [out_rank, in_rank, out[4], in[4], axes[4]].

buf_alloc(n_bytes)

Allocate a zero-initialised device buffer of n_bytes. Returns {:ok, ref}.

buf_byte_size(ref)

Buffer size in bytes (returns integer, never crashes on a valid resource).

buf_download(ref)

Read a device buffer back to a host binary. Returns {:ok, binary}.

buf_upload(data)

Allocate a device buffer + upload data to it. Returns {:ok, ref}. The ref is a Rustler resource that owns the underlying Arc<Buffer> — when the BEAM GCs it, vulkano's Drop runs and the GPU memory is freed.

buf_upload_into(ref, data)

Overwrite an existing device buffer with new host data. Returns :ok or {:error, :size_mismatch} when sizes disagree.

cast(out, a, n, spv_path)

Elementwise dtype cast (in binding 0 -> out binding 1). The shader defines the source/dest types; n = element count. Used for f32<->f64 on the GPU.

concat_buffers(inputs)

Concatenate N device buffers into a single fresh buffer via vkCmdCopyBuffer (no shader). Inputs are copied in list order into the destination starting at offset 0; total output size = Σ inputs[i].n_bytes. Returns {:ok, output_ref}.

Tier 2 step 1 of SHAPE_C_PLAN.md — keeps the result on the device so downstream ops don't pay the download+upload round trip that the host-fallback path imposed.

conv_gemm(out, col, kernel, n, cout, o_total, k, spv_path)

Conv GEMM: out{N,Cout,O_total} = im2col(A){M,K} · kernel{Cout,K}, written in canonical output layout. f64.

conv_im2col(col, in, params, n, cin, o_total, k_total, k, spv_path)

im2col unfold for conv (spatial rank <= 3, groups == 1), f64. Fills the column matrix col (M×K = NO_total × CinK_total) from in. params is a 21-int buffer of per-dim [D,O,K,stride,pad_lo,input_dil,kernel_dil].

device_name()

Physical Vulkan device name + type, e.g. {:ok, "NVIDIA GeForce GT 650M", "DiscreteGpu"}. Used to label benchmark/parity reports per host.

device_supports_f64()

Whether the physical device advertises shaderFloat64{:ok, boolean}. The _f64.spv shaders and any generated f64 kernel require it; on a device without it, pipeline creation for those fails at dispatch time, so callers must gate on this and take a host fallback instead.

dispatch_generated(out, in_refs, n, spv_path)

Generic JIT-shader dispatch (thrust 3 fusion compiler). Runs a runtime- generated shader: inputs at bindings 0..k-1 (in the order of in_refs), output at binding k, push {n = element count}. One dispatch replaces a whole fused elementwise chain generated by Nx.Vulkan.Codegen.

dispatch_generated_reduce(out, in_refs, outer, reduce_size, inner, spv_path)

Generic JIT fused-reduce dispatch (thrust 3). Runtime-generated shader that fuses an elementwise chain into a reduction: inputs at bindings 0..k-1, output at k, push {outer, reduce_size, inner}. One invocation per output slot; the reduce op (sum f64-acc / max / min) is baked into the generated shader.

fft(out, in, n, logn, batch, is_complex, inverse, bitrev_spv, stage_spv)

Radix-2 Cooley-Tukey FFT (power-of-two, last-axis, batched) in f64.

out is a complex buffer of batch*n*16 bytes (interleaved re/im f64). in is either real f64 (batch*n*8 bytes, is_complex = 0) or complex (batch*n*16, is_complex = 1). logn = log2(n). inverse = 1 applies the 1/n normalisation and the conjugate twiddle. Two shaders: bit-reversed load then log2(n) butterfly stages.

leapfrog_chain_synth(q, p, extras, push, k, spv_path)

leapfrog_chain_synth_batch(q, p, extras, push, k, spv_path)

leapfrog_chain_synth_f64(q, p, extras, push, k, spv_path)

Boundary-cast f64 leapfrog chain dispatch.

Same dispatch contract but all binaries are little-endian f64 (8 bytes per element). Push block is 24+ bytes (eps is f64 at byte offset 16). SPV at spv_path must be the f64-compiled synth shader (uses GL_ARB_gpu_shader_fp64 for storage; transcendentals via emitter-generated double(log(float(x))) wrappers).

Returns {:ok, {q_chain_bin, p_chain_bin, grad_chain_bin, logp_chain_bin}} as little-endian f64 binaries — each q/p/grad is K * d * 8 bytes; logp is K * 8 bytes.

See docs/EXMC_VULKAN_DOS_AND_DONTS.md for why this two-variant split exists (GLSL.std.450 has no f64 transcendentals at the Khronos spec layer; boundary-cast is the workaround).

matmul32(out, a, b, m, n, k, spv_path)

Register-blocked matmul dispatch (32-wide output tiles) for the *_rb32 shaders. Benchmark-only — the register-blocked kernels regressed on Kepler; used by examples/matmul_rb_race.exs to evaluate them on other GPUs. See F32_PLAN.md.

matmul(out, a, b, m, n, k, spv_path)

2D matmul. C = A · B where A is M×K row-major, B is K×N row-major, C is M×N row-major. All f32. Buffers: a (mk4), b (kn4), out (mn4).

reduce_axis(out, a, outer, reduce_size, inner, op_code, spv_path)

Per-axis reduction. op_code: 0=sum, 1=max, 2=min.

Input shape is interpreted as a virtual (outer, reduce_size, inner) tensor; output shape is (outer, inner) — i.e. the reduction collapses the middle axis. For full reductions use outer=1, reduce_size=n, inner=1.

Buffers: out has outer * inner elements; a has outer * reduce_size * inner elements.

reverse_nd(out, a, params, n, rank, spv_path)

Reverse along a set of axes, rank <= 4. Semantics match Nx.reverse(t, axes: axes).

Buffers: a, out (n * element_bytes each), params — an int32 buffer laid out as [rank, shape[4], rev[4]], where rev[d] is 1 when axis d is reversed. Input and output share a shape.

transpose_2d(out, a, m, n, spv_path)

2D transpose. Input A is M×N row-major; output is N×M row-major. Buffers: a (m*n*4 bytes), out (m*n*4 bytes).

transpose_nd(out, a, params, n, rank, spv_path)

Generic permuted transpose for rank <= 4. Semantics match Nx.transpose(t, axes: perm): output axis d is input axis perm[d].

Buffers: a, out (n * element_bytes each), params — an int32 buffer laid out as [rank, in[4], out[4], perm[4]] with shapes left-aligned and padded to 4.