Execution layer¶

How TrackGen emits device work: stream resolution, the execution policy shared by every producer (record into your capture, eager, or replay of the producer’s own CUDA graph; segment by segment for a pipeline with HOST stages), the device-side conditional node skip, bounded loops, the host operations of HOST stages and the capture errors; then the per-env seed and epoch state, and the torch interop. Calling a generator and Rules for recording into your graph show them in use; CUDA Graph — capture, replay, and buffer reuse explains the mechanics.

Launch contexts and runners¶

Execution layer: how track_gen emits device work.

LaunchContext pins a call to one device and one stream so the work can be recorded into a caller’s CUDA graph (Warp ScopedCapture, capture_begin(external=True), or torch.cuda.graph) and replayed correctly. GraphRunner is the execution policy of one producer (record into the caller’s capture / eager / replay of the producer’s own CUDA graph, its self-graph, a pipeline with HOST stages segment by segment). HostOps (ctx.host()) is how a HOST stage reaches the host. unconditional() keeps CUDA conditional nodes out of the captures recorded in its block. RebindAfterCaptureWarning flags a bound input replaced after a capture recorded the method that reads it.

class LaunchContext[source]¶

Bases: object

One device + one stream; the only way library code emits device work.

Construct with for_call() at a public API boundary and pass it down (ctx=) to every helper. Helpers accept ctx=None and resolve it with resolve(), which prefers the ambient context.

Besides launches, copies and fills it offers device-side control flow that stays recordable into any capture (skip_unless(), bounded_loop()) and, inside a HOST stage only, host operations on its stream (host()).

__init__(device, stream)[source]¶

Parameters:: stream (Stream | None)
Return type:: None

device¶

stream¶

classmethod for_call(device, stream=None, *, default=None)[source]¶

Context for a public call on device.

With no explicit stream, an ambient context on the same device (set by an enclosing scope()) is reused; otherwise the stream is resolved by resolve_stream().

Return type:: LaunchContext

classmethod current(device)[source]¶

Context pinned to Warp’s CURRENT stream on device (no adoption).

For code that must stay on the stream Warp allocates and zero-fills on (e.g. PerEnvSeededRNG construction, the facades’ backward-compatible _run()), regardless of ambient or torch state.

Return type:: LaunchContext

property capturing: bool¶: Whether this context’s stream is capturing. Informational only: enqueue code must not branch on it.

launch(kernel, dim, inputs=(), outputs=(), *, block_dim=256, max_blocks=0)[source]¶

wp.launch(kernel, dim, inputs, outputs) on this context’s device and stream (block_dim and max_blocks as in wp.launch).

Parameters:

inputs (Sequence)
outputs (Sequence)
block_dim (int)
max_blocks (int)

Return type:

None

copy(dst, src, *, dst_offset=0, src_offset=0, count=0)[source]¶

wp.copy(dst, src) on this context’s stream: count elements (0: the whole of src) from src[src_offset:] into dst[dst_offset:].

Parameters:

dst (array)
src (array)
dst_offset (int)
src_offset (int)
count (int)

Return type:

None

fill(arr, value)[source]¶

Fill arr with value on this stream (never arr.zero_()/ fill_(), whose memset targets Warp’s current stream).

Parameters:: arr (array)
Return type:: None

skip_unless(flag, body)[source]¶

Run body(self), skipped on device when flag[0] == 0 where that is possible; otherwise run it unconditionally. Returns True when body was recorded under a conditional node, False when it ran (or was recorded) unconditionally.

CONTRACT: body must be an exact no-op when flag[0] == 0 (it must leave every buffer bitwise unchanged), because whether the skip happens depends on the capture, not on the data:

inside a capture that Warp registered on this context’s stream (Warp ScopedCapture, capture_begin(external=True), or track_gen.interop.torch.registered(conditional=True) under torch.cuda.graph) on a device with conditional-node support (CUDA Toolkit and driver >= 12.4): body is recorded under a CUDA conditional node that reads flag[0] at replay time (wp.capture_if);
anywhere else (eager calls, the CPU device, an unregistered torch.cuda.graph capture or one registered without conditional=True, an unmarked Warp registration whose stream is torch’s current, capturing stream in this thread (a hand-made capture_begin(external=True) of torch’s capture, or your own ScopedCapture of torch’s current stream unless inside conditional()), a stream that only joined a capture registered on another stream, an older driver, inside unconditional()): body runs (or is recorded) unconditionally.

