`cuda.coop` Programming Guide#

cuda.coop lets threads cooperate inside a Python GPU kernel. You can load a tile, operate on its per-thread values, and write the result without leaving the kernel. You choose the participating group and the data each thread contributes.

The visualizations show how values move through these operations. Each explorer includes an example kernel and lets you step through the algorithm. The glossary defines terms and layouts; the FAQs explain common API choices.

This guide assumes you have written a CUDA kernel and know how threads, blocks, and device arrays work. The examples use Numba-CUDA-MLIR. The installation instructions describe the matching cuda-coop extra; the current backend requires numba-cuda-mlir>=0.5.0,<0.6. Check the validation scope for environment coverage. Keep a compiled kernel in its original CUDA context; see the device and context-lifetime limitation before reusing a dispatcher across devices or recreated contexts.

This guide describes the experimental Numba-CUDA-MLIR 0.5.x API. Operation support varies by group and backend. The examples below use supported block and warp operations.

Requirements#

Use Python 3.10 through 3.14 and numba-cuda-mlir>=0.5.0,<0.6. The installation instructions describe the cuda-coop extra matching your CUDA major version. Running the examples requires a CUDA-capable GPU. The developer guide’s validation scope describes the environments covered by CI. Before reusing compiled kernels, check the device and context-lifetime rules and disk-cache behavior.

A first kernel: copying a tile#

The cuda.coop.load() reference contains a complete, tested copy kernel, including allocation, launch, partial-tile handling, and an output check. All threads participate in Load and Store, while valid_items selects the valid prefix of the final tile. The examples below use these imports:

import numpy as np
from numba_cuda_mlir import cuda, types
from cuda import coop

Choosing the common or qualified API#

The common API is imported with:

from numba_cuda_mlir import cuda
from cuda import coop

cuda.coop expresses operations through a common vocabulary of groups, numeric values, ThreadData, and TempStorage. The kernel compiler uses its registered backend to implement those calls. Start here when these operations cover your kernel’s needs. Numba-CUDA-MLIR is the first backend; CUTLASS support is planned.

The qualified import selects the Numba-CUDA-MLIR API explicitly:

import cuda.coop.numba_mlir as numba_coop

Both imports can be used in the same program. This guide uses coop for common calls and numba_coop for qualified calls so the choice is visible. In a program that uses only the qualified API, importing it as coop is also fine.

The qualified API accepts Numba-specific payloads and adds controls to several operations:

Need	Common API	Numba-CUDA-MLIR-qualified API
Fixed data per thread	`ThreadData`; scalar inputs where the operation accepts them	Also accepts fixed local arrays in supported array operations; exposes `local` and `shared` memory namespaces
Load/Store algorithms and explicit scratch	String algorithm selectors and `TempStorage` on supported block calls	Same shared controls; qualifying the import is unnecessary for these

The common API defines backend-independent contracts. You must still check that the selected backend implements the requested group, dtype, and operation. The kernels here also contain Numba launch and indexing code; porting the complete kernel to another DSL involves those parts too.

Registering the compiler backend#

Register explicitly on the host before compiling when imports may occur in any order, as in a library or notebook:

from cuda import coop

coop.register("numba-cuda-mlir")

from numba_cuda_mlir import cuda

Repeated calls are safe. "numba_cuda_mlir" is also accepted. Importing Numba-CUDA-MLIR before cuda.coop activates the backend automatically; importing cuda.coop.numba_mlir as numba_coop does so explicitly and gives you the backend namespace. See backend registration for details.

These primitive calls belong inside kernels compiled by a compatible backend. Registration itself is a host-side operation.

Compilation and kernel reuse#

CUDA devices and context lifetime#

With Numba-CUDA-MLIR 0.5.0 through 0.5.4, select the CUDA device before a kernel’s first compilation and keep its dispatcher and configured launch callables in that CUDA context. Reusing them on another device or after destroying and recreating the context is not qualified: the compiler can reuse architecture, compiled overload, or launch state from the original context. This also applies to kernels that use cuda.coop.

The upstream context-isolation fix must be released and qualified with this integration before relying on that reuse. Switching devices does not require re-registering cuda.coop; registration installs compiler hooks and does not repair dispatcher context state.

