Writing a Stage

Before you start

• What you write: a Stage subclass of one of the open kinds below (a relax, width, elevation, validity, order or stations stage), optionally its Params class, and a reg.add(...) line in your package’s register(reg).

• What you need to know: the generator guide (Writing a Generator): Params, the StageBuilder, salts, capture classes and the hard rules of enqueue are the same for every stage, and only what differs is repeated here. Where your stage runs and what it reads: Pipeline stages and Stage kinds.

• How you test it: track_gen.testing.check_stage runs it in the default pipeline through the conformance checks (Pipeline stages: check_stage).

A stage is one step of the track pipeline or of the gate pipeline (see Pipeline stages and The gate pipeline). Besides the generators, six kinds of stage are open to plugins:

• relax reshapes the resampled centerline (cl.xy), as the built-in XPBD solver does;

• width writes the per-point half-width (cl.half_width) that XPBD, the band and the checks read;

• elevation writes the per-point altitude (cl.z), and the per-gate altitude (gates.z) when it runs in a gate pipeline;

• validity runs a check per env and sets one of the fail bits (valid.fail), on a track or on gates;

• order puts a gate generator’s anchors in course order (gates.xy), see Writing an ordering;

• stations places gate anchors on a track generator’s centerline (anchors.xy).

A gate generator is a stage too, of kind gate_generator, and the order kind is described with it: see Writing a Gate Generator.

A stage is a Stage subclass, the protocol a generator Stage uses (see The Stage contract), with read ports: it reads and writes channels. It ships the way a generator ships: in its own pip-installable package, found through the track_gen.plugins entry point, registered by the same register(reg) function, which may add generators and stages alike (Packaging a Plugin). It imports only public modules (track_gen.plugin, track_gen.exec, track_gen.ops), and track_gen.testing.check_stage checks it against the contract on this page (see Pipeline stages: check_stage).

Quick start: a relax stage

Taubin smoothing alternates a shrinking Laplacian step (lam > 0) and an inflating one (mu < 0), which smooths a closed polyline without shrinking it. It makes a relax stage with one static parameter (the number of passes, which fixes the launch count) and two dynamic fields (the step factors, per env, changeable at run time).

1. Declare the entry point, as for a generator (skip it if the package already has one; one register adds everything):

[project.entry-points."track_gen.plugins"]
acme = "acme_tracks:register"

2. Write the stage and register it (acme_tracks/taubin.py):

import dataclasses

import warp as wp

from track_gen import plugin


@wp.kernel
def _laplace_k(src: wp.array(dtype=wp.vec2f), count: wp.array(dtype=wp.int32),
               f: wp.array(dtype=wp.float32), S: int, dst: wp.array(dtype=wp.vec2f)):
    t = wp.tid()                         # one thread per point slot: t = e * S + i
    e = t // S
    i = t - e * S
    n = count[e]
    if i >= n:
        dst[t] = src[t]                  # carry the NaN padding: write every slot
        return
    base = e * S
    x = src[t]
    prev = src[base + (i + n - 1) % n]   # closed loop: wrap around the real points
    nxt = src[base + (i + 1) % n]
    dst[t] = x + f[e] * (0.5 * (prev + nxt) - x)


@dataclasses.dataclass(frozen=True, kw_only=True)
class TaubinParams(plugin.Params):
    passes: int = plugin.static(10, lo=1, doc="shrink/inflate pass pairs")
    lam: float = plugin.dynamic(0.5, lo=0.0, hi=1.0, doc="shrink factor")
    mu: float = plugin.dynamic(-0.53, lo=-1.0, hi=0.0, doc="inflate factor")


class Taubin(plugin.Stage):
    kind = "relax"
    name = "acme.taubin"
    Params = TaubinParams
    reads = {"xy": "cl.xy", "count": "cl.count"}
    writes = {"xy": "cl.xy"}           # same channel, same port: a rewrite

    def build(self, b, params):
        x, count = b.read("xy"), b.read("count")   # the version before this stage
        out = b.write("xy", wp.vec2f, per="cl")    # a NEW version, [E * N_max]
        tmp = b.temp_like("xy")                    # ping-pong partner, same shape
        lam, mu = b.param("lam"), b.param("mu")    # [E] float32, read by pointer
        S = b.layout.stride
        steps = [f for _ in range(params.passes) for f in (lam, mu)]

        def enqueue(ctx):
            src = x
            for k, f in enumerate(steps):        # x -> tmp -> out -> tmp -> ... -> out
                dst = tmp if k % 2 == 0 else out
                ctx.launch(_laplace_k, b.E * S, [src, count, f, S, dst])
                src = dst
        return enqueue


def register(reg):
    reg.add(Taubin)          # the class, not an instance

The stage rewrites cl.xy: it reads the channel and writes it through the same port. b.read("xy") is the version the previous stage wrote, and b.write("xy", ...) a new buffer that every later stage reads (see In-place rewrites). The stage never writes the buffer it reads. The steps alternate between tmp and out and their number is even, so the last one lands in out with no copy. passes is static: it fixes the number of launches, which a recorded graph replays forever.

3. Select it. Through a preset, by name, with params as a mapping or a TaubinParams:

from track_gen import TrackGenerator
from track_gen.pipeline import presets

device = "cuda:0"  # or "cpu"
spec = presets.track(num_envs=64, device=device,
                     relax=("acme.taubin", {"passes": 20}))
gen = TrackGenerator(spec)
track = gen.generate()

gen.params["relax"].set("lam", 0.6)     # every env, from the next generation on
track = gen.generate(advance=True)

Or through the flat config, whose relax_solver names any registered relax stage (it then takes TaubinParams.from_legacy(config), its defaults here):

from track_gen import TrackGenConfig

gen = TrackGenerator(TrackGenConfig(num_envs=64, device=device,
                                    relax_solver="acme.taubin"))

The rest of the pipeline is unchanged: the uniform resample, frames, offsets, elevation, validity and lift run on the smoothed loop, and gen.program.tap("cl.xy", after="resample") still holds the loop before it.

This example shows the mechanics, not a replacement for XPBD. Taubin smoothing enforces no separation between distant parts of the loop, so it fails the thickness check (fail bit 3) where relax=None does: at the default half_width=0.1 every bezier track does, and at half_width=0.03 (with n_max=1024) about a third pass. A relax stage you ship must leave a loop the band can be inflated around; Track.fail_mask tells you which check it fails. To separate the loop by the band’s own width, read cl.half_width (see Reading cl.half_width).