A graph holding a conditional node cannot be used as a CUDA child graph (e.g. wp.capture_if(c, on_true=that_graph)). wp.capture_if also ends and restarts the stream’s capture around the node, which CUDA refuses while a side stream that joined the capture has not been joined back (an unjoined fork invalidates the capture), and which gives the capture a new CUDA capture id: an external capturer that tracks its capture by id (torch’s caching allocator under torch.cuda.graph) loses track of it for the rest of the capture. Hence the opt-in under torch, and unconditional() to opt out of a Warp-registered capture.

Never takes Warp’s eager capture_if path, which reads the flag back on the host and synchronizes the stream.

Parameters:

flag (array) – [1] int32 array on this context’s device (e.g. written by track_gen.ops.any_set() earlier on the same stream).
body (Callable[[LaunchContext], None]) – body(ctx) emitting its work through ctx (this context).

Raises:

TypeError – flag is not a [1] int32 wp.array.
ValueError – flag is not on this context’s device (wp.capture_if would miss the capture and read the flag on the host).

Return type:

bool

bounded_loop(n_iters, body, *, active=None, first_check=0, check_every=1)[source]¶

Run body(ctx, it) for it in range(n_iters) (it a static Python int), stopping early once active[0] == 0 where the capture state allows it.

Iterations [0, min(first_check, n_iters)) always run. Then come the windows [lo, min(lo + check_every, n_iters)) for lo = first_check, first_check + check_every, ...: each is preceded by a check of active[0] and may be skipped, with every later one, once that check reads 0. Without active every iteration runs (a plain unrolled loop).

How a window check is lowered depends on where the call runs, never on data:

(a) a host poll (one 4-byte read) that stops the loop: the CPU device, eager, with no CPU graph capture open on it;
(b) a CUDA conditional node per window (wp.capture_if, via skip_unless()): inside a Warp-registered capture of this context’s stream with conditional-node support (CUDA >= 12.4), e.g. a producer’s own graph, a Warp ScopedCapture or track_gen.interop.torch.registered( conditional=True);
(c) every window runs (or is recorded) unconditionally: eager calls on CUDA (they never host-sync), plain torch.cuda.graph captures, older drivers, CPU graph captures, recordings inside unconditional().

CONTRACT (checked by the conformance kit, bounded.absorbing): active is a [1] int32 array on this context’s device that the stage writes on the device before the first check (its enqueue, before the loop, when first_check == 0; otherwise usually the iteration lo - 1 that precedes each window), in every generation; once a check reads active[0] == 0, every remaining iteration leaves everything read after the loop bitwise unchanged (an ABSORBING freeze); the body emits only through the context it is given and mutates no Python state read after the loop. The three lowerings then give bitwise identical results; how many iterations actually execute is unspecified (n_iters is a bound). A stage whose loop passes active declares capture = CaptureClass.BOUNDED.

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

Parameters:

n_iters (int) – The iteration bound (>= 0).
body (Callable[[LaunchContext, int], None]) – body(ctx, it), emitting its work through ctx.
active (array | None) – [1] int32 flag on this context’s device (None: no checks).
first_check (int) – The first iteration that may be skipped (>= 0).
check_every (int) – Iterations per window (>= 1).

Raises:

TypeError – a non-int count (bool included), or active not a [1] int32 wp.array.
ValueError – a negative count, check_every < 1, active on another device, or active passed inside a stage that declares capture = CaptureClass.STATIC.

Return type:

None

host()[source]¶

The HostOps of the HOST stage running on this context: readbacks, uploads and waits on this context’s stream through the stage’s build-time pinned host buffers.

Raises:

CaptureError – not inside the enqueue of a pipeline stage on this device.
ValueError – inside the enqueue of a stage that does not declare capture = CaptureClass.HOST (the error a wrong declared capture class gives, as for bounded_loop(active=...) in a STATIC stage and StageBuilder.host_buffer), or of a legacy generator spec.
NotCapturableError – this context’s stream is capturing, or a CPU graph capture is open on the device (host work would be recorded, or read stale data).

Return type:

HostOps

sync()[source]¶

Debug-only host sync on this stream (TRACKGEN_DEBUG_SYNC=1).

Return type:: None

scope()[source]¶

Make this the ambient context for nested ctx=None helpers and legacy plugins (context-variable scoped: thread- and task-local).

Return type:: Iterator[LaunchContext]

class GraphRunner[source]¶