Disk caching#

With Numba-CUDA-MLIR 0.5.x, activating the cuda.coop backend disables the compiler’s cache=True disk cache for all kernels in that process. Compiled kernels still have an in-memory cache. The provider cache controlled by CUDA_COOP_ENABLE_CACHE is separate; see compiler activation and reuse for its configuration and lifetime.

Groups: which threads cooperate#

A group defines the participants in an operation. this_block() uses all threads in the current block; this_warp() uses a complete physical warp of 32 threads. this_warp().group_by(8) partitions it into four consecutive logical warps of eight threads. The compiler obtains the block shape from the kernel launch; the group factories take no size arguments.

Load and Store support blocks, physical warps, and logical warp widths of 1, 2, 4, 8, 16, or 32. Warp operations require a block size divisible by 32. The partition width is a compile-time constant. Every member of a participating group must reach the same primitive invocation.

The common API also declares other hierarchy descriptors. Their presence does not imply executable operations or queries on those scopes in this backend.

Participation and synchronization#

Every required participant must reach the same primitive invocation. For a block operation, a branch around the call must be uniform across the block. For a logical-warp operation, it must be uniform within that logical warp. Complete sibling logical groups can follow different paths.

In particular, putting a block primitive inside if index < count breaks participation on a partial tile. Keep the primitive outside the per-element condition. For a partial Load, supply valid_items and oob_default, or initialize invalid slots after the Load and before reading them. An early return by some block threads has the same problem if the remaining threads later execute a block primitive.

Use cuda.syncthreads() for explicit block synchronization in Numba kernels. Every thread in the block must reach the barrier.

An operation’s scratch-reuse barrier protects its temporary storage. Automatic synchronization inserts a trailing barrier after each call that consumes scratch. It does not prove that arbitrary user control flow is safe; every member must still reach the primitive and its barrier. Arrange synchronization for your own shared-memory communication as well. Constructing a group or a ThreadData object does not synchronize threads.

`ThreadData`: the part of a tile owned by one thread#

coop.ThreadData(items_per_thread) gives each thread that many slots. With 128 threads, the group owns 128 * items_per_thread values. Each thread indexes its own slots with items[i]. Each thread accesses only its own slots.

ThreadDataLike names the common payload interface used in API signatures. It describes the item count, dtype, and indexed reads and writes. Use ThreadData() to construct a payload for the active compiler backend. Other payload representations require support from that backend.

The item count must be a positive compile-time integer. Pass items_per_thread as a kernel argument; the compiler specializes the kernel for each supplied value. You can also use items.items_per_thread as a loop bound:

for i in range(items.items_per_thread):
    items[i] = types.int32(items[i] * 2)

Initialize every slot before reading it. Constructing ThreadData does not fill it with zeros. A full Load initializes the entire payload; a partial Load leaves out-of-bounds slots unspecified unless oob_default is provided. Initialize those slots after the Load or supply oob_default before reading them.

Blocked and striped order#

The glossary defines blocked and striped ownership and compares their layouts.

A primitive needs to know how the per-thread slots correspond to the group’s tile. Here is a small layout illustration with four threads and two items each. The numbers are positions in the tile; this illustration does not prescribe a supported CUDA launch size.

Thread rank	Blocked: `items[0], items[1]`	Striped: `items[0], items[1]`
0	0, 1	0, 4
1	2, 3	1, 5
2	4, 5	2, 6
3	6, 7	3, 7

For group size G, items per thread K, thread rank t, and local item index i, blocked order uses tile position t * K + i. Striped order uses t + i * G.

The payload does not carry a runtime layout tag that corrects mismatched operations. The algorithms you call determine the interpretation. Pairing a striped Load with a direct Store without a conversion permutes the data.

Compare the Load and Store visualizations to see how each algorithm maps memory positions to per-thread slots.

Dtypes and storage#

The current numeric payload types are signed and unsigned integers of 8, 16, 32, or 64 bits, and 32- or 64-bit floating point. Boolean, half precision, complex, and structured payloads are outside this contract.