The relax contract: bead identity. A relax stage moves points and never re-indexes them: point i of its output is point i of its input, moved. The pipeline relies on it: the channels carried from the generator (the half-width, a provided altitude, a plugin’s per-point attributes) keep their point index through relax, with no copy, and the uniform resample after it re-interpolates them with the relaxed loop’s arc length (see Carried channels). A stage that resamples the loop internally, or shifts its start, keeps the shape but hands every carried value to another point: it must not be combined with carried channels. The conformance kit does not check this (XPBD’s own drift along the loop is as large as an off-by-one shift on some envs), so state it in your stage’s documentation if it does not hold.

The open kinds

kind	slot (preset keyword)	writes	usual reads	flat-config shortcut
relax	relax, between width and uniform	cl.xy (a rewrite)	cl.xy, cl.count, cl.half_width	relax_solver
width	width, between resample and relax	cl.half_width	cl.count, cl.xy, a generator’s cl.<package>.<name> attribute	none
elevation	elevation, after offset (gates: after gate_relax, and after gate_arclen for a non-flat profile)	cl.z (gates: gates.z, rewired by the planner)	seed, cl.count, cl.arclen2, cl.length2	z_profile (both configs)
validity	checks, after elevation (gates: after gate_frames)	valid.fail	valid.fail and anything written before it	none
order	order (gates), after the gate generator	gates.xy, gates.count	anchors.xy, anchors.count, seed	GateGenConfig.gate_ordering
stations	stations (gates), after a centerline generator	anchors.xy, anchors.count (and anchors.z, anchors.<package>.<name>)	raw.xy, gen.valid (or cl.xy, cl.count, cl.truncated)	none

A stage may also write plugin channels of its domain: cl.<package>.<name> for the four track kinds, gates.<package>.<name> for a stage on the gate slots, and anchors.<package>.<name> for a stations stage (see Ports and channels). Each kind’s built-ins are registered under plain names (relax:xpbd, relax:none, width:constant, width:from_channel, elevation:flat / uniform / random_walk / noise, validity:default, validity:gate_default, order:raw / ccw / random_pairs, stations:raw / relaxed); registry.available("elevation") lists them with the installed plugins. The randomize kind (randomize.uniform, which draws other stages’ dynamic fields, see Per-env randomization: randomize=) is TrackGen’s own and not open to plugins in 0.x.

A plugin stations stage computes what it writes itself: it gets no carried channels from the centerline in this release (the built-in stations do, see 3D station spacing). One that writes anchors.z is treated as providing the altitude (carry.z, order.raw), and its anchors.<package>.<name> channels reach the gates through carry.attributes. presets.gates(spacing_metric=...) fills a spacing_metric field of its Params when it has one; without the field, "xyz" raises ValueError.

Ports and channels

reads and writes map a local port (the name your build uses) to a channel (the pipeline’s name for a buffer). The channels are those of the catalogue (see Pipeline stages) and your own namespaced ones, such as cl.acme.rank: a domain prefix, then at least two dot-separated identifiers. The prefix is the point domain (Plugin channels): cl. for the track kinds, gates. for a stage on the gate slots, anchors. for a stations stage (raw. is the generators’). A user can rewire a port to another channel of the same dtype and size with StageRef(..., ports={"xy": "cl.acme.smoothed"}).

Required read ports. A reads value may name a domain instead of a channel, "cl.*" or "gates.*": the stage reads a channel the user picks, and the StageRef must wire it, StageRef("width", "from_channel", ports={"src": "cl.acme.hw"}). An unwired one raises PlanError at build, naming the port and the form to write (Required read ports); a write port cannot be one. registry.describe(kind, name).reads shows "src": "cl.*". Check a stage that has one with check_stage(..., ports={"src": ...}) (Pipeline stages: check_stage).

registry.add checks the declaration, and raises TypeError or ValueError (the plugin is then quarantined under warn) when:

• the class is an instance, abstract, or of a kind that is not open (resample, frames, offset, lift: “not open to plugins”);

• output is set (it is for generators), provides is not empty, topology is not "closed" or "any", api_version is above API_VERSION, capture is not a CaptureClass, or the Params class breaks the rules of Parameters: Params;

• a channel is neither in the catalogue nor a namespaced plugin channel, or a plugin channel is not of the stage’s domain (a raw.* channel is for generators, an anchors.* one for gate generators and stations stages);

• two write ports name one channel, or one channel is read through one port and written through another (a rewrite reads and writes through one port), or a port in both maps names two channels;

• the stage writes a catalogue channel its kind may not write: a relax stage writes cl.xy only, a width stage cl.half_width, an elevation stage cl.z (or gates.z, the channel the planner rewires it to in a gate pipeline), a validity stage valid.fail (plus plugin channels of their domain for all four; seed never), an order stage gates.xy and gates.count, a stations stage anchors.xy, anchors.count and optionally anchors.z. For example Stage Taubin (kind 'relax') writes 'cl.kappa' (port 'k'): a relax stage may write ['cl.xy'] and plugin channels 'cl.<package>.<name>' or 'gates.<package>.<name>';

• a validity stage does not write valid.fail, an order stage does not write both gates.xy and gates.count, or a stations stage both anchors.xy and anchors.count;

• a stage that is not a generator declares a non-empty preferred (a generator attribute: a gate generator’s "order" / "normalize", a centerline generator’s "relax"), or overrides declare() (read for centerline generators only in 0.x: Stage X.declare() is read for centerline generators only in 0.x, X being the Stage class name).

A non-Stage spec (a GeneratorSpec-like object) registered under one of these kinds raises TypeError: pipeline stages are Stage subclasses only.

The planner checks the rest when it builds a pipeline (each as a PlanError naming the stage, the port and the channel): the rules above again after ports rewiring, every required read port is wired, every read channel has an earlier writer (a width stage that reads cl.kappa must come after frames), a second writer of a channel reads it (a rewrite), and no stage reads a cl.* plugin channel that a stage wrote on the other side of the uniform resample (the values sit on other points, see Carried channels). b.write checks that a later writer of a plugin channel keeps the dtype and per its first writer gave it (TypeError). The rewrite model, with its five rules, is In-place rewrites; the one rule you write code for is: write your whole output into the buffer b.write returns, and never write a buffer b.read returned.

