`cuda.stf._experimental` API Reference#

Warning

cuda.stf._experimental is experimental. The API is subject to change without notice.

The core context, logical-data, task, and place types are implemented as a compiled extension; they are covered in the narrative guide and the C++ CUDASTF documentation. Because the extension is compiled (and mocked during documentation builds), autodoc cannot introspect these types, so their public surface is documented explicitly below. The pure-Python helper layers follow.

Contexts#

class cuda.stf._experimental.context(use_graph=False, *, stream=None, handle=None)#

Owns a Sequential Task Flow graph. All logical data and tasks belong to one context. Pass use_graph=True for the CUDA-graph backend, stream= a caller-owned CUDA stream (any __cuda_stream__ object or raw pointer) to emit work on top of it, and handle= an async_resources to share stream pools / cached graphs across contexts. Prefer with context() as ctx: so finalize() runs on exit.

logical_data(buf, dplace=None, name=None)#: Register an existing buffer (NumPy array, CUDA Array Interface object, or Python buffer) as logical data. Registration placement defaults to data_place.host(). The source must be C-contiguous; a read-only source rejects write()/rw().

logical_data_empty(shape, dtype=None, name=None, *, no_export=False)#: Allocate uninitialized logical data of the given shape/dtype.

logical_data_full( shape, fill_value, dtype=None, where=None, exec_place=None, name=None, )#: Allocate and fill logical data with a constant. Supports any 1/2/4/8-byte element type without optional third-party packages.

logical_data_zeros(shape, dtype=None, **kwargs)#

logical_data_ones(shape, dtype=None, **kwargs)#: Convenience wrappers over logical_data_full().

token()#: Create a token (buffer-less logical data) for pure ordering dependencies.

task(*args)#: Open a task. Positional args are dependencies (ld.read() / write() / rw()) and at most one exec_place. Use as a context manager.

cuda_kernel(*args)#: Like task() but for kernels described directly to STF as CUDA-graph nodes.

host_launch(*deps, fn, args=None, symbol=None)#: Schedule a Python callback with dependency tracking. Non-token dependencies are materialized as NumPy arrays and passed positionally to fn; token dependencies are ordering-only and are not materialized or passed. Exceptions raised by fn are captured and re-raised by a blocking wait()/finalize() or by check_errors().

wait(ld)#: Block until ld is available and return a host NumPy copy. Re-raises any pending host-callback exception.

fence()#: Return a raw CUstream (as int) that completes when all pending tasks finish, without destroying the context.

check_errors()#: Re-raise the first pending host-callback exception, if any (and clear it). Use with caller-stream contexts after synchronizing the stream, since their finalize() is asynchronous.

finalize()#: Run the graph and release resources. Blocking for a default context; asynchronous (non-blocking on the caller stream) for a context created with stream=.

place_resources#: Borrowed exec_place_resources owned by this context. Do not use past finalize().

class cuda.stf._experimental.stackable_context#

Nestable context supporting graph_scope() / while_loop() / repeat(count) scopes and record-once graphs. Mirrors context for logical_data*, task, host_launch, token, fence and check_errors. finalize() is only legal at root: every graph_scope / while_loop / repeat / LaunchableGraph must be closed first, or finalize() raises.

graph_scope()#

while_loop()#

repeat(count)#

Return context managers for a nested graph, a conditional while loop, or a fixed-count repeat scope. The object yielded by while_loop() exposes:

continue_while(...)

Set the loop’s continuation condition; call exactly once per body, after the body tasks. Accepts a single comparison in legacy form, continue_while(ld, op, threshold), or a condition expression of cond leaves combined with & (continue while all hold) or | (continue while any holds), optionally negated with ~:

loop.continue_while((lres > tol_sq) & (liter < max_iter))

One combiner applies per condition (up to 8 terms); mixed nesting such as (a & b) | c raises NotImplementedError.

cond_handle: Raw cudaGraphConditionalHandle as uint64_t, for a custom condition kernel that calls cudaGraphSetConditional() directly. Advanced: such a kernel cannot be written with Numba or PyTorch – cudaGraphSetConditional() is a device-runtime function, so the kernel must be compiled with relocatable device code and linked against cudadevrt (e.g. via NVRTC + nvJitLink). Prefer continue_while(...).