Bases: object

Execution policy (record / eager / self-graph) for one producer.

Parameters:

device – the producer’s Warp device.
capturable – whether the enqueue body obeys the graph-capture contract. False runs eagerly outside a capture and raises NotCapturableError (before any launch) inside one.
legacy_stream – the body contains a legacy-ABI plugin (no accepts_ctx) whose own launches go to Warp’s CURRENT stream (see the module docstring for what that changes). Recording into a caller’s capture requires the capture stream to be that current stream (LegacyPluginStreamError otherwise).
default_stream – stream used when a call passes none and no ambient context applies (see resolve_stream()).
owner – object anchored on a caller’s Warp graph when the body is recorded into it, so the buffers the graph references outlive the producer. Held weakly by the runner.
name – label used in error messages.

A body passed to run() may also be a segmented body (one part per Program segment, built by the Program): see the SEGMENTED row of the module docstring.

A runner is single-writer: running the same producer concurrently from two threads is a race on its output buffers (use one producer per concurrent caller). Different producers are independent: each self-graph is captured on the producer’s own private stream, so no process-wide lock is taken (legacy-plugin producers, and calls on Warp’s current stream while one of them may still self-capture, excepted).

__init__(device, *, capturable, legacy_stream=False, default_stream=None, owner=None, name='', _host_ids=(), _legacy_capture=None)[source]¶

Parameters:

capturable (bool)
legacy_stream (bool)
default_stream (Stream | None)
name (str)

Return type:

None

property graph: Graph | None¶: The default-key self-graph, once captured (first CUDA call with key=None outside a capture); None on CPU, for a non-capturable producer (also a segmented one: its parts have one graph each), or before that call.

graph_for(key=None)[source]¶

The whole-body self-graph cached for key (see run()), if captured; None for the key of a segmented body (see segment_graphs_for()).

Return type:: Graph | None

property graph_keys: tuple¶: The keys of every self-graph captured so far (None = default), segmented bodies included (their graphs: segment_graphs_for()).

segment_graphs_for(key=None)[source]¶

The per-part self-graphs of the segmented body cached for key, one per part in order (None at a HOST part), or None before they were captured (and for a body that is not segmented).

Return type:: tuple | None

property device¶: The producer’s Warp device (a warp.context.Device).

property capturable: bool¶: Whether the body may be recorded into a caller’s capture (the constructor’s capturable); False runs eagerly and raises NotCapturableError inside a capture.

property legacy_stream: bool¶: Whether the body holds a legacy-ABI plugin that launches on Warp’s current stream (the constructor’s legacy_stream).

property default_stream: Stream | None¶: The stream used when a call passes none and no ambient context applies (None on CPU, or when the constructor got none).

property name: str¶: The label used in error messages (the constructor’s name).

property prepared: bool¶: Whether prepare() ran (one real pass outside any capture).

context(stream=None)[source]¶

The launch context of one call (explicit stream > ambient > default_stream > torch capture stream > Warp’s current stream).

Return type:: LaunchContext

ensure_not_capturing()[source]¶

ensure_not_capturing() for this runner’s device and default stream: call it before the producer’s first allocation.

Return type:: LaunchContext

serial(ctx)[source]¶

The lock a call on ctx must hold while it decides its mode and runs (a context manager; a no-op for most calls).

runtime._CAPTURE_LOCK (reentrant) for a producer that self-captures legacy-plugin code on Warp’s current stream (CUDA), and for any call resolving to that stream (or of a legacy producer) while a legacy producer on the device may still self-capture there. Callers that issue their own launches ahead of run() (e.g. a reseed) hold it around both. An eager run of legacy code holds the lock on its own (see the module docstring).

Parameters:: ctx (LaunchContext)

prepare(enqueue)[source]¶

Run enqueue once, for real, outside any capture, and block until it finished.

Launching for real loads every module / kernel overload / block-dim variant the hot path uses (force_module_load and load_module miss launch-time overloads), so later recordings never load modules while capturing. The pass leaves a valid result in the producer’s buffers.

When the resolved stream is not Warp’s current stream (a default_stream or an ambient scope), it first waits on the current stream, where the producer’s buffers were just allocated and zeroed. The host then blocks on the resolved stream (a cold, construction-time cost), so the result is complete for any later stream or readback.

Every LaunchContext.bounded_loop of the pass runs all its windows (forced unrolled lowering), so every kernel a window can emit is loaded.