The StageBuilder in a pipeline stage

The builder is the one The StageBuilder describes, with these differences:

Attribute or method	In a pipeline stage
stage_id	The slot id ("relax", "elevation", or the id a StageRef gives). It keys gen.params, the scratch names and the salts.
layout.stride	N_max in a track pipeline: a per-point channel is [E * N_max], env e’s slots are e * N_max to e * N_max + N_max - 1, and the first count[e] are real. In a gate pipeline, G = max_gates for a stage on the gate slots and N_max for one with a cl.* port (see The gate pipeline).
stride_of(port), per_of(port), dtype_of(port)	The row stride (1 for an [E] channel), the per and the dtype of the channel the port is wired to, after ports rewiring: size reads from them when a stage reads two domains (an order stage reads A-strided anchors and writes G-strided gates), and check a required read port’s channel with them (width.from_channel wants a per-point float32).
seed	The [E] effective seed, as for a generator. Declaring a seed read port (reads = {"seed": "seed"}) gives the same buffer.
read(port)	The current version of reads[port]: written by an earlier stage, read-only for this one.
write(port, dtype=None, *, per=None, interp=None, export=False)	A new buffer for writes[port], zero-initialised once, read by every later stage. per is "env" ([E]) or "cl" ([E * N_max]), and for the gate catalogue "anchors" ([E * A]) or "gates" ([E * G]). For a catalogue channel, or a plugin channel an earlier stage wrote, omit both (they are the channel’s) or pass matching ones (TypeError otherwise); a plugin channel’s first writer fixes them (dtype required, per defaults to "env" for cl.* / gates.*, and is "anchors" for anchors.*). export=True (or the key itself) makes a per-point plugin channel a Track.attributes / GateSequence.attributes entry, keyed by its name without the domain prefix (Per-point attributes in the result): it needs per="cl" / "gates" (ValueError on the default per="env"). interp= is for a generator’s raw.* channels only (Per-point attributes); a stage’s channels are not resampled (ValueError). A rewired port is checked against the channel the stage declared and sized by the one it is wired to. Call it once per write port, in build.
temp_like(port)	A private zero-initialised buffer shaped like a write port’s channel (call write first for a plugin channel): a ping-pong partner. It works on anchors.* and gates.* ports too.
scratch, param, input, salt	As for a generator, namespaced by the stage id: scratch:<id>/<name>, and salt(stream) hashes "<id>/<stream>", so a stage moved to another slot id draws another random stream. input(name, dtype, default=0) is a per="env" [E] input (per="raw" is for centerline generators). Name an input with a dotted path in your own namespace ("acme.target_speed"): input names share one table per Program, "regen_mask" is reserved and centerline.* is the ingestion namespace (Ingesting your own centerlines).
fail_bit(name=None)	Validity stages only: the index (16 to 30) of a new fail bit this stage owns (see A validity check).
host_buffer(dtype, shape)	HOST stages only (ValueError otherwise): a host staging buffer for ctx.host(), pinned on cuda (see Writing a host-driven (HOST) stage).

build runs once, in pipeline order, after every earlier stage’s build. When it returns, every declared write port must have been written (PlanError otherwise). enqueue(ctx) follows the generator’s hard rules: device work through ctx and track_gen.ops only, no allocation, no host sync, no branch on device data, a launch sequence fixed at build (an iteration that stops early is a ctx.bounded_loop, see Writing an iterative (BOUNDED) stage). Per env, its output must depend only on that env’s inputs (its rows of the channels it reads, its seed, its dynamic parameters and inputs) and the static parameters.

View safety. The buffers b.write returns for the result’s channels (the final valid.fail, an exported attribute, …) may be views into a larger caller-owned result (out=big.view(e0, e1), Mixed batches). A kernel is view-safe when it indexes only [0, shape[0]) from wp.tid() and the per-env layout (e * stride + i, with stride from b.layout or b.stride_of), and reaches every buffer through its arguments: the bounds-correctness an owned array needs anyway. Never derive an index from an array’s address or assume a buffer starts an allocation. The conformance kit’s output.view check builds every stage with out= the middle of a poisoned larger result and fails a write outside it.

Topology. topology = "closed" (the default) says that the stage assumes closed loops (index (i + 1) % count wraps); "any" that it works on open curves too. 0.x accepts closed loops only: BatchSpec(closed=False) is refused before any stage is bound, naming the stages that declare "closed" (Topology), so a stage’s Params.validate and build never see b.layout.closed == False yet.

Writing a width stage

A width stage writes cl.half_width, the band’s half-width, at every slot. It runs between resample and relax, on the constant-spacing points (the beads): XPBD separates the loop with it, the uniform resample carries it to the final points (it stays attached to its bead through relax), offset builds the band from it and the checks hold the band to it (see Width). Its real points are the first count[e] of each env’s row.

A writer. This one narrows the road waves times per lap, by a per-env fraction depth (with the quick start’s imports):

@wp.kernel
def _ripple_k(count: wp.array(dtype=wp.int32), hw: float, waves: int,
              depth: wp.array(dtype=wp.float32), S: int,
              w: wp.array(dtype=wp.float32)):
    t = wp.tid()
    e = t // S
    i = t - e * S
    n = count[e]
    if i >= n:
        w[t] = hw                    # padding: the whole output is written
        return
    u = float(i) / float(n)          # the beads are equally spaced: an arc fraction
    w[t] = hw * (1.0 - depth[e] * 0.5 * (1.0 - wp.cos(6.2831853 * float(waves) * u)))


@dataclasses.dataclass(frozen=True, kw_only=True)
class RippleParams(plugin.Params):
    half_width: float = plugin.static(0.1, doc="nominal half-width")
    waves: int = plugin.static(4, lo=1, doc="narrowings per lap")
    depth: float = plugin.dynamic(0.3, lo=0.0, hi=0.9, doc="narrowing fraction")


class Ripple(plugin.Stage):
    kind = "width"
    name = "acme.ripple"
    Params = RippleParams
    reads = {"count": "cl.count"}
    writes = {"half_width": "cl.half_width"}

    def build(self, b, params):
        count = b.read("count")
        w = b.write("half_width", wp.float32, per="cl")
        depth, S = b.param("depth"), b.layout.stride
        hw, waves = float(params.half_width), int(params.waves)

        def enqueue(ctx):
            ctx.launch(_ripple_k, b.E * S, [count, hw, waves, depth, S, w])
        return enqueue