class cuda.stf._experimental.cond(ld, op, threshold)#

One while-loop continuation term: continue while ld <op> threshold, where ld is a 1-element logical data of a stackable context (float32 / float64 / int32 / int64), op is one of ">", "<", ">=", "<=" and threshold is a host-side real scalar. This is the canonical constructor; the ordering operators on stackable logical data (lres > tol_sq, …) are sugar that lowers onto it. Leaves combine with & / | and negate with ~; expressions have no Python truth value, so and / or / not raise TypeError. == / != on logical data keep identity semantics.

launchable_graph_scope()#: Return a context manager that instantiates the nested graph into a reusable cudaGraphExec_t launchable multiple times within the scope.

pop_prologue_shared()#: Return a storable LaunchableGraph for a graph built after push().

Logical data and dependencies#

class cuda.stf._experimental.logical_data#

A registered or allocated buffer tracked by a context. Created through the context logical_data* / token factories, not directly.

read(dplace=None)#

write(dplace=None)#

rw(dplace=None)#: Build a dep for a task/host_launch. Dependency placement defaults to data_place.affine(). write()/rw() raise on read-only sources.

dtype#

shape#

symbol#

readonly#: Metadata; readonly is True when the backing source forbids writes.

empty_like()#: Create a new logical data with the same shape/dtype metadata.

class cuda.stf._experimental.dep#: The result of ld.read()/write()/rw(). Also produced by the module-level read() / write() / rw() helpers.

cuda.stf._experimental.read(ld, dplace=None)#

cuda.stf._experimental.write(ld, dplace=None)#

cuda.stf._experimental.rw(ld, dplace=None)#: Functional forms of the dependency builders.

class cuda.stf._experimental.AccessMode#: IntFlag of access modes: NONE, READ, WRITE, RW.

Tasks, kernels, and streams#

class cuda.stf._experimental.task#

Returned by context.task(...); used as a context manager. Buffer/stream accessors are valid only while the task is active.

get_arg(index)#: Raw device pointer (int) for dependency index. Raises for token arguments.

get_arg_cai(index)#

args_cai()#: CUDA Array Interface view(s) for the non-token arguments. The views advertise no stream (CAI stream is None); launch your own work on the task stream(s) (stream_ptr() for scalar tasks, get_stream_at_index() / get_stream_ptrs() for grids), which STF has already ordered behind the data’s producers. Tokens are skipped.

stream_ptr()#: The task’s CudaStream.

get_grid_dims()#

get_stream_at_index(place_index)#

get_stream_ptrs()#: Grid-task helpers: grid shape (x, y, z, t) or None, the per-place stream at a linear index, and the list of all place streams.

set_exec_place(exec_place)#

set_symbol(name)#

class cuda.stf._experimental.cuda_kernel#: Returned by context.cuda_kernel(...). Adds launch(kernel, grid, block, args, shmem=0) on top of the task accessors, describing a kernel as a native CUDA-graph node.

class cuda.stf._experimental.CudaStream#: int subclass wrapping a raw CUstream; implements __cuda_stream__ and exposes .ptr.

Places, grids, and resources#

class cuda.stf._experimental.exec_place#

Where a task runs. Construct with device(), host(), current_device(), green_ctx(), or from_context().

static device(dev_id)#

static host()#

static current_device()#

static green_ctx(view, use_green_ctx_data_place=False)#

static from_context(ctx, dev_id=-1)#: Build a place from a device, the host, the current device, a green-context view, or an external CUcontext. Green-context and external-context places retain the objects they reference.

kind#

dims#

size#

backing_context#: kind is "host"/"device"; dims/size describe grids; backing_context is the external object backing a from_context place (else None), retained for the place’s lifetime.

set_affine_data_place(dplace)#

affine_data_place#

pick_stream(resources, for_computation=True)#

get_place(idx)#

class cuda.stf._experimental.exec_place_grid#

A grid of execution places (subclass of exec_place).

static from_devices(device_ids)#