Leave the constructor’s element type unspecified for normal use. Load infers it from the source, and Store can establish it from the destination. Supported typed assignments also provide the element type; use a typed scalar when your computation needs a particular width or precision. Inference follows the backend’s supported producers and assignments. Conflicting type requirements are errors. See element-type inference for cases that need additional information.

Load writes into the payload supplied by the caller. Transpose Store algorithms may rearrange their input payload in place, as in CUB. Copy values before Store if they are needed later. Both operations return None.

Numba can promote integer arithmetic. Store requires an exact match to the destination dtype, so cast computed values when necessary, as in the types.int32 assignment above. Accumulation dtype also matters: an integer sum can overflow, and parallel floating-point sums can differ from a sequential CPU sum because of evaluation order.

The compiler can keep payload elements in registers. Indexing patterns, address-taking, and register pressure influence the final placement. Increasing the item count increases the amount of live data per thread.

The optional alignment keyword requests a minimum power-of-two alignment in bytes when the compiler materializes the payload. For example, coop.ThreadData(items_per_thread, alignment=16) requests at least 16-byte alignment. The compiler may strengthen it. This setting applies to payload storage; alignment of the input and output arrays remains a separate property.

Load, operate, store#

Load and Store accept one-dimensional contiguous arrays. offset counts elements from the array’s beginning. valid_items counts elements in the valid prefix of the selected group’s tile. Both must be uniform within that group.

For a block, a tile holds block_threads * items_per_thread elements. For a physical or logical warp, it holds warp_width * items_per_thread elements. Clamp the valid count to that range before calling Load or Store. An out-of-range runtime count causes a device trap and invalidates the CUDA context.

Load leaves invalid slots unspecified unless you pass oob_default, even if those slots were initialized before Load. Store leaves destination elements outside the valid prefix untouched. Supply an operation-appropriate identity when processing padded data: zero for sum, one for multiplication, and a suitable upper or lower bound for minimum or maximum. For a runtime oob_default, use exactly the payload dtype and the same value across the group.

Algorithm choices#

Algorithm	Payload order	Memory access and rearrangement	Groups
`"direct"`	Blocked	Each thread accesses its consecutive items	Block and Warp
`"striped"`	Striped	Neighboring threads access neighboring elements at each item index	Block and Warp
`"vectorize"`	Blocked	Uses vector accesses where the specialization and alignment allow	Block and Warp
`"transpose"`	Blocked	Uses striped accesses and shared scratch to rearrange values	Block and Warp
`"warp_transpose"`	Blocked	Rearranges through warp-sized portions of the block tile	Block
`"warp_transpose_timesliced"`	Blocked	Reuses scratch across those portions	Block

The two block warp-transpose variants require a block size divisible by 32. Use string selectors. Begin with direct for a simple kernel; measure alternatives with the actual dtype, item count, and surrounding work. Coalesced accesses can justify rearrangement costs, but the best choice depends on the kernel.

Warp tile addresses#