def register(reg):
    reg.add(Ripple)

Select it with presets.track(..., width="acme.ripple") (there is no flat-config shortcut for widths). Its half_width field is the nominal width, and the preset fills it: presets.track(..., width="acme.ripple", half_width=0.15) runs RippleParams(half_width=0.15), and params that set another value than the preset’s half_width raise ValueError naming it. The planner’s agreement rule binds every stage that has a half_width field to one value (the nominal width; the per-point width is cl.half_width). The default spacing is 0.6 * half_width, as for width.constant.

XPBD now separates the narrow stretches by their own width, the band narrows with them, and a stretch narrowed below w_floor fails fail bit 4 on the final points.

A rewrite. A width stage that reads cl.half_width and writes it through the same port rewrites the width an earlier stage wrote. This one scales it by a per-env factor:

@wp.kernel
def _rescale_k(w_in: wp.array(dtype=wp.float32), count: wp.array(dtype=wp.int32),
               k: wp.array(dtype=wp.float32), S: int,
               w_out: wp.array(dtype=wp.float32)):
    t = wp.tid()
    e = t // S
    if t - e * S >= count[e]:
        w_out[t] = wp.nan               # NaN padding
        return
    w_out[t] = k[e] * w_in[t]


@dataclasses.dataclass(frozen=True, kw_only=True)
class RescaleParams(plugin.Params):
    factor: float = plugin.dynamic(1.0, lo=0.0, doc="width factor")


class Rescale(plugin.Stage):
    kind = "width"
    name = "acme.rescale"
    Params = RescaleParams
    reads = {"half_width": "cl.half_width", "count": "cl.count"}
    writes = {"half_width": "cl.half_width"}     # same channel, same port: a rewrite

    def build(self, b, params):
        w_in, count = b.read("half_width"), b.read("count")
        w_out = b.write("half_width", wp.float32, per="cl")
        k, S = b.param("factor"), b.layout.stride

        def enqueue(ctx):
            ctx.launch(_rescale_k, b.E * S, [w_in, count, k, S, w_out])
        return enqueue


def register(reg):
    reg.add(Rescale)

Insert it after the width slot:

from track_gen.pipeline import StageRef

spec = presets.track(num_envs=64, device=device).insert_after(
    "width", StageRef("width", "acme.rescale", {"factor": 0.8}, id="rescale"))
gen = TrackGenerator(spec)

The rewritten width is no longer width.constant’s constant, so XPBD, offset and validity run their per-point kernels: the outputs are close to, not byte-equal to, those of presets.track(half_width=0.08, spacing=0.06) (the rescaled spec keeps the default spacing of half_width=0.1; presets.track(half_width=0.08) alone also shortens the spacing to 0.048). Only width.constant selects the constant-width kernels (a private declaration of the built-in).

From a channel. The built-in width.from_channel makes any per-point float32 channel the width, scaled and clamped per env (min(max(scale * src, lo), max(hi, lo)): lo <= hi is checked at build, a run-time hi below lo clamps to lo; NaN kept). Its src is a required read port: wire it to the channel. With a generator that writes a per-point raw.acme.hw attribute (the generator guide’s acme.ribbon, carried to cl.acme.hw, see Per-point attributes):

spec = presets.track(
    num_envs=64, device="cuda:0", generator="acme.ribbon",
    width=StageRef("width", "from_channel", {"lo": 0.0, "hi": 0.3},
                   ports={"src": "cl.acme.hw"}))

A generator that knows its half-widths can also provide them (raw.half_width, Providing altitude and width): the width slot then goes away, with no stage at all.

Geometry-dependent widths. A width computed from the final geometry (a road that widens in tight turns, from cl.kappa) cannot run before relax: frames computes the curvature after the uniform resample. Keep a nominal width in the width slot, for XPBD, and insert the geometry-dependent stage after frames as a rewrite of cl.half_width (spec.insert_after("frames", StageRef("width", "acme.curvature_scaled", id="width_kappa"))). XPBD did not see that width, so the thickness check may reject a loop that XPBD separated for the nominal one. A stage that writes cl.half_width after offset raises PlanError: the band and Track.half_width come from offset’s read, and validity would check another width than the band.

Test a width stage with check_stage("width", "acme.ripple") (and ports= for a stage with a required read port); width.values checks that every real point gets a finite half-width >= 0 (see Pipeline stages: check_stage).

Reading cl.half_width

A relax or validity stage that needs the band’s width reads cl.half_width like any channel (reads = {..., "half_width": "cl.half_width"}). A relax stage reads the beads’ width, the one the uniform resample carries forward; a check after uniform reads the final points’ width, the band the Track has. Read real points only (i < count[e]): the padding of the version a relax stage reads is whatever the width stage wrote there (the constant for width.constant, NaN for a carried raw.half_width, a plugin’s own choice otherwise), and that of the uniform resample’s version is NaN, or the constant at constant width. Track.half_width is the NaN-padded field.

The built-in XPBD separates two beads i and j to r_i + r_j, with r_i = w_i * (1 + margin), and keeps every bead’s curvature radius above r_i; the per-point thickness check asks |p_i - p_j| >= (1 - relax_tol) (w_i + w_j) of every non-neighbouring pair. A plugin that reads the array is right at every width. The built-ins also keep scalar-width kernels for a constant width (a host scalar computed in double precision, which the per-point kernels’ float32 reads do not reproduce bit for bit), chosen from width.constant’s private declaration (Width); that choice is not open to plugins, so a plugin width that writes a constant value runs the per-point kernels.

An elevation stage

An elevation stage writes cl.z, one altitude per point, from the plan-view arc length (cl.arclen2, cumulative, and cl.length2, the lap length) that frames writes before it. The lift then builds level cross-sections at that altitude, and the 3D tangent and arc length. This one draws a random phase per env and makes waves sinusoidal hills per lap with a per-env amplitude:

from track_gen.ops.funcs import hash32


