`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. Explicit registration also works when automatic registration is disabled with CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION=1.

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 = 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

Examples

Construct a payload from the kernel’s item-count argument, fill its items with squared indices, and store the resulting blocked tiles:

import numpy as np
from numba_cuda_mlir import cuda, types

import cuda.coop.numba_mlir as numba_coop  # noqa: F401 -- activate the backend
from cuda import coop

@cuda.jit
def square_indices(destination, items_per_thread):
    block = coop.this_block()
    items = coop.ThreadData(items_per_thread, alignment=16)
    offset = cuda.blockIdx.x * cuda.blockDim.x * items_per_thread
    for i in range(items.items_per_thread):
        index = offset + cuda.threadIdx.x * items_per_thread + i
        items[i] = types.int32(index * index)
    coop.store(block, destination, items, offset=offset)

for items_per_thread in (1, 4):
    destination = cuda.device_array(
        2 * 64 * items_per_thread, dtype=np.int32
    )
    square_indices[2, 64](destination, items_per_thread)
    expected = np.arange(2 * 64 * items_per_thread, dtype=np.int32) ** 2
    np.testing.assert_array_equal(destination.copy_to_host(), expected)

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 launch-time shared-memory 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

Examples

Reuse one descriptor for transpose Load and Store. The loop processes two independent tiles. Explicit auto_sync=True enables barriers between operations and between iterations.

import numpy as np
from numba_cuda_mlir import cuda

import cuda.coop.numba_mlir as numba_coop  # noqa: F401 -- activate the backend
from cuda import coop

@cuda.jit
def copy_tiles(source, destination, items_per_thread):
    block = coop.this_block()
    scratch = coop.TempStorage(alignment=16, auto_sync=True)
    items = coop.ThreadData(items_per_thread)
    for tile in range(2):
        offset = tile * cuda.blockDim.x * items_per_thread
        coop.load(
            block,
            source,
            items,
            offset=offset,
            algorithm="transpose",
            temp_storage=scratch,
        )
        coop.store(
            block,
            destination,
            items,
            offset=offset,
            algorithm="transpose",
            temp_storage=scratch,
        )

for items_per_thread in (1, 4):
    values = (np.arange(2 * 128 * items_per_thread) % 7).astype(np.int32)
    source = cuda.to_device(values)
    destination = cuda.device_array_like(source)
    copy_tiles[1, 128](source, destination, items_per_thread)
    np.testing.assert_array_equal(destination.copy_to_host(), values)

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 before summing it. 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

See also

cub::BlockLoad, cub::WarpLoad: C++ block and warp Load primitives.

Examples

Copy an array with Numba-CUDA-MLIR, using 128 threads and a kernel argument for the values per thread. Each block loads up to 128 * items_per_thread elements. The last block pads missing values with zero and stores only the valid prefix.

import numpy as np
from numba_cuda_mlir import cuda

import cuda.coop.numba_mlir as numba_coop  # noqa: F401 -- activate the backend
from cuda import coop

@cuda.jit
def copy_tiles(source, destination, items_per_thread):
    block = coop.this_block()
    items = coop.ThreadData(items_per_thread)
    tile_size = cuda.blockDim.x * items_per_thread
    offset = cuda.blockIdx.x * tile_size
    valid = min(max(source.size - offset, 0), tile_size)
    coop.load(
        block,
        source,
        items,
        valid_items=valid,
        oob_default=0,
        offset=offset,
    )
    coop.store(
        block,
        destination,
        items,
        valid_items=valid,
        offset=offset,
    )

for items_per_thread in (1, 4):
    expected = np.arange(1000, dtype=np.int32)
    source = cuda.to_device(expected)
    destination = cuda.device_array_like(source)
    tile_size = 128 * items_per_thread
    blocks = (expected.size + tile_size - 1) // tile_size
    copy_tiles[blocks, 128](source, destination, items_per_thread)
    np.testing.assert_array_equal(destination.copy_to_host(), expected)

The qualified import activates the Numba-CUDA-MLIR backend even if another module imported cuda.coop first. Use the qualified cuda.coop.<backend> API for backend-specific behavior.

See also

Load visualization

cuda.coop.store( group: BlockGroup | WarpGroup, destination: object, value: _CommonNumericT | CommonThreadDataLike[_CommonNumericT], /, *, algorithm: BlockLoadStoreAlgorithm | WarpLoadStoreAlgorithm = 'direct', valid_items: ValidItems | None = None, offset: IntegerValue | None = None, temp_storage: TempStorageLike | None = None, ) → None#

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

Parameters:

group (cuda.coop.ThreadGroup) – Participating threads; see thread groups. Supports blocks and physical or logical warps. Warp stores require an enclosing block size divisible by 32.
destination (array) – Writable one-dimensional contiguous array in device-accessible memory, with the same element dtype as value. It must contain every element selected by offset and valid_items.
value (numeric scalar or cuda.coop.ThreadDataLike) – This thread’s value or readable payload. Initialize every item that will be stored. The tile contains group_size * items_per_thread elements, with one item per thread for a scalar. As in CUB, transpose algorithms may rearrange the payload in place. Do not rely on its contents after Store; copy values before the call if they are needed later.
algorithm (str, optional) – Compile-time store algorithm, default "direct". "direct" expects blocked values; "striped" expects striped values. "vectorize" and "transpose" also expect blocked values and use vector accesses or shared-memory rearrangement, respectively. Blocks additionally support "warp_transpose" and "warp_transpose_timesliced", both requiring a block size divisible by 32. See data layouts before pairing different Load and Store algorithms.
valid_items (int or integer scalar, optional) – Number of valid elements in the group’s tile, uniform across the group and between zero and the tile size, inclusive. None stores the entire tile. Elements outside the valid prefix are not written.
offset (int or integer scalar, optional) – Nonnegative offset in elements from the start of destination, uniform across the group. None means zero. Supply each block’s origin explicitly. Warp stores also add the within-block group origin automatically, using the same addressing rule as cuda.coop.load().
temp_storage (cuda.coop.TempStorageLike, optional) – Scratch descriptor for block transpose-family algorithms. None uses automatic scratch. Direct, striped, and vectorized stores need no shared scratch. Warp stores require None.

Returns:

Writes to destination and may rearrange the input payload.

Return type:

None

See also

cub::BlockStore, cub::WarpStore: C++ block and warp Store primitives.

Examples

Store a partial tile at an element offset. The untouched prefix and suffix keep their sentinel values. The input payload is not used after the transpose Store.

import numpy as np
from numba_cuda_mlir import cuda, types

import cuda.coop.numba_mlir as numba_coop  # noqa: F401 -- activate the backend
from cuda import coop

@cuda.jit
def store_prefix(destination, items_per_thread):
    block = coop.this_block()
    rank = cuda.threadIdx.x
    items = coop.ThreadData(items_per_thread)
    for item in range(items_per_thread):
        items[item] = types.int32(items_per_thread * rank + item)
    coop.store(
        block,
        destination,
        items,
        algorithm="transpose",
        valid_items=cuda.blockDim.x * items_per_thread - 5,
        offset=5,
    )

for items_per_thread in (1, 4):
    destination = cuda.to_device(
        np.full(128 * items_per_thread + 7, -1, dtype=np.int32)
    )
    store_prefix[1, 128](destination, items_per_thread)
    expected = np.full(128 * items_per_thread + 7, -1, dtype=np.int32)
    expected[5 : 128 * items_per_thread] = np.arange(
        128 * items_per_thread - 5, dtype=np.int32
    )
    np.testing.assert_array_equal(destination.copy_to_host(), expected)

See participation and synchronization for control-flow requirements at primitive calls.

See also

Store visualization

Numba-CUDA-MLIR-qualified API#

Use this module for the extensions below. Shared parameters and behavior follow the Common API.

import cuda.coop.numba_mlir as coop

Qualified calls also accept fixed-size, one-dimensional local arrays where the operation accepts per-thread payloads. local and shared expose Numba-CUDA-MLIR’s memory namespaces.

Payloads and temporary storage#

cuda.coop.numba_mlir.ThreadData(items_per_thread, dtype=None, *, alignment=None)#

Create a fixed-size Numba local array for cooperative operations.

The parameters, uninitialized contents, and example follow cuda.coop.ThreadData(). This qualified constructor returns the active Numba-CUDA-MLIR runtime’s local-array representation. The compiler infers the element type from supported context, including Load’s source, Store’s destination, and typed indexed assignments. Leave the constructor type unspecified for those uses. See element-type inference for cases that need additional type information.

Prefer a kernel parameter named items_per_thread and construct the payload with ThreadData(items_per_thread). The compiler specializes the kernel for each supplied count, which remains fixed during execution.

alignment requests a minimum alignment in bytes. This backend also enforces the compiler’s pointer-alignment requirement. See per-thread payloads for indexing and lifetime.

class cuda.coop.numba_mlir.TempStorage( size_in_bytes=None, *, alignment=None, auto_sync=False, sharing='shared', )#

Shared-memory requirements for cooperative operations in one kernel.

Parameters, defaults, synchronization rules, and the executable reuse example follow cuda.coop.TempStorage(). This qualified descriptor exposes size_in_bytes, alignment, auto_sync, and sharing for the Numba-CUDA-MLIR planner. auto_sync=None becomes False. Set auto_sync=True to request automatic reuse barriers.

Only supported block algorithms accept an explicit descriptor. The planner determines capacity and alignment from its uses; its contents are opaque to user code. See temporary storage for shared versus exclusive slices and manual reuse synchronization.

Memory operations#

cuda.coop.numba_mlir.load( group: BlockGroup | WarpGroup, source: object, output: ThreadDataLike[_CommonNumericT] | numpy.ndarray, /, *, 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 block or warp tile with the Numba-CUDA-MLIR backend.

Parameters, algorithm choices, tile addressing, and the None return follow cuda.coop.load(). In this backend, output may also be a supported fixed-size Numba local array. It must be writable; an untyped local array can infer its dtype from the source.

See per-thread payloads and temporary storage for allocation rules. The executable example in cuda.coop.load() activates this backend explicitly and shows a guarded final tile.

See also

cub::BlockLoad, cub::WarpLoad: C++ Load primitives used for these group scopes.

Examples

Copy a partial final tile using a Numba local array.

import numpy as np
from numba_cuda_mlir import cuda, types

import cuda.coop.numba_mlir as coop

@cuda.jit
def copy_tiles(source, destination, items_per_thread):
    block = coop.this_block()
    # Qualified calls also accept a fixed-size Numba local array.
    items = cuda.local.array(items_per_thread, dtype=types.int32)
    tile_size = cuda.blockDim.x * items_per_thread
    offset = cuda.blockIdx.x * tile_size
    valid = min(source.size - offset, tile_size)
    coop.load(
        block,
        source,
        items,
        offset=offset,
        valid_items=valid,
        oob_default=0,
    )
    coop.store(block, destination, items, offset=offset, valid_items=valid)

for items_per_thread in (1, 4):
    values = np.arange(1000, dtype=np.int32)
    source = cuda.to_device(values)
    destination = cuda.device_array_like(source)
    tile_size = 64 * items_per_thread
    blocks = (values.size + tile_size - 1) // tile_size
    copy_tiles[blocks, 64](source, destination, items_per_thread)
    np.testing.assert_array_equal(destination.copy_to_host(), values)

See also

Load visualization

cuda.coop.numba_mlir.store( group: BlockGroup | WarpGroup, destination: object, value: _CommonNumericT | CommonThreadDataLike[_CommonNumericT] | numpy.ndarray, /, *, algorithm: BlockLoadStoreAlgorithm | WarpLoadStoreAlgorithm = 'direct', valid_items: ValidItems | None = None, offset: IntegerValue | None = None, temp_storage: TempStorageLike | None = None, ) → None#

Store a block or warp tile with the Numba-CUDA-MLIR backend.

Parameters, algorithm choices, tile addressing, and the None return follow cuda.coop.store(). value may be a numeric scalar or a supported fixed-size Numba local array. Transpose algorithms may rearrange that input in place; copy values before Store if they are needed later.

See per-thread payloads and temporary storage for allocation rules, and cuda.coop.store() for an executable partial-tile example.

See also

cub::BlockStore, cub::WarpStore: C++ Store primitives used for these group scopes.

Examples

Store only the valid prefix of the final tile.

import numpy as np
from numba_cuda_mlir import cuda, types

import cuda.coop.numba_mlir as coop

@cuda.jit
def copy_tiles(source, destination, items_per_thread):
    block = coop.this_block()
    # Qualified calls also accept a fixed-size Numba local array.
    items = cuda.local.array(items_per_thread, dtype=types.int32)
    tile_size = cuda.blockDim.x * items_per_thread
    offset = cuda.blockIdx.x * tile_size
    valid = min(source.size - offset, tile_size)
    coop.load(
        block,
        source,
        items,
        offset=offset,
        valid_items=valid,
        oob_default=0,
    )
    coop.store(block, destination, items, offset=offset, valid_items=valid)

for items_per_thread in (1, 4):
    values = np.arange(1000, dtype=np.int32)
    source = cuda.to_device(values)
    destination = cuda.device_array_like(source)
    tile_size = 64 * items_per_thread
    blocks = (values.size + tile_size - 1) // tile_size
    copy_tiles[blocks, 64](source, destination, items_per_thread)
    np.testing.assert_array_equal(destination.copy_to_host(), values)

See also

Store visualization

cuda.coop API Reference#

Common API#

Backend registration#

Thread groups#

Payloads and temporary storage#

Memory operations#

Numba-CUDA-MLIR-qualified API#

Payloads and temporary storage#

Memory operations#

`cuda.coop` API Reference#