Descriptions

A kernel describes the buffers it needs as data, holding no memory, and lets the caller decide where they come from. The atom is Undef — an element type and a shape — which alloc materializes against a Space of the caller's choosing.

A kernel bundles its Undefs into named descriptions, dispatched on the kernel function and split by lifetime:

outputs — the arrays a pass produces. They belong to the caller, so they are materialized from the unframed space and outlive the scratchspace block.
checkpoints — buffers a forward pass saves and its backward pass reads (a norm's Rstd, attention's M/L). They bridge forward → backward.
scratch — buffers a single pass needs only while it runs. They live and die within a frame; the one-step form materializes against a Frame only.

Pol.Undef — Type

Undef(T, dims)
Undef(T, dims...)
Undef(x, [T=eltype(x)])
Undef(x, [dims=size(x)])

A description of the element type and size of an array, holding no memory. Array{T}(undef, dims) allocates the array; Undef{T}(dims) only describes one, for alloc to allocate against a space of the caller's choosing.

source

Pol.alloc — Method

alloc(space, u::Undef)
alloc(space, spec::NamedTuple)

Materialize an Undef description — or a whole NamedTuple spec of them, name for name — from space. Specs nest: a NamedTuple value materializes recursively, so buffers bundled under one name stay bundled. A spec value that is already an array passes through untouched — merge an owned array over its Undef before the alloc and that buffer is never materialized, which is how a caller aliases one output in place (e.g. S′ = S) while the rest allocate.

source

Pol.outputs — Function

outputs(f, args...; kwargs...) -> NamedTuple
outputs(space, f, args...; kwargs...)

The arrays a pass produces — the destinations, first in the mutating signature — described as Undefs. Implementations dispatch on the pass function. A leading Space materializes the description on the spot: outputs(space, f, args...) is alloc(space, outputs(f, args...)). Outputs belong to the caller: materialize them from the unframed space, before the scratchspace block they must outlive.

source

Pol.checkpoints — Function

checkpoints(f, args...; kwargs...) -> NamedTuple
checkpoints(space, f, args...; kwargs...)

The buffers a primitive saves in its forward pass and reads in its backward pass (a norm's Rstd, attention's M/L), described as Undefs: the primitive states eltypes and shapes, allocates nothing. Implementations dispatch on the forward function. A primitive that saves nothing has no method. The described names are keyword arguments the primitive already takes; a leading Space materializes them — splat the result into the call.

Distinct from scratch, which lives and dies within a scope.

source

Pol.scratch — Function

scratch(f, args...; kwargs...) -> NamedTuple
scratch(frame::Frame, f, args...; kwargs...)

The buffers a single pass needs only while it runs (workspace, partial-reduction buffers), described as Undefs. Dispatched per pass function — rms_norm! and ∇rms_norm! each describe their own; a pass that needs none has no method. The described names are keyword arguments the pass already takes; materialize them against the Frame a scratchspace block passes — scratch(frame, f, args...) is alloc(frame, scratch(f, args...)) — and splat into the call. The one-step form takes a Frame only: scratch lifetime is frame lifetime.

Distinct from checkpoints, which bridges forward → backward.

source