`cuda.coop` API Reference#

Warning

cuda.coop is an experimental API and is subject to change.

Common API#

The primitive functions below are compiler markers; register is a host-side configuration function. The installed .pyi files are authoritative for overload and result typing.

Backend registration#

See registering a backend.

cuda.coop.register( backend: Literal['numba-cuda-mlir', 'numba_cuda_mlir'], ) → None#

Register cooperative primitives with the selected compiler backend.

Call this on the host before compiling a kernel, including when cuda.coop was imported before the compiler. Repeated calls are safe.

Parameters:

backend – Compiler backend to register. "numba_cuda_mlir" is an alias for "numba-cuda-mlir".

Raises:

ValueError – The backend name is unsupported.
ImportError – The backend adapter or its required compiler is unavailable or incompatible. Other backend initialization errors propagate unchanged.

Thread groups#

See thread groups and participation and synchronization for the shared execution model. ThreadGroup and ThreadHierarchy describe which threads cooperate. A descriptor’s availability does not imply that every primitive supports that group.

cuda.coop.this_thread() → ThreadGroup#

Describe the calling thread as a one-thread group.

Returns:: A symbolic descriptor. Load and Store do not accept one-thread groups. See thread groups.
Return type:: cuda.coop.ThreadGroup

cuda.coop.this_warp() → ThreadGroup#

Describe the calling thread’s physical warp.

Returns:: A group of 32 consecutive threads in the block’s linear thread order. Use group_by() to form smaller logical warps.
Return type:: cuda.coop.ThreadGroup

Notes

Warp Load and Store require a block size divisible by 32. Constructing this descriptor does not turn a partial final warp into a complete group. See participation requirements.

cuda.coop.this_block() → ThreadGroup#

Describe all threads in the calling thread’s CUDA block.

Returns:: A descriptor whose size comes from the kernel launch’s block dimensions. For multidimensional blocks, thread ranks are linearized with the x coordinate varying fastest.
Return type:: cuda.coop.ThreadGroup

Notes

The factory takes no size argument and does not synchronize the block. See thread groups and the primitive’s participation requirements.

cuda.coop.this_cluster() → ThreadGroup#

Describe the thread-block cluster containing the calling thread.

Returns:: A symbolic descriptor. Cluster groups are not Load or Store targets.
Return type:: cuda.coop.ThreadGroup

Notes

Construction does not create a cluster or enable cluster scheduling. See thread groups.

cuda.coop.this_grid() → ThreadGroup#

Describe all threads in the current kernel grid.

Returns:: A symbolic descriptor. Grid groups are not Load or Store targets.
Return type:: cuda.coop.ThreadGroup

Notes

Construction does not request a cooperative launch or synchronize the grid. See thread groups.

class cuda.coop.ThreadGroup( kind: str, hierarchy: ThreadHierarchy = <factory>, parent: ThreadGroup | None = None, mapping: GroupByMapping | None = None, source: str = 'explicit', )#

Describe the participants in a cooperative operation.

Use the this_* factories to describe groups in the current kernel launch. Constructing a descriptor does not synchronize threads or launch a kernel. Load and Store support blocks, physical warps, and logical warps; see thread groups.

Group descriptors can also be constructed in ordinary Python. Runtime rank, size, membership, and synchronization queries are not available in this API layer.

group_by( count: int, *, exhaustive: bool = True, ) → _ThreadGroupT#

Partition a physical warp by threads or a block by warps.

Parameters:

count (int) – Positive compile-time number of units in each subgroup. For a warp parent, the unit is one thread; for a block parent, it is one physical warp. Thus this_warp().group_by(8) describes eight lanes, while this_block().group_by(2) describes 64 threads.
exhaustive (bool, optional) – Compile-time flag, default True. An exhaustive partition must divide the parent’s unit count exactly. False permits a remainder outside the complete groups. Each primitive still determines which partitions it supports.

Returns:

A descriptor for the subgroup containing the calling thread. Nested partitions are unsupported. Load and Store support logical warp widths of 1, 2, 4, 8, 16, or 32; mapped groups of physical warps are not Load or Store targets.

Return type:

cuda.coop.ThreadGroup

Examples

Descriptors can be inspected without compiling or launching a kernel:

from cuda import coop

group = coop.this_warp().group_by(8)
assert group.kind == "threads_within_warp"
assert group.static_size == 8

class cuda.coop.ThreadHierarchy#

Describe the launch dimensions known during group planning.

Public construction describes the current kernel launch with unresolved dimensions. Backends supply exact extents from LaunchFacts when they resolve a group. Callers cannot pass independent launch dimensions to the public constructor.

block_dim#

Threads along each axis of one block, once known.

Type:: tuple of int or None

cluster_dim#

Blocks along each axis of one cluster, once needed and known.

Type:: tuple of int or None

grid_dim#

Clusters along each axis of the resolved grid. On a non-cluster launch, each cluster contains one block.

Type:: tuple of int or None

implicit#

Whether this descriptor still refers to the current launch without explicit dimensions. A resolved hierarchy can leave higher levels unknown if the selected group does not need them.

Type:: bool

class cuda.coop.Hierarchy#: Alias for ThreadHierarchy.

Payloads and temporary storage#

ThreadData constructs the per-thread payload supplied to a primitive; thread-group descriptors describe the participating threads.

cuda.coop.ThreadData( items_per_thread: int, dtype: object = None, *, alignment: int | None = None, ) → ThreadDataLike[Any]#

Construct a fixed-size payload owned by the calling thread.

Each thread has its own slots. See per-thread payloads for their relationship to a group tile and the blocked and striped layouts.

Parameters:

items_per_thread (int) – Positive compile-time number of items owned by each thread. This extent is fixed for the lifetime of the payload and is available inside the kernel as items.items_per_thread. Pass it as a compile-time kernel argument using the backend’s supported specialization mechanism.
dtype (dtype-like, optional) – Optional numeric element type. Leave unspecified for normal use: a supported producer such as cuda.coop.load() supplies it. Backend-supported typed assignments can also establish the type. Supply it explicitly when the available type information cannot identify the intended representation. All items have the same type; conflicting typed values are an error.
alignment (int, optional) – Compile-time minimum storage alignment in bytes, expressed as a positive power of two. None lets the compiler choose. The request applies when payload storage is materialized; it does not assert alignment of a Load source or Store destination.

Returns:

Writable per-thread payload with indexed reads and writes. Its contents are uninitialized; write every item before reading it. Construction does not synchronize threads.

Return type:

cuda.coop.ThreadDataLike

class cuda.coop.ThreadDataLike(*args, **kwargs)#

Mutable fixed-size per-thread payload understood by supported backends.

See per-thread payloads for construction, item access, and dtype requirements.

cuda.coop.TempStorage( size_in_bytes: int | None = None, *, alignment: int | None = None, auto_sync: bool | None = False, sharing: str = 'shared', ) → TempStorageLike#

Describe shared scratch for supported cooperative block operations.

Construct the descriptor inside the kernel and pass it as temp_storage to operations that accept explicit block scratch. See temporary storage for supported operations, allocation lifetime, and synchronization requirements.

Parameters:

size_in_bytes (int, optional) – Positive compile-time capacity in bytes. None lets the compiler determine the capacity from all uses. An explicit capacity must be large enough for those operations; undersized storage is rejected.
alignment (int, optional) – Compile-time minimum alignment in bytes, expressed as a positive power of two. None lets the compiler choose. The allocation satisfies both this request and the operations’ alignment needs.
auto_sync (bool, optional) – Whether to insert a trailing barrier after each scratch-using call. Defaults to False; None also disables automatic reuse synchronization. The caller must synchronize before reusing the scratch, including on the next iteration of a loop. Pass True to request automatic reuse barriers.
sharing ({"shared", "exclusive"}, optional) – Compile-time allocation policy, default "shared". Calls using the same descriptor can reuse one scratch slice. "exclusive" gives distinct call sites separate slices, using more shared memory to avoid barriers needed solely for cross-call scratch reuse when auto_sync=False. Repeated execution of one call site still reuses its slice and requires synchronization before reuse. Omitting temp_storage leaves layout and reuse barriers to the compiler.

Returns:

Compiler-recognized scratch descriptor. The storage contents are opaque; keep application data in cuda.coop.ThreadData() or application-owned arrays.

Return type:

cuda.coop.TempStorageLike

class cuda.coop.TempStorageLike(*args, **kwargs)#

Explicit scratch descriptor understood by supported backends.

See temporary storage for construction, allocation sharing, and synchronization.

Memory operations#

cuda.coop.load( group: BlockGroup | WarpGroup, source: object, output: ThreadDataLike[_CommonNumericT], /, *, algorithm: BlockLoadStoreAlgorithm | WarpLoadStoreAlgorithm = 'direct', valid_items: ValidItems | None = None, oob_default: _CommonNumericT | float | None = None, offset: IntegerValue | None = None, temp_storage: TempStorageLike | None = None, ) → None#

Load a group tile from memory into per-thread values.

Parameters:

group (cuda.coop.ThreadGroup) – Participating threads; see thread groups. Supports blocks and physical or logical warps. Warp loads require an enclosing block size divisible by 32.
source (array) – One-dimensional contiguous source array in device-accessible memory. Its element dtype must match output; an untyped ThreadData infers its dtype from this array. The array must contain all elements selected by offset and valid_items.
output (cuda.coop.ThreadDataLike) – Writable per-thread payload. Load populates this payload in place. The group’s tile contains group_size * items_per_thread elements.
algorithm (str, optional) – Compile-time load algorithm, default "direct". "direct" gives each thread consecutive elements (blocked order); "striped" gives neighboring threads neighboring elements at each item index. "vectorize" uses vector accesses when possible, and "transpose" uses shared scratch to rearrange striped accesses into blocked order. Both return blocked order. Blocks also support "warp_transpose" and "warp_transpose_timesliced", which return blocked order and require a block size divisible by 32.
valid_items (int or integer scalar, optional) – Number of valid elements in the group’s tile, shared by all threads in that group. Supply a value between zero and the tile size, inclusive. None loads the full tile. Slots beyond this valid prefix have unspecified values unless oob_default is given.
oob_default (numeric scalar, optional) – Value written to slots beyond valid_items. Requires an explicit valid_items count. For example, use zero to pad a partial tile. A runtime value must have the payload dtype and be uniform across the group. With None, those slots are unspecified, even if initialized before the Load; assign them before reading them.
offset (int or integer scalar, optional) – Nonnegative offset in elements from the start of source, uniform across the group. None means zero. For block tiles, supply the block’s starting offset explicitly. For Warp tiles, the backend adds (linear_thread_rank // group_size) * tile_size automatically; do not include that within-block group offset a second time.
temp_storage (cuda.coop.TempStorageLike, optional) – Scratch descriptor for block transpose-family algorithms. None uses automatic scratch. Direct, striped, and vectorized loads need no shared scratch. Warp loads require None.

Returns:

The call populates output in place.

Return type:

None

cuda.coop API Reference#

Common API#

Backend registration#

Thread groups#

Payloads and temporary storage#

Memory operations#

`cuda.coop` API Reference#