vkML 0.1.0

Shaders and the GLSL layer

24 compute shaders, one shared preamble, and the conventions that let one module serve many operators.

Descriptor-less binding

Buffers arrive as 64-bit device addresses in push constants, not through descriptor sets. Using bufferDeviceAddress deletes descriptor pools, set layouts, per-dispatch vkUpdateDescriptorSets and the pool-growth logic that ggml-vulkan needs — roughly 200–500 lines of the most error-prone code in a Vulkan backend, plus real per-dispatch CPU cost.

scalarBlockLayout means these structs lay out identically in GLSL and C++, so a push block is declared once per operator and mirrored by a plain C++ struct with no std140 padding rules to get wrong.

One module, many pipelines

The operation is a specialisation constant, so the driver folds the switch away at pipeline creation and each variant is as tight as a dedicated shader. That is why unary.comp serves 20 operators and binary.comp serves 13 without any of them paying for the others.

Comparisons live in binary.comp alongside arithmetic because everything except the final store is identical — same broadcast indexing, same operand layout, same bounds check. Only the destination element type differs, and that is decided at pipeline creation too.

Why f32→f16 narrowing is done in software

A load from an f16 buffer is a hardware widening, which is exact and needs no help. A store goes through f32_to_f16_bits, an integer-domain round-to-nearest-even routine, rather than float16_t(value).

⚠ Warning

SPIR-V leaves OpFConvert's rounding mode implementation-defined. RADV rounds to nearest even; AMD's Windows compiler rounds toward zero. The same program therefore produced different f16 results on the two, and the cross-backend oracle — which compares bit for bit against the CPU's round-to-nearest-even routine — failed on Windows.

Determinism across drivers is a project invariant, so the fix could not be a tolerance. It could have been the RoundingModeRTE execution mode from VK_KHR_shader_float_controls, but that needs a device capability, a fallback for devices without it, and trust that the driver honours it — three things not verifiable from inside the shader. The routine used instead contains no floating-point operation whose rounding a driver could choose: integer shifts and comparisons only.

It is bit-for-bit the same function as vkml::fp32_to_fp16 in src/core/dtype.cpp, which reaches the same result by a different route, and the two are checked against each other over the whole f32 exponent range.

f16 is storage, never arithmetic

Both conversion helpers sit at the memory boundary and everything between them is float. The dtype is a specialisation constant at every call site, so the branch is folded away at pipeline creation and an f32 kernel compiles to exactly what it did before f16 existed.

This is deliberately written to look like the CPU backend's widen, because it implements the same half of the numerical contract.

Rank 4, mirrored from the C++ side

kMaxDims = 4 keeps a three-operand shape/stride block at 96 bytes, comfortably inside the budget; rank 8 would need 192 and force the metadata into a uniform buffer, adding an indirection to every kernel.

Vector loads are not automatically one instruction

Tile loading uses 4-wide access with buffer_reference_align = 4 rather than 16. scalarBlockLayout permits a vec4 at 4-byte alignment, so a tile row not starting on a 16-byte boundary is still legal — but the driver splits the access when it cannot prove alignment, which is the correct fallback and also means a vec4 load is not automatically a single instruction.

Shared helpers

HelperWhat it does
global_index() Flat invocation index, folding y into x. An identity while y holds a single group, so the ordinary case is unchanged.
global_group_index() The same at workgroup granularity, for kernels that index by group — reductions, softmax, GEMV and the three GEMM variants.
offset_from() Byte offset of a logical index under arbitrary strides. Broadcasting is free: a stride of 0 contributes nothing.
f32_to_f16_bits() Software round-to-nearest-even narrowing, as above.

The shaders

Generated from the tree — adding a shader adds a row, and changing one changes its numbers. "Ops" counts the OP_ constants a module dispatches on; a dash means the shader serves one operation.

ShaderLinesPushSpecSharedBarriersOps
binary.comp132124 B3013
cast.comp5720 B20
cat.comp86112 B20
col2im.comp10960 B90
fill.comp3620 B10
gemm_db.comp306108 B10yes1
gemm_naive.comp102108 B20
gemm_reg.comp276108 B10yes2
gemm_split_k_reduce.comp8624 B20
gemm_tiled.comp150108 B3yes2
gemv.comp203108 B6yes5
im2col.comp8760 B70
index_select.comp50104 B10
max_pool2d.comp18160 B110
probe_mlp.comp8120 B2yes1
probe_private_array.comp9712 B40
rand.comp7028 B00
reduce.comp28688 B4yes26
scaled_add.comp6784 B20
scatter_add.comp71104 B10
slice_backward.comp70100 B00
softmax.comp153120 B4yes7
tri.comp5396 B30
unary.comp29292 B3020
where.comp64116 B20

Compilation

Shaders are compiled to SPIR-V at build time by glslc or glslangValidator from the Vulkan SDK, and the result is embedded in the binary — so a built vkML has no runtime dependency on a shader compiler and cannot fail at first dispatch because a .comp file moved.

vkML — Vulkan-first machine learning in C++20. Apache-2.0. Signatures on this page are generated from the installed module.