Kernel metaprogramming utilities

PTX.Utils collects the metaprogramming idioms register-resident kernels repeat: loops whose induction value must be a compile-time constant (Val arguments, per-iteration variable names, constant tuple indexing), and wide tuple folds that must not become closures (an ntuple ... do body past a handful of registers escapes the inliner and turns into a real device call). Not exported; access as using PTX.Utils: @unroll, strided_reduce.

Note the division of labor with LLVM's own unroller: a loopinfo hint (KernelAbstractions' @unroll, llvm.loop.unroll.* metadata) asks the optimizer to unroll and cannot make the induction variable a constant — @unroll expands at macro time and can, which is what barrier-slot Vals and distinct loop-carried phase registers require.

PTX.UtilsModule

Parse-time metaprogramming helpers for register-resident device code.

Two utilities, both existing because GPU kernels routinely need loop bodies where the induction value is a compile-time constant — Val(i) arguments, distinct per-iteration variable names (an indexed collection of loop-carried state demotes to local memory), tuple getindex with a constant — and where large ntuple(...) do closures escape the inliner and become real device CALLs:

  • @unroll: statement-level loop unrolling at macro-expansion time. The LLVM-hint kind of unrolling (loopinfo metadata) cannot provide any of the above; this is the guaranteed, constant-substituting kind.
  • strided_reduce: @generated strided tuple reduction with a pinned pairwise tree — the closure-free replacement for hand-written reduction trees.

Pure leaf module: no PTX dependencies, host-testable.

source
PTX.Utils.@unrollMacro
@unroll for x in <literal range or tuple>
    body
end

Expand the loop at macro-expansion time: the body is repeated once per iteration with x replaced by that iteration's value, so x is a compile-time constant inside each copy — Val(x), constant tuple indexing, and constant shifts all work.

Two substitution rules, following Base.Cartesian's conventions:

  • x as a standalone symbol becomes the iteration value (for tuple iterators, the element expression verbatim).
  • An identifier with a literal _x suffix gets the suffix replaced by the value: ph_xph_0, ph_1, ... This is how per-iteration names are minted — the mechanism that keeps loop-carried state in distinct registers instead of an indexable (hence local-memory) collection.

The iterator must be a parse-time literal: a range of integer literals (0:3, 0:2:6) or a tuple, optionally destructured — for (s, tma) in ((0, tma_K), (1, tma_V)) binds s to a constant and tma to the spliced expression. Anything else is refused at expansion time.

The expanded bodies are spliced into the enclosing scope (like Base.Cartesian.@nexprs, unlike a real for): assignments made in the body persist after the loop, which is what per-iteration phase variables need.

source
@unroll CAP for x in start:stop        body end
@unroll CAP for x in start:step:stop   body end

The capped form: for trip counts that are compile-time constants but not parse-time literals (Val-derived tile parameters — the norm kernels' @nexprs 16 + folded-guard idiom, generated).

CAP (a literal positive integer) is the maximum number of iterations; start and step must be literal integers, stop may be any expression. The macro expands CAP guarded copies — copy i substitutes the literal value start + i·step and executes only when that value is ≤ stop (≥ stop for a negative step). The guard compares against the spliced stop expression, so 0:(D8 - 1) and 1:D8 both do exactly what they say; there is no privileged start value. stop is evaluated once, before the first copy.

Semantics: identical to for x in start:step:stop (with the enclosing- scope splice of the uncapped form) provided the range has at most CAP iterations — iterations beyond the cap are silently absent, so the cap is a caller-asserted ceiling; check it where the values are chosen (host side), the way the norm kernels assert v4_iters ≤ 16. When stop folds to a specialization-time constant the guards vanish; a genuinely runtime stop leaves a predicated unroll, which is still correct.

source
PTX.Utils.strided_reduceFunction
strided_reduce(op, w::NTuple{N}, ::Val{S}[, ::Val{F}]) -> NTuple{S}

Fold an N-tuple into S strided accumulators: out[a] = op(w[a], w[a+S], w[a+2S], ...) for a in 1:S, with N divisible by S.

Fully unrolled by @generated expansion — no closure, nothing for the inliner to drop — so the operands stay in registers. Each accumulator reduces through a balanced tree of up-to-F-ary op calls (default F = 2, pairwise), and the shape is part of the contract: for non-associative op (floating-point +) the pairwise result is bit-reproducible against any other pairwise implementation, at tree depth ⌈log₂(N/S)⌉ rather than the N/S - 1 of a left fold.

F = 3 groups leaves in threes — op(op(l₁,l₂,l₃), op(l₄,l₅,l₆), ...) — which NVPTX fuses into the 3-input max/min instructions on sm_100+ for max/min reductions. Only reach for it when op is exactly associative (max, min, integer +): for floating-point addition a fanout change is an association change, i.e. different numerics.

source