Raises:: CaptureError – if the resolved stream or Warp’s current stream is capturing, or torch is capturing in the calling thread.
Return type:: None

run(enqueue, *, stream=None, policy='auto', sync=False, key=None)[source]¶

Run enqueue once under the execution policy; returns the call’s resolved LaunchContext.

A non-capturable runner refuses a caller’s capture (RECORD) and a CPU graph capture (wp.ScopedCapture("cpu")) with NotCapturableError, before any launch (see check_record()). Outside a capture it runs a plain body eagerly; a segmented body (one part per Program segment) runs its HOST parts eagerly and replays the other parts’ self-graphs on CUDA (SEGMENTED, see the module docstring).

Parameters:

enqueue – the producer’s body, enqueue(ctx) -> None (or a segmented body built by a Program).
stream – optional wp.Stream (or "torch") for this call.
policy (str) – "auto" (record / self-graph / eager as above) or "eager" (record when capturing, otherwise run the body directly; never self-graph).
sync (bool) – block the host on the resolved stream afterwards; skipped when recording, and while torch is capturing in the calling thread.
key – hashable self-graph cache key identifying WHICH body enqueue is (e.g. (advance, mask.ptr)). The self-graph of a key is captured on its first SELF-GRAPH call and replayed by every later call with an equal key, so equal keys must mean identical bodies (same launches, same buffer pointers). None (default) is the producer’s default body (graph). A non-None key requires prepare() to have run (a keyed body may not be idempotent, so it is never used as a warm-up pass). Ignored when recording or eager. A legacy-plugin producer runs non-None keys eagerly (it self-captures on Warp’s shared current stream, once). At most MAX_KEYED_GRAPHS (16) non-None keys get a self-graph; calls with a further new key run eagerly (with a one-time RuntimeWarning).

Return type:

LaunchContext

check_record(ctx)[source]¶

Raise where run() would refuse ctx before any launch: ctx is capturing and this producer cannot be recorded there, or a CPU graph capture (wp.ScopedCapture("cpu")) is open on the device and the producer is not capturable. A no-op otherwise. Lets a caller validate before issuing its own launches ahead of run() (e.g. a reseed).

Raises:

NotCapturableError – the producer is not capturable (it names its HOST stages).
LegacyPluginStreamError – a legacy-plugin producer and the capture stream is not Warp’s current stream.

Parameters:

ctx (LaunchContext)

Return type:

None

class HostOps[source]¶

Bases: object

Host operations of a HOST stage, bound to one call’s stream (ctx.host()).

Every copy runs on ctx.stream through a pinned host buffer the stage allocated at build (StageBuilder.host_buffer), and every wait blocks on that stream only, so a HOST stage never allocates per call and never waits on the whole device. On the CPU device copies are synchronous and sync() does nothing (the same code path runs).

Obtain it with h = ctx.host() inside the enqueue of a stage declaring capture = CaptureClass.HOST:

def build(self, b, p):
    seeds_h = b.host_buffer(wp.int32, b.E)
    xy_h = b.host_buffer(wp.vec2f, b.E * p.points)
    xy = b.write("xy", wp.vec2f, per="raw")

    def enqueue(ctx):
        h = ctx.host()
        seeds = h.readback(b.seed, via=seeds_h)        # numpy view of seeds_h
        pts = my_host_algorithm(seeds)                  # [E * points, 2] float32
        h.upload(xy, pts, via=xy_h)
    return enqueue

__init__(ctx, slot)[source]¶

Parameters:: slot (_SlotInfo)
Return type:: None

property stage_id: str¶: The id of the HOST stage these operations belong to.

sync()[source]¶

Block the host until the work queued on ctx.stream finished (CPU: no-op). A stream synchronize in relaxed capture mode, never a device-wide one.

Return type:: None

read_i32(arr, index=0)[source]¶

arr[index] of an int32 device array, through the slot’s build-time pinned [1] buffer (copy on ctx.stream, then sync()).

Raises:

TypeError – arr is not an int32 wp.array, or index not an int.
ValueError – arr on another device, not contiguous, or index out of range.

Parameters:

arr (array)
index (int)

Return type:

int

readback(src, *, via)[source]¶

Copy src (a device array) into the host buffer via on ctx.stream, wait, and return a numpy VIEW of via shaped like src (vector dtypes add a trailing component axis). The view is valid until via is used again.

Raises:

TypeError – a non-array, or via of another dtype.
ValueError – src on another device or not contiguous, via not a host buffer (not pinned, on CUDA) or too small.

Parameters:

src (array)
via (array)

Return type:

ndarray

upload(dst, data, *, via)[source]¶

Write data (anything numpy converts to via’s dtype, dst.size elements: e.g. an [n, 2] float array for a vec2f dst) into the host buffer via, copy it to the device array dst on ctx.stream and wait, so via is reusable on return.

Raises:

TypeError – a non-array, or via of another dtype.
ValueError – dst on another device or not contiguous, via not a host buffer (not pinned, on CUDA) or too small, data of the wrong size.

Parameters:

dst (array)
via (array)

Return type:

None

resolve_stream(device, stream=None, default=None)[source]¶

Resolve the single stream a public call runs on (see module docstring).

Parameters:

device – a Warp device or device string.
stream – None, a wp.Stream, or "torch".
default – stream to use when stream is None and no ambient context applies (e.g. a stream bound at construction).

Return type:

Stream | None

resolve(ctx, device)[source]¶

ctx if given, else LaunchContext.for_call() on device.

Parameters:: ctx (LaunchContext | None)
Return type:: LaunchContext

is_capturing(stream)[source]¶

True when stream is capturing (any capture: Warp, torch, external).

Parameters:: stream (Stream | None)
Return type:: bool

unconditional(enabled=True)[source]¶

Record no CUDA conditional node in this block (thread- and task-local).

Inside the block LaunchContext.skip_unless() records its body unconditionally in every capture, so Course.regenerate(mask=...) and every LaunchContext.bounded_loop() window (the repulsive generator, a plugin’s BOUNDED stage) take lowering (c) (every window recorded), with the same results. Use it around a recording into a Warp-registered capture (wp.ScopedCapture, wp.capture_begin) that must not get conditional nodes:

a graph you will use as a CUDA child graph (wp.capture_if(c, on_true=graph), a graph node of another graph): a graph holding a conditional node cannot be one;
a capture whose stream has an outstanding fork (a side stream that joined the capture and has not been joined back): wp.capture_if ends and restarts the stream’s capture around each node, which CUDA refuses with an unjoined fork (cudaErrorStreamCaptureUnjoined), invalidating your capture.

A producer’s own graph captured in the block (its first CUDA call outside a capture, e.g. gen.graph) is recorded unconditionally too, and stays so. Eager calls are unchanged. unconditional(False) restores the default inside an enclosing block.

Parameters:: enabled (bool)
Return type:: Iterator[None]

conditional(enabled=True)[source]¶

Let your own Warp capture of torch’s current stream record CUDA conditional nodes (thread- and task-local).

LaunchContext.skip_unless() (a masked Course.regenerate, every LaunchContext.bounded_loop() window) records no conditional node in a Warp registration nobody opted in whose stream is torch’s current stream while that stream is capturing: Warp does not record who began the capture, and a hand-made capture_begin(external=True) of a running torch.cuda.graph must not get one (wp.capture_if renews the capture id torch’s caching allocator tracks). That also catches YOUR OWN wp.ScopedCapture of torch’s current stream (e.g. under torch.cuda.stream(t) + wp.ScopedStream(wp.stream_from_torch(t))), whose windows are then recorded unconditionally: same bytes, no device-side skip. Inside this block such a capture records conditional nodes; use it only when torch is NOT the capturer (under torch.cuda.graph use track_gen.interop.torch.registered(conditional=True) instead). unconditional() wins over it; the library’s own graphs (a producer’s graph) are unaffected (Warp captures them, marked). conditional(False) restores the default inside an enclosing block.

Parameters:: enabled (bool)
Return type:: Iterator[None]

Capture errors¶

Their causes and fixes are summarized in Errors and warnings.

exception CaptureError[source]¶

Bases: RuntimeError

A call was made in a capture state it does not support.

For example, constructing a generator (which runs one real warm-up pass) while the resolved stream is being captured into a CUDA graph.

exception NotCapturableError[source]¶

Bases: CaptureError

A non-capturable producer was asked to record into a caller’s capture (a CUDA graph capture, or a CPU graph capture: wp.ScopedCapture("cpu")).

Raised BEFORE any device work is issued, so the caller’s capture is left untouched. Outside a capture the same producer runs: its HOST stages eagerly, the rest of a pipeline from its own graphs (segmented execution, see GraphRunner).

name¶: The producer’s name (e.g. the generator name).

