Nx.Vulkan.Compiler (nx_vulkan v0.2.0)

Copy Markdown View Source

Nx.Defn.Compiler for the Vulkan backend — thrust 3, the fusion compiler.

EXLA's structural advantage over an eager backend is whole-graph compilation: it fuses a chain of elementwise ops into a single kernel instead of dispatching each op separately (each with its own launch + intermediate buffer). This compiler does the same for the cases it supports.

At __compile__ time it traces the defn to an Nx.Defn.Expr tree. If the single output is a same-shape elementwise chain (see Nx.Vulkan.Codegen), it generates one GLSL shader for the whole chain, compiles it once (cached by source hash), and returns a function that uploads the inputs, issues ONE dispatch_generated call, and hands back a GPU-resident result. Anything it can't fuse — tuple outputs, batched dot, broadcasting between differing shapes, integer dtypes — falls through to Nx.Defn.Evaluator, so results are always correct; the worst case is "no fusion, same as eager."

Dtypes

{:f, 32} and {:f, 64} both fuse; the generated shaders are parameterised on the element type and every buffer is sized by it. A graph must be of ONE float dtype throughout — mixed precision would need a cast stage this compiler does not emit, so it falls back.

Two f64 caveats:

  • No f64 transcendentals exist in SPIR-V (GLSL.std.450 defines exp/log/pow/tanh for 16- and 32-bit floats only), so an f64 tree containing one falls back whole-graph rather than silently boundary-casting through f32. f64 arithmetic, sqrt, abs, sign, min/max, floor/ceil/round all fuse — GLSL.std.450 covers 64-bit for those.
  • f64 fusion requires the device to advertise shaderFloat64 (Nx.Vulkan.Device.f64?/0, override NXV_F64=0|1). glslangValidator compiles fp64 GLSL regardless of the device, so the failure would otherwise surface only at dispatch.

f64 matmul is ~1/32 rate on consumer NVIDIA cards, so f64 fusion is about coverage, not speed.

Usage

Nx.Defn.jit(&my_fun/2, compiler: Nx.Vulkan.Compiler).(a, b)

Set NXV_FUSE_DEBUG=1 to log which path each defn takes.

Reductions

An elementwise chain feeding a reduction (sum/reduce_max/reduce_min) is fused into a single parallel workgroup-per-slot shared-memory tree reduce (Codegen.emit_fused_reduce + dispatch_generated_reduce), which grid-strides over output slots so one launch handles any slot count. It beats even the eager path, whose own reduce_axis is one-thread-per-slot serial. Enabled by default for a contiguous reduce (inner_stride == 1) with FEW output slots — full reductions and small-output reductions — which win across the fleet: ~8-27x over eager on the GT 650M and ~2.8-6.7x on the RTX 3060 Ti (the case where EXLA had out-run the eager backend). The many-slot wide-reduce regime is grid- stride-capable and wins on the weak Kepler eager path (~4.4x) but REGRESSES on the much stronger Ampere eager path (0.44x), so it is hardware-dependent: it is auto-enabled only on GPUs Nx.Vulkan.Device classifies :weak (integrated / software / older low-end discrete), and stays off on strong GPUs. Force it on any GPU with NXV_FUSE_REDUCE=1. Non-contiguous, short-axis and mid-slot reductions fall back to the already-parallel eager path — no regressions. =0 disables all reduce fusion. See reduce_beneficial?/3 and Nx.Vulkan.Device.