@wp.kernel
def _hills_k(seed: wp.array(dtype=wp.int32), salt: wp.uint32,
             cum: wp.array(dtype=wp.float32), perim: wp.array(dtype=wp.float32),
             count: wp.array(dtype=wp.int32), amp: wp.array(dtype=wp.float32),
             waves: int, base: float, S: int, z: wp.array(dtype=wp.float32)):
    t = wp.tid()
    e = t // S
    i = t - e * S
    if i >= count[e]:
        z[t] = 0.0                   # padding, as the built-in profiles write it
        return
    key = hash32(wp.uint32(seed[e]) ^ salt) & wp.uint32(0x7FFFFFFF)
    state = wp.rand_init(wp.int32(key))
    phase = 6.2831853 * wp.randf(state)          # the same draw for every point of e
    u = cum[t] / wp.max(perim[e], 1.0e-9)        # lap fraction in [0, 1)
    z[t] = base + amp[e] * wp.sin(6.2831853 * float(waves) * u + phase)


@dataclasses.dataclass(frozen=True, kw_only=True)
class HillsParams(plugin.Params):
    amplitude: float = plugin.dynamic(0.2, lo=0.0, hi=5.0, doc="hill height")
    waves: int = plugin.static(3, lo=1, doc="hills per lap")
    base: float = plugin.static(0.0, doc="mean altitude")

    @classmethod
    def from_legacy(cls, cfg):
        return cls(base=float(cfg.z_base))      # TrackGenConfig(z_profile=...) users


class SinusoidHills(plugin.Stage):
    kind = "elevation"
    name = "acme.sinusoid_hills"
    Params = HillsParams
    reads = {"cum": "cl.arclen2", "perim": "cl.length2", "count": "cl.count"}
    writes = {"z": "cl.z"}

    def build(self, b, params):
        cum, perim, count = b.read("cum"), b.read("perim"), b.read("count")
        z = b.write("z", wp.float32, per="cl")
        amp, salt, S = b.param("amplitude"), b.salt("phase"), b.layout.stride
        waves, base = int(params.waves), float(params.base)

        def enqueue(ctx):
            ctx.launch(_hills_k, b.E * S,
                       [b.seed, salt, cum, perim, count, amp, waves, base, S, z])
        return enqueue


def register(reg):
    reg.add(SinusoidHills)

Select it with presets.track(..., elevation="acme.sinusoid_hills"), or with TrackGenConfig(z_profile="acme.sinusoid_hills"): z_profile names any registered elevation stage, and a plugin’s Params come from its from_legacy(config), which may map the flat z_* fields (here z_base). Either way the pipeline lifts with lift.elevated, since the elevation is not elevation.flat, and validity.default’s grade check (z_valid_grade, fail bit 6) reads the new cl.z.

The random stream is keyed on the effective seed and the stage’s salt (Random streams: b.salt), so env e’s hills change when its seed or epoch changes and depend on nothing else.

On gates. The same class runs in a gate pipeline, with no change: presets.gates(..., elevation="acme.sinusoid_hills") or GateGenConfig(z_profile="acme.sinusoid_hills"). The planner rewires its ports to the gate channels (cum -> gates.cum, perim -> gates.perim, count -> gates.count, z -> gates.z, and a plugin channel cl.acme.slope -> gates.acme.slope; see Elevation on gates), b.layout.stride is then max_gates, and b.write("z", wp.float32, per="cl") returns a [E * max_gates] buffer: per is checked against the channel the stage declares (a plugin channel written per="cl", the form export=True needs on a track, becomes a per-gate gates.<pkg>.<name>). The salt is keyed by the slot id, "elevation" in both pipelines. A stage that sizes buffers itself should use b.stride_of(port) rather than assume N_max. Through GateGenConfig(z_profile=...), from_legacy(config) receives the GateGenConfig: only the z_* fields are shared with TrackGenConfig, so read any other field with getattr(config, name, default) (a from_legacy that reads a TrackGenConfig-only field raises ValueError naming it). Test it on gates with check_stage("elevation", "acme.sinusoid_hills", pipeline="gates") (Pipeline stages: check_stage).

A validity check

A validity stage reads valid.fail, ORs in the bit it owns where its check fails, and writes the result to a new version. b.fail_bit() assigns the bit: 16 for the first plugin check of a pipeline, 17 for the next, in pipeline order. Pass it to the kernel as a host scalar:

@wp.kernel
def _kmax_k(kappa: wp.array(dtype=wp.float32), count: wp.array(dtype=wp.int32),
            k_max: wp.array(dtype=wp.float32), S: int, bit: int,
            fail_in: wp.array(dtype=wp.int32), fail_out: wp.array(dtype=wp.int32)):
    e = wp.tid()                     # one thread per env
    m = float(0.0)
    for i in range(count[e]):
        m = wp.max(m, wp.abs(kappa[e * S + i]))
    f = fail_in[e]                   # the bits of the earlier checks
    if m > k_max[e]:
        f = f | (1 << bit)
    fail_out[e] = f                  # every env, failed or not


@dataclasses.dataclass(frozen=True, kw_only=True)
class KappaParams(plugin.Params):
    kappa_max: float = plugin.dynamic(10.0, lo=0.0, doc="largest |curvature| allowed")


class MaxCurvature(plugin.Stage):
    kind = "validity"
    name = "acme.max_curvature"
    Params = KappaParams
    reads = {"kappa": "cl.kappa", "count": "cl.count", "fail": "valid.fail"}
    writes = {"fail": "valid.fail"}

    def build(self, b, params):
        kappa, count = b.read("kappa"), b.read("count")
        fail_in = b.read("fail")
        fail_out = b.write("fail", wp.int32)
        bit = b.fail_bit()                       # 16 for the first plugin check
        k_max, S = b.param("kappa_max"), b.layout.stride

        def enqueue(ctx):
            ctx.launch(_kmax_k, b.E, [kappa, count, k_max, S, bit, fail_in, fail_out])
        return enqueue


def register(reg):
    reg.add(MaxCurvature)

Add it after the built-in checks:

from track_gen.pipeline import StageRef, presets

spec = presets.track(num_envs=64, device="cuda:0",
                     checks=("default", ("acme.max_curvature", {"kappa_max": 8.0})))
gen = TrackGenerator(spec)
gen.program.fail_bits      # {'generator': 0, ..., 'truncation': 7, 'acme.max_curvature': 16}

# or into any spec, under an id of your choice:
spec = spec.insert_after("validity", StageRef("validity", "acme.max_curvature", id="kmax"))

Track.valid is fail_mask == 0, so the check rejects envs like the built-in ones, and Track.fail_mask tells which check rejected each env. The rules of fail bits:

• Call b.fail_bit(name=None) in build, once per independent check the stage runs. name keys the bit in Program.fail_bits / Track.fail_bits and defaults to the stage id ("acme.max_curvature" for a check added by name in checks, "kmax" above). A stage of another kind gets ValueError.

• A validity stage owns at least one bit (PlanError otherwise), sets only its own bits and keeps the others: fail_out = fail_in | own. check_stage’s validity.bits check enforces it.

• A pipeline holds at most 15 plugin bits (16 to 30); one more, or a name already taken, raises PlanError. Bits 0 to 15 belong to TrackGen’s checks (see Fail bits).

• Write every env. An env that passes gets fail_in[e] unchanged.

A check must read valid.fail. One that writes it without reading it would start the bitmask in place of validity.default, dropping the generator, turning, thickness and truncation checks (a truncated loop would pass): the planner refuses such a pipeline (PlanError) until the split built-in checks arrive, and validity.default comes first in checks.

With on_invalid="poison" the commit empties every env with a bit set, yours included (see The commit).

On gates. A gate check is a validity stage that reads gate channels (gates.position, gates.count, …; Gate channels) with b.layout.stride == max_gates, under the same rules: it reads and rewrites valid.fail, owns bits from 16 on, and follows validity.gate_default, which owns bits 0 to 5 (Gate fail bits). Add it with presets.gates(..., checks=("gate_default", "acme.max_height")); its bit then shows in GateSequence.fail_mask and GateSequence.fail_bits. check_stage runs a check that reads a gates.* or anchors.* channel in the default gate pipeline.

A gate check may also publish what it measured, per gate, as a gate plugin channel: add a write port on gates.acme.clearance and write it with b.write("clearance", wp.float32, per="gates", export=True); GateSequence.attributes["acme.clearance"] is then that buffer (NaN past count[e]: write the padding), held, restored and poisoned by the commit with the other fields (Gate plugin channels). A track check does the same with a cl.acme.<name> channel and per="cl", written after uniform (on the final points), into Track.attributes.

Runtime parameters

A stage’s dynamic parameters work as a generator’s (see Runtime parameters and inputs), keyed by the stage id: gen.params lists "generator" and every stage with a dynamic field, gen.params["relax"]["lam"] is the live [E] array, and gen.params["relax"].set("lam", 0.6) writes it, on the resolved stream, recordable inside a caller’s graph. The new value is used from the next generation. b.param returns the stage’s working copy, copied from the live array at the start of every generation (for the masked envs only, under a masked regeneration; how the copy works); read it in enqueue, never write it, and read each env’s own entry (param[e]): the kit’s params.per_env check sets the field per env and compares every env with a scalar run. The copy runs at the start of segment 0, so in a Program with a HOST stage a write recorded after segment 0 applies at the next generation (Capture classes and segments). The built-in elevation stages work the same way: their z_min / z_max (and z_max_step, z_noise_amplitude) are dynamic, gen.params["elevation"]["z_max"] (Per-env altitude: randomization and curricula), and so are the gate stages’ normalize.size and gate_frames.gate_width (Per-env bar width and course size). gen.params[stage_id].in_effect(field) returns the values the stage read for each env’s current rows (its working copy), and bounds(field) the declared range (the tuning guide).

Only static fields reach build and capture_class (reading a dynamic one raises AttributeError): a host value derived from a dynamic field would be baked into the launches and any recorded graph.

Per-env draws. A dynamic float32 field of your stage can be drawn per env, in the graph, at every generation, by the built-in randomize.uniform stage: presets.track(..., relax="acme.taubin", randomize={"relax.lam": (0.4, 0.6)}) (Per-env randomization: randomize=). The draw is lo + u * (hi - lo) keyed on the env’s effective seed, clamped to the field’s declared lo / hi: declare bounds on every field that is safe to randomize. Nothing changes in your stage: b.param(field) returns the drawn [E] array, filled before your stage runs in the same generation (it is not copied, and gen.params[...].set on it raises). int32 fields cannot be drawn in 0.x.

Capture classes

As for a generator (Capture classes), and the same in a gate pipeline:

• STATIC (default): a fixed launch sequence, recorded and replayed like the built-ins.

• BOUNDED: an iteration with a static bound that may stop early, through ctx.bounded_loop. Recordable by anything: the generator’s own graph, or a Warp-registered capture (wp.ScopedCapture, wp.capture_begin, registered(conditional=True)) on CUDA 12.4+, skips the converged windows on device (a conditional node); other captures (a plain torch.cuda.graph, inside track_gen.exec.unconditional()) and eager cuda calls run the full bound (see Writing an iterative (BOUNDED) stage).

• HOST: host readbacks or host control flow. The stage runs eagerly on every call, as its own segment of the Program: the stages before and after it still replay from graphs. A caller cannot record the Program (NotCapturableError before any launch), and PipelineSpec(require_caller_capturable=True) refuses it at build (see Writing a host-driven (HOST) stage).

The Program’s capture class is the highest of its stages’ (STATIC < BOUNDED < HOST); Program.segments shows how the stages are grouped (see Pipeline stages).

Writing an iterative (BOUNDED) stage

A stage that iterates until each env has converged keeps a static launch sequence by writing its loop as ctx.bounded_loop(n_iters, body, active=..., first_check=..., check_every=...): n_iters bounds the iteration count, body(ctx, it) emits one iteration (it is a plain int), and active, a [1] int32 flag the body writes on device, says whether any env still needs the next window of check_every iterations. Where the call runs decides whether the loop really stops early (see Iterative stages (BOUNDED)): an eager cpu run polls active and stops, the generator’s own graph or another Warp-registered capture with conditional nodes skips the remaining windows on device, every other run executes the full bound. The outputs must be the same bytes in every case.

This relax stage moves each env toward a fixed point (a Laplacian smoothing step plus a pull back to its input loop, an iteration that contracts) and freezes an env once its largest step falls under tol:

import dataclasses

import warp as wp

from track_gen import ops, plugin