stage_ids¶: The ids of the HOST stages that make it non-capturable (() when unknown, e.g. a bare GraphRunner(capturable=False)).

__init__(name='', detail='', *, stage_ids=())[source]¶

Parameters:

name (str)
detail (str)

Return type:

None

exception LegacyPluginStreamError[source]¶

Bases: CaptureError

A legacy-ABI plugin would launch on the wrong stream under capture.

Legacy generator plugins (registered without accepts_ctx=True) issue their own wp.launch calls on Warp’s CURRENT stream. Recording one into a caller’s capture is therefore only correct when the capture stream IS Warp’s current stream.

__init__(name='', detail='')[source]¶

Parameters:

name (str)
detail (str)

Return type:

None

exception RebindAfterCaptureWarning[source]¶

Bases: UserWarning

A bound input array was replaced after a hot method that reads it was recorded into a capture.

The recorded graph keeps reading the arrays that were bound when it was recorded, so after a rebind its replays silently read the old buffers. Emitted (once per bind call) by Course.bind, Component.bind and the utilities’ bind / bind_inputs / bind_window when a pointer changes after their query() / update() ran while its stream (or a CPU graph capture) was capturing; rebinding the same arrays is silent. The documented recipe “rebind, then record again” also warns once: filter this category when you re-record on purpose:

import warnings

import track_gen.exec

warnings.simplefilter("ignore", track_gen.exec.RebindAfterCaptureWarning)

Seeds and epochs¶

Per-env seeds and device-side epochs, the pair that keys every generation; see Seeds, epochs and the generation counter.

Per-env seeds and epochs: the device-side state that keys every generation.

A SeedState holds two [E] int32 arrays:

seeds: the per-env base seeds. Either an array the state OWNS (built from None / an int, or copied with copy=True), or a caller array BOUND by pointer (a wp.array, or a PerEnvSeededRNG’s seeds_warp). SeedState never writes bound seeds, and neither do the generators: write new seeds into that array yourself (they are read at every generation, also by every graph replay), or use an owned state and SeedState.set(). The rng itself does write its seeds_warp: PerEnvSeededRNG.set_seeds_warp and Course.generate(seeds=...) overwrite it in place, and it IS the caller’s array when the rng (or CourseConfig.seeds) was given a wp.array.
epoch: owned, zero-initialised, and only ever changed on device (SeedState.advance(), SeedState.reset_epoch(), SeedState.set()), so it can be driven from inside a CUDA graph.

Every generation first folds the two into an EFFECTIVE seed per env (_effective_seed_k, the first launch of every generator body):

effective = seed                                       if epoch == 0
          = hash32(u32(seed) ^ hash32(u32(epoch) * 0x9E3779B9)) & 0x7FFFFFFF
                                                       otherwise

Epoch 0 leaves the seed unchanged, so a never-advanced env generates exactly what its base seed alone gives. A MASKED generation (effective(out, mask=m)) re-derives the effective seed of the envs with m[e] != 0 only: every other env keeps the effective seed its current output was built from, so it is recomputed to identical values whatever happened to its (seed, epoch) since. Determinism: the pair (seed, epoch) fully reproduces an env. A generator whose seed for that env is effective_seed(seed, epoch) at epoch 0 produces the identical track or gate sequence (effective_seed() is the bit-exact host twin of the kernel).

class SeedState[source]¶

Bases: object

Per-env base seeds plus a device-side epoch counter (see the module docstring).

Parameters:

num_envs – number of environments E (>= 1).
device – Warp device of the state.
seeds – None (owned 0 + arange(E)), an int k (owned k + arange(E) with int32 wrap-around, the same expansion as PerEnvSeededRNG(seeds=k)), an [E] int32 wp.array on device (BOUND by pointer, never written by SeedState), or a PerEnvSeededRNG (binds its seeds_warp by pointer, so rng.set_seeds_warp(...) keeps reseeding generation). The rng must therefore live on device.
copy – with an array or rng seeds, take an OWNED copy instead of binding (then set() is allowed).
default_stream – optional wp.Stream the methods use when a call passes neither stream= nor ctx= (and no ambient context applies), as track_gen.exec.resolve_stream()’s default. The facades pass their constructor stream=, so gen.seeds.advance() is ordered with gen.generate().

Construction allocates (and, for an owned int expansion, uploads) on Warp’s current stream: build it outside any CUDA graph capture. Every other method is one or two launches on the call’s stream: allocation-free, sync-free and capturable.