static create(places, grid_dims=None, mapper=None)#: Build a grid from device ordinals, or from explicit places with an optional grid_dims shape (validated for rank, positivity, and product) and a mapper partition function. The grid retains its sub-places.

class cuda.stf._experimental.data_place#

Where logical data lives. Construct with device(), host(), managed(), affine(), current_device(), green_ctx(), or composite().

static device(dev_id)#

static host()#

static managed()#

static affine()#

static current_device()#

static green_ctx(view)#

static composite(grid, mapper)#: composite retains both the grid and the ctypes mapper closure it references.

kind#

device_id#

allocate(nbytes, stream=None)#

deallocate(ptr, nbytes, stream=None)#

class cuda.stf._experimental.exec_place_resources#: Per-place stream-pool registry. Construct standalone (exec_place_resources()) or borrow context.place_resources.

class cuda.stf._experimental.async_resources#: Shareable async_resources_handle. Reuse one across contexts to amortize graph-instantiation and share stream pools; it must outlive every context it is passed to.

class cuda.stf._experimental.LaunchableGraph#

Storable, shared-ownership handle for a re-launchable stackable graph, returned by stackable_context.pop_prologue_shared().

launch()#

reset()#

valid#

exec_graph#

stream#

graph#: launch() replays the graph; reset() drops the shared reference (running pop_epilogue when it was the last one) and is idempotent; valid reports whether the handle still refers to a live graph. Assigning the handle to another variable aliases the same object – there is no handle-duplication API – so resetting one resets all aliases.

Green-context places#

cuda.stf._experimental.green_places.green_places( sms_per_place: int, n_places: int | None = None, device_id: int = 0, coscheduled_sm_count: int = 0, ) → list[cuda.stf._experimental._stf_bindings.exec_place]#

Partition a device’s SMs into green contexts and return one STF place per partition.

Parameters:

sms_per_place – Number of SMs per place. Rounded up to the device’s minimum partition size by the driver.
n_places – Number of places to create, or None to create as many as the device’s SM count allows.
device_id – The device to partition.
coscheduled_sm_count – Optional co-scheduling constraint forwarded to SMResourceOptions.

Returns:

A list of exec_place, each backed by a cuda.core green Context. The contexts are kept alive by the places (see exec_place.from_context()); places also expose them via the read-only place.backing_context property for interop (e.g. warp.map_cuda_device).

Note

The split is performed by carving groups off the device’s SM resource through cuda.core; if fewer than n_places groups fit, a RuntimeError is raised. The caller’s current CUDA context is saved on entry and restored on return, so partitioning a device does not leave a different context current for the caller.

Record-once task graphs#

Use task_graph() to create a record-once task graph. It returns a TaskGraph object, which is the context manager and launch handle for the recorded graph.

cuda.stf._experimental.task_graph.task_graph() → TaskGraph#

Create a single-record, many-launch CUDASTF task graph.

Returns:: The object used as the recording context manager and launch handle.
Return type:: TaskGraph

class cuda.stf._experimental.task_graph.TaskGraph#

Object returned by task_graph().

A TaskGraph records a CUDASTF task DAG once and launches the recorded graph many times. User code should normally create instances with task_graph() rather than calling this class directly.

property raw: Any#: Return the underlying launchable graph after recording.

property graph: int#: Raw cudaGraph_t as a plain Python int.

property exec_graph: int#: Raw cudaGraphExec_t as a plain Python int.

property stream: int#: Raw cudaStream_t as a plain Python int.

launch() → None#: Launch the recorded graph once.

reset() → None#: Release the recorded graph and prevent future launches.

finalize() → None#: Release any recorded graph and finalize the owned context.

Device allocations#

Lightweight device array backed by data_place.allocate().

Implements BOTH device-memory interchange protocols – they are complementary, and a consumer picks by construction:

__cuda_array_interface__ (CAI v3): a description of the memory. The importer retains nothing; the DeviceArray (or something holding it) must outlive every borrowed view. This is the borrowed / zero-copy path – cuda.compute algorithms, Numba, torch.as_tensor.
__dlpack__ / __dlpack_device__ (DLPack): an ownership-carrying export. The capsule holds the owning array alive, and the consumer’s deleter releases it when the imported tensor’s storage dies – e.g. torch.from_dlpack gives a tensor whose lifetime carries the allocation, with the DeviceArray finalizer remaining the single deallocation point.

