Pol

Memory management primitives: protocols for kernels to describe the buffers they need as data, one verb that materializes descriptions into a space of the caller's choosing, and a bump-allocator with frame-scoped lifetimes.

Installation

using Pkg
Registry.add(url="https://registry.jool.space")
Pkg.add("Pol")

Quick start

using Pol

# a kernel takes its buffers as plain keyword arguments and describes them:
softmax!(y, x; tmp) = (tmp .= exp.(x); y .= tmp ./ sum(tmp, dims=1))
Pol.outputs(::typeof(softmax!), x) = (; y = Pol.Undef(x))
Pol.scratch(::typeof(softmax!), x) = (; tmp = Pol.Undef(x))

x, y = rand(1000), zeros(1000)

# materialize from an arena — the call is its own frame: carve, run, retract
arena = Arena(Vector{UInt8}(undef, 2^20))

scratchspace(arena) do frame
    softmax!(y, x; scratch(frame, softmax!, x)...)
end

# the functional form: output from the unframed space (caller lifetime),
# scratch inside the frame (dies with the block)
function softmax(x; space = Similar(x))
    (; y) = outputs(space, softmax!, x)
    scratchspace(space) do frame
        softmax!(y, x; scratch(frame, softmax!, x)...)
    end
    return y
end

That functional shell is what Allocating writes for you — with a NamedTuple-out verb softmax!((; y), x; tmp) it is one line, callable explicitly or against the ambient space:

const softmax = Allocating(softmax!)

y = softmax(arena, x).y      # outputs from the arena, scratch framed
withspace(arena) do
    y = softmax(x).y         # same call, space from scope
end

Arena allocations are aligned arrays (256-byte by default; see Arena(slab; alignment)) aliasing a single preallocated byte slab — for a plain Vector{UInt8} slab they are unsafe_wrapped Arrays, not GC-rooted. They are only valid until the scratchspace block they were carved in closes. Arenas are task-local and never grow; allocating past the end throws.

Manual

• Spaces — the alloc verb, the Space it materializes against, the GC-owned Similar default, and the Arena bump allocator.
• Scratchspaces — scratchspace blocks and frame-scoped lifetimes.
• Descriptions — how a kernel describes the buffers it needs as Undefs: its outputs, checkpoints, and scratch.
• Verbs — the Allocating form of a mutating verb.
• Shadows — pairing primals with gradient buffers for reverse-mode kernels.

The full symbol index lives in the API reference.