Raises:

TypeError – seeds is not one of the accepted kinds.
ValueError – num_envs < 1, or a seed array that is not [E] int32 on device.

__init__(num_envs, device, seeds=None, *, copy=False, default_stream=None)[source]¶

Parameters:

num_envs (int)
copy (bool)
default_stream (Stream | None)

Return type:

None

property seeds: array¶: [E] int32 base seeds (a caller array when bound).

property epoch: array¶: [E] int32 per-env epoch (owned; changed only on device).

property bound: bool¶: Whether seeds is a caller array bound by pointer (never written by SeedState; set() raises).

property num_envs: int¶: The number of environments E (the length of seeds and epoch).

property device: str¶: The Warp device of the state, as a string ("cpu", "cuda:0").

check_mask(mask, name='mask')[source]¶

Validate a contiguous [E] int32 mask on this state’s device (host-side). Contiguous, because the producers cache one CUDA graph of their own per mask POINTER: two views with one pointer and different strides would share it.

Parameters:: name (str)
Return type:: None

advance(mask=None, *, stream=None, ctx=None)[source]¶

epoch[e] += 1 for every env (mask=None) or where mask[e] != 0.

One launch on the call’s stream; capturable (inside a caller’s graph every replay advances again).

Parameters:

mask (array | None) – optional [E] int32 mask on this device, read on device.
stream – optional wp.Stream (or "torch"); resolved per track_gen.exec.
ctx – a LaunchContext to use instead of stream.

Return type:

None

reset_epoch(mask=None, *, stream=None, ctx=None)[source]¶

epoch[e] = 0 for every env (mask=None) or where mask[e] != 0 (one launch; capturable).

Parameters:: mask (array | None)
Return type:: None

set(seeds, ids=None, *, stream=None, ctx=None)[source]¶

Write new base seeds into an OWNED state and zero those envs’ epochs.

ids=None: seeds is [E] int32 and replaces every env’s seed. Otherwise ids is [n] int32 env indices and seeds the [n] int32 new seeds (seeds[i] goes to env ids[i]; duplicate ids race). One launch (or a copy plus a launch); capturable.

Raises:

RuntimeError – the seeds are bound to a caller array (SeedState never writes bound seeds): write that array yourself, or build the state with copy=True.
ValueError – wrong shape, dtype or device.

Parameters:

seeds (array)
ids (array | None)

Return type:

None

effective(out, mask=None, *, stream=None, ctx=None)[source]¶

Write the effective per-env seeds (see the module docstring) into the [E] int32 out: for every env, or (mask) only where mask[e] != 0, leaving the other entries of out untouched. One launch; capturable. This is the first launch of every generator body (after the epoch advance, if any).

Parameters:

out (array)
mask (array | None)

Return type:

None

effective_seed(seed, epoch)[source]¶

Host twin of the effective-seed kernel: the seed env generation actually uses for base seed seed at epoch epoch (both taken as int32).

effective_seed(s, 0) == s; for epoch != 0 the result is a non-negative int32 hash of (s, epoch). A generator whose rng seed for an env is effective_seed(s, k) (at epoch 0) reproduces that env’s batch at (s, k) exactly.

Parameters:

seed (int)
epoch (int)

Return type:

int

Torch interop¶

Registering a torch.cuda.graph capture with Warp, and tying producers’ lifetime to a torch graph; see CUDA graphs in a sim loop. Importing track_gen does not import torch.

PyTorch interop for recording track_gen work into torch.cuda.graph captures.

A plain torch.cuda.graph(g) capture is invisible to Warp: library calls made inside it are still recorded into g (they adopt torch’s capture stream, see track_gen.exec), but two Warp features need the capture to be REGISTERED with Warp:

lifetime anchoring: a producer recorded into a registered capture is attached to the Warp graph object of that capture;
device-side conditional skips (track_gen.exec.LaunchContext.skip_unless(), used by Course.regenerate(mask=...) and by every window of an iterative BOUNDED stage such as the repulsive generator), opt-in with conditional=True: otherwise the guarded body is recorded unconditionally (always correct, never skipped). A hand-made Warp registration of torch’s capture (wp.capture_begin(stream=..., external=True) on torch’s capture stream) gets no conditional node either: opt in through registered(). (Your own wp.ScopedCapture of torch’s current stream, when torch is not capturing, opts in with track_gen.exec.conditional().)

