JIT GLSL codegen from Nx.Defn.Expr trees — the heart of the thrust-3
fusion compiler.
Nx.Vulkan.Compiler traces a defn to an expression tree; this module
turns a fully-fusable elementwise subtree into a single GLSL compute shader
(one thread per element, the whole chain inlined) and compiles it to SPIR-V,
cached by source hash. One generated shader + one dispatch_generated call
replaces the N per-op dispatches the eager backend would issue.
Scope: same-shape elementwise chains over unary and binary arithmetic ops plus
scalar constants, at {:f, 32} or {:f, 64} — the emitters are parameterised
on the element type (glsl_scalar/1), so one code path emits float or
double SSBOs, temps and literals. A tree must be of one dtype throughout.
At f64 the transcendentals are excluded (@f64_unsafe_ops): GLSL.std.450 has
no 64-bit exp/log/pow/tanh. Comparisons (which change dtype to u8) are not
fused — Nx.Vulkan.Compiler falls back to the Evaluator for those, so
behaviour stays correct.
Ported from the dropped 577baf9/9a9e3ad codegen, retargeted to Nx 0.13 and
the current Nx.Vulkan.NativeV dispatch primitives.
Summary
Functions
Compile GLSL to a cached .spv, keyed by source hash. Returns
{:ok, spv_path} or {:error, reason}. Reuses an existing .spv on hit,
so a given fused kernel is compiled by glslangValidator exactly once.
Emit a GLSL compute shader for a fusable elementwise expression.
Emit a GLSL shader that fuses an elementwise inner chain into a reduction
over the (outer, reduce_size, inner_stride) view, using a parallel
workgroup-per-slot shared-memory tree reduce: each output slot gets one
workgroup of 256 threads that stride the reduce axis, accumulate into
a shared array, then tree-reduce to a single value. sum accumulates in f64
to match BinaryBackend.
Like emit_fused_reduce/3 but for a reduction region whose leaf inputs are
given explicitly as inputs ({{:param, pidx} | {:stage, node_id}, shape}) —
so the reduced inner chain may read earlier stages' output buffers, not just
parameters. Used by the multi-stage split to materialise a reduce as a stage
(e.g. mean(x) in an x - mean(x) layernorm graph). Returns {glsl, %{n_inputs: k}}.
Emit an elementwise shader for a fusion region whose leaf inputs are given by
inputs — an ordered list of {{:param, pidx} | {:stage, node_id}, shape}.
Parameters and stage-materialised buffers are both loaded from input bindings
(broadcast-aware); constants inline. Used by both the single-region compile
and the multi-stage split (where a leaf may be a prior stage's output buffer).
Returns {glsl, %{n_inputs: k}}.
True when the whole expression tree is a same-shape elementwise chain
(unary/binary ops, parameters, scalar constants) of a single float dtype
({:f, 32} or {:f, 64}) that emit_elementwise/1 can compile into one
shader. Every interior node must share the root's dtype — a mixed-precision
tree needs an as_type cast, which is not a fusable op anyway.
Set of ops that can be fused into one elementwise shader (excludes params/constants).
Like fusable_op?/1 but dtype-aware: at {:f, 64} the transcendentals are
excluded (no f64 GLSL.std.450 equivalents — see @f64_unsafe_ops).
True if op is a reduction this module can fuse an elementwise inner into.
Functions
Compile GLSL to a cached .spv, keyed by source hash. Returns
{:ok, spv_path} or {:error, reason}. Reuses an existing .spv on hit,
so a given fused kernel is compiled by glslangValidator exactly once.
Emit a GLSL compute shader for a fusable elementwise expression.
Returns {glsl, %{param_order: [param_index, ...], n_inputs: k}} where
param_order[b] is the runtime argument index bound to input binding b.
Emit a GLSL shader that fuses an elementwise inner chain into a reduction
over the (outer, reduce_size, inner_stride) view, using a parallel
workgroup-per-slot shared-memory tree reduce: each output slot gets one
workgroup of 256 threads that stride the reduce axis, accumulate into
a shared array, then tree-reduce to a single value. sum accumulates in f64
to match BinaryBackend.
This is 256x more parallel than a serial per-slot loop and beats even
the eager path (whose reduce_axis is itself one-thread-per-slot). It is only
valid when the number of slots fits the one-dimensional workgroup-count limit
(maxComputeWorkGroupCount[0], typically 65535) — the caller gates on that.
Dispatch outer*inner workgroups (one per slot), NOT ceil(slots/256).
scale (a number or nil) applies a final / scale to each output slot — this
is how mean fuses: divide(sum(...), n) becomes a fused sum scaled by 1/n.
Returns {glsl, %{param_order: [...], n_inputs: k}}.
Like emit_fused_reduce/3 but for a reduction region whose leaf inputs are
given explicitly as inputs ({{:param, pidx} | {:stage, node_id}, shape}) —
so the reduced inner chain may read earlier stages' output buffers, not just
parameters. Used by the multi-stage split to materialise a reduce as a stage
(e.g. mean(x) in an x - mean(x) layernorm graph). Returns {glsl, %{n_inputs: k}}.
Emit an elementwise shader for a fusion region whose leaf inputs are given by
inputs — an ordered list of {{:param, pidx} | {:stage, node_id}, shape}.
Parameters and stage-materialised buffers are both loaded from input bindings
(broadcast-aware); constants inline. Used by both the single-region compile
and the multi-stage split (where a leaf may be a prior stage's output buffer).
Returns {glsl, %{n_inputs: k}}.
True when the whole expression tree is a same-shape elementwise chain
(unary/binary ops, parameters, scalar constants) of a single float dtype
({:f, 32} or {:f, 64}) that emit_elementwise/1 can compile into one
shader. Every interior node must share the root's dtype — a mixed-precision
tree needs an as_type cast, which is not a fusable op anyway.
Set of ops that can be fused into one elementwise shader (excludes params/constants).
Like fusable_op?/1 but dtype-aware: at {:f, 64} the transcendentals are
excluded (no f64 GLSL.std.450 equivalents — see @f64_unsafe_ops).
True if op is a reduction this module can fuse an elementwise inner into.