@wp.kernel
def _step_k(x: wp.array(dtype=wp.vec2f), x0: wp.array(dtype=wp.vec2f),
            count: wp.array(dtype=wp.int32), S: int, lam: float, pull: float,
            running: wp.array(dtype=wp.int32), tmp: wp.array(dtype=wp.vec2f),
            moved: wp.array(dtype=wp.float32)):
    t = wp.tid()
    e = t // S
    i = t - e * S
    n = count[e]
    if running[e] == 0 or i >= n:        # a frozen env is never touched again
        return
    base = e * S
    p = x[t]
    mid = 0.5 * (x[base + (i + n - 1) % n] + x[base + (i + 1) % n])
    d = lam * (mid - p) + pull * (x0[t] - p)
    tmp[t] = p + d
    wp.atomic_max(moved, e, wp.length(d))   # max: independent of the thread order


@wp.kernel
def _apply_k(tmp: wp.array(dtype=wp.vec2f), count: wp.array(dtype=wp.int32), S: int,
             running: wp.array(dtype=wp.int32), x: wp.array(dtype=wp.vec2f)):
    t = wp.tid()
    e = t // S
    if running[e] != 0 and t - e * S < count[e]:
        x[t] = tmp[t]


@wp.kernel
def _freeze_k(moved: wp.array(dtype=wp.float32), tol: float,
              running: wp.array(dtype=wp.int32)):
    e = wp.tid()
    if moved[e] < tol:
        running[e] = 0                   # absorbing: nothing sets it back to 1


@dataclasses.dataclass(frozen=True, kw_only=True)
class SettleParams(plugin.Params):
    max_iters: int = plugin.static(64, lo=1, doc="iteration bound")
    check_every: int = plugin.static(8, lo=1, doc="iterations per window")
    lam: float = plugin.static(0.3, lo=0.0, hi=0.5, doc="smoothing factor")
    pull: float = plugin.static(0.2, lo=0.0, hi=0.5, doc="pull back to the input")
    tol: float = plugin.static(1e-5, lo=0.0, doc="largest step of a settled env")


class Settle(plugin.Stage):
    kind = "relax"
    name = "acme.settle"
    capture = plugin.CaptureClass.BOUNDED   # its loop may stop early
    Params = SettleParams
    reads = {"xy": "cl.xy", "count": "cl.count"}
    writes = {"xy": "cl.xy"}

    def build(self, b, params):
        x, count = b.read("xy"), b.read("count")
        out = b.write("xy", wp.vec2f, per="cl")
        tmp = b.temp_like("xy")
        running = b.scratch("running", wp.int32, b.E)    # 1 while the env moves
        moved = b.scratch("moved", wp.float32, b.E)      # its largest step
        active = b.scratch("active", wp.int32, 1)        # any env still running
        S, E, K = b.layout.stride, b.E, params.check_every
        lam, pull, tol = float(params.lam), float(params.pull), float(params.tol)

        def body(ctx, it):
            ctx.fill(moved, 0.0)
            ctx.launch(_step_k, E * S,
                       [out, x, count, S, lam, pull, running, tmp, moved])
            ctx.launch(_apply_k, E * S, [tmp, count, S, running, out])
            if (it + 1) % K == 0:                   # the iteration before a window
                ctx.launch(_freeze_k, E, [moved, tol, running])
                ops.any_set(ctx, running, active)

        def enqueue(ctx):
            ctx.copy(out, x)                        # the padding too
            ctx.fill(running, 1)
            ctx.bounded_loop(params.max_iters, body, active=active,
                             first_check=K, check_every=K)
        return enqueue


def register(reg):
    reg.add(Settle)

For eight bezier tracks (seeds 0 to 7) on cpu, the loop stops after three of its seven windows, with the same bytes as all 64 iterations give. Like the Taubin quick start, the stage shows the mechanics: it enforces no separation, so at the default half_width its tracks fail the thickness check, as they do with relax=None.

The checklist of a BOUNDED stage:

• Declare it. capture = plugin.CaptureClass.BOUNDED (or a capture_class that returns it). A stage declared STATIC that calls bounded_loop(active=...) raises ValueError at construction (... declare capture = CaptureClass.BOUNDED). active=None (a plain static unroll) is fine in any stage.

• Make the freeze absorbing (once frozen, an env never changes again). Once active reads 0, every remaining iteration must leave the stage’s outputs (and anything it reads after the loop) bitwise unchanged: a per-env flag that only ever goes from running to frozen, and every kernel that writes an output returning early for a frozen env. Scratch the loop alone uses (moved, active here) may differ.

• Refresh active on the iteration before each window. The windows start at first_check, first_check + check_every, …; the check before the window at lo reads what iteration lo - 1 wrote. Here first_check = check_every = K and the refresh runs when (it + 1) % K == 0. Write it before the first check, in every generation: a check must never read last generation’s value.

• Keep the body pure in it. Branch only on it and on static params; mutate no Python state that code after the loop reads (the loop may stop at any window), and emit only through the ctx the body is given.

• Keep each window device-only. Under lowering (b) a window is the body of a CUDA conditional node, which may hold kernel launches, memsets and device-to-device copies only. Launch kernels and use ctx.copy / ctx.fill between device arrays on ctx’s stream; no copy to or from pageable host memory, no event, no other stream (no wait_stream fork or join). CUDA accepts a copy from pinned, device-mapped host memory in the body, but keep windows to device memory for portability. Anything else fails inside wp.capture_if while the caller’s capture is being recorded, so the iterations before the window are already recorded. check_stage(..., devices=["cuda:0"]) (bounded.lowering) records every window that way.

• Mind the cost. The generator’s own graph replays a conditional node per window; a caller’s plain torch.cuda.graph (and any eager cuda call) runs the full n_iters, so pick a bound you can afford, with the per-iteration work of a frozen env reduced to early returns.

check_stage("relax", "acme.settle") checks it like any stage, plus bounded.absorbing (the host-poll lowering against the unrolled one, with active poisoned) and, on CUDA, bounded.lowering (see Testing a Plugin). A loop that keeps moving a converged env also fails batch_independence: under an early stop it moves until the WHOLE batch has converged, so its output depends on the other envs.

ops.any_set(ctx, mask, flag) (any env set) and ops.all_set(ctx, mask, flag) (every env set) reduce an [E] int32 mask into the [1] flag in two launches, deterministically. repulsive is the built-in example: its stall check freezes an env on device, and its growth loop is a bounded_loop over that freeze (see Repulsive Generator).

Writing a host-driven (HOST) stage

Some work needs the host: a NumPy or SciPy routine, a readback that decides what runs next. Declare capture = CaptureClass.HOST. The stage then runs eagerly on every call, outside every capture, as its own segment: the stages around it still replay from graphs on cuda (see HOST stages and segments). It reaches the host through ctx.host(), a HostOps, with staging buffers it allocates at build with b.host_buffer(dtype, shape) (pinned on cuda, so nothing allocates per call).