registered() registers torch’s capture with Warp for the duration of a with block; retain() ties Python objects to the torch graph’s lifetime (here course is a bound Course and done your [E] int32 episode-end mask):

import torch
import track_gen.interop.torch

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    with track_gen.interop.torch.registered(graph=g):
        course.step()
        course.regenerate(mask=done, advance=True)
track_gen.interop.torch.retain(g, course)   # keep the course's buffers alive with g

This module imports torch lazily (only when a function is called).

registered(stream=None, *, graph=None, capture_mode=None, conditional=False)[source]¶

Use INSIDE torch.cuda.graph(g). Resolves the capture stream to the library’s cached Warp wrapper of that handle (never a throwaway one: dropping a wp.stream_from_torch wrapper unregisters the handle process-wide), then:

if Warp already has a capture registered on that handle (e.g. an enclosing registered() or Warp ScopedCapture), yields that registered stream and does nothing else (graph= and conditional= are then ignored: the enclosing registration owns the Warp-side graph object);
if the stream only JOINED a capture Warp registered on another stream (a side stream forked from it), yields the stream and registers nothing (a second registration of the same capture would break the first one’s teardown); the conditional skip then does not apply on this stream;
otherwise opens wp.ScopedCapture(stream=s, external=True) on it (which does not begin a new CUDA capture, only registers torch’s) and closes it with the SAME stream object at the end of the block (Warp looks an open capture up by stream object identity).

Inside the block, library calls adopt the stream as usual and additionally get lifetime anchoring and, with conditional=True, conditional skips (LaunchContext.skip_unless, CUDA 12.4+). The registration ends with the block; the torch capture continues.

conditional=True caveat: a conditional node (wp.capture_if) ends and restarts the CUDA capture of the stream, which gives it a new capture id, and torch’s caching allocator recognises its graph’s capture by the id taken at capture_begin. After the first conditional node, torch therefore no longer routes the allocations of later torch ops in the same torch.cuda.graph to the graph’s private pool (they either fail to capture or take blocks from the general pool that eager code may reuse while the graph still writes them), and other torch features that look the capture up by its id may fail. The first conditional node comes with a masked Course.regenerate(mask=...) or with ANY generation that runs a BOUNDED stage (the repulsive generator, a plugin’s ctx.bounded_loop with active=: one conditional node per window), not only with masked regeneration. So with conditional=True, record those calls AFTER every torch op of the graph that allocates (or preallocate their outputs and use out= / in-place ops), up to the end of the torch.cuda.graph block, not just of this one. Ending the capture also fails while a side stream that joined it has not been joined back. The graph also cannot be used as a CUDA child graph. Without conditional=True the masked regeneration and every BOUNDED window are recorded unconditionally (exact no-ops when skipped would be, costing their full work per replay) and torch’s capture is never touched.

Parameters:

stream – None (torch’s current stream on the current CUDA device, which is the capture stream inside torch.cuda.graph), a torch.cuda.Stream, or a wp.Stream.
graph – optional torch.cuda.CUDAGraph being captured. When given, the Warp-side graph object of the registration (it holds the producers anchored during the block and the Warp modules the recorded kernels use) is retain()-ed on it, so they live as long as graph. Without it they are released when the block exits: keep the producers alive yourself, e.g. retain(g, course).
capture_mode – the mode torch began the capture with (torch.cuda.graph(..., capture_error_mode=...)): "global" (default, torch’s default), "thread_local", "relaxed", or a wp.CaptureMode.
conditional (bool) – allow LaunchContext.skip_unless to record CUDA conditional nodes into this registration (a masked Course.regenerate, every window of a BOUNDED stage; default False; see the caveat above).

Yields:

The registered wp.Stream (pass it as stream= if you like; library calls inside the block resolve to it anyway).

Raises:

track_gen.exec.CaptureError – the stream is not capturing (not inside torch.cuda.graph).

Return type:

Iterator[Stream]

retain(torch_graph, *owners)[source]¶

Keep owners alive for as long as torch_graph is.

A CUDA graph stores raw device pointers; it does not keep the buffers it reads and writes alive. Replaying torch_graph after a producer (Course, TrackGenerator, …) recorded into it was garbage-collected reads freed memory. retain(g, course) stores strong references in an attribute of the torch graph object (g._track_gen_owners), so they are released together with g. Idempotent per object; None owners are ignored.

Raises:: TypeError – torch_graph cannot carry attributes (no __dict__).
Return type:: None