Add Numba-CUDA-MLIR support for cuda.coop Load and Store

[tpn]

Why this is needed

The common API in #11182 needs a compiler integration to execute inside GPU
kernels. This adds Numba-CUDA-MLIR lowering for the complete Block, physical
Warp, and logical Warp Load/Store surface.

Install the extra for the matching CUDA major, then register before compiling:

python -m pip install "cuda-coop[numba-cuda-mlir-cu13]"

from cuda import coop

coop.register("numba-cuda-mlir")
from numba_cuda_mlir import cuda

Registration works with either import order, is safe to repeat, and also
accepts "numba_cuda_mlir". Numba-first automatic registration remains
available. import cuda.coop.numba_mlir as numba_coop registers the backend
and exposes its memory namespaces. The base wheel includes this integration.
The CUDA 12/13 extras install the same wheel and add the required backend
dependencies; the base installation declares no Python package dependencies.

Load/Store behavior

Algorithms	Temporary storage	Default reuse synchronization
direct, striped, vectorize	None	None
Block transpose, warp_transpose, warp_transpose_timesliced	Per-block	Trailing block barrier
Warp transpose	Per physical/logical group	Trailing barrier with that group's mask

Load mutates its supplied ThreadData and returns None. Transpose Store
preserves its input payload. Partial Load preserves invalid slots unless
oob_default is supplied. Each Warp group addresses its own consecutive tile;
offset is applied afterward and valid_items is relative to that tile.

Block operations support caller-supplied TempStorage; Warp operations use
compiler-owned storage. Shared and exclusive allocations both synchronize
by default. auto_sync=False transfers reuse synchronization to the caller.
Runtime controls must be uniform within the selected group, and callers must
keep valid_items within the tile bounds. Invalid runtime bounds trap before
narrowing to the CUB ABI.

The adapter validates dtypes, alignment, supported group forms, descriptor
provenance, and shared-memory compatibility before provider generation.
Collective helpers must inline into their kernels. The guarded compiler
integration supports Numba-CUDA-MLIR 0.5.x. With this compiler, user dynamic
shared arrays cannot coexist with cooperative backing; user static arrays
require static cooperative backing. Storage-free operations introduce no
such conflict.

This PR includes compile/runtime tests, Load/Store examples, and GPU-hidden
compiler plus GPU runtime CI stages. Stacked on #11182; subsequent algorithm
families start with Exchange and Shuffle in #11208.

Validation

Fresh installed-wheel checks at ecec1c57ed97:

• Host contracts and packaging (full suites): 434 passed on Python 3.14.
• Numba backend unit/compiler tests (focused, GPUs hidden): 67 passed on Python 3.14.
• GPU runtime tests (focused): 19 passed on Python 3.14.

Wheel validation and source, installed-package, and bundled-header comparisons passed.

Remote CI passed at ecec1c57ed97: GitHub Actions run. All exact-head checks and statuses, including pre-commit, are successful or conditionally skipped.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[Copilot]

🟡 Changes recommended

A critical issue remains with forwarding temp_storage to storage-free factories.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Completes execution support for all six CUB Block Load/Store algorithms across portable and qualified APIs, including storage planning, payload preservation, synchronization, diagnostics, and expanded tests.

Changes:

• Adds storage-free and storage-bearing provider routing.
• Preserves transpose-store payloads and partial-load invalid slots.
• Expands compiler, runtime, contract, documentation, and diagnostic coverage.

File summaries

File	Summary
python/cuda_coop/tests/contracts/core/test_core_group_load_store.py	Tests algorithm identities and storage contracts.
python/cuda_coop/tests/contracts/core/test_core_block_load_store.py	Tests algorithm specializations and preservation behavior.
python/cuda_coop/tests/backends/numba_mlir/unit/test_data_movement.py	Tests planning and rewrite behavior.
python/cuda_coop/tests/backends/numba_mlir/runtime/test_block_load_store.py	Adds runtime coverage for all algorithms.
python/cuda_coop/tests/backends/numba_mlir/compile/test_storage_diagnostics.py	Tests shared-memory diagnostics.
python/cuda_coop/tests/backends/numba_mlir/compile/test_block_load_store_compile.py	Tests provider compilation and caching.
python/cuda_coop/README.md	Documents algorithms and storage behavior.
python/cuda_coop/cuda/coop/numba_mlir/_lowering/_load_store.py	Registers algorithm-specific providers.
python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planner.py	Supports keyword local-array shapes.
python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_load_store.py	Plans storage and provider lowering.
python/cuda_coop/cuda/coop/_core/group/load_store.py	Expands algorithm and storage contracts.
python/cuda_coop/cuda/coop/_core/block/load_store.py	Adds CUB specializations and preservation wrappers.
docs/python/coop.rst	Documents execution and storage behavior.
docs/python/coop_api.rst	Updates qualified API documentation.

Review details

• Files reviewed: 14/14 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[github-actions]

🥳 CI Workflow Results

🟩 Finished in 3h 20m: Pass: 100%/558 | Total: 7d 04h | Max: 2h 16m | Hits: 99%/855781

See results here.