Warp Load and Store automatically add the group’s origin within the block. For width G and K items per thread, this origin is (linear_thread_rank // G) * G * K. The explicit offset is added after that origin. In a multi-block traversal, pass the block’s global tile origin as offset.

The valid count still belongs to each individual group. Compute it using both the block origin and the group’s origin:

@cuda.jit
def copy_warp_tiles(source, destination, count, items_per_thread):
    group = coop.this_warp().group_by(8)
    items = coop.ThreadData(items_per_thread)
    block_origin = cuda.blockIdx.x * cuda.blockDim.x * items_per_thread
    group_origin = (cuda.threadIdx.x // 8) * 8 * items_per_thread
    valid = min(
        max(count - block_origin - group_origin, 0), 8 * items_per_thread
    )

    coop.load(
        group,
        source,
        items,
        offset=block_origin,
        valid_items=valid,
        oob_default=0,
    )
    coop.store(
        group, destination, items, offset=block_origin, valid_items=valid
    )

for items_per_thread in (1, 4):
    source = np.arange(531, dtype=np.int32)
    destination = np.full_like(source, -1)
    blocks = (source.size + 128 * items_per_thread - 1) // (
        128 * items_per_thread
    )
    copy_warp_tiles[blocks, 128](
        source, destination, source.size, items_per_thread
    )
    cuda.synchronize()
    np.testing.assert_array_equal(destination, source)

Adding group_origin to offset here would count it twice. This automatic origin applies to Warp Load and Store; ordinary array indexing in a kernel uses exactly the index you write.

Device helpers#

Use @cuda.jit(device=True) to put cooperative operations in a reusable helper. Pass a group and a ThreadData payload as arguments just as you would pass them to an operation in the kernel. The group still names the caller’s participating threads; each thread passes its own payload. Changes to that payload’s items are visible to the caller.

Cooperative helpers must be inlined into the calling kernel before planning. Numba-CUDA-MLIR defaults device helpers to inline="always"; the example spells it out on load_into. inline="never" prevents this frontend inlining and is unsupported for a helper containing cooperative operations. Inlining during later optimization or device linking comes too late to supply the group and payload information this planner needs.

This example uses nested helpers. load_tile creates a payload, passes it and the group to load_into, and returns the initialized payload to the kernel. Load infers the element type from source:

import numpy as np
from numba_cuda_mlir import cuda

from cuda import coop

@cuda.jit(device=True, inline="always")
def load_into(group, source, items):
    coop.load(group, source, items)

@cuda.jit(device=True)
def load_tile(group, source, items_per_thread):
    items = coop.ThreadData(items_per_thread)
    load_into(group, source, items)
    return items

@cuda.jit
def copy_tile(source, destination, items_per_thread):
    block = coop.this_block()
    items = load_tile(block, source, items_per_thread)
    coop.store(block, destination, items)

for items_per_thread in (1, 4):
    expected = np.arange(128 * items_per_thread, dtype=np.int32)
    source = cuda.to_device(expected)
    destination = cuda.device_array_like(source)
    copy_tile[1, 128](source, destination, items_per_thread)
    np.testing.assert_array_equal(destination.copy_to_host(), expected)

After inlining, the planner can see the helpers’ operations, group arguments, and payload uses together. The calling kernel supplies the compiler context and launch dimensions. An inlined helper can also create or return a group. coop.this_block() inside the helper refers to the same block as the caller. Group descriptors are resolved during compilation and have no separate runtime representation for an out-of-line device call.

Every required group member must reach each cooperative call inside the helper, including through nested calls or loops. Moving an operation into a helper does not change its participation or synchronization requirements. Cooperative operations inside standalone operation callbacks remain unsupported; those callbacks may perform ordinary device computation.

An ordinary inline="never" helper can index and modify a payload whose element type is already established. It receives a local array, so items.items_per_thread is unavailable there. Continue using the original payload in the caller: returning the array from an out-of-line helper loses the payload metadata needed by subsequent common cooperative calls. Inline helpers that need that metadata or return a payload for those calls.

Ordinary scalar helpers can also remain out of line. Their return values may need a typed assignment or another supported operation to establish a ThreadData element type, even when the helper has an explicit signature. This is separate from cooperative helper inlining; see element-type inference for workarounds.

Keep cooperative group mappings, operation selectors, and payload or storage shapes known at compile time. The current planner does not support literal_unroll values that determine these properties. Write those calls with compile-time constants. Ordinary runtime loops with fixed cooperative shapes, and unrelated uses of literal_unroll, remain supported.

`TempStorage`: scratch used during a primitive#

Some algorithms exchange intermediate values through shared memory. By default, cuda.coop allocates the scratch they need and inserts the required reuse barrier. Start with that behavior.

TempStorageLike names the common interface for explicit scratch descriptors in API signatures. Construct one inside the kernel with TempStorage(). Its size_in_bytes, alignment, auto_sync, and sharing properties control capacity, alignment, synchronization, and allocation sharing. The active compiler backend must recognize the descriptor.

An explicit descriptor lets several supported block calls reuse an allocation. Its contents are opaque; keep application values in ThreadData or your own arrays.

Calls	Scratch behavior in the current backend
Direct, striped, or vectorize Load/Store	No shared scratch or reuse barrier
Block transpose-family Load/Store	Automatic scratch, or an explicit `TempStorage`
Warp transpose Load/Store	Automatic scratch per group; explicit descriptors are rejected

Capacity, alignment, and lifetime#

TempStorage() defaults to sharing="shared" and auto_sync=False, leaving reuse synchronization to the caller. size_in_bytes=None and alignment=None let the compiler determine the requirements. An explicit capacity must be large enough for the operations using it. An explicit alignment requests a minimum positive power of two in bytes; the compiler can strengthen it. Only size_in_bytes may be positional. The other options are keyword-only.

sharing="exclusive" gives distinct call sites separate slices. With auto_sync=False, those calls need no barrier solely to reuse each other’s scratch. This uses more shared memory than sharing a slice. A loop reuses each call site’s slice on its next iteration, so that reuse still needs a barrier. sharing and auto_sync are independent controls; application data dependencies may require additional synchronization. See the exclusive-storage FAQ.

Omitted storage also participates in the backend’s shared-memory plan and launch accounting. The compiler may reuse its scratch across compatible calls and inserts the required reuse barriers.

With the default auto_sync=False, the kernel must provide reuse barriers. Put an explicit block barrier before a later call reuses the same scratch, including between loop iterations.

Set auto_sync=True to insert automatic trailing barriers for scratch reuse. Without an explicit descriptor, the compiler synchronizes scratch automatically. An unrelated memory access between calls does not establish a block barrier. The compiler rejects merging multiple manually synchronized TempStorage constructors into one descriptor, including conditional definitions. Use one constructor and keep its synchronization explicit. This restriction reflects what the planner can establish; it does not mean that every rejected program necessarily races.

Scratch lasts for the kernel’s execution on that block. It cannot carry state between blocks or kernel launches. Keep persistent application state in a separate payload.

When the combined scratch requirement exceeds the default static shared-memory limit, the backend can use dynamic shared memory, subject to the GPU’s opt-in limit. Numba-CUDA-MLIR automatically includes those required bytes in the launch configuration. You do not need to copy a compiler-reported byte count into the launch yourself. Requirements above the device limit are rejected. The required dynamic byte count is a minimum: launch-time shared bytes are not added to it as a separate allocation. Dynamic cooperative backing supports alignment requirements up to 16 bytes.

With the currently supported compiler, do not combine user dynamic or runtime-sized cuda.shared.array allocations with cooperative scratch. User static shared arrays may coexist with static cooperative backing, but are rejected when the cooperative backing becomes dynamic, including an implicit oversized allocation. These checks apply after helper inlining. Keep both allocations static within the device limit, use global memory for the user buffer, or separate the work into kernels. Storage-free operations do not create this conflict. The compatibility restrictions remain until a released compiler with the shared-memory fix has passed the coexistence tests.

Extra shared memory can reduce resident blocks per multiprocessor. The default inferred allocation and unsized shared descriptor are sufficient for the kernels above; use an explicit capacity when you have a reason to reserve that amount of shared memory.

Checking and tuning a kernel#

Check results before comparing algorithms. Useful cases include one full tile, several tiles, a single valid element in the final tile, and an empty input handled on the host. Layout conversions are easier to inspect with distinct input values. Compare computed values against a CPU reference with a suitable accumulation dtype and floating-point tolerance.

Keep launch dimensions, logical-warp widths, payload extents, and algorithm choices consistent with the code. The compiler specializes group operations using these facts. A runtime tile offset or valid count can change from one call to the next; the size of a ThreadData payload must be known at compile time.

Warm up the kernel before timing it, use device-resident arrays, and account for asynchronous execution with CUDA events or explicit synchronization. Then vary one choice at a time: threads per block, items per thread, or a Load/Store algorithm. More items per thread can amortize primitive work while increasing register pressure. Scratch-heavy choices consume shared memory and can reduce occupancy. Measure the complete kernel, including conversions and synchronization.

If compilation fails, check the group/operation combination, payload dtype, static parameters, and import order first. The API reference records exact signatures; the overview collects operation restrictions and configuration. For generated-source diagnostics and the compiler integration, see the Developer Overview.

cuda.coop Programming Guide#