This generator draws each env’s loop on the host with NumPy, from the env’s effective seed:

import dataclasses

import numpy as np
import warp as wp

from track_gen import ops, plugin


@dataclasses.dataclass(frozen=True, kw_only=True)
class BlobParams(plugin.Params):
    points: int = plugin.static(64, lo=8, doc="points per loop")


class Blob(plugin.Stage):
    kind = "generator"
    name = "acme.blob"
    output = "centerline"
    capture = plugin.CaptureClass.HOST      # host RNG: runs eagerly
    Params = BlobParams
    writes = {"xy": "raw.xy", "valid": "gen.valid"}

    def capacity(self, params, layout):
        return {"raw": params.points}

    def build(self, b, params):
        n, E = params.points, b.E
        xy = b.write("xy", wp.vec2f, per="raw")         # [E * n] on the device
        valid = b.write("valid", wp.int32)              # [E]
        seed_h = b.host_buffer(wp.int32, E)             # host staging buffers
        xy_h = b.host_buffer(wp.vec2f, E * n)
        theta = np.linspace(0.0, 2.0 * np.pi, n, endpoint=False)
        radius = 0.5 * ops.constants.TARGET_EXTENT * b.scale

        def enqueue(ctx):
            h = ctx.host()
            seeds = h.readback(b.seed, via=seed_h)      # a NumPy view of seed_h
            pts = np.empty((E, n, 2), np.float32)
            for e, s in enumerate(seeds):
                rng = np.random.default_rng(int(s) & 0x7FFFFFFF)
                lobes, phase = rng.integers(2, 6), rng.uniform(0.0, 2.0 * np.pi)
                r = radius * (1.0 + 0.2 * np.sin(lobes * theta + phase))
                pts[e, :, 0], pts[e, :, 1] = r * np.cos(theta), r * np.sin(theta)
            h.upload(xy, pts.reshape(E * n, 2), via=xy_h)   # returns once copied
            ctx.fill(valid, 1)                          # device work goes through ctx
        return enqueue


def register(reg):
    reg.add(Blob)

HostOps has four operations, each on the call’s stream:

• read_i32(arr, index=0): one element of an int32 device array, through a one-element host buffer the Program allocates for the stage.

• readback(src, via=...): copies src into via (a host_buffer of the same dtype, at least as large), waits, and returns a NumPy view of via shaped like src, valid until via is used again.

• upload(dst, data, via=...): copies data into via, then via into dst, and waits, so via is free again on return.

• sync(): waits until the call’s stream has drained (a no-op on cpu).

The rules of a HOST stage:

• ctx.host() works only inside a HOST stage’s enqueue: a STATIC or BOUNDED stage that calls it raises ValueError on its first run, naming the stage and its declared class (as bounded_loop(active=...) in a STATIC stage does), and so does b.host_buffer in another stage.

• Never under a capture: a caller’s capture of the Program raises NotCapturableError naming the HOST stages before any launch, and ctx.host() raises it too.

• Reach the host only through ctx.host(), with via= buffers from b.host_buffer (on cuda a via that is not pinned raises ValueError). A direct .numpy(), wp.synchronize*() or a per-call allocation reads on the wrong stream, may allocate, and can invalidate another thread’s CUDA capture: the conformance kit fails hot_path.host_sync / hot_path.allocation for them, while the HostOps operations pass (they are safe while other threads capture on the device).

• The rest of the contract holds: per-env determinism in the seed, no state carried from one generation to the next, every output slot written on every call.

Users who need to record the whole pipeline into their own graph build it with PipelineSpec(..., require_caller_capturable=True), which refuses HOST stages at build with a PlanError naming them. check_generator("acme.blob") (or check_stage for a pipeline stage) runs the HOST variants of the checks: eager determinism and statelessness, the segmented graphs against an eager run, and the refusals (see the HOST checks).

Selecting a stage

register(reg) adds stage classes with reg.add(MyStage), next to any generators; names, duplicates, loading and quarantine follow the rules of Packaging a Plugin (dotted names, Names). registry.describe(kind, name) reports a stage’s reads, writes, capture, capturable and params_schema. A stage is selected by:

• presets.track(..., relax=..., width=..., elevation=..., checks=(...)), with a name, (name, params) or a StageRef (see Presets: presets.track);

• presets.gates(..., generator=..., order=..., stations=..., elevation=..., checks=(...)), with the same forms (see Gate presets: presets.gates);

• the edit operations of any PipelineSpec (replace, insert_after, …);

• for relax and elevation, the flat TrackGenConfig fields relax_solver and z_profile (see Pipeline stages: relax_solver and z_profile); for gates, the GateGenConfig fields generator, gate_ordering and z_profile.

Testing a stage

track_gen.testing.check_stage runs a stage inside the default track pipeline (an order or stations stage, or a check that reads gate channels, inside the default gate pipeline) through the conformance checks: shapes and NaN padding of the Track (or GateSequence), determinism, batch independence, statelessness, the hot-path guards, caller capture on CUDA, the bounded.* checks for a stage that runs a bounded_loop and host.* for a HOST stage, and the stage-specific checks: stage.isolation (the stage writes none of its inputs and recomputes its output from them), for a validity stage validity.bits, and for a width stage width.values:

from track_gen.testing import check_stage

check_stage("relax", "acme.taubin", {"passes": 4}).raise_for_failures()
check_stage("width", "acme.ripple").raise_for_failures()
check_stage("width", "from_channel", ports={"src": "cl.acme.hw"},
            upstream={"generator": "acme.ribbon"}).raise_for_failures()
check_stage("elevation", "acme.sinusoid_hills").raise_for_failures()
check_stage("validity", "acme.max_curvature",
            upstream={"elevation": "acme.sinusoid_hills"}).raise_for_failures()
check_stage("order", "acme.reverse").raise_for_failures()

See Pipeline stages: check_stage for the checks and the arguments (check_stage("order", "acme.reverse") checks the ordering of Writing an ordering). Test the stage’s own properties too: that a tap of its output matches a NumPy reference on a small batch (Taps), that a dynamic parameter set with gen.params[...].set takes effect, that a check sets its bit exactly where it should (Testing your plugin’s own properties shows the patterns).