Also provides copy_to_host / copy_to_device helpers that mirror the Numba DeviceNDArray API.

class cuda.stf._experimental.device_array.DeviceArray(size: int, dtype, dplace: data_place, stream=None)#

1-D device array allocated through a data_place.

Parameters:

size (int) – Number of elements.
dtype (numpy dtype-like) – Element type.
dplace (data_place) – The data place that owns the allocation.
stream (optional) – CUDA stream for stream-ordered allocation.

__init__( size: int, dtype, dplace: data_place, stream=None, )#

static from_host( host_array: np.ndarray, dplace: data_place, stream=None, ) → DeviceArray#: Allocate on dplace and copy host_array to the device.

__dlpack__( *, stream=None, max_version=None, dl_device=None, copy=None, )#

Export as a "dltensor" capsule (ownership-carrying).

The capsule keeps the OWNING array (a view’s root) alive; the consumer’s deleter drops that reference when the imported tensor’s storage dies, and the DeviceArray finalizer – the single deallocation point – then frees the memory once no other reference remains. Complements __cuda_array_interface__, which describes the same memory but transfers no ownership.

max_version is accepted and answered with an unversioned capsule (permitted by the spec; understood by all consumers).

property dtype: dtype#

property shape#

property size: int#

property nbytes: int#

property data_place: data_place#: The data_place backing this array.

copy_to_host() → ndarray#: Synchronous device-to-host copy. Returns a new NumPy array.

copy_to_device(host_array: ndarray) → None#

Copy host_array into this device buffer (synchronous H2D).

The source must match this buffer’s byte size exactly – including for empty buffers and sliced views. A size mismatch is a programming error (it would otherwise silently leave part of the buffer untouched) and raises instead of performing a partial copy.

Path discovery#

Locate the CUDASTF C development headers and shared library.

These helpers let external C/CUDA projects compile and link against the same STF C ABI that the Python bindings use. Importing this module is cheap: it does not load the STF extension (_stf_bindings_impl) or preload CUDA libraries, so it is safe to use from build scripts.

The stf include path contains the C STF and cudax headers. When cuda-cccl is installed, get_include_paths() also returns its libcudacxx, CUB, and Thrust include paths.

cuda.stf._experimental.paths.iter_site_roots()#

Yield unique candidate roots under which an installed cuda package may live.

Scans sys.path plus the interpreter’s site directories. The site directories are required for pip build isolation, which strips the venv site-packages from sys.path while the package remains installed there (sys.prefix still points at the venv, so site.getsitepackages() recovers it). getsitepackages is missing in some virtualenv setups, so it is probed defensively.

class cuda.stf._experimental.paths.IncludePaths( cuda: 'Optional[Path]', libcudacxx: 'Optional[Path]', cub: 'Optional[Path]', thrust: 'Optional[Path]', stf: 'Optional[Path]', )#

cuda: Path | None#

libcudacxx: Path | None#

cub: Path | None#

thrust: Path | None#

stf: Path | None#

as_tuple()#

__init__( cuda: Path | None, libcudacxx: Path | None, cub: Path | None, thrust: Path | None, stf: Path | None, ) → None#

cuda.stf._experimental.paths.get_stf_include_dir() → Path#: Return cuda-stf’s own include root (cudax + C STF headers).

cuda.stf._experimental.paths.get_include_paths() → IncludePaths#

Return the include paths needed to compile against the STF C/C++ API.

The stf field contains the C STF and cudax headers. The libcudacxx, cub, and thrust fields contain cuda-cccl’s include root when available and are None otherwise.

cuda.stf._experimental.paths.get_library_dir() → Path#: Return the directory containing the STF C shared library.

cuda.stf._experimental.paths.get_library_path() → Path#: Return the full path to the STF C shared library.

Numba interop#

Numba interop helpers for cuda.stf._experimental.

This module provides:

get_arg_numba() and numba_arguments() – low-level converters from STF CAI objects to Numba CUDA device arrays.
numba_task() – context manager that opens an STF task and yields its arguments as Numba device arrays plus the task stream pointer.
jit() – an ergonomic @jit decorator that lets a Numba kernel be invoked directly with STF dep arguments. The first call compiles the underlying Numba kernel; subsequent calls reuse the cached compilation.

Numba is imported lazily inside each function. Importing this module does not require Numba to be installed; calling a function that uses Numba without numba-cuda available raises ImportError with an installation hint.

cuda.stf._experimental.interop.numba.get_arg_numba(task, index)#

Return one task argument as a Numba device array.

task.get_arg_cai(index) returns an stf_cai exposing the __cuda_array_interface__ protocol.

cuda.stf._experimental.interop.numba.jit(*jit_args, **jit_kwargs)#

STF-aware @jit decorator wrapping numba.cuda.jit.

A decorated function can be invoked as kernel[grid, block](*args) where arguments that are STF dep objects are transparently converted into Numba device arrays inside an STF task. The Numba compilation happens at first call.

Examples

Bare decorator:

@jit
def axpy(a, x, y):
    ...

With Numba cuda.jit arguments:

@jit(fastmath=True)
def kernel(...):
    ...

Then:

axpy[grid, block](2.0, lX.read(), lY.rw())

cuda.stf._experimental.interop.numba.numba_arguments(task)#

Return all task buffer arguments as Numba device arrays.

Same shape as task.args_cai(): None, a single array, or a tuple of arrays.

cuda.stf._experimental.interop.numba.numba_task(ctx, *args, symbol=None)#

Context manager: ctx.task(*args) yielding (numba_arrays, stream).

numba_arrays is a tuple of Numba CUDA device arrays (one per non-token dep), converted from each stf_cai via the CUDA Array Interface.

stream is the STF task’s stream pointer and implements the __cuda_stream__ protocol, so it can be passed as stream= to cuda.compute algorithms.

Example

>>> from cuda.stf._experimental.interop.numba import numba_task
>>> with numba_task(ctx, lA.read(), lB.read(), lC.rw()) as (args, stream):
...     cuda.compute.binary_transform(
...         args[0], args[1], args[2], OpKind.PLUS, N, stream=stream
...     )

class cuda.stf._experimental.interop.numba.stf_kernel_decorator(pyfunc, jit_args, jit_kwargs)#

Decorator-class wrapper around a Numba CUDA kernel for STF.

Created by jit(); not intended for direct instantiation. Indexing (kernel[grid, block, ...]) returns a fresh _stf_bound_kernel so launch configuration is never shared mutable state; only the compiled kernel is cached (and shared) here.

__init__(pyfunc, jit_args, jit_kwargs)#

PyTorch interop#

PyTorch interop helpers for cuda.stf._experimental.

This module provides:

tensor_arg() and tensor_arguments() – convert one or all STF task arguments to torch.Tensor views via the CUDA Array Interface.
pytorch_task() – context manager that opens an STF task, makes the task stream the current PyTorch CUDA stream, and yields the task arguments as torch.Tensor views.

PyTorch is imported lazily inside each function. Importing this module does not require PyTorch to be installed; calling a function that uses PyTorch without it raises ImportError with an installation hint.

cuda.stf._experimental.interop.pytorch.pytorch_task(ctx, *args)#

Context manager: ctx.task(*args) with PyTorch stream + tensor conversion.

Yields the tensor(s) from task.args_cai() converted to torch.Tensor as a tuple. The STF task stream is also made the current PyTorch CUDA stream for the duration of the with block.

Example

>>> from cuda.stf._experimental.interop.pytorch import pytorch_task
>>> with pytorch_task(ctx, lX.read(), lY.rw()) as (x_tensor, y_tensor):
...     y_tensor[:] = x_tensor * 2

cuda.stf._experimental.interop.pytorch.tensor_arg(task, index)#

Return one task argument as a torch.Tensor.

task.get_arg_cai(index) returns an stf_cai exposing the __cuda_array_interface__ protocol.

cuda.stf._experimental.interop.pytorch.tensor_arguments(task)#

Return all task buffer arguments as torch.Tensor views.

Same shape as task.args_cai(): None, a single tensor, or a tuple of tensors.

cuda.stf._experimental API Reference#