Infer launch bounds for exact cuda.coop block specializations

[tpn]

Why this is needed

Cooperative planning knows a kernel's exact block size, but the device compiler
previously received no corresponding launch bound. A large block could fail to
launch when its compiled register requirements exceeded the available resources.

Behavior

Infer scalar launch_bounds=x*y*z in compile-local options when launch_bounds
is absent or None and max_registers is not explicitly set. Preserve explicit
integer/tuple bounds, register limits, dispatcher options, exact dimensions,
and specialization identity. Reject an exact launch exceeding an explicit
maximum with an attributable error.

This supplies launch-resource information; it does not infer occupancy or
cluster bounds. Lower register use may introduce spilling.

Stacked on #11320, following the full primitive and documentation stack. The
launch-bound implementation patch is unchanged by the restack.

Validation

Fresh installed-wheel checks at 216928d0eac1:

• Host contracts and packaging (full suites): 891 passed on Python 3.10; 891 passed on Python 3.14.
• Numba backend unit/compiler tests (full suites, GPUs hidden): 1,415 passed on Python 3.14.
• GPU runtime tests (full suites): 795 passed, 3 skipped on Python 3.14.
• Standalone GPU examples: 7 executions passed.
• Mamba with nested device helpers: all 44 cases passed against an independent float64 reference (rtol=3e-4, atol=3e-5).

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

Remote CI passed at 216928d0eac1: 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]

🔵 Needs a closer look

The review requests stronger verification of explicit max_registers preservation and a documentation wording correction.

Pull request overview

Adds inferred CUDA launch bounds for exact cooperative block specializations while preserving explicit resource options and specialization identity.

Changes:

• Infers scalar bounds from exact 1D/2D/3D block dimensions.
• Preserves explicit launch and register limits and reports oversized launches.
• Adds compile/runtime coverage and programming-guide documentation.

File summaries

File	Description
python/cuda_coop/tests/backends/numba_mlir/runtime/test_scan.py	Validates large 1D/2D/3D scans.
python/cuda_coop/tests/backends/numba_mlir/compile/test_group_hierarchy_compile.py	Tests bounds, options, specializations, and diagnostics.
python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planner.py	Infers and validates launch bounds.
docs/python/coop/programming_guide.rst	Documents launch-resource behavior.

Review details

Suppressed comments (2)

docs/python/coop/programming_guide.rst:1069

• important: The new test passes launch_bounds=None, and the implementation treats that as unset and infers a bound, but this wording says that supplying launch_bounds explicitly suppresses inference. Please distinguish non-None bounds (and non-None max_registers) so the documented behavior matches the tested behavior.

A kernel specialized for an exact cooperative launch block receives an inferred
``launch_bounds`` equal to that block's thread count, unless the kernel supplies
``launch_bounds`` or ``max_registers`` explicitly. This lets the compiler account

python/cuda_coop/tests/backends/numba_mlir/compile/test_group_hierarchy_compile.py:233

• important: The max_registers case currently only asserts that no .maxntid directive was added. That would still pass if the explicit register cap were dropped, so this test does not cover the advertised preservation; also assert that the compiled target options (or PTX) retain max_registers=64.

    if maximum is None:
        assert ".maxntid" not in ptx
    else:
        assert f".maxntid {maximum}, 1, 1" in ptx

• Files reviewed: 4/4 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[github-actions]

🥳 CI Workflow Results

🟩 Finished in 34m 44s: Pass: 100%/8 | Total: 1h 27m | Max: 28m 29s

See results here.