From b7343f53de9d3d610d8884d630a93c86fa1314da Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 27 Oct 2026 12:36:17 -0700 Subject: [PATCH 1/2] Remove the existing cuda.coop implementation Prepare for a CUTLASS-focused first release by removing the shared foundation and Numba-CUDA-MLIR backend. Remove their packaging, tests, documentation, CI jobs, and editor configuration together. Signed-off-by: Trent Nelson --- .pre-commit-config.yaml | 14 - .vscode/launch.json | 34 - AGENTS.md | 4 - ci/build_cuda_coop_python.sh | 53 - ci/matrix.yaml | 25 +- ci/project_files_and_dependencies.yaml | 30 +- ci/test/inspect_changes/core_dirty.output | 2 +- .../cuda_coop_python.dirty_files | 4 - .../inspect_changes/cuda_coop_python.output | 2 - .../cuda_coop_python_cudax.dirty_files | 1 - .../cuda_coop_python_cudax.output | 2 - .../cuda_coop_python_license.dirty_files | 4 - .../cuda_coop_python_license.output | 2 - .../cuda_coop_python_pyproject.dirty_files | 1 - .../cuda_coop_python_pyproject.output | 2 - ...oop_python_pyproject_near_miss.dirty_files | 2 - ...uda_coop_python_pyproject_near_miss.output | 2 - .../cuda_coop_python_windows_ci.dirty_files | 1 - .../cuda_coop_python_windows_ci.output | 2 - .../inspect_changes/ignored_only.dirty_files | 1 - .../inspect_changes/libcudacxx_both.output | 2 +- .../libcudacxx_public_only.output | 2 +- .../inspect_changes/libcudacxx_thrust.output | 2 +- ci/test_cuda_coop_python.sh | 226 -- ci/validate_cuda_coop_wheel.py | 273 -- ci/windows/build_cuda_coop_python.ps1 | 123 - docs/_ext/coop_visualization.py | 103 - docs/_static/coop-api.css | 12 - docs/_static/coop-explorer.js | 332 -- docs/_static/coop-load.css | 241 -- docs/_static/coop-load.js | 397 --- docs/_static/coop-movement.js | 113 - docs/conf.py | 38 +- docs/python/api_reference.rst | 1 - docs/python/coop.rst | 127 - docs/python/coop/concepts.rst | 401 --- docs/python/coop/configuration.rst | 143 - docs/python/coop/debugger_walkthrough.py | 43 - docs/python/coop/developer_overview.rst | 1686 ---------- docs/python/coop/faqs.rst | 189 -- docs/python/coop/glossary.rst | 156 - docs/python/coop/programming_guide.rst | 609 ---- docs/python/coop/source_dumps/direct.cpp.txt | 33 - .../coop/source_dumps/transpose.cpp.txt | 48 - .../compiler-phases-classes-mobile.svg | 261 -- .../compiler-phases-classes.svg | 257 -- .../coop/visualizations/compiler-phases.svg | 64 - .../compiler-swimlane-classes-mobile.svg | 354 -- .../compiler-swimlane-classes.svg | 354 -- .../coop/visualizations/compiler-swimlane.svg | 89 - docs/python/coop/visualizations/index.rst | 18 - .../coop/visualizations/load-transpose.svg | 80 - docs/python/coop/visualizations/load.rst | 84 - .../coop/visualizations/store-transpose.svg | 86 - docs/python/coop/visualizations/store.rst | 90 - docs/python/coop_api.rst | 116 - docs/python/index.rst | 4 - python/cuda_coop/AGENTS.md | 81 - python/cuda_coop/CMakeLists.txt | 178 - python/cuda_coop/LICENSE | 1 - python/cuda_coop/README.md | 364 -- python/cuda_coop/cuda/coop/__init__.py | 50 - python/cuda_coop/cuda/coop/__init__.pyi | 44 - python/cuda_coop/cuda/coop/_core/__init__.py | 191 -- python/cuda_coop/cuda/coop/_core/_adapter.py | 128 - .../cuda_coop/cuda/coop/_core/_algorithm.py | 269 -- .../cuda/coop/_core/_auto_registration.py | 258 -- python/cuda_coop/cuda/coop/_core/_bindings.py | 434 --- .../cuda_coop/cuda/coop/_core/_source_dump.py | 93 - python/cuda_coop/cuda/coop/_core/_symbols.py | 292 -- python/cuda_coop/cuda/coop/_core/_types.py | 404 --- .../cuda_coop/cuda/coop/_core/api/__init__.py | 55 - .../cuda/coop/_core/api/__init__.pyi | 34 - .../cuda/coop/_core/api/_dispatch.py | 166 - .../cuda_coop/cuda/coop/_core/api/_payload.py | 90 - .../cuda/coop/_core/api/load_store.py | 224 -- .../cuda/coop/_core/api/load_store.pyi | 110 - .../cuda/coop/_core/api/temp_storage.py | 84 - .../cuda/coop/_core/api/temp_storage.pyi | 16 - .../cuda/coop/_core/api/thread_data.py | 77 - .../cuda/coop/_core/api/thread_data.pyi | 31 - .../cuda/coop/_core/api/thread_group.py | 135 - .../cuda/coop/_core/api/thread_group.pyi | 110 - .../cuda/coop/_core/block/__init__.py | 34 - .../cuda/coop/_core/block/_common.py | 57 - .../cuda/coop/_core/block/load_store.py | 607 ---- .../cuda_coop/cuda/coop/_core/dtype_policy.py | 98 - .../cuda/coop/_core/group/__init__.py | 61 - .../cuda/coop/_core/group/_dispatch.py | 279 -- .../_core/group/_execution_requirements.py | 340 -- .../cuda_coop/cuda/coop/_core/group/_model.py | 947 ------ .../cuda/coop/_core/group/_resolution.py | 304 -- .../cuda/coop/_core/group/load_store.py | 612 ---- python/cuda_coop/cuda/coop/_core/launch.py | 369 -- .../cuda_coop/cuda/coop/_core/thread_group.py | 1054 ------ .../cuda/coop/_core/warp/__init__.py | 27 - .../cuda/coop/_core/warp/load_store.py | 643 ---- .../cuda_coop/cuda/coop/_headers/__init__.py | 334 -- .../cuda_coop/cuda/coop/_headers/_identity.py | 393 --- .../cuda_coop/cuda/coop/_headers/_toolkit.py | 632 ---- python/cuda_coop/cuda/coop/_registration.py | 56 - python/cuda_coop/cuda/coop/_typing.py | 104 - python/cuda_coop/cuda/coop/_typing.pyi | 167 - .../cuda/coop/numba_mlir/__init__.py | 74 - .../cuda/coop/numba_mlir/__init__.pyi | 37 - .../coop/numba_mlir/_compiler/__init__.py | 5 - .../coop/numba_mlir/_compiler/_activation.py | 221 -- .../coop/numba_mlir/_compiler/_artifacts.py | 66 - .../coop/numba_mlir/_compiler/_caching.py | 397 --- .../_compiler/_descriptor_provenance.py | 253 -- .../numba_mlir/_compiler/_group_errors.py | 321 -- .../numba_mlir/_compiler/_group_load_store.py | 891 ----- .../numba_mlir/_compiler/_group_planner.py | 2192 ------------ .../_compiler/_group_planner_support.py | 121 - .../numba_mlir/_compiler/_group_planning.py | 1143 ------- .../numba_mlir/_compiler/_group_rewriting.py | 90 - .../cuda/coop/numba_mlir/_compiler/_layout.py | 128 - .../_compiler/_load_store_algorithms.py | 59 - .../_compiler/_numba_mlir_compat.py | 225 -- .../cuda/coop/numba_mlir/_compiler/_nvrtc.py | 739 ---- .../coop/numba_mlir/_compiler/_operations.py | 784 ----- .../coop/numba_mlir/_compiler/_parameters.py | 471 --- .../coop/numba_mlir/_compiler/_planner.py | 99 - .../coop/numba_mlir/_compiler/_rewrite.py | 1006 ------ .../_compiler/_rewrite_arguments.py | 396 --- .../_compiler/_rewrite_group_metadata.py | 130 - .../_compiler/_rewrite_invocables.py | 304 -- .../numba_mlir/_compiler/_rewrite_launch.py | 256 -- .../_compiler/_rewrite_load_store.py | 819 ----- .../numba_mlir/_compiler/_rewrite_payload.py | 223 -- .../_compiler/_rewrite_provenance.py | 2165 ------------ .../numba_mlir/_compiler/_rewrite_storage.py | 1830 ---------- .../numba_mlir/_compiler/_rewrite_support.py | 414 --- .../_compiler/_scalar_provenance.py | 336 -- .../cuda/coop/numba_mlir/_group/__init__.py | 3 - .../coop/numba_mlir/_group/_load_store.py | 153 - .../coop/numba_mlir/_group/_load_store.pyi | 98 - .../cuda/coop/numba_mlir/_group/_marker.py | 26 - .../coop/numba_mlir/_lowering/__init__.py | 15 - .../cuda/coop/numba_mlir/_lowering/_core.py | 684 ---- .../coop/numba_mlir/_lowering/_load_store.py | 765 ----- .../cuda/coop/numba_mlir/_semantic.py | 62 - .../cuda/coop/numba_mlir/_temp_storage.py | 79 - .../cuda/coop/numba_mlir/_temp_storage.pyi | 25 - .../cuda/coop/numba_mlir/_thread_data.py | 120 - .../cuda/coop/numba_mlir/_thread_data.pyi | 68 - .../cuda/coop/numba_mlir/_thread_group.py | 100 - .../cuda/coop/numba_mlir/_thread_group.pyi | 85 - .../cuda_coop/cuda/coop/numba_mlir/_types.py | 2978 ----------------- .../cuda_coop/cuda/coop/numba_mlir/py.typed | 0 python/cuda_coop/cuda/coop/py.typed | 0 .../examples/numba_mlir/block_load.py | 53 - .../examples/numba_mlir/block_store.py | 56 - .../examples/numba_mlir/source_dumps.py | 86 - python/cuda_coop/pyproject.toml | 176 - python/cuda_coop/tests/__init__.py | 3 - python/cuda_coop/tests/_group_planning.py | 29 - .../tests/backends/numba_mlir/__init__.py | 3 - .../backends/numba_mlir/compile/__init__.py | 3 - .../compile/test_block_load_store_compile.py | 975 ------ .../test_installed_activation_compile.py | 216 -- .../compile/test_load_store_target_compile.py | 182 - .../compile/test_scalar_store_typing.py | 136 - .../compile/test_storage_diagnostics.py | 452 --- .../compile/test_warp_load_store_compile.py | 637 ---- .../backends/numba_mlir/runtime/__init__.py | 3 - .../runtime/test_block_load_store.py | 1830 ---------- .../runtime/test_device_helpers_example.py | 50 - .../numba_mlir/runtime/test_load_example.py | 60 - .../test_programming_guide_examples.py | 60 - .../test_qualified_load_store_examples.py | 54 - .../runtime/test_storage_examples.py | 87 - .../numba_mlir/runtime/test_store_example.py | 55 - .../runtime/test_warp_load_store.py | 2011 ----------- .../backends/numba_mlir/unit/__init__.py | 3 - .../backends/numba_mlir/unit/test_caching.py | 266 -- .../numba_mlir/unit/test_data_movement.py | 2919 ---------------- .../unit/test_family_registration.py | 827 ----- .../numba_mlir/unit/test_group_hierarchy.py | 291 -- .../unit/test_group_lowering_plan.py | 955 ------ .../unit/test_invocable_identity.py | 104 - .../unit/test_invocable_lifecycle.py | 106 - .../unit/test_launch_dimension_rewrites.py | 359 -- .../backends/numba_mlir/unit/test_nvrtc.py | 733 ---- .../numba_mlir/unit/test_parameter_names.py | 176 - .../numba_mlir/unit/test_public_foundation.py | 325 -- .../unit/test_runtime_activation.py | 484 --- .../numba_mlir/unit/test_storage_rewrites.py | 2423 -------------- .../numba_mlir/unit/test_thread_data.py | 121 - .../numba_mlir/unit/test_thread_groups.py | 96 - .../cuda_coop/tests/packaging/test_typing.py | 129 - .../tests/packaging/typing/common_consumer.py | 76 - .../packaging/typing/invalid_consumer.py | 134 - .../tests/packaging/typing/numba_consumer.py | 119 - python/cuda_coop/tests/test_bindings.py | 53 - .../cuda_coop/tests/test_block_load_store.py | 202 -- .../cuda_coop/tests/test_group_load_store.py | 295 -- .../cuda_coop/tests/test_group_resolution.py | 38 - python/cuda_coop/tests/test_header_bundle.py | 309 -- .../cuda_coop/tests/test_header_identity.py | 118 - python/cuda_coop/tests/test_headers.py | 180 - python/cuda_coop/tests/test_registration.py | 51 - python/cuda_coop/tests/test_source_dump.py | 51 - python/cuda_coop/tests/test_thread_group.py | 127 - python/cuda_coop/tests/test_toolkit.py | 522 --- python/cuda_coop/tests/test_types.py | 127 - .../cuda_coop/tests/test_warp_load_store.py | 93 - .../tests/typing/pyright-unused.json | 13 - 208 files changed, 9 insertions(+), 59871 deletions(-) delete mode 100755 ci/build_cuda_coop_python.sh delete mode 100644 ci/test/inspect_changes/cuda_coop_python.dirty_files delete mode 100644 ci/test/inspect_changes/cuda_coop_python.output delete mode 100644 ci/test/inspect_changes/cuda_coop_python_cudax.dirty_files delete mode 100644 ci/test/inspect_changes/cuda_coop_python_cudax.output delete mode 100644 ci/test/inspect_changes/cuda_coop_python_license.dirty_files delete mode 100644 ci/test/inspect_changes/cuda_coop_python_license.output delete mode 100644 ci/test/inspect_changes/cuda_coop_python_pyproject.dirty_files delete mode 100644 ci/test/inspect_changes/cuda_coop_python_pyproject.output delete mode 100644 ci/test/inspect_changes/cuda_coop_python_pyproject_near_miss.dirty_files delete mode 100644 ci/test/inspect_changes/cuda_coop_python_pyproject_near_miss.output delete mode 100644 ci/test/inspect_changes/cuda_coop_python_windows_ci.dirty_files delete mode 100644 ci/test/inspect_changes/cuda_coop_python_windows_ci.output delete mode 100755 ci/test_cuda_coop_python.sh delete mode 100644 ci/validate_cuda_coop_wheel.py delete mode 100644 ci/windows/build_cuda_coop_python.ps1 delete mode 100644 docs/_ext/coop_visualization.py delete mode 100644 docs/_static/coop-api.css delete mode 100644 docs/_static/coop-explorer.js delete mode 100644 docs/_static/coop-load.css delete mode 100644 docs/_static/coop-load.js delete mode 100644 docs/_static/coop-movement.js delete mode 100644 docs/python/coop.rst delete mode 100644 docs/python/coop/concepts.rst delete mode 100644 docs/python/coop/configuration.rst delete mode 100644 docs/python/coop/debugger_walkthrough.py delete mode 100644 docs/python/coop/developer_overview.rst delete mode 100644 docs/python/coop/faqs.rst delete mode 100644 docs/python/coop/glossary.rst delete mode 100644 docs/python/coop/programming_guide.rst delete mode 100644 docs/python/coop/source_dumps/direct.cpp.txt delete mode 100644 docs/python/coop/source_dumps/transpose.cpp.txt delete mode 100644 docs/python/coop/visualizations/compiler-phases-classes-mobile.svg delete mode 100644 docs/python/coop/visualizations/compiler-phases-classes.svg delete mode 100644 docs/python/coop/visualizations/compiler-phases.svg delete mode 100644 docs/python/coop/visualizations/compiler-swimlane-classes-mobile.svg delete mode 100644 docs/python/coop/visualizations/compiler-swimlane-classes.svg delete mode 100644 docs/python/coop/visualizations/compiler-swimlane.svg delete mode 100644 docs/python/coop/visualizations/index.rst delete mode 100644 docs/python/coop/visualizations/load-transpose.svg delete mode 100644 docs/python/coop/visualizations/load.rst delete mode 100644 docs/python/coop/visualizations/store-transpose.svg delete mode 100644 docs/python/coop/visualizations/store.rst delete mode 100644 docs/python/coop_api.rst delete mode 100644 python/cuda_coop/AGENTS.md delete mode 100644 python/cuda_coop/CMakeLists.txt delete mode 100644 python/cuda_coop/LICENSE delete mode 100644 python/cuda_coop/README.md delete mode 100644 python/cuda_coop/cuda/coop/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/__init__.pyi delete mode 100644 python/cuda_coop/cuda/coop/_core/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/_core/_adapter.py delete mode 100644 python/cuda_coop/cuda/coop/_core/_algorithm.py delete mode 100644 python/cuda_coop/cuda/coop/_core/_auto_registration.py delete mode 100644 python/cuda_coop/cuda/coop/_core/_bindings.py delete mode 100644 python/cuda_coop/cuda/coop/_core/_source_dump.py delete mode 100644 python/cuda_coop/cuda/coop/_core/_symbols.py delete mode 100644 python/cuda_coop/cuda/coop/_core/_types.py delete mode 100644 python/cuda_coop/cuda/coop/_core/api/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/_core/api/__init__.pyi delete mode 100644 python/cuda_coop/cuda/coop/_core/api/_dispatch.py delete mode 100644 python/cuda_coop/cuda/coop/_core/api/_payload.py delete mode 100644 python/cuda_coop/cuda/coop/_core/api/load_store.py delete mode 100644 python/cuda_coop/cuda/coop/_core/api/load_store.pyi delete mode 100644 python/cuda_coop/cuda/coop/_core/api/temp_storage.py delete mode 100644 python/cuda_coop/cuda/coop/_core/api/temp_storage.pyi delete mode 100644 python/cuda_coop/cuda/coop/_core/api/thread_data.py delete mode 100644 python/cuda_coop/cuda/coop/_core/api/thread_data.pyi delete mode 100644 python/cuda_coop/cuda/coop/_core/api/thread_group.py delete mode 100644 python/cuda_coop/cuda/coop/_core/api/thread_group.pyi delete mode 100644 python/cuda_coop/cuda/coop/_core/block/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/_core/block/_common.py delete mode 100644 python/cuda_coop/cuda/coop/_core/block/load_store.py delete mode 100644 python/cuda_coop/cuda/coop/_core/dtype_policy.py delete mode 100644 python/cuda_coop/cuda/coop/_core/group/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/_core/group/_dispatch.py delete mode 100644 python/cuda_coop/cuda/coop/_core/group/_execution_requirements.py delete mode 100644 python/cuda_coop/cuda/coop/_core/group/_model.py delete mode 100644 python/cuda_coop/cuda/coop/_core/group/_resolution.py delete mode 100644 python/cuda_coop/cuda/coop/_core/group/load_store.py delete mode 100644 python/cuda_coop/cuda/coop/_core/launch.py delete mode 100644 python/cuda_coop/cuda/coop/_core/thread_group.py delete mode 100644 python/cuda_coop/cuda/coop/_core/warp/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/_core/warp/load_store.py delete mode 100644 python/cuda_coop/cuda/coop/_headers/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/_headers/_identity.py delete mode 100644 python/cuda_coop/cuda/coop/_headers/_toolkit.py delete mode 100644 python/cuda_coop/cuda/coop/_registration.py delete mode 100644 python/cuda_coop/cuda/coop/_typing.py delete mode 100644 python/cuda_coop/cuda/coop/_typing.pyi delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/__init__.pyi delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_activation.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_artifacts.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_caching.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_descriptor_provenance.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_errors.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_load_store.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planner.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planner_support.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planning.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_rewriting.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_layout.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_load_store_algorithms.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_numba_mlir_compat.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_nvrtc.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_operations.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_parameters.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_planner.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_arguments.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_group_metadata.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_invocables.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_launch.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_load_store.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_payload.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_provenance.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_storage.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_support.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_compiler/_scalar_provenance.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_group/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_group/_load_store.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_group/_load_store.pyi delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_group/_marker.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_lowering/__init__.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_lowering/_core.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_lowering/_load_store.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_semantic.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_temp_storage.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_temp_storage.pyi delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_thread_data.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_thread_data.pyi delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_thread_group.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_thread_group.pyi delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/_types.py delete mode 100644 python/cuda_coop/cuda/coop/numba_mlir/py.typed delete mode 100644 python/cuda_coop/cuda/coop/py.typed delete mode 100644 python/cuda_coop/examples/numba_mlir/block_load.py delete mode 100644 python/cuda_coop/examples/numba_mlir/block_store.py delete mode 100644 python/cuda_coop/examples/numba_mlir/source_dumps.py delete mode 100644 python/cuda_coop/pyproject.toml delete mode 100644 python/cuda_coop/tests/__init__.py delete mode 100644 python/cuda_coop/tests/_group_planning.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/__init__.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/compile/__init__.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/compile/test_block_load_store_compile.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/compile/test_installed_activation_compile.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/compile/test_load_store_target_compile.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/compile/test_scalar_store_typing.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/compile/test_storage_diagnostics.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/compile/test_warp_load_store_compile.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/runtime/__init__.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/runtime/test_block_load_store.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/runtime/test_device_helpers_example.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/runtime/test_load_example.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/runtime/test_programming_guide_examples.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/runtime/test_qualified_load_store_examples.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/runtime/test_storage_examples.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/runtime/test_store_example.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/runtime/test_warp_load_store.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/__init__.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_caching.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_data_movement.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_family_registration.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_group_hierarchy.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_group_lowering_plan.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_invocable_identity.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_invocable_lifecycle.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_launch_dimension_rewrites.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_nvrtc.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_parameter_names.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_public_foundation.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_runtime_activation.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_storage_rewrites.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_thread_data.py delete mode 100644 python/cuda_coop/tests/backends/numba_mlir/unit/test_thread_groups.py delete mode 100644 python/cuda_coop/tests/packaging/test_typing.py delete mode 100644 python/cuda_coop/tests/packaging/typing/common_consumer.py delete mode 100644 python/cuda_coop/tests/packaging/typing/invalid_consumer.py delete mode 100644 python/cuda_coop/tests/packaging/typing/numba_consumer.py delete mode 100644 python/cuda_coop/tests/test_bindings.py delete mode 100644 python/cuda_coop/tests/test_block_load_store.py delete mode 100644 python/cuda_coop/tests/test_group_load_store.py delete mode 100644 python/cuda_coop/tests/test_group_resolution.py delete mode 100644 python/cuda_coop/tests/test_header_bundle.py delete mode 100644 python/cuda_coop/tests/test_header_identity.py delete mode 100644 python/cuda_coop/tests/test_headers.py delete mode 100644 python/cuda_coop/tests/test_registration.py delete mode 100644 python/cuda_coop/tests/test_source_dump.py delete mode 100644 python/cuda_coop/tests/test_thread_group.py delete mode 100644 python/cuda_coop/tests/test_toolkit.py delete mode 100644 python/cuda_coop/tests/test_types.py delete mode 100644 python/cuda_coop/tests/test_warp_load_store.py delete mode 100644 python/cuda_coop/tests/typing/pyright-unused.json diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index be6d0d6ef37f..55d5c4b5bdae 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -79,21 +79,7 @@ repos: rev: v0.15.20 hooks: - id: ruff # linter - exclude: ^python/cuda_coop/ - id: ruff-format # formatter - exclude: ^python/cuda_coop/ - - # Use the newer Ruff rules for `cuda.coop`. - - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.16.9 - hooks: - - id: ruff-check - name: ruff check (cuda.coop) - files: ^python/cuda_coop/ - - id: ruff-format - name: ruff format (cuda.coop) - files: ^python/cuda_coop/ - types_or: [python, pyi, jupyter] # CMake formatting - repo: https://github.com/BlankSpruce/gersemi diff --git a/.vscode/launch.json b/.vscode/launch.json index 09b6daed1d14..af87418fd753 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -1,39 +1,5 @@ { "configurations": [ - { - "console": "integratedTerminal", - "cwd": "${workspaceFolder}", - "env": { - "CUDA_COOP_CCCL_ROOT": "${workspaceFolder}", - "CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION": "0", - "CUDA_COOP_ENABLE_CACHE": "0", - "CUDA_COOP_SOURCE_DUMP_DIR": "${workspaceFolder}/build/coop-debug-sources", - "CUDA_VISIBLE_DEVICES": "0", - "PYTHONPATH": "${workspaceFolder}/python/cuda_coop" - }, - "justMyCode": false, - "name": "cuda.coop: Debugger walkthrough", - "program": "${workspaceFolder}/docs/python/coop/debugger_walkthrough.py", - "request": "launch", - "type": "debugpy" - }, - { - "console": "integratedTerminal", - "cwd": "${workspaceFolder}", - "env": { - "CUDA_COOP_CCCL_ROOT": "${workspaceFolder}", - "CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION": "0", - "CUDA_COOP_ENABLE_CACHE": "0", - "CUDA_COOP_SOURCE_DUMP_DIR": "${workspaceFolder}/build/coop-debug-sources", - "CUDA_VISIBLE_DEVICES": "0", - "PYTHONPATH": "${workspaceFolder}/python/cuda_coop" - }, - "justMyCode": false, - "name": "cuda.coop: Debug active Python file", - "program": "${file}", - "request": "launch", - "type": "debugpy" - }, { "cwd": "${workspaceFolder}", "initCommands": [ diff --git a/AGENTS.md b/AGENTS.md index 2f991e0403ce..3fb68488059f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,10 +2,6 @@ This document provides guidelines for building, testing, and contributing to the CCCL repository. It is primarily written for agentic AIs, but the information is also useful for CCCL developers. -For `cuda.coop` work, including `python/cuda_coop/**` and -`docs/python/coop*`, read [python/cuda_coop/AGENTS.md](python/cuda_coop/AGENTS.md) -and apply its API, behavior, and example conventions. - --- ## Overview diff --git a/ci/build_cuda_coop_python.sh b/ci/build_cuda_coop_python.sh deleted file mode 100755 index ae628c15015c..000000000000 --- a/ci/build_cuda_coop_python.sh +++ /dev/null @@ -1,53 +0,0 @@ -#!/usr/bin/env bash -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -# Build the universal cuda-coop wheel and check its archive contents. -# GitHub Actions uploads the checked wheel for the installed-package tests. - -set -euo pipefail - -ci_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -repo_root="$(cd "$ci_dir/.." && pwd)" - -usage="Usage: $0 -py-version [additional options...]" - -# shellcheck source=ci/util/python/common_arg_parser.sh -source "$ci_dir/util/python/common_arg_parser.sh" -parse_python_args "$@" -require_py_version "$usage" || exit 1 - -# shellcheck source=ci/pyenv_helper.sh -source "$ci_dir/pyenv_helper.sh" -if [[ "${CCCL_CUDA_COOP_PYENV_READY:-0}" != "1" ]]; then - setup_python_env "${py_version}" ".cccl-coop-build-venv" -fi - -python -m pip install build - -cd "$repo_root" -mkdir -p wheelhouse -rm -f wheelhouse/cuda_coop-*.whl -python -m build --wheel --outdir wheelhouse python/cuda_coop - -mapfile -t wheels < <(find wheelhouse -maxdepth 1 -name 'cuda_coop-*.whl' -print | sort) -if [[ ${#wheels[@]} -ne 1 ]]; then - echo "Error: expected exactly one cuda-coop wheel, found ${#wheels[@]}:" >&2 - printf ' %s\n' "${wheels[@]}" >&2 - exit 1 -fi - -if [[ "${wheels[0]}" != *-py3-none-any.whl ]]; then - echo "Error: cuda-coop must produce a universal py3-none-any wheel: ${wheels[0]}" >&2 - exit 1 -fi - -python "$ci_dir/validate_cuda_coop_wheel.py" "${wheels[0]}" - -echo "Built cuda-coop wheel: ${wheels[0]}" - -if [[ -n "${GITHUB_ACTIONS:-}" ]]; then - wheel_artifact_name="$(CCCL_WHEEL_KIND=coop "$ci_dir/util/workflow/get_wheel_artifact_name.sh")" - "$ci_dir/util/artifacts/upload.sh" "$wheel_artifact_name" 'wheelhouse/cuda_coop-.*\.whl' -fi diff --git a/ci/matrix.yaml b/ci/matrix.yaml index 00df78927f55..e5486893b8d7 100644 --- a/ci/matrix.yaml +++ b/ci/matrix.yaml @@ -214,12 +214,6 @@ workflows: - {jobs: ['test_py_compute_minimal'], project: 'python_tsan', ctk: '13.X', py_version: '3.14t', gpu: 'l4', cxx: 'gcc13'} # cuda.cccl.headers: CPU-only (gpu:false), arch-insensitive - {jobs: ['test_headers'], project: 'python', ctk: ['12.X', '13.X'], py_version: '3.14', cxx: ['gcc13', 'msvc2022'], py_ctk_mode: ['pinned', 'latest', 'sysctk']} - # cuda.coop: installed-wheel host contracts, GPU-free compilation, one - # representative runtime lane, and Windows universal-wheel endpoints. - - {jobs: ['test_py_coop_host'], project: 'cuda_coop_python', ctk: '13.X', py_version: ['3.10', '3.14'], cxx: 'gcc13'} - - {jobs: ['test_py_coop_numba_mlir_compile'], project: 'cuda_coop_python', ctk: '13.X', py_version: '3.14', cxx: 'gcc13'} - - {jobs: ['test_py_coop_numba_mlir_runtime'], project: 'cuda_coop_python', ctk: '13.X', py_version: '3.14', gpu: 'l4', cxx: 'gcc13'} - - {jobs: ['build'], project: 'cuda_coop_python', ctk: '13.X', py_version: ['3.10', '3.14'], cxx: 'msvc2022'} # CCCL packaging: - {jobs: ['test'], project: 'packaging', ctk: '12.0', cxx: ['gcc10', 'clang14'], gpu: 't4', args: '-min-cmake'} - {jobs: ['test'], project: 'packaging', ctk: '12.X', cxx: ['gcc10', 'clang14'], gpu: 't4'} @@ -286,10 +280,6 @@ workflows: - {project: 'python', jobs: ['test'], ctk: '13.X', py_version: '3.14', gpu: 'l4', cxx: ['gcc13', 'msvc2022']} - {project: 'python', jobs: ['test_headers'], ctk: '13.X', py_version: '3.14', cxx: ['gcc13', 'msvc2022']} - {project: 'python', jobs: ['test_py_stf'], ctk: '13.X', py_version: '3.14', gpu: 'l4', cxx: 'gcc13'} - # Test the installed `cuda-coop` wheel when its dependencies change. - - {project: 'cuda_coop_python', jobs: ['test_py_coop_host'], ctk: '13.X', py_version: '3.14', cxx: 'gcc13'} - - {project: 'cuda_coop_python', jobs: ['test_py_coop_numba_mlir_compile'], ctk: '13.X', py_version: '3.14', cxx: 'gcc13'} - - {project: 'cuda_coop_python', jobs: ['test_py_coop_numba_mlir_runtime'], ctk: '13.X', py_version: '3.14', gpu: 'l4', cxx: 'gcc13'} # Packaging / install - {project: 'packaging', jobs: ['test'], ctk: '13.X', cxx: ['gcc', 'clang'], gpu: 't4', sm: 'gpu'} - {project: 'packaging', jobs: ['test'], args: '-min-cmake', gpu: 't4', sm: 'gpu'} @@ -437,13 +427,6 @@ workflows: - {jobs: ['test_py_compute_minimal'], project: 'python', ctk: ['12.0', '12.X', '13.0', '13.X'], py_version: '3.14t', gpu: 'l4', cxx: ['gcc13', 'msvc2022'], py_ctk_mode: 'sysctk'} # cuda.cccl.headers (CPU-only): all CTK x source x OS, py endpoints 3.10 + 3.14 - {jobs: ['test_headers'], project: 'python', ctk: ['12.0', '12.X', '13.0', '13.X'], py_version: ['3.10', '3.14'], cxx: ['gcc13', 'msvc2022'], py_ctk_mode: ['pinned', 'latest', 'sysctk']} - # cuda.coop: run the older supported CUDA major in `nightly`. The CUDA 13 - # compiler/runtime counterparts run in `pull_request`, while the pure wheel - # is checked at both Python endpoints on Linux and Windows in both. - - {jobs: ['test_py_coop_host'], project: 'cuda_coop_python', ctk: '13.X', py_version: ['3.10', '3.14'], cxx: 'gcc13'} - - {jobs: ['test_py_coop_numba_mlir_compile'], project: 'cuda_coop_python', ctk: '12.X', py_version: '3.14', cxx: 'gcc13'} - - {jobs: ['test_py_coop_numba_mlir_runtime'], project: 'cuda_coop_python', ctk: '12.X', py_version: '3.14', gpu: 'l4', cxx: 'gcc13'} - - {jobs: ['build'], project: 'cuda_coop_python', ctk: '13.X', py_version: ['3.10', '3.14'], cxx: 'msvc2022'} # CCCL packaging: - {jobs: ['test'], project: 'packaging', ctk: '12.0', cxx: ['gcc10', 'clang14'], gpu: 't4', args: '-min-cmake'} - {jobs: ['test'], project: 'packaging', ctk: '12.X', cxx: ['gcc10', 'clang14'], gpu: 'rtx2080'} @@ -607,8 +590,7 @@ workflows: # deliberately absent (like the TSan lanes below): the experimental # cuda-stf wheel must not be published, and the combined STF producer's # cuda_cccl wheel would collide with the release one in dist/. STF - # coverage lives in the pull_request/nightly/weekly workflows. cuda.coop - # is also absent until cuda-coop has a registered release target. + # coverage lives in the pull_request/nightly/weekly workflows. - {jobs: ['test'], project: 'python', ctk: ['12.0', '12.X', '13.0', '13.X'], py_version: ['3.10', '3.11', '3.12', '3.13', '3.14', '3.14t'], gpu: 'l4', cxx: 'gcc13'} - {jobs: ['test'], project: 'python', ctk: ['12.X', '13.0', '13.X'], py_version: ['3.10', '3.11', '3.12', '3.13', '3.14', '3.14t'], gpu: 'l4', cxx: 'msvc2022'} # Windows + CTK 12.0: compute only, no examples (NVBug 6706273; see the nightly comment). @@ -856,9 +838,6 @@ jobs: test_py_compute_minimal: { name: "Test cuda.compute minimal", gpu: true, needs: 'build_py_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_compute_minimal'} } test_py_examples: { name: "Test cuda.cccl.examples", gpu: true, needs: 'build_py_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_cccl_examples'} } test_py_stf: { name: "Test cuda.stf._experimental", gpu: true, needs: 'build_py_stf_wheel', force_producer_ctk: "pybuild", invoke: { prefix: 'test_cuda_stf'} } - test_py_coop_host: { name: "Test cuda.coop (host)", gpu: false, needs: 'build', force_producer_ctk: "pybuild", invoke: { prefix: 'test', args: '-stage contracts'} } - test_py_coop_numba_mlir_compile: { name: "Compile cuda.coop Numba-CUDA-MLIR (GPU hidden)", gpu: false, needs: 'build', force_producer_ctk: "pybuild", invoke: { prefix: 'test', args: '-stage numba-mlir-compile'} } - test_py_coop_numba_mlir_runtime: { name: "Test cuda.coop Numba-CUDA-MLIR (L4)", gpu: true, needs: 'build', force_producer_ctk: "pybuild", invoke: { prefix: 'test', args: '-stage numba-mlir-runtime'} } # Run jobs for 'target' project (ci/util/build_and_test_targets.sh): run_cpu: { gpu: false } @@ -942,8 +921,6 @@ projects: job_map: build: ['build_py_wheel'] test: ['test_py_compute_minimal'] - cuda_coop_python: - name: "Python (cuda.coop)" cccl_c_parallel: name: 'CCCL C Parallel' stds: [20] diff --git a/ci/project_files_and_dependencies.yaml b/ci/project_files_and_dependencies.yaml index 39064a2e8694..d2b7a763e596 100644 --- a/ci/project_files_and_dependencies.yaml +++ b/ci/project_files_and_dependencies.yaml @@ -54,7 +54,6 @@ projects: lite_dependencies: [c2h] full_dependencies: [libcudacxx_public] include_regexes: ["libcudacxx/"] - exclude_regexes: ["^libcudacxx/LICENSE\\.TXT$"] exclude_project_files: [libcudacxx_public] cub_public: @@ -69,7 +68,6 @@ projects: lite_dependencies: [c2h, nvbench_helper] full_dependencies: [cub_public] include_regexes: ["cub/"] - exclude_regexes: ["^cub/LICENSE\\.TXT$"] exclude_project_files: [cub_public] thrust_public: @@ -84,7 +82,6 @@ projects: lite_dependencies: [nvbench_helper] full_dependencies: [thrust_public] include_regexes: ["thrust/"] - exclude_regexes: ["^thrust/LICENSE$"] exclude_project_files: [thrust_public] cudax_public: @@ -99,7 +96,6 @@ projects: lite_dependencies: [c2h, nvbench_helper] full_dependencies: [cudax_public] include_regexes: ["cudax/"] - exclude_regexes: ["^cudax/LICENSE\\.TXT$"] exclude_project_files: [cudax_public] cccl_c_parallel_public: @@ -176,29 +172,6 @@ projects: - "pyproject.toml" exclude_regexes: [] - cuda_coop_python: - name: "cuda.coop Python" - matrix_project: "cuda_coop_python" - # The wheel carries these three public header trees at one CCCL revision. - # Header changes therefore exercise the focused installed-wheel lane. - lite_dependencies: - - libcudacxx_public - - cub_public - - thrust_public - full_dependencies: [] - include_regexes: - - "python/cuda_coop/" - - "ci/build_cuda_coop_python\\.sh" - - "ci/test_cuda_coop_python\\.sh" - - "ci/validate_cuda_coop_wheel\\.py" - - "ci/windows/build_cuda_coop_python\\.ps1" - - "^LICENSE$" - - "^cub/LICENSE\\.TXT$" - - "^libcudacxx/LICENSE\\.TXT$" - - "^thrust/LICENSE$" - - "^pyproject\\.toml$" - exclude_regexes: [] - packaging: name: "CCCL Packaging" matrix_project: "packaging" @@ -274,8 +247,6 @@ projects: # Files matching any of these regexes will be ignored globally. ignore_regexes: - - '^cudax/LICENSE\.TXT$' - - '^libcudacxx/test/LICENSE\.TXT$' - '.+\.md$' - '\.branch_notes/' - '^\.agent/' @@ -334,3 +305,4 @@ ignore_regexes: - 'ci/rapids/' - '^ci/util/pre-commit/' - 'docs/' + - 'LICENSE' diff --git a/ci/test/inspect_changes/core_dirty.output b/ci/test/inspect_changes/core_dirty.output index cac815153ad5..456e3158f677 100644 --- a/ci/test/inspect_changes/core_dirty.output +++ b/ci/test/inspect_changes/core_dirty.output @@ -1,2 +1,2 @@ -FULL_BUILD=libcudacxx cub thrust cudax cccl_c_parallel cccl_c_parallel_v2 python_v2 python_tsan cccl_c_stf python cuda_coop_python packaging stdpar nvbench_helper nvrtcc tidy +FULL_BUILD=libcudacxx cub thrust cudax cccl_c_parallel cccl_c_parallel_v2 python_v2 python_tsan cccl_c_stf python packaging stdpar nvbench_helper nvrtcc tidy LITE_BUILD= diff --git a/ci/test/inspect_changes/cuda_coop_python.dirty_files b/ci/test/inspect_changes/cuda_coop_python.dirty_files deleted file mode 100644 index 24ee129f18b5..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python.dirty_files +++ /dev/null @@ -1,4 +0,0 @@ -python/cuda_coop/cuda/coop/__init__.py -ci/build_cuda_coop_python.sh -ci/validate_cuda_coop_wheel.py -ci/test_cuda_coop_python.sh diff --git a/ci/test/inspect_changes/cuda_coop_python.output b/ci/test/inspect_changes/cuda_coop_python.output deleted file mode 100644 index 25b5a5aa9040..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python.output +++ /dev/null @@ -1,2 +0,0 @@ -FULL_BUILD=cuda_coop_python -LITE_BUILD= diff --git a/ci/test/inspect_changes/cuda_coop_python_cudax.dirty_files b/ci/test/inspect_changes/cuda_coop_python_cudax.dirty_files deleted file mode 100644 index 9b8672c0cb6b..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python_cudax.dirty_files +++ /dev/null @@ -1 +0,0 @@ -cudax/include/cuda/experimental/coop/group diff --git a/ci/test/inspect_changes/cuda_coop_python_cudax.output b/ci/test/inspect_changes/cuda_coop_python_cudax.output deleted file mode 100644 index bd358baf03b2..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python_cudax.output +++ /dev/null @@ -1,2 +0,0 @@ -FULL_BUILD=cudax tidy -LITE_BUILD=cccl_c_stf python packaging diff --git a/ci/test/inspect_changes/cuda_coop_python_license.dirty_files b/ci/test/inspect_changes/cuda_coop_python_license.dirty_files deleted file mode 100644 index f03844f8d445..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python_license.dirty_files +++ /dev/null @@ -1,4 +0,0 @@ -LICENSE -cub/LICENSE.TXT -libcudacxx/LICENSE.TXT -thrust/LICENSE diff --git a/ci/test/inspect_changes/cuda_coop_python_license.output b/ci/test/inspect_changes/cuda_coop_python_license.output deleted file mode 100644 index 25b5a5aa9040..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python_license.output +++ /dev/null @@ -1,2 +0,0 @@ -FULL_BUILD=cuda_coop_python -LITE_BUILD= diff --git a/ci/test/inspect_changes/cuda_coop_python_pyproject.dirty_files b/ci/test/inspect_changes/cuda_coop_python_pyproject.dirty_files deleted file mode 100644 index cc736a365426..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python_pyproject.dirty_files +++ /dev/null @@ -1 +0,0 @@ -pyproject.toml diff --git a/ci/test/inspect_changes/cuda_coop_python_pyproject.output b/ci/test/inspect_changes/cuda_coop_python_pyproject.output deleted file mode 100644 index 86050dd10ed5..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python_pyproject.output +++ /dev/null @@ -1,2 +0,0 @@ -FULL_BUILD=python_v2 python_tsan python cuda_coop_python -LITE_BUILD= diff --git a/ci/test/inspect_changes/cuda_coop_python_pyproject_near_miss.dirty_files b/ci/test/inspect_changes/cuda_coop_python_pyproject_near_miss.dirty_files deleted file mode 100644 index 4c1b6afb44b3..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python_pyproject_near_miss.dirty_files +++ /dev/null @@ -1,2 +0,0 @@ -pyprojectXtoml -pyproject.toml.bak diff --git a/ci/test/inspect_changes/cuda_coop_python_pyproject_near_miss.output b/ci/test/inspect_changes/cuda_coop_python_pyproject_near_miss.output deleted file mode 100644 index e14e27c732a0..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python_pyproject_near_miss.output +++ /dev/null @@ -1,2 +0,0 @@ -FULL_BUILD=python_v2 python_tsan python -LITE_BUILD= diff --git a/ci/test/inspect_changes/cuda_coop_python_windows_ci.dirty_files b/ci/test/inspect_changes/cuda_coop_python_windows_ci.dirty_files deleted file mode 100644 index 94e38eb85601..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python_windows_ci.dirty_files +++ /dev/null @@ -1 +0,0 @@ -ci/windows/build_cuda_coop_python.ps1 diff --git a/ci/test/inspect_changes/cuda_coop_python_windows_ci.output b/ci/test/inspect_changes/cuda_coop_python_windows_ci.output deleted file mode 100644 index 25b5a5aa9040..000000000000 --- a/ci/test/inspect_changes/cuda_coop_python_windows_ci.output +++ /dev/null @@ -1,2 +0,0 @@ -FULL_BUILD=cuda_coop_python -LITE_BUILD= diff --git a/ci/test/inspect_changes/ignored_only.dirty_files b/ci/test/inspect_changes/ignored_only.dirty_files index 3f4f99c929ce..7d2fa473a692 100644 --- a/ci/test/inspect_changes/ignored_only.dirty_files +++ b/ci/test/inspect_changes/ignored_only.dirty_files @@ -1,2 +1 @@ docs/index.rst -cudax/LICENSE.TXT diff --git a/ci/test/inspect_changes/libcudacxx_both.output b/ci/test/inspect_changes/libcudacxx_both.output index 1cacd88a3835..eddc77079559 100644 --- a/ci/test/inspect_changes/libcudacxx_both.output +++ b/ci/test/inspect_changes/libcudacxx_both.output @@ -1,2 +1,2 @@ FULL_BUILD=libcudacxx tidy -LITE_BUILD=cub thrust cudax cccl_c_parallel cccl_c_parallel_v2 python_v2 python_tsan cccl_c_stf python cuda_coop_python packaging stdpar nvbench_helper +LITE_BUILD=cub thrust cudax cccl_c_parallel cccl_c_parallel_v2 python_v2 python_tsan cccl_c_stf python packaging stdpar nvbench_helper diff --git a/ci/test/inspect_changes/libcudacxx_public_only.output b/ci/test/inspect_changes/libcudacxx_public_only.output index 1cacd88a3835..eddc77079559 100644 --- a/ci/test/inspect_changes/libcudacxx_public_only.output +++ b/ci/test/inspect_changes/libcudacxx_public_only.output @@ -1,2 +1,2 @@ FULL_BUILD=libcudacxx tidy -LITE_BUILD=cub thrust cudax cccl_c_parallel cccl_c_parallel_v2 python_v2 python_tsan cccl_c_stf python cuda_coop_python packaging stdpar nvbench_helper +LITE_BUILD=cub thrust cudax cccl_c_parallel cccl_c_parallel_v2 python_v2 python_tsan cccl_c_stf python packaging stdpar nvbench_helper diff --git a/ci/test/inspect_changes/libcudacxx_thrust.output b/ci/test/inspect_changes/libcudacxx_thrust.output index 0ad3962639bb..00f10a620f0f 100644 --- a/ci/test/inspect_changes/libcudacxx_thrust.output +++ b/ci/test/inspect_changes/libcudacxx_thrust.output @@ -1,2 +1,2 @@ FULL_BUILD=libcudacxx thrust tidy -LITE_BUILD=cub cudax cccl_c_parallel cccl_c_parallel_v2 python_v2 python_tsan cccl_c_stf python cuda_coop_python packaging stdpar nvbench_helper +LITE_BUILD=cub cudax cccl_c_parallel cccl_c_parallel_v2 python_v2 python_tsan cccl_c_stf python packaging stdpar nvbench_helper diff --git a/ci/test_cuda_coop_python.sh b/ci/test_cuda_coop_python.sh deleted file mode 100755 index b0d71a7ac61e..000000000000 --- a/ci/test_cuda_coop_python.sh +++ /dev/null @@ -1,226 +0,0 @@ -#!/usr/bin/env bash -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -# Test the installed wheel in one stage: host contracts, device compilation, -# or GPU execution. The compiler stage hides devices to check that providers -# can compile with an explicit target and no running kernel. - -set -euo pipefail - -ci_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -repo_root="$(cd "$ci_dir/.." && pwd)" - -usage="Usage: $0 -py-version [-stage contracts|numba-mlir-compile|numba-mlir-runtime]" - -# shellcheck source=ci/util/python/common_arg_parser.sh -source "$ci_dir/util/python/common_arg_parser.sh" -parse_python_args "$@" -require_py_version "$usage" || exit 1 - -stage="contracts" -while [[ $# -gt 0 ]]; do - case "$1" in - -stage | --stage) - if [[ $# -lt 2 || -z "$2" ]]; then - echo "Error: $1 requires a value" >&2 - exit 1 - fi - stage="$2" - shift 2 - ;; - -stage=* | --stage=*) - stage="${1#*=}" - if [[ -z "$stage" ]]; then - echo "Error: $1 requires a value" >&2 - exit 1 - fi - shift - ;; - -py-version | -ctk-mode) - if [[ $# -lt 2 || -z "$2" ]]; then - echo "Error: $1 requires a value" >&2 - exit 1 - fi - shift 2 - ;; - -py-version=* | -ctk-mode=*) - shift - ;; - *) - echo "Error: unknown argument '$1'" >&2 - echo "$usage" >&2 - exit 1 - ;; - esac -done - -needs_cuda_toolkit=false -case "$stage" in - contracts) - ;; - numba-mlir-compile) - needs_cuda_toolkit=true - ;; - numba-mlir-runtime) - needs_cuda_toolkit=true - ;; - *) - echo "Error: unknown cuda.coop test stage '$stage'" >&2 - echo "$usage" >&2 - exit 1 - ;; -esac - -if [[ "$stage" == "numba-mlir-compile" ]]; then - # Hide every device for the complete compiler-contract stage, including - # import and installed-wheel isolation probes. - export CUDA_VISIBLE_DEVICES="" -fi - -# shellcheck source=ci/pyenv_helper.sh -source "$ci_dir/pyenv_helper.sh" - -if [[ "$needs_cuda_toolkit" == true ]]; then - if ! command -v nvcc >/dev/null 2>&1; then - echo "Error: cuda.coop stage '$stage' requires nvcc on PATH" >&2 - exit 1 - fi - pin_cuda_toolkit "${ctk_mode}" -fi - -setup_python_env "${py_version}" ".cccl-coop-test-venv" - -if [[ -n "${GITHUB_ACTIONS:-}" ]]; then - wheel_artifact_name="$(CCCL_WHEEL_KIND=coop "$ci_dir/util/workflow/get_wheel_artifact_name.sh")" - "$ci_dir/util/artifacts/download.sh" "$wheel_artifact_name" "$repo_root/" -else - CCCL_CUDA_COOP_PYENV_READY=1 \ - "$ci_dir/build_cuda_coop_python.sh" -py-version "${py_version}" -fi - -mapfile -t wheels < <(find "$repo_root/wheelhouse" -maxdepth 1 -name 'cuda_coop-*.whl' -print | sort) -if [[ ${#wheels[@]} -ne 1 ]]; then - echo "Error: expected exactly one cuda-coop wheel, found ${#wheels[@]}:" >&2 - printf ' %s\n' "${wheels[@]}" >&2 - exit 1 -fi - -case "$stage" in - contracts) - python -m pip install "${wheels[0]}[test]" - ;; - numba-mlir-compile | numba-mlir-runtime) - python -m pip install \ - "${wheels[0]}[test,numba-cuda-mlir-cu${cuda_major_version}]" - ;; - *) - echo "Error: unhandled cuda.coop test stage '$stage'" >&2 - exit 1 - ;; -esac - -python -m pip check -# Isolated mode prevents the checkout or PYTHONPATH from hiding wheel defects. -python -I - <<'PY' -import importlib.metadata -import sys -from pathlib import Path - -import cuda -from cuda import coop -from cuda.coop._headers import resolve_include_paths - -distribution = importlib.metadata.distribution("cuda-coop") -expected_root = Path( - distribution.locate_file("cuda/coop/__init__.py") -).resolve() -assert Path(coop.__file__).resolve() == expected_root -assert getattr(cuda, "__file__", None) is None -assert coop.this_block().kind == "block" - -paths = resolve_include_paths( - start=Path(sys.prefix), - required_headers=( - "cub/block/block_load.cuh", - "cub/block/block_store.cuh", - "thrust/detail/raw_pointer_cast.h", - "cuda/std/cstdint", - ), -) -assert paths.origin == "cuda-coop wheel header bundle" -PY - -tests_root="$repo_root/python/cuda_coop/tests" - -case "$stage" in - contracts) - # Check shared helpers without requiring optional compiler backends. - python -m pyright --warnings --project "$repo_root/python/cuda_coop/tests/typing/pyright-unused.json" \ - --pythonpath "$(command -v python)" - cd "$tests_root" - python -m pytest -v test_*.py packaging/ - ;; - numba-mlir-compile) - python -m pyright --project "$repo_root/python/cuda_coop" \ - --pythonpath "$(command -v python)" --warnings - # The compile contract is deliberately GPU-free. Tests may replace only - # the backend's current-device query with a fixed compute capability; NVRTC - # and nvJitLink remain real. - python -I - <<'PY' -import importlib.metadata - -from numba_cuda_mlir import cuda - -if cuda.is_available(): - raise SystemExit("GPU-hidden cuda.coop compile stage can access a CUDA device") -print(f"numba-cuda-mlir={importlib.metadata.version('numba-cuda-mlir')}") -PY - cd "$tests_root" - python -m pytest -v \ - backends/numba_mlir/unit/ \ - backends/numba_mlir/compile/ - ;; - numba-mlir-runtime) - python -I - <<'PY' -import importlib.metadata - -from numba_cuda_mlir import cuda - -if not cuda.is_available(): - raise SystemExit("numba-cuda-mlir cannot access an NVIDIA GPU") -print(f"numba-cuda-mlir={importlib.metadata.version('numba-cuda-mlir')}") -PY - nvidia-smi \ - --query-gpu=name,compute_cap,driver_version \ - --format=csv,noheader - cd "$tests_root" - python -m pytest -v backends/numba_mlir/runtime/ - - mapfile -t examples < <( - find "$repo_root/python/cuda_coop/examples/numba_mlir" \ - -maxdepth 1 -name '*.py' ! -name '__init__.py' -print | sort - ) - if [[ ${#examples[@]} -eq 0 ]]; then - echo "Error: no cuda.coop Numba-CUDA-MLIR examples were found" >&2 - exit 1 - fi - for example in "${examples[@]}"; do - case "${example##*/}" in - source_dumps.py) - for source_dump_mode in direct transpose; do - python -I "$example" "$source_dump_mode" - done - ;; - *) - python -I "$example" - ;; - esac - done - ;; - *) - echo "Error: unhandled cuda.coop test stage '$stage'" >&2 - exit 1 - ;; -esac diff --git a/ci/validate_cuda_coop_wheel.py b/ci/validate_cuda_coop_wheel.py deleted file mode 100644 index c9a3d8dd3f87..000000000000 --- a/ci/validate_cuda_coop_wheel.py +++ /dev/null @@ -1,273 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check a ``cuda-coop`` wheel archive at the end of a build. - -Linux and Windows build jobs use these checks to enforce the package layout, -bundled headers, licenses, and universal wheel metadata. The checks read the -archive without installing it or importing its contents. Installed-package -checks run separately in the CI scripts. -""" - -from __future__ import annotations - -import email.parser -import json -import re -import sys -import zipfile -from pathlib import Path, PurePosixPath - -_REQUIRED_PACKAGE_FILES = { - "cuda/coop/__init__.py", - "cuda/coop/__init__.pyi", - "cuda/coop/_typing.pyi", - "cuda/coop/py.typed", - "cuda/coop/_core/_auto_registration.py", - "cuda/coop/_core/api/__init__.py", - "cuda/coop/_core/api/__init__.pyi", - "cuda/coop/_core/api/load_store.py", - "cuda/coop/_core/api/load_store.pyi", - "cuda/coop/_core/api/temp_storage.py", - "cuda/coop/_core/api/temp_storage.pyi", - "cuda/coop/_core/api/thread_data.py", - "cuda/coop/_core/api/thread_data.pyi", - "cuda/coop/_core/api/thread_group.py", - "cuda/coop/_core/api/thread_group.pyi", - "cuda/coop/_core/block/load_store.py", - "cuda/coop/_core/group/load_store.py", - "cuda/coop/_headers/_identity.py", - "cuda/coop/_headers/_toolkit.py", - "cuda/coop/numba_mlir/__init__.py", - "cuda/coop/numba_mlir/__init__.pyi", - "cuda/coop/numba_mlir/_compiler/_activation.py", - "cuda/coop/numba_mlir/_compiler/_artifacts.py", - "cuda/coop/numba_mlir/_compiler/_caching.py", - "cuda/coop/numba_mlir/_compiler/_group_load_store.py", - "cuda/coop/numba_mlir/_compiler/_group_planner.py", - "cuda/coop/numba_mlir/_compiler/_nvrtc.py", - "cuda/coop/numba_mlir/_compiler/_rewrite.py", - "cuda/coop/numba_mlir/_group/_load_store.py", - "cuda/coop/numba_mlir/_group/_load_store.pyi", - "cuda/coop/numba_mlir/_lowering/_load_store.py", - "cuda/coop/numba_mlir/_temp_storage.py", - "cuda/coop/numba_mlir/_temp_storage.pyi", - "cuda/coop/numba_mlir/_thread_data.py", - "cuda/coop/numba_mlir/_thread_data.pyi", - "cuda/coop/numba_mlir/_thread_group.py", - "cuda/coop/numba_mlir/_thread_group.pyi", - "cuda/coop/numba_mlir/py.typed", -} - -_REQUIRED_HEADER_FILES = { - "cuda/coop/_headers/cccl-bundle-provenance.json", - "cuda/coop/_headers/include/cub/version.cuh", - "cuda/coop/_headers/include/cub/block/block_load.cuh", - "cuda/coop/_headers/include/cub/block/block_store.cuh", - "cuda/coop/_headers/include/thrust/detail/raw_pointer_cast.h", - "cuda/coop/_headers/include/cuda/std/cstdint", - "cuda/coop/_headers/include/nv/target", -} - -_REQUIRED_LICENSES = { - "LICENSE", - "cub/LICENSE.TXT", - "libcudacxx/LICENSE.TXT", - "thrust/LICENSE", -} - -_OBSOLETE_LAYOUT_COMPONENTS = {"_block", "_dsl", "_internal", "_warp"} -_NATIVE_SUFFIXES = {".a", ".dll", ".dylib", ".exe", ".lib", ".pyd", ".so"} -_FORBIDDEN_PACKAGE_FILES = { - "cuda/coop/_aot_cli.py", - "cuda/coop/numba_mlir/_enums.py", - "cuda/coop/numba_mlir/_enums.pyi", - "cuda/coop/_core/api/reduce.py", - "cuda/coop/_core/api/reduce.pyi", - "cuda/coop/_core/api/scan.py", - "cuda/coop/_core/api/scan.pyi", - "cuda/coop/_core/block/reduce.py", - "cuda/coop/_core/block/scan.py", - "cuda/coop/_core/group/reduce.py", - "cuda/coop/_core/group/scan.py", - "cuda/coop/numba_mlir/_dataclass.py", - "cuda/coop/numba_mlir/_stateful_function.py", - "cuda/coop/numba_mlir/_group/_reduce.py", - "cuda/coop/numba_mlir/_group/_scan.py", - "cuda/coop/numba_mlir/_lowering/_reduce.py", - "cuda/coop/numba_mlir/_lowering/_scan.py", - "cuda/coop/numba_mlir/_lowering/_thread_group.py", - "cuda/coop/numba_mlir/_compiler/_rewrite_reduce.py", - "cuda/coop/numba_mlir/_compiler/_rewrite_scan.py", -} - - -def _one_member(names: set[str], suffix: str) -> str: - matches = sorted(name for name in names if name.endswith(suffix)) - if len(matches) != 1: - raise SystemExit( - f"cuda-coop wheel must contain exactly one {suffix}: {matches}" - ) - return matches[0] - - -def _validate_metadata(archive: zipfile.ZipFile, names: set[str]) -> None: - """Check the distribution name, supported Python versions, and wheel tag. - - These declarations control how installers select the wheel. Keep them - consistent with a pure-Python package shared by all supported platforms. - """ - metadata_name = _one_member(names, ".dist-info/METADATA") - wheel_name = _one_member(names, ".dist-info/WHEEL") - - metadata = email.parser.Parser().parsestr( - archive.read(metadata_name).decode("utf-8") - ) - if metadata.get("Name") != "cuda-coop": - raise SystemExit( - "cuda-coop wheel metadata has an unexpected project name: " - f"{metadata.get('Name')!r}" - ) - if metadata.get("Requires-Python") != ">=3.10": - raise SystemExit( - "cuda-coop wheel metadata must declare Requires-Python: >=3.10; got " - f"{metadata.get('Requires-Python')!r}" - ) - expected_extras = { - "numba-cuda-mlir-cu12", - "numba-cuda-mlir-cu13", - "test", - } - extras = set(metadata.get_all("Provides-Extra", [])) - if extras != expected_extras: - raise SystemExit( - "cuda-coop wheel metadata has unexpected extras: " - f"expected {sorted(expected_extras)}, got {sorted(extras)}" - ) - if any( - "cutlass" in requirement.lower() - for requirement in metadata.get_all("Requires-Dist", []) - ): - raise SystemExit("cuda-coop wheel must not depend on CUTLASS") - - wheel_metadata = email.parser.Parser().parsestr( - archive.read(wheel_name).decode("utf-8") - ) - if wheel_metadata.get("Root-Is-Purelib", "").lower() != "true": - raise SystemExit("cuda-coop wheel must declare Root-Is-Purelib: true") - if wheel_metadata.get_all("Tag", []) != ["py3-none-any"]: - raise SystemExit( - "cuda-coop wheel must declare only Tag: py3-none-any; got " - f"{wheel_metadata.get_all('Tag', [])}" - ) - - -def _validate_provenance(archive: zipfile.ZipFile) -> None: - """Check the recorded header revision's schema and token syntax. - - The build can record ``unknown`` when it cannot identify the source. - This check validates the record; it does not verify a source checkout. - """ - member = "cuda/coop/_headers/cccl-bundle-provenance.json" - try: - provenance = json.loads(archive.read(member)) - except (json.JSONDecodeError, UnicodeDecodeError) as exc: - raise SystemExit(f"cuda-coop header provenance is invalid JSON: {exc}") from exc - if set(provenance) != {"cccl_source_commit"}: - raise SystemExit( - "cuda-coop header provenance must contain only cccl_source_commit" - ) - revision = provenance["cccl_source_commit"] - if ( - not isinstance(revision, str) - or re.fullmatch(r"[0-9A-Za-z._+-]+", revision) is None - ): - raise SystemExit("cuda-coop header provenance has an invalid source revision") - - -def validate(wheel: str | Path) -> None: - """Check a wheel and stop on a packaging contract violation. - - Require the shared API and its header bundle, reject native binaries and - excluded implementations, and preserve the shared ``cuda`` namespace. - A failed check raises ``SystemExit`` with a packaging diagnostic. - File and ZIP errors propagate to the caller. - """ - wheel_path = Path(wheel) - if not wheel_path.name.endswith("-py3-none-any.whl"): - raise SystemExit( - f"cuda-coop must produce a universal py3-none-any wheel: {wheel_path.name}" - ) - - with zipfile.ZipFile(wheel_path) as archive: - names = set(archive.namelist()) - _validate_metadata(archive, names) - - missing = (_REQUIRED_PACKAGE_FILES | _REQUIRED_HEADER_FILES) - names - if missing: - raise SystemExit( - f"cuda-coop wheel is missing required files: {sorted(missing)}" - ) - - if "cuda/__init__.py" in names: - raise SystemExit( - "cuda-coop wheel must not contain cuda/__init__.py; " - "it would break the PEP 420 cuda namespace" - ) - if any(name.startswith("cuda/coop/cutlass/") for name in names): - raise SystemExit("cuda-coop wheel must not contain a CUTLASS backend") - forbidden = sorted(_FORBIDDEN_PACKAGE_FILES & names) - if forbidden: - raise SystemExit( - f"cuda-coop wheel contains excluded implementations: {forbidden}" - ) - - obsolete = sorted( - name - for name in names - if name.startswith("cuda/coop/") - and not name.startswith("cuda/coop/_headers/include/") - and _OBSOLETE_LAYOUT_COMPONENTS.intersection(PurePosixPath(name).parts) - ) - if obsolete: - raise SystemExit( - f"cuda-coop wheel contains obsolete package layouts: {obsolete}" - ) - - native = sorted( - name - for name in names - if any( - suffix.lower() in _NATIVE_SUFFIXES - for suffix in PurePosixPath(name).suffixes - ) - ) - if native: - raise SystemExit( - f"cuda-coop wheel must not contain native binaries: {native}" - ) - - license_members = { - name.split(".dist-info/licenses/", 1)[1] - for name in names - if ".dist-info/licenses/" in name - } - missing_licenses = _REQUIRED_LICENSES - license_members - if missing_licenses: - raise SystemExit( - "cuda-coop wheel is missing license payloads: " - f"{sorted(missing_licenses)}" - ) - - _validate_provenance(archive) - - -def main() -> None: - if len(sys.argv) != 2: - raise SystemExit(f"Usage: {sys.argv[0]} ") - validate(sys.argv[1]) - - -if __name__ == "__main__": - main() diff --git a/ci/windows/build_cuda_coop_python.ps1 b/ci/windows/build_cuda_coop_python.ps1 deleted file mode 100644 index d6a46b4c3849..000000000000 --- a/ci/windows/build_cuda_coop_python.ps1 +++ /dev/null @@ -1,123 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -<# -.SYNOPSIS - Build and validate the pure-Python cuda-coop wheel on Windows. - -.DESCRIPTION - Check the archive, then install it in a temporary environment. Verify - that imports use the installed wheel and resolve its bundled headers. - These checks do not compile or launch GPU kernels. - -.PARAMETER PyVersion - Python version used to build and install the universal wheel. -#> - -[CmdletBinding()] -Param( - [Parameter(Mandatory = $true)] - [Alias("py-version")] - [ValidatePattern("^\d+\.\d+t?$")] - [string]$PyVersion -) - -$ErrorActionPreference = "Stop" - -Import-Module "$PSScriptRoot/build_common_python.psm1" -Force - -$RepoRoot = (Resolve-Path (Join-Path $PSScriptRoot "../..")).Path -$PackageRoot = Join-Path $RepoRoot "python/cuda_coop" -$Wheelhouse = Join-Path $RepoRoot "wheelhouse" -$PythonExe = Get-Python -Version $PyVersion - -New-Item -ItemType Directory -Path $Wheelhouse -Force | Out-Null -Get-ChildItem -LiteralPath $Wheelhouse -Filter "cuda_coop-*.whl" -File | - Remove-Item -Force - -& $PythonExe -m pip install build -if ($LASTEXITCODE -ne 0) { - throw "Failed to install the Python build frontend." -} - -& $PythonExe -m build --wheel --outdir $Wheelhouse $PackageRoot -if ($LASTEXITCODE -ne 0) { - throw "cuda-coop wheel build failed." -} - -$Wheels = @(Get-ChildItem -LiteralPath $Wheelhouse -Filter "cuda_coop-*.whl" -File) -if ($Wheels.Count -ne 1) { - throw "Expected one cuda-coop wheel under $Wheelhouse; found $($Wheels.Count)." -} -if ($Wheels[0].Name -notmatch "-py3-none-any\.whl$") { - throw "Expected a py3-none-any wheel; found $($Wheels[0].Name)." -} - -$Validator = Join-Path $RepoRoot "ci/validate_cuda_coop_wheel.py" -& $PythonExe $Validator $Wheels[0].FullName -if ($LASTEXITCODE -ne 0) { - throw "cuda-coop wheel validation failed." -} - -$SmokeEnv = Join-Path ([System.IO.Path]::GetTempPath()) ( - "cuda-coop-smoke-" + [guid]::NewGuid().ToString("N") -) -try { - & $PythonExe -m venv $SmokeEnv - if ($LASTEXITCODE -ne 0) { - throw "Failed to create the cuda-coop wheel smoke-test environment." - } - - $SmokePython = Join-Path $SmokeEnv "Scripts/python.exe" - & $SmokePython -m pip install --disable-pip-version-check $Wheels[0].FullName - if ($LASTEXITCODE -ne 0) { - throw "Failed to install the cuda-coop wheel." - } - & $SmokePython -m pip check - if ($LASTEXITCODE -ne 0) { - throw "The installed cuda-coop wheel has inconsistent dependencies." - } - - $SmokeCommand = @' -import importlib.metadata -import sys -from pathlib import Path - -import cuda -from cuda import coop -from cuda.coop._headers import resolve_include_paths - -distribution = importlib.metadata.distribution("cuda-coop") -expected_root = Path( - distribution.locate_file("cuda/coop/__init__.py") -).resolve() -assert Path(coop.__file__).resolve() == expected_root -assert getattr(cuda, "__file__", None) is None -assert coop.__version__ != "0+unknown" -assert coop.this_block().kind == "block" - -paths = resolve_include_paths( - start=Path(sys.prefix), - required_headers=( - "cub/block/block_load.cuh", - "cub/block/block_store.cuh", - "thrust/detail/raw_pointer_cast.h", - "cuda/std/cstdint", - ), -) -assert paths.origin == "cuda-coop wheel header bundle" -'@ - $SmokeScript = Join-Path $SmokeEnv "validate_install.py" - Set-Content -LiteralPath $SmokeScript -Value $SmokeCommand -Encoding UTF8 - & $SmokePython -I $SmokeScript - if ($LASTEXITCODE -ne 0) { - throw "Installed cuda-coop wheel smoke test failed." - } -} -finally { - if (Test-Path -LiteralPath $SmokeEnv) { - Remove-Item -LiteralPath $SmokeEnv -Recurse -Force - } -} - -Write-Host "Built cuda-coop wheel: $($Wheels[0].FullName)" diff --git a/docs/_ext/coop_visualization.py b/docs/_ext/coop_visualization.py deleted file mode 100644 index 0be593eeacf7..000000000000 --- a/docs/_ext/coop_visualization.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Embed cooperative explorers while retaining readable static content. - -The directive keeps its RST fallback for non-HTML output and browsers where -JavaScript is unavailable. Page hooks load assets only where they are used. -""" - -from docutils import nodes -from sphinx.util.docutils import SphinxDirective - -_EXPLORERS = { - "load": "coop-load.js", - "store": "coop-movement.js", -} - -_VISUALIZATION_TITLES = {} - -_API_VISUALIZATIONS = { - "load": "load", - "store": "store", -} - - -def add_api_visualization_link(app, what, name, obj, options, lines): - """Append the matching explorer link to common and qualified API docs.""" - module, _, function = name.rpartition(".") - if what != "function" or module not in {"cuda.coop", "cuda.coop.numba_mlir"}: - return - visualization = _API_VISUALIZATIONS.get(function) - if visualization is not None: - title = _VISUALIZATION_TITLES.get(visualization, visualization.title()) - lines.extend( - [ - "", - ".. seealso::", - "", - ( - f" :doc:`{title} visualization " - f"`" - ), - "", - ] - ) - - -class CoopVisualization(SphinxDirective): - """Pair an HTML explorer host with the directive's static RST content.""" - - required_arguments = 1 - has_content = True - - def run(self): - """Validate the name and retain the fallback in the document tree.""" - name = self.arguments[0] - if name not in _EXPLORERS: - raise self.error(f"Unknown cooperative visualization: {name!r}.") - self.assert_has_content() - container = nodes.container(classes=["coop-visualization"]) - container["coop_explorer"] = name - attribute = ( - "data-coop-load" if name == "load" else f'data-coop-explorer="{name}"' - ) - container += nodes.raw("", f"
", format="html") - fallback = nodes.container(classes=["coop-visualization-fallback"]) - self.state.nested_parse(self.content, self.content_offset, fallback) - container += fallback - return [container] - - -def add_visualization_assets(app, pagename, templatename, context, doctree): - """Load styles and scripts for the explorers recorded on this page.""" - if pagename == "python/coop_api": - app.add_css_file("coop-api.css") - if doctree is None: - return - explorers = { - node["coop_explorer"] - for node in doctree.findall(nodes.container) - if "coop_explorer" in node - } - if explorers: - # Assets are relative to Sphinx's output root, including dirhtml builds. - app.add_css_file("coop-load.css") - if explorers - {"load"}: - app.add_js_file("coop-explorer.js", defer="defer") - for script in sorted({_EXPLORERS[name] for name in explorers}): - app.add_js_file(script, defer="defer") - - -def setup(app): - app.setup_extension("sphinx.ext.autodoc") - app.add_directive("coop-visualization", CoopVisualization) - app.connect("autodoc-process-docstring", add_api_visualization_link) - app.connect("html-page-context", add_visualization_assets) - return { - "version": "8", - "env_version": 7, - "parallel_read_safe": True, - "parallel_write_safe": True, - } diff --git a/docs/_static/coop-api.css b/docs/_static/coop-api.css deleted file mode 100644 index 866aa852caf0..000000000000 --- a/docs/_static/coop-api.css +++ /dev/null @@ -1,12 +0,0 @@ -/* Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. - * - * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - */ - -#cuda-coop-api-reference .sig-name, -#cuda-coop-api-reference .sig-prename, -#cuda-coop-api-reference .sig-param, -#cuda-coop-api-reference code.xref .pre { - overflow-wrap: anywhere; - white-space: normal; -} diff --git a/docs/_static/coop-explorer.js b/docs/_static/coop-explorer.js deleted file mode 100644 index c1bd29f8e4bf..000000000000 --- a/docs/_static/coop-explorer.js +++ /dev/null @@ -1,332 +0,0 @@ -// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -// -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -// Render the rows, phases, and value tokens supplied by each explorer. -// Algorithm models own data movement; this renderer owns controls, URL state, -// keyboard selection, and animation between phases. -(() => { - "use strict"; - - const colors = ["#76b900", "#00a9ce", "#f6b21a", "#d6538f", "#9d80d7", "#e86f18", "#2a9d8f", "#8198a7"]; - const definitions = new Map(); - const options = (source, state) => (typeof source === "function" ? source(state) : source) - .map((item) => typeof item === "object" ? { ...item, value: String(item.value) } : { value: String(item), label: String(item) }); - - function svg_element(tag, attributes = {}, label) { - const element = document.createElementNS("http://www.w3.org/2000/svg", tag); - for (const [name, value] of Object.entries(attributes)) element.setAttribute(name, value); - if (label !== undefined) element.textContent = label; - return element; - } - - function mount(host, name, definition, instance) { - const reduced_motion = window.matchMedia("(prefers-reduced-motion: reduce)"); - const prefix = `coop-${name}-${instance}`; - const controls = definition.controls || []; - const state = {}; - let model; - let phase = 0; - let playing = !reduced_motion.matches; - let visible = true; - let timer = null; - let selected = null; - let geometry; - let token_elements = new Map(); - const root = document.createElement("section"); - root.className = "coop-load coop-explorer"; - root.setAttribute("aria-label", definition.title); - root.innerHTML = ` -
-

-
- - - - -
-
-
-
-

-
-
- -
-

-
-

Selected value

-

Result

-
-
-

`; - const query = (selector) => root.querySelector(selector); - const svg = query("svg"); - query("[data-title]").textContent = definition.title; - query("[data-eyebrow]").textContent = definition.eyebrow; - - function algorithms() { - return typeof definition.algorithms === "function" ? definition.algorithms(state) : definition.algorithms; - } - - function normalize() { - for (const control of controls) { - const choices = options(control.choices, state); - if (!choices.some((choice) => choice.value === state[control.id])) { - state[control.id] = choices.some((choice) => choice.value === String(control.value)) ? String(control.value) : choices[0].value; - } - } - const choices = algorithms(); - if (!choices.some((choice) => choice.id === state.algorithm)) { - state.algorithm = choices.some((choice) => choice.id === definition.defaultAlgorithm) ? definition.defaultAlgorithm : choices[0].id; - } - } - - function read_url() { - const params = new URL(window.location.href).searchParams; - for (const control of controls) state[control.id] = params.get(`${name}-${control.id}`) ?? String(control.value); - state.algorithm = params.get(`${name}-algorithm`) || definition.defaultAlgorithm; - normalize(); - phase = 0; - playing = !reduced_motion.matches; - } - - function write_url() { - const url = new URL(window.location.href); - for (const [key, value] of Object.entries(state)) url.searchParams.set(`${name}-${key}`, value); - try { window.history.replaceState(window.history.state, "", url); } catch { /* Local file previews can forbid history updates. */ } - } - - function position(row_id, index) { - const row = model.rows.find((row) => row.id === row_id); - const row_index = model.rows.indexOf(row); - const start = row.align === "start" ? 20 : (geometry.width - row.count * 42 + 10) / 2; - return [start + index * 42, 76 + row_index * 120]; - } - - function rebuild() { - normalize(); - model = definition.build(state); - phase = Math.min(phase, model.phases.length - 1); - query(".coop-load-options").replaceChildren(...algorithms().map((algorithm) => { - const button = document.createElement("button"); - button.type = "button"; - button.dataset.algorithm = algorithm.id; - button.setAttribute("aria-pressed", algorithm.id === state.algorithm); - const title = document.createElement("strong"); - title.textContent = algorithm.label; - const tag = document.createElement("span"); - tag.textContent = algorithm.tag || ""; - button.append(title, tag); - return button; - })); - query(".coop-load-settings").replaceChildren(...controls.filter((control) => !control.hidden?.(state)).map((control) => { - const label = document.createElement("label"); - label.htmlFor = `${prefix}-${control.id}`; - label.textContent = control.label; - const select = document.createElement("select"); - select.id = label.htmlFor; - select.dataset.control = control.id; - select.append(...options(control.choices, state).map((choice) => new Option(choice.label, choice.value))); - select.value = state[control.id]; - label.append(select); - return label; - })); - query(".coop-load-phases").replaceChildren(...model.phases.map((item, index) => { - const button = document.createElement("button"); - button.type = "button"; - button.dataset.phase = index; - button.textContent = `${index + 1}. ${item.label}`; - return button; - })); - query("[data-detail]").textContent = model.detail; - query("[data-summary]").textContent = model.summary; - query("[data-notes]").replaceChildren(...(model.notes || []).map((note) => { - const paragraph = document.createElement("p"); - paragraph.textContent = note; - return paragraph; - })); - query("[data-caption]").textContent = model.caption || "Teaching model: eight threads and four lanes per displayed warp; physical CUDA warps have 32 lanes. Stages show data dependencies and ownership, not instruction timing. Select a value to inspect it; arrow keys move between values."; - geometry = { width: Math.max(850, Math.max(...model.rows.map((row) => row.count)) * 42 + 40) }; - svg.replaceChildren(); - svg.setAttribute("viewBox", `0 0 ${geometry.width} ${model.rows.length * 120 + 10}`); - svg.style.minWidth = `${geometry.width * 0.8}px`; - svg.append(svg_element("title", { id: `${prefix}-title` }, definition.title)); - svg.append(svg_element("desc", { id: `${prefix}-description` }, model.detail)); - model.rows.forEach((row, index) => { - const y = 32 + index * 120; - svg.append(svg_element("text", { x: 20, y, class: "coop-load-axis" }, row.label)); - for (const group of row.groups || []) { - const [x] = position(row.id, group.start); - svg.append(svg_element("rect", { x: x - 5, y: y + 20, width: group.count * 42 - 1, height: 70, rx: 7, class: "coop-load-card" })); - svg.append(svg_element("text", { x, y: y + 35, class: "coop-load-thread-label" }, group.label)); - } - for (let cell = 0; cell < row.count; ++cell) { - const [x, cell_y] = position(row.id, cell); - svg.append(svg_element("rect", { x, y: cell_y, width: 32, height: 32, rx: 4, class: "coop-load-slot", "data-row": row.id, "data-index": cell })); - if (!row.groups?.length) svg.append(svg_element("text", { x: x + 16, y: cell_y - 8, class: "coop-load-address" }, cell)); - } - }); - svg.append(svg_element("g", { "data-links": "", "aria-hidden": "true" })); - token_elements = new Map(); - render_frame(); - } - - // Keep token elements by ID so a value retains its identity and animates - // from its previous position when the phase changes. - function render_frame() { - const frame = model.phases[phase]; - const live_ids = new Set(frame.tokens.map((token) => String(token.id))); - if (!live_ids.has(selected)) selected = frame.tokens.length ? String(frame.tokens[0].id) : null; - for (const [id, element] of token_elements) { - if (!live_ids.has(id)) { - element.style.opacity = "0"; - element.setAttribute("tabindex", "-1"); - element.setAttribute("aria-hidden", "true"); - element.style.pointerEvents = "none"; - } - } - for (const token of frame.tokens) { - const id = String(token.id); - let element = token_elements.get(id); - if (!element) { - element = svg_element("g", { class: "coop-load-token coop-explorer-token", role: "button", "data-token-id": id }); - element.append(svg_element("rect", { x: -3, y: -3, width: 38, height: 38, rx: 7, class: "coop-load-token-outline" })); - element.append(svg_element("rect", { width: 32, height: 32, rx: 4, class: "coop-load-token-body" })); - element.append(svg_element("text", { x: 16, y: 21, "text-anchor": "middle" })); - const origin = token.from || token; - const [x, y] = position(origin.row, origin.index); - element.style.transform = `translate(${x}px, ${y}px)`; - svg.append(element); - token_elements.set(id, element); - // Establish an entry point before transitioning a newly formed partial. - element.getBoundingClientRect(); - } - const [x, y] = position(token.row, token.index); - element.style.transform = `translate(${x}px, ${y}px)`; - element.style.opacity = token.muted ? "0.45" : "1"; - element.style.pointerEvents = "auto"; - element.style.setProperty("--coop-token-color", colors[(token.color ?? token.index) % colors.length]); - element.querySelector("text").textContent = token.label; - element.dataset.row = token.row; - element.dataset.index = token.index; - element.dataset.value = token.value ?? token.label; - element.dataset.selected = id === selected; - element.setAttribute("tabindex", id === selected ? "0" : "-1"); - element.setAttribute("aria-pressed", id === selected); - element.setAttribute("aria-label", token.detail || String(token.label)); - element.removeAttribute("aria-hidden"); - } - const links = query("[data-links]"); - links.replaceChildren(...(frame.links || []).map((link) => { - const [x1, y1] = position(link.from.row, link.from.index); - const [x2, y2] = position(link.to.row, link.to.index); - return svg_element("path", { d: `M ${x1 + 16} ${y1 + 16} L ${x2 + 16} ${y2 + 16}`, class: "coop-load-path" }); - })); - const token = frame.tokens.find((token) => String(token.id) === selected); - query("[data-inspection]").textContent = token ? token.detail || String(token.label) : "No defined values at this stage."; - for (const button of query(".coop-load-phases").children) { - if (Number(button.dataset.phase) === phase) button.setAttribute("aria-current", "step"); - else button.removeAttribute("aria-current"); - } - query("[data-status]").setAttribute("aria-live", playing ? "off" : "polite"); - query("[data-status]").textContent = `Step ${phase + 1} of ${model.phases.length}. ${frame.description}`; - query("[data-action=play]").textContent = playing ? "Pause" : "Play"; - query("[data-action=play]").setAttribute("aria-label", playing ? "Pause animation" : "Play animation"); - root.dataset.phase = phase; - root.dataset.algorithm = state.algorithm; - } - - // Run playback only while the user has enabled it and can see the page. - function sync_timer() { - window.clearInterval(timer); - if (playing && visible && !document.hidden) { - timer = window.setInterval(() => { phase = (phase + 1) % model.phases.length; render_frame(); }, 1600); - } - } - - root.addEventListener("click", (event) => { - const token = event.target.closest("[data-token-id]"); - if (token) { selected = token.dataset.tokenId; render_frame(); return; } - const button = event.target.closest("button"); - if (!button) return; - if (button.dataset.algorithm) { - const algorithm = button.dataset.algorithm; - state.algorithm = algorithm; - phase = 0; - rebuild(); - query(`[data-algorithm="${algorithm}"]`)?.focus(); - write_url(); - } else if (button.dataset.phase !== undefined) { - phase = Number(button.dataset.phase); - playing = false; - } else if (button.dataset.action === "play") playing = !playing; - else { - playing = false; - phase = button.dataset.action === "reset" ? 0 : (phase + (button.dataset.action === "next" ? 1 : model.phases.length - 1)) % model.phases.length; - } - render_frame(); - sync_timer(); - }); - root.addEventListener("change", (event) => { - const id = event.target.dataset.control; - if (!id) return; - state[id] = event.target.value; - phase = 0; - rebuild(); - query(`[data-control="${id}"]`)?.focus(); - write_url(); - sync_timer(); - }); - root.addEventListener("keydown", (event) => { - const element = event.target.closest("[data-token-id]"); - if (!element) return; - const ids = model.phases[phase].tokens.map((token) => String(token.id)); - let index = ids.indexOf(element.dataset.tokenId); - if (["ArrowRight", "ArrowDown"].includes(event.key)) index = (index + 1) % ids.length; - else if (["ArrowLeft", "ArrowUp"].includes(event.key)) index = (index + ids.length - 1) % ids.length; - else if (event.key === "Home") index = 0; - else if (event.key === "End") index = ids.length - 1; - else if (!["Enter", " "].includes(event.key)) return; - event.preventDefault(); - selected = ids[index]; - render_frame(); - token_elements.get(selected).focus(); - }); - - read_url(); - rebuild(); - host.replaceChildren(root); - host.closest(".coop-visualization")?.classList.add("coop-load-ready"); - window.addEventListener("popstate", () => { read_url(); rebuild(); sync_timer(); }); - document.addEventListener("visibilitychange", sync_timer); - reduced_motion.addEventListener("change", () => { - if (reduced_motion.matches) { playing = false; render_frame(); sync_timer(); } - }); - if (typeof IntersectionObserver === "function") { - const observer = new IntersectionObserver(([entry]) => { visible = entry.isIntersecting; sync_timer(); }); - observer.observe(host); - } - sync_timer(); - } - - function initialize(name) { - document.querySelectorAll(`[data-coop-explorer="${name}"]`).forEach((host, index) => { - if (host.dataset.mounted) return; - try { - mount(host, name, definitions.get(name), index); - host.dataset.mounted = "true"; - } catch (error) { - console.warn(`The ${name} explorer could not be initialized.`, error); - } - }); - } - - window.CoopExplorer = { - register(name, definition) { - definitions.set(name, definition); - if (document.readyState === "loading") document.addEventListener("DOMContentLoaded", () => initialize(name)); - else initialize(name); - }, - }; -})(); diff --git a/docs/_static/coop-load.css b/docs/_static/coop-load.css deleted file mode 100644 index 2f5a20772cac..000000000000 --- a/docs/_static/coop-load.css +++ /dev/null @@ -1,241 +0,0 @@ -/* Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. - * - * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - * - * Adapted from python/cuda_coop/docs/fern/fern/main.css at - * cccl-mirror trentn/dev/cuda-coop 5dba3d36b6eaae48b967d6fa48f9d15e98136000. - */ - -/* Match the developer overview's wide layout. Sphinx includes this stylesheet - * only on pages that embed the explorer. - */ -body { overflow-x: clip; } -.bd-page-width { max-width: 120rem; } -.bd-header .logo__title { - max-width: calc(100vw - 9rem); - overflow: hidden; - text-overflow: ellipsis; - white-space: nowrap; -} -.bd-sidebar-primary, -.bd-sidebar-secondary { flex-basis: var(--pst-sidebar-secondary); } -.bd-main .bd-content .bd-article-container { - max-width: none; - min-width: 0; -} - -.coop-visualization { - min-width: 0; - max-width: 100%; -} - -.coop-load { - --coop-accent: #527f00; - --coop-panel: var(--pst-color-surface, #f3f5f7); - --coop-background: var(--pst-color-background, #fff); - --coop-text: var(--pst-color-text-base, #222); - --coop-muted: var(--pst-color-text-muted, #52606d); - --coop-border: var(--pst-color-border, #cbd2d8); - color: var(--coop-text); - background: var(--coop-background); - border: 1px solid var(--coop-border); - border-radius: 12px; - padding: 20px; - margin: 1.5rem 0; - min-width: 0; - overflow: hidden; -} - -html[data-theme="dark"] .coop-load { - --coop-accent: #93ce32; -} - -.coop-load-heading { - display: flex; - flex-wrap: wrap; - align-items: center; - justify-content: space-between; - gap: 12px; - margin-bottom: 16px; -} - -.coop-load .coop-load-eyebrow { - color: var(--coop-accent); - font-size: 0.8rem; - font-weight: 700; - margin: 0 0 5px; - text-transform: uppercase; - letter-spacing: 0.06em; -} - -.coop-load h2 { - font-size: 1.35rem; - margin: 0; -} - -.coop-load button, -.coop-load select { - font: inherit; - color: var(--coop-text); - background: var(--coop-background); - border: 1px solid var(--coop-border); - border-radius: 6px; - padding: 7px 10px; -} - -.coop-load button { cursor: pointer; } - -.coop-load button:hover { - border-color: var(--coop-accent); -} - -.coop-load button:focus-visible, -.coop-load select:focus-visible, -.coop-load-scroll:focus-visible { - outline: 3px solid var(--coop-accent); - outline-offset: 3px; -} - -.coop-load-transport { - display: flex; - flex-wrap: wrap; - gap: 5px; - font-size: 0.8rem; -} - -.coop-load-transport [data-action="play"] { - color: var(--coop-background); - background: var(--coop-accent); - border-color: var(--coop-accent); - font-weight: 700; - min-width: 60px; -} - -.coop-load-options { - display: grid; - grid-template-columns: repeat(3, minmax(0, 1fr)); - gap: 8px; -} - -.coop-load-options button { - min-width: 0; - text-align: left; - padding: 10px; - background: var(--coop-panel); - line-height: 1.35; -} - -.coop-load-options strong, -.coop-load-options span { - display: block; - overflow-wrap: anywhere; -} - -.coop-load-options strong { font-size: 0.87rem; } -.coop-load-options span { font-size: 0.72rem; margin-top: 5px; color: var(--coop-muted); } - -.coop-load-options [aria-pressed="true"] { - border-color: var(--coop-accent); - box-shadow: inset 3px 0 var(--coop-accent); -} - -.coop-load-settings { - display: flex; - flex-wrap: wrap; - align-items: flex-end; - gap: 12px 16px; - margin-top: 18px; - font-size: 0.8rem; -} - -.coop-load-settings label { display: grid; gap: 5px; } -.coop-load-settings [data-tile-size] { color: var(--coop-muted); padding-bottom: 9px; } -.coop-load .coop-load-detail { margin: 16px 0; font-size: 0.88rem; } - -.coop-load-phases { - display: flex; - flex-wrap: wrap; - gap: 5px; - margin-bottom: 10px; -} - -.coop-load-phases button { - flex: 1 1 100px; - font-size: 0.75rem; - padding: 8px 6px; -} - -.coop-load-phases [aria-current="step"] { - border-color: var(--coop-accent); - color: var(--coop-accent); - font-weight: 700; - box-shadow: inset 0 -2px var(--coop-accent); -} - -.coop-load-scroll { - width: 100%; - overflow-x: auto; - border: 1px solid var(--coop-border); - border-radius: 8px; - background: var(--coop-panel); -} - -.coop-load-diagram { display: block; width: 100%; } -.coop-load-diagram text { font-family: var(--pst-font-family-base, sans-serif); } -.coop-load-axis { fill: var(--coop-text); font-size: 13px; font-weight: 600; } -.coop-load-address { fill: var(--coop-muted); font-size: 10px; text-anchor: middle; } -.coop-load-card { fill: var(--coop-background); stroke: var(--coop-border); } -.coop-load-thread-label { fill: var(--coop-muted); font-size: 12px; } -.coop-load-warp-label { fill: var(--coop-muted); font-size: 11px; } -.coop-load-slot { fill: var(--coop-background); stroke: var(--coop-border); } -.coop-load-shared-slot { stroke-dasharray: 3 2; } -.coop-load-vector { fill: none; stroke: var(--coop-accent); stroke-width: 2px; } -.coop-load-path { fill: none; stroke: var(--coop-accent); stroke-width: 2px; stroke-dasharray: 5 5; opacity: 0.65; } -.coop-load-token { cursor: pointer; transition: transform 650ms ease; } -.coop-load-token-body { fill: var(--coop-token-color); } -.coop-load-token text { fill: #111; font-weight: 700; font-size: 12px; pointer-events: none; } -.coop-load-token-outline { fill: none; stroke: var(--coop-text); stroke-width: 2px; opacity: 0; } -.coop-load-token[data-selected="true"] .coop-load-token-outline, -.coop-load-token:focus .coop-load-token-outline { opacity: 1; } -.coop-load-token:focus { outline: none; } -.coop-explorer-token { transition: transform 650ms ease, opacity 350ms ease; } - -.coop-load .coop-load-status { - font-size: 0.85rem; - margin: 12px 0; - border-left: 3px solid var(--coop-accent); - padding-left: 10px; - min-height: 2.8em; -} - -.coop-load-inspection { - display: grid; - grid-template-columns: 1fr 1fr; - gap: 10px; -} - -.coop-load-inspection > div { padding: 12px; border-radius: 7px; background: var(--coop-panel); } -.coop-load-inspection h3 { font-size: 0.8rem; margin: 0 0 6px; color: var(--coop-muted); } -.coop-load-inspection p { font-size: 0.85rem; margin: 0; overflow-wrap: anywhere; } -.coop-load-notes { font-size: 0.8rem; margin: 14px 0; } -.coop-load-notes p { margin: 6px 0; } -.coop-load .coop-load-caption { font-size: 0.75rem; color: var(--coop-muted); margin-bottom: 0; } - -@media screen { - .coop-load-ready .coop-visualization-fallback { display: none; } -} - -@media (max-width: 600px) { - .coop-load { padding: 12px; } - .coop-load-options { grid-template-columns: repeat(2, minmax(0, 1fr)); } - .coop-load-inspection { grid-template-columns: 1fr; } -} - -@media (prefers-reduced-motion: reduce) { - .coop-load-token, .coop-explorer-token { transition: none; } -} - -@media print { - [data-coop-load], [data-coop-explorer] { display: none; } - .coop-visualization-fallback { display: block; } -} diff --git a/docs/_static/coop-load.js b/docs/_static/coop-load.js deleted file mode 100644 index 45cbfabc8238..000000000000 --- a/docs/_static/coop-load.js +++ /dev/null @@ -1,397 +0,0 @@ -// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -// -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -// Adapted from python/cuda_coop/docs/fern/fern/components/CooperativeDataMotion.tsx -// at cccl-mirror trentn/dev/cuda-coop 5dba3d36b6eaae48b967d6fa48f9d15e98136000. - -(() => { - "use strict"; - - // Use a small teaching tile; these sizes do not describe a CUDA launch. - // Each value keeps its tile index as it moves between owners. - const threads = 8; - const warp_threads = 4; - const colors = ["#76b900", "#00a9ce", "#f6b21a", "#d6538f", "#9d80d7", "#e86f18", "#2a9d8f", "#8198a7"]; - const algorithms = { - direct: { - label: "Direct", tag: "Blocked scalar loads", access: "blocked", result: "blocked", - detail: "Each thread loads consecutive items. At each scalar issue step, adjacent threads access addresses separated by the items-per-thread count.", - scratch: "No exchange scratch.", caution: "Larger per-thread counts can reduce memory-transaction utilization.", - }, - striped: { - label: "Striped", tag: "Striped loads and registers", access: "striped", result: "striped", - detail: "Adjacent threads load adjacent addresses at each issue step. Values remain distributed in stripes across the threads.", - scratch: "No exchange scratch.", caution: "A following operation that expects blocked ownership needs an explicit conversion.", - }, - vectorize: { - label: "Vectorize", tag: "Blocked vector candidates", access: "blocked", result: "blocked", - detail: "Each thread owns consecutive items that can be grouped into vector loads. The final register ownership matches Direct.", - scratch: "No exchange scratch.", caution: "Vectorization depends on the type, pointer, alignment, and item count. The diagram does not promise a particular instruction.", - }, - transpose: { - label: "Transpose", tag: "Striped loads, blocked registers", access: "striped", result: "blocked", exchange: true, - detail: "Threads first load striped values. A block-wide shared-memory exchange redistributes them into consecutive per-thread registers.", - scratch: "Shared scratch for the block exchange; actual storage can include padding.", - caution: "The exchange adds shared-memory access and synchronization.", - }, - warp_transpose: { - label: "Warp transpose", tag: "Warp-striped loads and exchange", access: "warp-striped", result: "blocked", exchange: true, - detail: "Each warp loads its own tile in stripes, then exchanges values within that warp to produce blocked registers.", - scratch: "Separate shared scratch for every warp in the block; padding is omitted here.", - caution: "The real block size must be a multiple of the 32-lane physical warp.", - }, - warp_transpose_timesliced: { - label: "Warp transpose, timesliced", tag: "Warps reuse shared scratch", access: "warp-striped", result: "blocked", exchange: true, timesliced: true, - detail: "All warps load warp-striped values into registers first. Warps then take turns using the same shared scratch for their exchange.", - scratch: "One warp's shared scratch, reused across exchange rounds; padding is omitted here.", - caution: "Requires complete physical warps. Less scratch comes with additional serialized exchange rounds.", - }, - }; - - function ownership(layout, value, items) { - if (layout === "blocked") return [Math.floor(value / items), value % items]; - if (layout === "striped") return [value % threads, Math.floor(value / threads)]; - const warp_size = warp_threads * items; - return [Math.floor(value / warp_size) * warp_threads + value % warp_threads, Math.floor((value % warp_size) / warp_threads)]; - } - - function svg_element(tag, attributes = {}, text = null) { - const element = document.createElementNS("http://www.w3.org/2000/svg", tag); - for (const [name, value] of Object.entries(attributes)) element.setAttribute(name, value); - if (text !== null) element.textContent = text; - return element; - } - - function mount(host, index) { - const wrapper = host.closest(".coop-visualization"); - const reduced_motion = window.matchMedia("(prefers-reduced-motion: reduce)"); - const state = { algorithm: "transpose", items: 2, value: 5, thread: 0, phase: 0, playing: !reduced_motion.matches }; - const prefix = `coop-load-${index}`; - let phases = []; - let geometry; - let tokens = []; - let timer = null; - let visible = true; - const root = document.createElement("section"); - root.className = "coop-load"; - root.setAttribute("aria-label", "Cooperative load explorer"); - root.innerHTML = ` -
-

Memory to registers

Follow a cooperative load

-
- - - - -
-
-
-
- - - - -
-

-
-
- -
-

-
-

Selected value

-

Result registers

-
-

-

Teaching model: eight threads, with four lanes per displayed warp. Physical CUDA warps have 32 lanes. Colors identify each value's original blocked group. The phases illustrate ownership, not instruction timing or transaction counts. Click a value to follow its path; arrow keys move between values.

`; - - const query = (selector) => root.querySelector(selector); - const svg = query("svg.coop-load-diagram"); - const option_grid = query(".coop-load-options"); - for (const [name, option] of Object.entries(algorithms)) { - const button = document.createElement("button"); - button.type = "button"; - button.dataset.algorithm = name; - const title = document.createElement("strong"); - title.textContent = option.label; - const tag = document.createElement("span"); - tag.textContent = option.tag; - button.append(title, tag); - option_grid.append(button); - } - for (let thread = 0; thread < threads; ++thread) query("[data-thread]").add(new Option(`T${thread}`, thread)); - - function read_url() { - const params = new URL(window.location.href).searchParams; - const algorithm = params.get("load-algorithm"); - state.algorithm = Object.hasOwn(algorithms, algorithm) ? algorithm : "transpose"; - const items = Number(params.get("load-items")); - state.items = [1, 2, 4].includes(items) ? items : 2; - for (const [name, fallback, limit] of [["value", 5, threads * state.items], ["thread", 0, threads]]) { - const raw = params.get(`load-${name}`); - const value = raw === null || raw.trim() === "" ? NaN : Number(raw); - state[name] = Number.isInteger(value) && value >= 0 && value < limit ? value : fallback; - } - state.phase = 0; - state.playing = !reduced_motion.matches; - } - - function write_url() { - const url = new URL(window.location.href); - for (const name of ["algorithm", "items", "value", "thread"]) url.searchParams.set(`load-${name}`, state[name]); - // A locally opened HTML build can disallow history updates; the explorer still works. - try { window.history.replaceState(window.history.state, "", url); } catch { /* Keep local controls usable. */ } - } - - function register_point(layout, value, y) { - const [thread, slot] = ownership(layout, value, state.items); - return [geometry.register_x + thread * geometry.thread_step + 8 + slot * 34, y]; - } - - function point(value, location) { - if (location === "memory") return [geometry.memory_x + value * 34, 56]; - if (location === "read") return register_point(algorithms[state.algorithm].access, value, 178); - if (location === "scratch") { - const slot = algorithms[state.algorithm].timesliced ? value % (warp_threads * state.items) : value; - return [geometry.scratch_x + slot * 34, 306]; - } - return register_point(algorithms[state.algorithm].result, value, geometry.output_y); - } - - function add_text(x, y, content, class_name) { - svg.append(svg_element("text", { x, y, class: class_name }, content)); - } - - function draw_registers(y, title, layout) { - add_text(20, y - 44, title, "coop-load-axis"); - for (let thread = 0; thread < threads; ++thread) { - const x = geometry.register_x + thread * geometry.thread_step; - svg.append(svg_element("rect", { x, y: y - 26, width: geometry.card_width, height: 70, rx: 7, class: "coop-load-card" })); - add_text(x + 8, y - 10, `T${thread}`, "coop-load-thread-label"); - } - for (let value = 0; value < threads * state.items; ++value) { - const [x] = register_point(layout, value, y); - svg.append(svg_element("rect", { x, y, width: 28, height: 28, rx: 4, class: "coop-load-slot" })); - } - } - - // Rebuild layout and exchange phases after algorithm or extent changes. - // Animation phases show ownership, not GPU instruction timing. - function build_scene() { - const option = algorithms[state.algorithm]; - const total = threads * state.items; - geometry = { card_width: state.items * 34 + 10, thread_step: state.items * 34 + 24 }; - geometry.width = Math.max(850, threads * geometry.thread_step + 32); - geometry.register_x = (geometry.width - (threads * geometry.thread_step - 14)) / 2; - geometry.memory_x = (geometry.width - (total * 34 - 6)) / 2; - geometry.scratch_x = (geometry.width - ((option.timesliced ? warp_threads * state.items : total) * 34 - 6)) / 2; - geometry.output_y = option.exchange ? 438 : 306; - phases = [ - { label: "Memory", positions: ["memory", "memory"], description: "The full input tile starts in consecutive memory locations." }, - { label: "Load into registers", positions: ["read", "read"], description: `All displayed threads have loaded their ${option.access} values into registers.` }, - ]; - if (option.timesliced) { - phases.push( - { label: "Warp 0 exchange", positions: ["scratch", "read"], description: "Displayed warp 0 uses shared scratch. Warp 1 keeps its already-loaded values in registers." }, - { label: "Warp 1 exchange", positions: ["output", "scratch"], description: "Warp 0 has completed. Displayed warp 1 reuses the same shared scratch for its exchange." }, - ); - } else if (option.exchange) { - phases.push({ label: "Shared exchange", positions: ["scratch", "scratch"], description: "Loaded values pass through shared scratch at their logical positions. Storage padding is omitted." }); - } - phases.push({ label: "Result registers", positions: ["output", "output"], description: `${option.result === "blocked" ? "Consecutive" : "Striped"} per-thread values are ready for the next operation.` }); - query(".coop-load-phases").replaceChildren(...phases.map((phase, index) => { - const button = document.createElement("button"); - button.type = "button"; - button.dataset.phase = index; - button.textContent = `${index + 1}. ${phase.label}`; - return button; - })); - query("[data-value]").replaceChildren(...Array.from({ length: total }, (_, value) => new Option(value, value))); - svg.replaceChildren(); - svg.setAttribute("viewBox", `0 0 ${geometry.width} ${geometry.output_y + 60}`); - svg.style.minWidth = `${geometry.width * 0.8}px`; - svg.append(svg_element("title", { id: `${prefix}-title` }, `${option.label}: memory to register ownership`)); - svg.append(svg_element("desc", { id: `${prefix}-description` }, `${total} values across eight illustrative threads, ${state.items} per thread. The next paragraphs describe the selected value and result thread.`)); - add_text(20, 24, "Global memory · logical item index", "coop-load-axis"); - for (let value = 0; value < total; ++value) { - const [x, y] = point(value, "memory"); - svg.append(svg_element("rect", { x, y, width: 28, height: 28, rx: 4, class: "coop-load-slot" })); - add_text(x + 14, y - 6, value, "coop-load-address"); - } - if (state.algorithm === "vectorize" && state.items > 1) { - for (let thread = 0; thread < threads; ++thread) { - svg.append(svg_element("rect", { x: geometry.memory_x + thread * state.items * 34 - 3, y: 52, width: state.items * 34, height: 36, rx: 5, class: "coop-load-vector" })); - } - } - draw_registers(178, `Loaded registers · ${option.access} ownership`, option.access); - if (option.exchange) { - add_text(20, 263, option.timesliced ? "Shared scratch · reused by one warp at a time" : "Shared scratch · logical positions (padding omitted)", "coop-load-axis"); - const count = option.timesliced ? warp_threads * state.items : total; - for (let slot = 0; slot < count; ++slot) { - const x = geometry.scratch_x + slot * 34; - svg.append(svg_element("rect", { x, y: 306, width: 28, height: 28, rx: 4, class: "coop-load-slot coop-load-shared-slot" })); - add_text(x + 14, 298, slot, "coop-load-address"); - } - } - draw_registers(geometry.output_y, `Result registers · ${option.result} ownership`, option.result); - if (option.access === "warp-striped") { - for (let warp = 0; warp < 2; ++warp) { - add_text(geometry.register_x + warp * warp_threads * geometry.thread_step + 8, 229, `Displayed warp ${warp} · T${warp * 4}–T${warp * 4 + 3}`, "coop-load-warp-label"); - } - } - svg.append(svg_element("path", { class: "coop-load-path", "aria-hidden": "true" })); - tokens = Array.from({ length: total }, (_, value) => { - const [owner_thread, owner_slot] = ownership(option.result, value, state.items); - const [access_thread, access_slot] = ownership(option.access, value, state.items); - const token = svg_element("g", { - role: "button", "data-value-index": value, "data-owner-thread": owner_thread, - "data-owner-slot": owner_slot, "data-access-thread": access_thread, "data-access-slot": access_slot, - "aria-label": `Value ${value}: loaded by T${access_thread}, slot ${access_slot}; result T${owner_thread}, slot ${owner_slot}`, - class: "coop-load-token", - }); - token.style.setProperty("--coop-token-color", colors[Math.floor(value / state.items)]); - token.append(svg_element("rect", { x: -3, y: -3, width: 34, height: 34, rx: 7, class: "coop-load-token-outline" })); - token.append(svg_element("rect", { width: 28, height: 28, rx: 4, class: "coop-load-token-body" })); - token.append(svg_element("text", { x: 14, y: 19, "text-anchor": "middle" }, value)); - svg.append(token); - return token; - }); - query("[data-detail]").textContent = option.detail; - query("[data-scratch]").textContent = option.scratch; - query("[data-caution]").textContent = state.algorithm === "vectorize" && state.items === 1 ? "One item per thread has no multi-item vector bundle in this illustration. All three ownership layouts coincide." : option.caution; - query("[data-tile-size]").textContent = `8 threads × ${state.items} = ${total} values`; - for (const button of option_grid.children) button.setAttribute("aria-pressed", button.dataset.algorithm === state.algorithm); - } - - function render_frame() { - const option = algorithms[state.algorithm]; - const phase = phases[state.phase]; - for (const button of query(".coop-load-phases").children) { - if (Number(button.dataset.phase) === state.phase) button.setAttribute("aria-current", "step"); - else button.removeAttribute("aria-current"); - } - tokens.forEach((token, value) => { - const warp = Math.floor(value / (warp_threads * state.items)); - const location = phase.positions[warp]; - const [x, y] = point(value, location); - token.style.transform = `translate(${x}px, ${y}px)`; - token.dataset.location = location; - token.dataset.selected = value === state.value; - token.setAttribute("aria-pressed", value === state.value); - token.setAttribute("tabindex", value === state.value ? "0" : "-1"); - }); - const route = ["memory", "read", ...(option.exchange ? ["scratch"] : []), "output"].map((location) => point(state.value, location)); - query(".coop-load-path").setAttribute("d", route.map(([x, y], index) => `${index ? "L" : "M"} ${x + 14} ${y + 14}`).join(" ")); - const [access_thread, access_slot] = ownership(option.access, state.value, state.items); - const [owner_thread, owner_slot] = ownership(option.result, state.value, state.items); - query("[data-inspection]").textContent = `Value ${state.value}: memory[${state.value}] → T${access_thread}, slot ${access_slot}${option.exchange ? " → shared scratch" : ""} → T${owner_thread}, slot ${owner_slot}.`; - const values = Array.from({ length: state.items }, (_, slot) => tokens.findIndex((token) => Number(token.dataset.ownerThread) === state.thread && Number(token.dataset.ownerSlot) === slot)); - query("[data-thread-inspection]").textContent = `T${state.thread} receives [${values.join(", ")}] in slots 0${state.items > 1 ? `–${state.items - 1}` : ""}.`; - query("[data-status]").setAttribute("aria-live", state.playing ? "off" : "polite"); - query("[data-status]").textContent = `Step ${state.phase + 1} of ${phases.length}. ${phase.description}`; - query("[data-items]").value = state.items; - query("[data-value]").value = state.value; - query("[data-thread]").value = state.thread; - query("[data-action=play]").textContent = state.playing ? "Pause" : "Play"; - query("[data-action=play]").setAttribute("aria-label", state.playing ? "Pause animation" : "Play animation"); - } - - function sync_timer() { - window.clearInterval(timer); - timer = null; - if (state.playing && visible && !document.hidden) { - timer = window.setInterval(() => { - state.phase = (state.phase + 1) % phases.length; - render_frame(); - }, 1600); - } - } - - function choose_value(value, focus = false) { - state.value = value; - state.thread = ownership(algorithms[state.algorithm].result, value, state.items)[0]; - render_frame(); - write_url(); - if (focus) tokens[value].focus(); - } - - root.addEventListener("click", (event) => { - const token = event.target.closest("[data-value-index]"); - if (token) { choose_value(Number(token.dataset.valueIndex)); return; } - const button = event.target.closest("button"); - if (!button) return; - if (button.dataset.algorithm) { - state.algorithm = button.dataset.algorithm; - state.phase = 0; - build_scene(); - write_url(); - } else if (button.dataset.phase !== undefined) { - state.phase = Number(button.dataset.phase); - state.playing = false; - } else if (button.dataset.action === "play") { - state.playing = !state.playing; - } else { - state.playing = false; - if (button.dataset.action === "reset") state.phase = 0; - else state.phase = (state.phase + (button.dataset.action === "next" ? 1 : phases.length - 1)) % phases.length; - } - render_frame(); - sync_timer(); - }); - root.addEventListener("change", (event) => { - if (event.target.matches("[data-items]")) { - state.items = Number(event.target.value); - state.value = Math.min(state.value, threads * state.items - 1); - state.phase = 0; - build_scene(); - } else if (event.target.matches("[data-value]")) { - choose_value(Number(event.target.value)); - return; - } else if (event.target.matches("[data-thread]")) state.thread = Number(event.target.value); - render_frame(); - sync_timer(); - write_url(); - }); - root.addEventListener("keydown", (event) => { - const token = event.target.closest("[data-value-index]"); - if (!token) return; - const value = Number(token.dataset.valueIndex); - let next = value; - if (["ArrowRight", "ArrowDown"].includes(event.key)) next = (value + 1) % tokens.length; - else if (["ArrowLeft", "ArrowUp"].includes(event.key)) next = (value + tokens.length - 1) % tokens.length; - else if (event.key === "Home") next = 0; - else if (event.key === "End") next = tokens.length - 1; - else if (event.key !== "Enter" && event.key !== " ") return; - event.preventDefault(); - choose_value(next, true); - }); - - read_url(); - build_scene(); - render_frame(); - host.replaceChildren(root); - wrapper?.classList.add("coop-load-ready"); - window.addEventListener("popstate", () => { read_url(); build_scene(); render_frame(); sync_timer(); }); - document.addEventListener("visibilitychange", sync_timer); - reduced_motion.addEventListener("change", () => { - if (reduced_motion.matches) { state.playing = false; render_frame(); sync_timer(); } - }); - if (typeof IntersectionObserver === "function") { - const observer = new IntersectionObserver(([entry]) => { visible = entry.isIntersecting; sync_timer(); }); - observer.observe(host); - } - sync_timer(); - } - - function initialize() { - document.querySelectorAll("[data-coop-load]").forEach((host, index) => { - try { mount(host, index); } catch (error) { - // The adjacent static explanation remains available when enhancement fails. - console.warn("The cooperative load explorer could not be initialized.", error); - } - }); - } - if (document.readyState === "loading") document.addEventListener("DOMContentLoaded", initialize); - else initialize(); -})(); diff --git a/docs/_static/coop-movement.js b/docs/_static/coop-movement.js deleted file mode 100644 index ee387459128f..000000000000 --- a/docs/_static/coop-movement.js +++ /dev/null @@ -1,113 +0,0 @@ -// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -// -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -// Ownership models adapted from CooperativeDataMotion.tsx in -// python/cuda_coop/docs/fern/fern/components at cccl-mirror -// trentn/dev/cuda-coop 5dba3d36b6eaae48b967d6fa48f9d15e98136000. - -(() => { - "use strict"; - - const threads = 8; - const warp_threads = 4; - const item_control = {id: "items", label: "Items per thread", value: "2", choices: ["1", "2", "4"]}; - - function owner(layout, value, items) { - if (layout === "blocked") return [Math.floor(value / items), value % items]; - if (layout === "striped") return [value % threads, Math.floor(value / threads)]; - const warp_size = warp_threads * items; - return [Math.floor(value / warp_size) * warp_threads + value % warp_threads, Math.floor((value % warp_size) / warp_threads)]; - } - - function register_index(layout, value, items) { - const [thread, slot] = owner(layout, value, items); - return thread * items + slot; - } - - function thread_row(id, label, items) { - return { - id, label, count: threads * items, - groups: Array.from({length: threads}, (_, thread) => ({start: thread * items, count: items, label: `T${thread}`})), - }; - } - - function owner_text(layout, value, items) { - const [thread, slot] = owner(layout, value, items); - return `T${thread}, slot ${slot}`; - } - - const stores = [ - {id: "direct", label: "Direct", tag: "Blocked scalar stores", input: "blocked", access: "blocked", - detail: "Each thread writes its consecutive values directly. Adjacent threads address locations separated by the items-per-thread count at each scalar issue step.", - note: "No exchange scratch. Larger blocked strides can reduce memory-transaction utilization."}, - {id: "striped", label: "Striped", tag: "Striped registers and stores", input: "striped", access: "striped", - detail: "Each thread already owns striped values. Adjacent threads write adjacent addresses at each item slot, without exchanging values.", - note: "No exchange scratch. The input payload must already have the intended striped ownership."}, - {id: "vectorize", label: "Vectorize", tag: "Blocked vector candidates", input: "blocked", access: "blocked", - detail: "Each thread writes its consecutive values using wider memory accesses when the pointer, type, alignment, and item count permit vectorization.", - note: "No exchange scratch. The diagram shows candidate bundles, not a guarantee of vector instructions."}, - {id: "transpose", label: "Transpose", tag: "Blocked input, striped stores", input: "blocked", access: "striped", exchange: true, - detail: "A shared-memory exchange converts blocked input into striped writer registers. Adjacent threads can then write adjacent memory locations.", - note: "The block exchange uses shared scratch and synchronization. Scratch padding is omitted from the diagram."}, - {id: "warp_transpose", label: "Warp transpose", tag: "Warp-striped writer registers", input: "blocked", access: "warp-striped", exchange: true, - detail: "Each warp exchanges its blocked values into a warp-striped arrangement, then writes its own contiguous memory tile.", - note: "All warps have exchange scratch. The real block size must be a multiple of the 32-lane physical warp."}, - {id: "warp_transpose_timesliced", label: "Warp transpose, timesliced", tag: "Warps reuse exchange scratch", input: "blocked", access: "warp-striped", exchange: true, timesliced: true, - detail: "Warps take turns exchanging through one shared scratch region. The rearranged registers then feed warp-striped stores.", - note: "One warp of exchange scratch is reused in serialized rounds. The real block must contain complete physical warps."}, - ]; - - // Describe each value's owners before and after the Store exchange. - // Stable token IDs let the shared renderer follow values across phases. - function build_store(state) { - const option = stores.find((entry) => entry.id === state.algorithm); - const items = Number(state.items); - const count = threads * items; - const rows = [thread_row("input", `Payload · ${option.input} input`, items)]; - if (option.exchange) rows.push({id: "scratch", label: option.timesliced ? "Shared scratch · one warp at a time" : "Shared scratch · logical positions, padding omitted", count: option.timesliced ? warp_threads * items : count}); - rows.push(thread_row("writers", `${option.access} writer registers`, items)); - rows.push({id: "memory", label: "Global memory · logical item index", count}); - - function tokens_for(locations) { - return Array.from({length: count}, (_, value) => { - const warp = Math.floor(value / (warp_threads * items)); - const row = locations[warp]; - let index = value; - if (row === "input") index = register_index(option.input, value, items); - if (row === "writers") index = register_index(option.access, value, items); - if (row === "scratch" && option.timesliced) index = value % (warp_threads * items); - return { - id: `v${value}`, label: String(value), value, row, index, - color: owner(option.input, value, items)[0], - detail: `Value ${value}: input ${owner_text(option.input, value, items)} → writer ${owner_text(option.access, value, items)} → memory[${value}].`, - }; - }); - } - - const phases = [{label: "Input registers", description: `The tile begins in ${option.input} ownership. Each displayed thread has ${items} item${items === 1 ? "" : "s"}.`, tokens: tokens_for(["input", "input"])}]; - if (option.timesliced) { - phases.push( - {label: "Warp 0 exchange", description: "Displayed warp 0 uses shared scratch. Warp 1 retains its blocked input registers.", tokens: tokens_for(["scratch", "input"])}, - {label: "Warp 1 exchange", description: "Warp 0 has warp-striped writer registers. Warp 1 reuses the same scratch for its exchange.", tokens: tokens_for(["writers", "scratch"])}, - ); - } else if (option.exchange) { - phases.push({label: "Shared exchange", description: "Values pass through scratch at their logical positions; readers then take them in the writer layout.", tokens: tokens_for(["scratch", "scratch"])}); - } - phases.push( - {label: "Writer registers", description: `${option.access} registers supply the memory writes. This is an ownership stage, not a barrier or exact instruction schedule.`, tokens: tokens_for(["writers", "writers"])}, - {label: "Memory", description: "The full tile is stored in consecutive memory locations. Each value reaches the address shown by its label.", tokens: tokens_for(["memory", "memory"])}, - ); - return { - detail: option.detail, rows, phases, - notes: [option.note, "Transpose Store algorithms may rearrange the input payload in place. Copy values before Store if they are needed later.", items === 1 ? "With one item per thread, blocked and striped ownership coincide; there is no multi-item vector bundle." : "The values and ownership are illustrative; access patterns do not specify a transaction count."], - summary: `Memory receives values 0–${count - 1} in order. ${option.exchange ? "The exchange changes which thread writes each value." : "Each thread writes directly from the illustrated input ownership."}`, - caption: "Eight illustrative threads; warp variants use two four-lane teaching warps. CUDA physical warps have 32 lanes. Timeslicing serializes exchange scratch use, not the subsequent global-memory store instruction stream.", - }; - } - - window.CoopExplorer.register("store", { - title: "Follow a cooperative store", eyebrow: "Registers to memory", defaultAlgorithm: "transpose", - algorithms: stores.map(({id, label, tag}) => ({id, label, tag})), controls: [item_control], build: build_store, - }); -})(); diff --git a/docs/conf.py b/docs/conf.py index f4bc43dbbbd0..48b1a2f71072 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -8,10 +8,9 @@ # Add extension directory to path sys.path.insert(0, os.path.abspath("_ext")) -# Add Python CCCL packages to path for autodoc. cuda-cccl, cuda-coop, and -# cuda-stf are separate distributions that contribute to the shared ``cuda`` -# namespace. -for _pkg in ("../python/cuda_cccl", "../python/cuda_coop", "../python/cuda_stf"): +# Add Python CCCL packages to path for autodoc. cuda-cccl and cuda-stf are +# separate distributions that both contribute to the shared ``cuda`` namespace. +for _pkg in ("../python/cuda_cccl", "../python/cuda_stf"): python_package_path = os.path.abspath(_pkg) if os.path.exists(python_package_path): sys.path.insert(0, python_package_path) @@ -54,7 +53,6 @@ # "exhale", # Disabled - causing build timeouts, API docs handled by breathe "sphinx_design", # For dropdown, card, and other directives "sphinx_copybutton", - "coop_visualization", # Interactive cooperative primitives with RST fallbacks "nbsphinx", # "rst_processor", # Disabled - breathe handles embed:rst natively "auto_api_generator", # Automatically generate API reference pages from Doxygen XML @@ -213,33 +211,6 @@ autodoc_type_hints = "description" autodoc_type_aliases = { "Operator": "Operator", - **{ - name: name - for name in ( - "BlockExchangeMode", - "BlockGroup", - "BlockLoadStoreAlgorithm", - "CommonNumericScalar", - "CommonShuffleMode", - "CommonThreadDataLike", - "IntegerValue", - "IntegralScalar", - "MemoryGroup", - "PortableNumericScalar", - "PortableThreadDataLike", - "ReduceAlgorithm", - "ReduceOperator", - "ReductionGroup", - "ScalarShuffleMode", - "ScanAlgorithm", - "ScanOperator", - "SignedIntegerScalar", - "ThreadDataLike", - "ValidItems", - "WarpGroup", - "WarpLoadStoreAlgorithm", - ) - }, } # Set Python domain primary for intersphinx @@ -247,9 +218,6 @@ # Mock imports for Python documentation - these modules may not be installed autodoc_mock_imports = [ - # Render real cooperative markers without initializing an optional compiler. - "numba_cuda_mlir", - "cuda.coop.numba_mlir._compiler._activation", "numba", "numba.core", "numba.core.cgutils", diff --git a/docs/python/api_reference.rst b/docs/python/api_reference.rst index 5942065de338..622c9113ddcf 100644 --- a/docs/python/api_reference.rst +++ b/docs/python/api_reference.rst @@ -5,5 +5,4 @@ API Reference :maxdepth: 1 compute_api - coop_api stf_api diff --git a/docs/python/coop.rst b/docs/python/coop.rst deleted file mode 100644 index f3a0f193e129..000000000000 --- a/docs/python/coop.rst +++ /dev/null @@ -1,127 +0,0 @@ -.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -.. -.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -.. _cccl-python-coop: - -``cuda.coop``: Cooperative Group Primitives -=========================================== - -.. toctree:: - :hidden: - :maxdepth: 2 - :caption: Shared concepts - - Overview - coop/concepts - coop/visualizations/index - coop/glossary - coop/faqs - coop/configuration - -.. toctree:: - :hidden: - :maxdepth: 2 - :caption: Numba-CUDA-MLIR - - coop/programming_guide - coop/developer_overview - -``cuda.coop`` brings CCCL's optimized cooperative algorithms to Python GPU -kernels. Use it when threads need to work together, such as loading a tile -of values or arranging data for the next stage of a computation. These -operations run inside a kernel, where you can combine them with your own -code and reuse algorithms maintained and tuned for NVIDIA GPUs. - -The Numba-CUDA-MLIR integration compiles cooperative calls inside Python -GPU kernels. You keep Numba's kernel syntax and launch conventions and -use ``cuda.coop`` for the cooperative parts of the work. - -Get started ------------ - -Start with the :doc:`Numba-CUDA-MLIR Programming Guide -` for complete kernels and host launch code. -Install the extra matching your CUDA major version: - -.. code-block:: console - - python -m pip install "cuda-coop[numba-cuda-mlir-cu13]" - # Use numba-cuda-mlir-cu12 with CUDA 12. - -The guide covers :ref:`requirements and kernel reuse -`. - -Inside a kernel ---------------- - -A cooperative operation acts on a *group* of threads, usually a warp or a -block. Each thread owns part of the group's data, held in ``ThreadData``. -Together, those parts form a tile. - -These calls copy a full tile from one array to another: - -.. code-block:: python - - # Inside a kernel, with `coop` imported from `cuda`: - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - - coop.load(block, source, items, offset=offset) - coop.store(block, destination, items, offset=offset) - -The kernel supplies the tile's ``offset`` and a compile-time -``items_per_thread`` count. With 128 threads and four items per thread, -the block processes 512 values. All threads in the group must reach the -cooperative calls, including when only part of the tile is valid. - -The :doc:`programming concepts ` explain participation, -data layouts, operation results, and temporary storage. - -Common API and compiler extensions ----------------------------------- - -Use ``from cuda import coop`` for the common API. It describes the -cooperative part of a kernel independently of its compiler. - -The qualified ``cuda.coop.numba_mlir`` namespace provides Numba-specific -operands and controls. The :ref:`Numba API comparison -` explains when to use it. - -Import your kernel compiler before ``cuda.coop`` to activate its -integration automatically. For varying import order, use explicit -:ref:`backend registration `. - -Explore further ---------------- - -The :doc:`interactive visualizations ` show -which values each thread owns and how an operation moves them. - -Use the :doc:`API reference ` for signatures and return values, -and :doc:`configuration ` for package and compiler -settings. - -To work on the integration, follow a kernel through the -:doc:`Numba-CUDA-MLIR Developer Guide `. - -.. raw:: html - - - - - - - - - - - - - - - - - - - diff --git a/docs/python/coop/concepts.rst b/docs/python/coop/concepts.rst deleted file mode 100644 index 6fe3b2b57caa..000000000000 --- a/docs/python/coop/concepts.rst +++ /dev/null @@ -1,401 +0,0 @@ -.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -.. -.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -.. _coop-programming-concepts: - -Programming concepts -==================== - -Cooperative operations let a group of threads work on data together. -This page explains the participating threads, their payloads and -layouts, and the rules for results and temporary storage. - -See the :doc:`Numba-CUDA-MLIR Programming Guide ` -for complete kernels and launch examples. - -.. _coop-backend-registration: - -Registering a backend ---------------------- - -When you cannot ensure import order, call :func:`cuda.coop.register` on the -host before compiling kernels to select the backend explicitly: - -.. code-block:: python - - from cuda import coop - - coop.register("numba-cuda-mlir") - - from numba_cuda_mlir import cuda - -Registration loads the backend and installs its compiler hooks, so this -works regardless of whether ``cuda.coop`` or Numba-CUDA-MLIR was imported -first. Repeated calls are safe and return ``None``. The spelling -``"numba_cuda_mlir"`` is also accepted. Registration requires the backend's -dependencies to be installed; it does not install packages. - -For convenience, importing ``cuda.coop`` after ``numba_cuda_mlir`` also -registers the backend automatically. A standalone ``cuda.coop`` import does -not discover or load optional compilers. Explicit registration is useful in -libraries and notebooks where another import may already have loaded -``cuda.coop``. - -Importing the backend namespace also registers it: - -.. code-block:: python - - import cuda.coop.numba_mlir as numba_coop - -Use ``numba_coop`` when mixing common and backend calls. If your program uses -only the backend namespace, you can import it as ``coop`` instead. See the -:ref:`namespace FAQ ` and -:ref:`API comparison `. - - -Kernel API ----------- - -The common root and qualified backend expose matching entry points: - -.. code-block:: python - - from numba_cuda_mlir import cuda, types - - from cuda import coop - - # Inside a Numba-CUDA-MLIR kernel with an items_per_thread argument: - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - tile_items = cuda.blockDim.x * items_per_thread - tile_offset = cuda.blockIdx.x * tile_items - valid_items = count - tile_offset - if valid_items < 0: - valid_items = 0 - elif valid_items > tile_items: - valid_items = tile_items - coop.load( - block, - source, - items, - valid_items=valid_items, - oob_default=0, - offset=tile_offset, - ) - coop.store( - block, - destination, - items, - valid_items=valid_items, - offset=tile_offset, - ) - -For a complete copy kernel, including allocation, launch, tail handling, and -an output check, see the example in :func:`cuda.coop.load`. The API reference -also includes tested examples for :func:`cuda.coop.store`, -:func:`cuda.coop.ThreadData`, and :func:`cuda.coop.TempStorage`. - -Use the qualified namespace when backend-specific types or controls are -required: - -.. code-block:: python - - import cuda.coop.numba_mlir as coop - -Both spellings are compiler markers. Calls must occur in a compatible compiler -context; they are not host-side data movement operations. - - -Groups and thread data ----------------------- - -:func:`cuda.coop.this_block` describes the current CUDA thread block, and -:func:`cuda.coop.this_warp` describes the current 32-thread physical warp. A -physical warp can be partitioned with ``this_warp().group_by(width)`` into -consecutive logical warps of 1, 2, 4, 8, 16, or 32 threads. Load and Store -support all three forms. The enclosing block must contain a multiple of 32 -threads, with no incomplete final physical warp. For a multidimensional block, -threads are linearized in x-major order. Every member of a participating group -must reach its collective; complete sibling logical groups may take different -control-flow paths. - -The common group vocabulary also includes thread, cluster, grid, and mapped -groups of physical warps, but those are not Load or Store targets. -``ThreadGroup`` objects are descriptor-only in this release. ``group_by`` is -compile-time vocabulary for describing a static partition. Runtime query, -membership, and synchronization methods such as ``rank``, ``count``, -``rank_as``, ``count_as``, ``sync``, ``sync_aligned``, and ``is_member`` are -not exposed. - - -Participation and synchronization ---------------------------------- - -Every member of a participating group must reach the same cooperative call. -A branch around a block operation must be uniform across the block; a branch -around a logical-warp operation must be uniform within that logical warp. -Complete sibling logical groups may follow different paths. Warp operations -require a block size divisible by 32, with no incomplete final physical warp. - -Do not put a block Load or Store inside a per-element ``if index < count`` -condition. Use ``valid_items`` to describe the valid prefix while all block -threads participate. An early return by some threads also violates -participation if the remaining threads later execute a block operation. - -A scratch-reuse barrier protects temporary storage. Its presence depends on -the algorithm and storage policy; arrange explicit synchronization wherever -application-owned shared memory requires it. A barrier does not make -divergent participation safe. - - -Per-thread payloads -------------------- - -``coop.ThreadData(items_per_thread)`` gives each participating thread a -fixed-size payload with that many items. Pass ``items_per_thread`` as a -kernel argument; Numba-CUDA-MLIR specializes the kernel for its value. -Common and qualified calls use the same inference rules: an untyped Load -output infers its dtype from the source, and Store combines the destination -dtype with payload writes. Load fills the supplied output in place and -returns ``None``. - -Both namespaces accept ``alignment`` as a compile-time positive power of two -in bytes. It specifies minimum alignment when the compiler materializes -payload storage; ``None`` lets the compiler choose. For example, -``coop.ThreadData(items_per_thread, alignment=16)`` requests at least -16-byte alignment. The backend may use stronger alignment, including for -requests smaller than its minimum allocation alignment. This option does not -assert alignment of source or destination arrays passed to Load or Store. - -Payload slots start uninitialized. Write every slot before reading it. A -partial Load leaves invalid slots unspecified unless you pass -``oob_default``. Initialize those slots after the Load or supply a default -before reading them. :class:`cuda.coop.ThreadDataLike` names the shared -payload interface in type signatures. Protocol compatibility alone does not -make an arbitrary Python object a supported kernel value. - -The payload's ``items_per_thread`` attribute is a compile-time integer and -can be used as a loop bound inside a kernel, including through payload -aliases. - -Supported payload types are signed and unsigned 8-, 16-, 32-, and 64-bit -integers plus 32- and 64-bit floating-point values. Boolean, 16-bit floating -point, complex, and mismatched payload types are rejected before NVRTC -compilation. - - -Load and Store semantics ------------------------- - -The signatures are: - -.. code-block:: python - - load( - group, source, output, /, *, - algorithm="direct", - valid_items=None, - oob_default=None, - offset=None, - temp_storage=None, - ) -> None - - store( - group, destination, value, /, *, - algorithm="direct", - valid_items=None, - offset=None, - temp_storage=None, - ) -> None - -``valid_items`` counts the valid prefix of the selected group tile, not the -number of valid items per thread. A Warp-group tile contains -``group_size * items_per_thread`` elements. Here ``group_size`` is 32 for -``this_warp()`` or the width passed to ``group_by``. The count must be uniform -within that group. With Load, invalid output slots are unspecified unless -``oob_default`` is supplied, even if initialized before Load. A runtime -default must also be uniform within the group. A default is valid only when -``valid_items`` is present. With Store, elements outside the valid prefix -are not written. - -.. warning:: - - ``valid_items`` must satisfy - ``0 <= valid_items <= group_size * items_per_thread``. Static values outside - that range are rejected while planning. Runtime values are checked rather - than saturated; do not rely on CUB's oversized-count behavior. An invalid - runtime value executes a deterministic device trap before narrowing to - CUB's integer parameter, and that trap poisons the current CUDA context. - Clamp grid-stride and tail counts as in the example above. Run intentional - failure probes in disposable processes. For a Warp-group call, a block-wide - remainder is not a valid count: subtract that group's tile origin and clamp - the result to ``[0, group_size * items_per_thread]``. - -``offset`` is an element offset into the source or destination. It is -independent of ``valid_items`` and is not measured in bytes. The value must -be uniform within each participating group; different groups may use -different offsets. Static offsets must be nonnegative; a runtime offset is a -caller-enforced nonnegative precondition. Source and destination arrays must -be one-dimensional and contiguous. Store accepts both scalar values and -multi-item ``ThreadData`` payloads. - -Runtime ``valid_items`` and ``offset`` accept signed integer types through 64 -bits and unsigned integer types through 32 bits. Boolean, floating-point, and -``uint64`` runtime values are rejected. A runtime ``oob_default`` is already -typed by the compiler and must exactly match the Load payload dtype. Ordinary -Python integer and floating-point literals are converted contextually and -range-checked against that dtype before provider generation. - -For a Warp group of width ``group_size``, the compiler first advances the -memory base by ``group_index * (group_size * items_per_thread)`` and then -applies the caller's ``offset``. The group index is the x-major linear -thread rank divided by the group size, so every physical or logical Warp -group in a block addresses a distinct tile. In a multi-block traversal, the -caller offset must also include the block's global tile origin. Do not add -the compiler-provided group origin again. Runtime offsets must leave enough -signed 64-bit range for the last group origin in the block; static offsets -are checked during planning. - -Store payloads must have exactly the destination dtype. Numba-CUDA-MLIR may -promote integer arithmetic even when its operands are 32-bit. Cast a computed -value explicitly before storing it: - -.. code-block:: python - - value = types.int32(source[cuda.threadIdx.x] + 1) - coop.store(block, destination, value, algorithm="direct") - - -Data layouts and algorithms -^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Both common and qualified entry points use the same string algorithm -vocabulary: ``direct``, ``striped``, ``vectorize``, ``transpose``, -``warp_transpose``, and ``warp_transpose_timesliced``. All six work with the -Numba-CUDA-MLIR backend. ``direct`` and ``vectorize`` use blocked ordering, so -each thread owns a contiguous segment of the tile. ``striped`` exposes striped -ordering, where item ``i`` for a thread is separated from its next item by the -block size. The three transpose algorithms use striped memory transactions but -present blocked ``ThreadData`` to the caller. The two warp-transpose variants -perform that reordering within each warp and require a block size divisible by -32. - -Physical and logical Warp Load and Store support ``direct``, ``striped``, -``vectorize``, and ``transpose``. Their layouts follow the same rules at the -selected group width: ``direct`` and ``vectorize`` expose blocked payloads, -``striped`` exposes a striped payload, and ``transpose`` uses striped memory -transactions while exposing a blocked payload. Common and qualified calls -use the same lowercase string selectors. Selectors are normalized to lowercase -underscore-delimited strings. Enum and integer selectors, including ``0``, are -rejected. - -For group size ``G``, items per thread ``K``, thread rank ``t``, and item -index ``i``, blocked order uses tile position ``t * K + i``; striped order -uses ``t + i * G``. The payload has no runtime layout tag that corrects a -mismatched Load/Store pair. - -Store consumes the arrangement associated with its selected algorithm. -Transpose Store algorithms may rearrange the input payload in place, following -CUB's behavior. Reload or reinitialize the payload before using its previous -arrangement again. - - -Temporary storage ------------------ - -Block Load and Store accept an optional ``TempStorage`` descriptor. - -Scratch belongs to one block during its kernel execution. A descriptor does -not carry data between blocks or kernel launches. - -For example: - -.. code-block:: python - - scratch = coop.TempStorage( - size_in_bytes=None, - alignment=None, - auto_sync=False, - sharing="shared", - ) - -Only ``size_in_bytes`` may be positional; ``alignment``, ``auto_sync``, and -``sharing`` are keyword-only. As with ``ThreadData``, ``alignment=None`` lets -the compiler choose. An explicit positive power of two sets minimum alignment -in bytes. The planner may strengthen it to satisfy every primitive using the -storage. Integer-like values implementing ``__index__`` are accepted. An -explicit ``size_in_bytes`` must still be large enough for the planned storage. - -For block, physical Warp, and logical Warp operations, ``direct``, -``striped``, and ``vectorize`` are storage-free. They default-construct the -CUB primitive, report zero temporary bytes, and emit no shared-memory -allocation, storage pointer, or synchronization barrier. For a block call, -an explicit descriptor, including an unsized descriptor, is validated as -compile-time vocabulary but does not change code generation for those -algorithms. Construct ``TempStorage`` inside the kernel; the current -Numba-CUDA-MLIR frontend does not resolve module-global storage descriptors. -A descriptor may be passed to a device function that Numba-CUDA-MLIR inlines -into the kernel, which is the default, but it cannot cross into a separately -compiled device function. - -The block ``transpose``, ``warp_transpose``, and -``warp_transpose_timesliced`` use CUB temporary storage. Without a -descriptor, the compiler allocates the specialization's exact storage and -inserts a block reuse barrier. An explicit descriptor selects the sharing -policy and may request capacity and alignment. The compiler selects static -or dynamic backing from the total requirement, as described below. The -provider remains authoritative for the required byte count and alignment, -and the backend validates the descriptor against the concrete lowering plan. - -Sharing selects only the slice layout: ``sharing="shared"`` overlaps every -call that passes the same descriptor on one region, while -``sharing="exclusive"`` gives each call site its own slice. A call site -inside a loop reuses its slice under either layout, so ``auto_sync`` is -independent of ``sharing`` and defaults to ``False`` for both. - -A descriptor names one region. Distinct descriptors and compiler-owned storage -never alias each other. With ``auto_sync=True``, the compiler appends -``cuda.syncthreads()`` immediately after every block call that consumes the -storage, including the last one. Warp calls do not accept a descriptor. The -compiler never inserts a barrier before a call. The trailing barrier orders -only reuse of the temporary storage, and storage-free calls get none. Arrange -explicit barriers for the kernel's own shared-memory traffic. With -``auto_sync=False`` (the default), the compiler adds no reuse barrier. The -caller issues ``cuda.syncthreads()`` between consecutive uses of the -descriptor, and a call site inside a loop counts as a reuse on every -iteration. Compiler-owned storage always synchronizes. - -The compiler stages every descriptor and every compiler-owned requirement of -a kernel into one shared-memory backing. Up to 48 KiB, it uses static shared -memory without querying the device. Larger requests, from explicit -``size_in_bytes`` or implicit requirements, trigger a query of the device's -default and opt-in limits. Backing above the default limit uses dynamic shared -memory, and the launch reserves its exact byte count within the opt-in limit. -Supported Numba-CUDA-MLIR releases do not separate static and dynamic shared -allocations reliably. A kernel using cooperative temporary storage must not -also declare a zero-sized or runtime-sized ``cuda.shared.array``. When -cooperative backing becomes dynamic, user static shared arrays are also -unsupported. Keep both user arrays and cooperative backing static, or move -the user data out of shared memory. Storage-free operations do not add this -restriction. - -With ``auto_sync=False``, a descriptor must originate from exactly one -constructor site. Selecting between multiple manual-sync constructors is -unsupported: the compiler cannot prove that caller barriers protect the -merged region, even when a particular program supplies sufficient barriers. - -Cooperative calls in device helpers must be inlined into the kernel; use -``@cuda.jit(device=True, inline="always")`` when selecting the helper's -policy explicitly. Standalone collective helpers and collectives inside -standalone callbacks are unsupported. Values from ``literal_unroll`` cannot -determine cooperative payload extents, group dimensions, selectors, or -descriptor constructor arguments. Use separate calls with explicit constants, -or an ordinary loop with one fixed cooperative shape. The restriction does -not apply to an unrelated ``literal_unroll`` loop. - -Warp ``transpose`` uses compiler-owned storage with one disjoint slice per -physical or logical group. The compiler inserts ``syncwarp`` with the exact -logical-group mask. Both the common and qualified APIs reject explicit -``TempStorage`` for every Warp Load and Store algorithm, including the -storage-free modes. diff --git a/docs/python/coop/configuration.rst b/docs/python/coop/configuration.rst deleted file mode 100644 index 06783677489a..000000000000 --- a/docs/python/coop/configuration.rst +++ /dev/null @@ -1,143 +0,0 @@ -.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -.. -.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -.. _coop-configuration: - -Configuration -============= - -Installation ------------- - -Install ``cuda-coop`` without adding Python package dependencies: - -.. code-block:: console - - python -m pip install cuda-coop - -The wheel includes the common API, every shipped DSL integration (including -``cuda.coop.numba_mlir``), type declarations, and a matching bundle of CUB, -Thrust, and libcu++ headers. The base install declares no Python -package dependencies. You can import ``cuda.coop`` without a compiler or GPU; -using an integration requires its backend dependencies to be installed. - -For Numba-CUDA-MLIR, install the extra matching your CUDA major version: - -.. code-block:: console - - python -m pip install "cuda-coop[numba-cuda-mlir-cu13]" - # Use numba-cuda-mlir-cu12 with CUDA 12. - -Both commands install the same ``cuda-coop`` wheel with the same DSL -integrations. The extra only adds the dependency requirements declared in -``pyproject.toml`` so pip installs the supported Numba-CUDA-MLIR stack for -the selected CUDA major version. The current integration requires -``numba-cuda-mlir>=0.5.0,<0.6``. -Installing an extra does not register a backend in a running Python process; -see :ref:`installation versus registration `. - -Installed-wheel compilation uses the bundled CCCL headers. Development from a -CCCL source checkout uses its matching headers. ``CUDA_COOP_CCCL_ROOT`` can -select another source checkout or ``cuda-coop`` header bundle. - - -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`` below is separate. - -Runtime environment variables ------------------------------ - -For the Boolean switches below, *truthy* means any value other than the -empty string, ``0``, ``false``, ``no``, or ``off``. For example, ``1``, -``true``, ``yes``, and ``on`` are all truthy. Values are case-insensitive, -and leading and trailing whitespace is ignored. An unset variable is false. - -``CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION`` - A truthy value disables automatic backend activation during - :mod:`cuda.coop` import. Explicit ``coop.register(...)`` and - backend imports still work. - -``CUDA_COOP_CCCL_ROOT`` - Selects a CCCL source checkout or a ``cuda-coop`` header bundle. An invalid - configured root is an error; resolution does not fall back to another CCCL - source. - -``CUDA_COOP_ENABLE_CACHE`` - A truthy value enables the persistent compiler cache. The value is read - when the backend cache module is imported. - -``XDG_CACHE_HOME`` - On Linux and other POSIX systems, sets the cache base directory; entries - are stored in ``/cccl``. Unset, empty, or relative values fall back - to ``~/.cache/cccl``. Read when the backend cache module is imported. - -``LOCALAPPDATA`` - On Windows, sets the cache base directory; entries are stored in - ``\cccl``. Unset, empty, or relative values fall back to - ``~\AppData\Local\cccl``. Read when the backend cache module is imported. - -``CUDA_COOP_SOURCE_DUMP_DIR`` - Writes generated CUDA source to this directory for compiler diagnostics. - Files use ``cuda_coop__.cu`` names so different backends can - share a directory. Set it before compiling; the Numba backend also writes - the source when its provider compilation cache is hit. Unset or empty - disables dumping. - -``CUDA_PATH`` - Supplies ``/include`` as a CUDA header candidate if - ``cuda-pathfinder`` does not resolve one. - -``CUDA_HOME`` - Supplies ``/include`` after ``CUDA_PATH`` under the same fallback - rule. - -``CUDA_ROOT`` - Supplies ``/include`` after ``CUDA_HOME`` under the same fallback - rule. - -On Linux and other POSIX systems, ``/usr/local/cuda/include`` is tried last. -Windows uses ``cuda-pathfinder`` or the configured toolkit roots above; it -does not try the Unix fallback. If no valid CUDA include directory is found, -compilation reports a header-resolution error. - -Build-time CMake variables --------------------------- - -``CUDA_COOP_INSTALL_HEADER_BUNDLE`` - Defaults to ``ON``. Installs the private CCCL header and CMake-package - bundle into the wheel. - -``CUDA_COOP_ALLOW_DIRTY_HEADER_BUNDLE`` - Defaults to ``OFF``. Allows a Git-worktree bundle when selected inputs are - changed or ``git status`` cannot verify them, and records its source - revision as ``unknown``. - -``CUDA_COOP_CCCL_SOURCE_REVISION`` - Defaults to empty. Supplies the revision token recorded instead of deriving - it from Git. A dirty or unverifiable Git worktree still records ``unknown``. - - -Compilation and headers ------------------------ - -``cuda-coop`` compiles providers against its configured CCCL root, the active -source checkout during in-tree development, or its installed header bundle, in -that order. It never substitutes the CUDA Toolkit's copy of CUB. CUDA headers, -NVRTC, ``nvrtc-builtins``, and nvJitLink must resolve to a compatible toolkit -root. The resulting compiler artifacts and caches include the launch -dimensions, dtype and item extent, storage ABI, compute capability, compiler -options, ordered header identity, and toolkit-library identity. - -Numba-CUDA-MLIR providers created during one compiler attempt share header and -toolkit identity when they select the same ordered include roots. Root selection -is checked for each new provider; header contents are fingerprinted once per -selection in that attempt. Keep headers and the toolkit installation unchanged -while compilation is in progress. A later attempt resolves current headers for -new providers, so edits between attempts do not require a cache-clearing call. -Existing provider contexts retain their original identity; recreate providers -and recompile kernels after changing headers. - -See :doc:`../coop_api` for the public API reference. diff --git a/docs/python/coop/debugger_walkthrough.py b/docs/python/coop/debugger_walkthrough.py deleted file mode 100644 index 753b15c05795..000000000000 --- a/docs/python/coop/debugger_walkthrough.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -# Import the kernel DSL before cuda.coop so automatic registration sees it. -# isort: off -import numpy as np -from numba_cuda_mlir import cuda - -from cuda import coop # First breakpoint: observe backend registration. -# isort: on - - -@cuda.jit -def copy_tile(source, destination, items_per_thread): - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - coop.load(block, source, items, algorithm="direct") - coop.store(block, destination, items, algorithm="direct") - - -def main(items_per_thread=2): - source = np.arange(128 * items_per_thread, dtype=np.int32) - destination = np.zeros_like(source) - - copy_tile[1, 128]( - source, destination, items_per_thread - ) # Compile, link, and launch. - cuda.synchronize() - np.testing.assert_array_equal(destination, source) - print("First launch: copy verified") - - destination.fill(-1) - copy_tile[1, 128]( - source, destination, items_per_thread - ) # Reuse the compiled kernel. - cuda.synchronize() - np.testing.assert_array_equal(destination, source) - print("Second launch: copy verified") - - -if __name__ == "__main__": - main() diff --git a/docs/python/coop/developer_overview.rst b/docs/python/coop/developer_overview.rst deleted file mode 100644 index 9e5e9724b1c7..000000000000 --- a/docs/python/coop/developer_overview.rst +++ /dev/null @@ -1,1686 +0,0 @@ -.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -.. -.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -.. _cuda.coop.developer_overview: - -``cuda.coop`` Developer Overview -================================ - -``cuda.coop`` makes CUB cooperative primitives callable inside a -Python GPU kernel. The Python compiler compiles the surrounding kernel; -``cuda.coop`` generates the C++ device functions for its primitive calls. -The two are linked together before the kernel runs. - -This overview follows a call through the Numba-CUDA-MLIR implementation. It -assumes some familiarity with CUDA threads, blocks, and shared memory. The -:doc:`Programming Guide ` covers writing kernels and -the :doc:`overview <../coop>` covers installation and supported operations; -the focus here is how the implementation works and where to change it. -For a hands-on tour, follow the :ref:`cuda.coop.debugger_walkthrough`. - -This overview describes the Numba-CUDA-MLIR 0.5.x integration. - -A tile copy ------------ - -Start with a kernel that copies 256 integers. We launch one block of 128 -threads, with two items per thread: - -.. code-block:: python - - import numpy as np - from numba_cuda_mlir import cuda - - from cuda import coop - - - @cuda.jit - def copy_tile(source, destination, items_per_thread): - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - coop.load(block, source, items, algorithm="direct") - coop.store(block, destination, items, algorithm="direct") - - - items_per_thread = 2 - source = np.arange(128 * items_per_thread, dtype=np.int32) - destination = np.zeros_like(source) - copy_tile[1, 128](source, destination, items_per_thread) - cuda.synchronize() - np.testing.assert_array_equal(destination, source) - -The example uses NumPy arrays, which Numba-CUDA-MLIR handles at the launch -boundary. An application can supply device arrays instead. The primitive -sees device pointers in either case. - -Each thread owns a separate ``items`` payload. With the direct algorithm, -thread 0 gets elements 0 and 1, thread 1 gets elements 2 and 3, and so on. -The Load fills that payload; the Store writes it back. This is a *blocked* -arrangement of the tile. ``ThreadData(items_per_thread)`` describes that -many values per thread, not values shared by the block. Numba-CUDA-MLIR -specializes the kernel for the ``items_per_thread`` argument. - -The :doc:`Load ` and :doc:`Store ` -visualizations show this ownership pattern and the exchanges used by other -algorithms. - -All 128 threads execute both calls. There is one kernel launch. Neither -``load`` nor ``store`` launches another kernel or returns to the host. - -.. _cuda.coop.calling_conventions: - -Positional operands and keyword-only options --------------------------------------------- - - -Primitive calls take the participating group first, followed by their data -operands. These arguments are positional-only. Options such as -``algorithm``, ``valid_items``, and ``offset`` are keyword-only: - -.. code-block:: python - - coop.load(block, source, items, algorithm="direct", valid_items=n) - coop.store(block, destination, items, algorithm="direct", valid_items=n) - -In the API reference, ``/`` marks the end of the positional-only arguments -and ``*`` introduces keyword-only parameters. Pass the group and data -operands in their documented order, and use names for optional controls. -This keeps the primitive call compact while making algorithm and partial-tile -choices visible. New optional controls can be added without changing the -operand order. - -``cuda.compute`` uses keyword-only arguments for device-wide algorithms, -as described in its :doc:`API conventions <../compute/index>`. Those calls -can contain several input/output arrays, item counts, offsets, and a stream. -For ``cuda.coop``, the group already names the participants and the smaller -operand list describes the data handled within the kernel. - - -Calling CUB from the kernel ---------------------------- - -For this fixed example, the C++ work is small. The Load can be expressed as: - -.. code-block:: c++ - - #include - - __device__ void load_tile(const int* source, int (&items)[2]) - { - using block_load = cub::BlockLoad; - block_load{}.Load(source, items); - } - -This function is called by every thread in the block. CUB uses the thread -index to decide which elements that thread should load. The block size, -item count, dtype, and algorithm are C++ template arguments. - -The Python compiler needs a callable device symbol with a known calling -convention. It cannot call a C++ template directly, and its array value is -not a C++ array reference. We therefore generate a wrapper that exposes the -specialization through the compiler's device calling convention (ABI). -For this Load, a simplified wrapper looks like: - -.. code-block:: c++ - - extern "C" __device__ int load_tile_abi( - void* result, void* source, void* items) - { - load_tile( - static_cast(source), - *reinterpret_cast(items)); - return 0; - } - -Here the integer return is the ABI status. The leading pointer is the -return-value slot and is unused for this void operation. The remaining -pointers address the input and this thread's output payload. For an -operation returning a scalar, the wrapper writes that scalar through the -return-value pointer. Scalar inputs can be passed by value; arrays and -scratch storage use pointers. - -These snippets explain the wrapper rather than reproduce its generated -spelling. The real wrapper names encode the specialization and compilation -identity, and the generated source includes the required support types. -``Algorithm._source_code()`` and ``Algorithm._emit_abi_wrapper()`` in -``numba_mlir/_types.py`` generate this code. - -The backend calls this implementation of the primitive a *provider*. -NVRTC compiles its generated C++ as relocatable device code using C++17 and -``-dlto``. The result is LTO-IR, the intermediate representation used for -link-time optimization. An ``Invocable`` records the callable signature, -link inputs, and any scratch requirements. Its compiler type uses -Numba-CUDA-MLIR's ``ExternFunction`` to represent the external call. - -Numba-CUDA-MLIR compiles the Python kernel and supplies the provider's -LTO-IR to the device link. Link-time optimization can inline the CUB wrapper -into the kernel and optimize across that call. Inlining and register use -still depend on the resulting code; LTO is not a guarantee about either. -Numba-CUDA-MLIR owns the final kernel compilation, loading, and launch. - -Compared with the :doc:`cuda.compute overview -<../compute/developer_overview>`, the same runtime compilation tools appear -at a different boundary. Here the generated C++ implements a device call -within a kernel supplied by the user. That means the group shape and the -kernel's other primitive calls matter to compilation. - -.. _cuda.coop.generated_shims: - -Kernels and their generated C++ -------------------------------- - -.. raw:: html - - - -The following pairs use source captured while compiling real kernels. Each -captured kernel used one block of 128 threads to copy 256 ``int32`` values, -with two values per thread. - -The C++ excerpts retain the emitted types, casts, and calls. Generated -identifiers have been shortened to names such as ``load_impl`` and -``load_abi``, and whitespace has been formatted to fit the page. The copy -excerpts show the no-offset Load helper and its ABI wrapper. Their full -translation units also contain Store and offset overloads. - -These excerpts were captured with Numba-CUDA-MLIR 0.5.1, targeting compute -capability 12.0. Generated names and emitted overloads can change with the -compiler, toolkit, and source checkout. - -Capturing the source yourself -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -From the CCCL repository root, in an environment with the Numba-CUDA-MLIR -dependencies installed: - -.. code-block:: bash - - export PYTHONPATH="$PWD/python/cuda_coop${PYTHONPATH:+:$PYTHONPATH}" - export CUDA_COOP_ENABLE_CACHE=0 - export CUDA_COOP_SOURCE_DUMP_DIR="$PWD/coop-source-dumps/direct" - python python/cuda_coop/examples/numba_mlir/source_dumps.py direct - - export CUDA_COOP_SOURCE_DUMP_DIR="$PWD/coop-source-dumps/transpose" - python python/cuda_coop/examples/numba_mlir/source_dumps.py transpose - -Each command launches the selected kernel, checks its result against NumPy, -and writes ``cuda_coop_numba_mlir_.cu`` under the selected directory. -Use ``CUDA_VISIBLE_DEVICES`` as well if you need to select a particular GPU. - -The script's copy examples use 64 threads and four items per thread. -The recorded copy captures use 128 threads and two items per thread. -Both arrangements copy 256 values, but their template arguments differ. - -The :github:`example script -` includes the imports, -launches, and result checks omitted from the panels below. - -Set the environment variables before starting Python. Unset or empty -``CUDA_COOP_SOURCE_DUMP_DIR`` disables dumping. The provider source is dumped -on provider-cache hits too. These commands disable that cache and use a -fresh process for each kernel so the captures are easy to associate with -their inputs. Reusing an already compiled kernel in the same process can -bypass provider generation entirely. - -The dump contains the C++ input to NVRTC. Numba compiles the surrounding -kernel separately, so its indexing, launches, and planner-inserted barriers -need to be inspected in the kernel's compiler output. Bundling can also put -several providers and overloads in one source file. A definition in the dump -does not by itself establish which overload the final kernel calls. - -Direct Load and Store -^^^^^^^^^^^^^^^^^^^^^ - -.. grid:: 1 1 2 2 - :gutter: 3 - :class-container: coop-shim-pair - - .. grid-item:: - - Python kernel - - .. literalinclude:: ../../../python/cuda_coop/examples/numba_mlir/source_dumps.py - :language: python - :start-after: # docs: start dump-direct - :end-before: # docs: end dump-direct - - .. grid-item:: - - Generated C++: Load excerpt - - .. literalinclude:: source_dumps/direct.cpp.txt - :language: cpp - :start-after: // excerpt-begin - -``128`` and ``2`` appear in the CUB template arguments. The ABI accepts -pointers, casts the payload pointer back to an array of two elements, and -calls ``Load``. Direct Load needs no scratch pointer. Its ``__ret`` slot is -unused because Load fills the caller's payload; the integer return value is -the ABI status. The Store wrapper uses the same pointer conversion and calls -``Store``. - -Transpose with a shared scratch descriptor -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -.. grid:: 1 1 2 2 - :gutter: 3 - :class-container: coop-shim-pair - - .. grid-item:: - - Python kernel - - .. literalinclude:: ../../../python/cuda_coop/examples/numba_mlir/source_dumps.py - :language: python - :start-after: # docs: start dump-transpose - :end-before: # docs: end dump-transpose - - .. grid-item:: - - Generated C++: Load excerpt - - .. literalinclude:: source_dumps/transpose.cpp.txt - :language: cpp - :start-after: // excerpt-begin - -The algorithm is now ``BLOCK_LOAD_TRANSPOSE``. CUB's ``TempStorage`` type -determines the required bytes and alignment. The ABI has an additional -pointer, ``temp_storage``, which it casts to that storage type before -calling the helper. The Python ``scratch`` descriptor causes the planner to -supply the allocation and reuse it across Load and Store. - -This pointer-taking helper has no block barrier of its own. The planner -inserts the barriers in the Python kernel's lowered code. The full source -also contains ``_alloc`` variants with local ``__shared__`` storage and -``__syncthreads()``; those are separate entry points. - - -Recovering the specialization ------------------------------ - -The fixed C++ example supplied all its template arguments by hand. In the -Python kernel, some of that information is in the call, some comes from -type inference, and some comes from the configured launch: - -.. list-table:: - :header-rows: 1 - :widths: 27 35 38 - - * - Information - - Source in the example - - Use - * - Group - - ``this_block()`` and ``[1, 128]`` - - Resolve a block with dimensions ``(128, 1, 1)``. - * - Payload element type - - Load's typed source array - - Infer the C++ element type and check the destination type. - * - Items per thread - - ``ThreadData(items_per_thread)`` with a launch argument of ``2`` - - Instantiate a fixed array extent of two. - * - Algorithm - - ``algorithm="direct"`` - - Select the CUB algorithm and its storage requirements. - * - Input and output addresses - - Kernel arguments - - Pass device pointers when the kernel executes. - * - Target architecture - - Compiler target - - Compile compatible device code for the surrounding kernel. - -``this_block()`` is a compile-time group descriptor. The planner resolves -it against the launch and removes the descriptor from the runtime code. -The user does not need to repeat the block size in the primitive call. -Launching the same kernel with a different block shape can require a -different specialization. - -The full shape matters. A block of ``(32, 4, 1)`` and a block of -``(128, 1, 1)`` both contain 128 threads, but their coordinate expressions -differ. Group indexing uses the x-major linear rank: - -.. code-block:: text - - threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z) - -Groups can also describe partitions. ``this_warp().group_by(8)`` describes -four consecutive groups of eight threads within a physical warp. The -partition width is a compile-time value. Each group's rank, memory tile, -scratch instance, and synchronization mask must agree on that partition. -For Warp Load/Store, the enclosing block must contain complete physical -warps, and every member of a participating logical group must reach the -call. - - -There is also a distinction between a static group size and a runtime -quantity measured within that group. A tail Load may use: - -.. code-block:: python - - # Inside a kernel, after clamping remaining to this block's tile size: - coop.load(block, source, items, valid_items=remaining, oob_default=0) - -The CUB specialization still describes the full tile. ``remaining`` says -how many elements in that tile are valid on this invocation. It can be a -runtime value, provided it is uniform across the group and lies between -zero and the tile capacity. - -The core represents optional arguments with ``ArgumentBinding``: - -* ``OMITTED`` selects the overload without that argument. -* ``STATIC`` carries a compile-time value into specialization and code - generation. -* ``RUNTIME`` keeps an operand in the device call. - -Those cases affect the ABI and cache identity. Treating a runtime value as -a constant would compile the wrong program; treating an omitted argument -as zero could select a different CUB overload. The planner rejects invalid -static controls before compilation. For bounded runtime controls such as -``valid_items``, the generated code checks the range before narrowing to -CUB's integer parameter. A failed check traps on the device. Callers must -also provide enough memory for the selected tile and offset. - -From Python syntax to an external call --------------------------------------- - -The integration registers one whole-function planner, -``CoopWholeFunctionPlanner`` in ``_compiler/_planner.py``. It inspects and -modifies the compiler's intermediate representation (IR) after device -helper inlining. Its ``run()`` method sequences two phases implemented by -separate mixins: - -#. ``_GroupPlanning._resolve_groups()`` finds group constructors, group - methods, and group-first primitive calls. It requests exact launch - metadata, resolves descriptors and payload information, and asks the - primitive family's core planner for a lowering plan. The family's Numba - implementation turns that plan into private provider calls carrying the - plan as compile-time metadata. -#. ``_CallRewriting._rewrite_calls()`` materializes those calls using - ``CoopSinglePhaseRewrite``. Load/Store analysis reuses the plan's dtype - and item count, checks agreement with factory inputs, and records the - payload facts needed for allocation. The rewrite replaces ``ThreadData`` - with fixed local arrays, builds invocables, supplies scratch pointers - where needed, and emits result handling and reuse barriers. It removes - the compile-time constructors from the runtime IR. - -When group resolution changes the IR, ``CoopWholeFunctionPlanner.run()`` -calls Numba's ``_planner_registry._repair_ir()`` before rewriting calls. -This rebuilds the definition lookup and control-flow analysis so the second -phase sees the newly introduced providers and their operands. Rewriting -also runs when group resolution makes no changes, since payload constructors -and private provider calls may remain. The mixins and rewrite helper have -no separate registration. The planner returns ``True`` if either phase -changed the IR; Numba then repairs the final IR before proceeding to later -planners and compiler passes. That result does not request another run of -the cooperative planner. - -Inlining is relevant to helper functions. A helper that receives a group and -calls ``coop.load`` can be planned after it is inlined into its kernel -caller. The descriptor then has the caller's launch context. A group -descriptor escaping into an arbitrary runtime object or a non-inlined device -call is not supported by this mechanism. Descriptor validation waits for -default helper inlining and recursively follows aliases and conditional -definitions. A surviving unsupported helper or descriptor escape is -diagnosed with its name. Standalone callbacks cannot contain primitives -because they lack the caller's cooperative launch context. - -Here, group descriptors are compile-time ``ThreadGroup`` and -``ThreadHierarchy`` values: they describe which threads cooperate. -``_GroupCallPlanner._validate_group_descriptor_references()`` checks their -remaining references after collecting replacements and before changing -block bodies. A descriptor may feed another group or hierarchy assignment -or a primitive call being replaced. Returning it or passing it to an -unrelated runtime call raises ``EscapingGroupDescriptorError``. ``ThreadData`` -constructs a per-thread payload array and follows the separate payload -analysis and materialization path. - -``literal_unroll`` values shaping cooperative groups, selectors, payloads, -or storage are unsupported. The planner diagnoses those uses and suggests -explicit calls with compile-time constants. Ordinary unrolling unrelated to -cooperative planning remains available. Supporting shaped unrolling would -require revisiting planner ordering; this implementation does not move -planning after SSA or unrolling. - -The common API primitives in ``_core/api/`` are compiler markers with shared -signatures and validation rules. Numba's planner recognizes their identity -and binds the call arguments. Reading the Python body alone does not show -the path that runs during kernel compilation. - -.. _cuda.coop.compiler_flow: - -Compiler phases and handoffs ----------------------------- - -The figures show a new kernel compilation in Numba-CUDA-MLIR 0.5.x. -Read each figure from top to bottom. These compiler stages run on the host. -The GPU executes the kernel after compilation and linking finish. - -The first figure locates cooperative planning within the compiler pipeline. -It groups small compiler passes by purpose. The green block is the single -registered ``cuda.coop`` planner. Select **Python classes** to see the classes -and methods that implement the same phases. - -.. tab-set:: - :class: coop-compiler-tabs - - .. tab-item:: Overview - - .. figure:: visualizations/compiler-phases.svg - :alt: Seven compilation phases. After device-helper inlining, the cooperative planner resolves calls and reuses planned payload facts during rewriting, before the remaining untyped passes and type inference. - :width: 100% - - Payload fact reuse stays within cooperative planning, before the - main type-inference pass. - - .. tab-item:: Python classes - - .. figure:: visualizations/compiler-phases-classes.svg - :alt: The seven phases with their Python classes. CoopWholeFunctionPlanner passes lowering plans into CoopSinglePhaseRewrite, which retains analyzed calls before MLIRTypeInference. MLIRBackend and MLIRLower prepare the kernel for compilation and linking. - :width: 100% - :figclass: coop-compiler-diagram-desktop - - The classes follow the same seven phases. - - .. figure:: visualizations/compiler-phases-classes-mobile.svg - :alt: The seven compiler phases with their Python classes, arranged vertically. CoopWholeFunctionPlanner supplies plans to CoopSinglePhaseRewrite before MLIRTypeInference, then MLIRBackend and MLIRLower lower the kernel. - :width: 100% - :figclass: coop-compiler-diagram-mobile - - The classes follow the same seven phases. - - Mixin methods run on their owning planner or rewriter. They are not - separate compiler passes. ``GroupLoweringPlan`` carries call semantics - into provider analysis; ``_RewriteMatch`` retains the analyzed call for - storage planning and replacement. - -Numba first builds *untyped IR*: statements and expressions without a full -map of their types. Argument types and some constants are already known. -Early rewrites fold constants and remove dead branches. Device-helper -inlining then makes supported helper calls visible in the caller's IR. - -``PostInlineWholeFunctionPlanners`` calls ``CoopWholeFunctionPlanner`` -at this point. The planner resolves groups, prepares providers, and rewrites -cooperative calls. It uses argument types, constants, and value definitions -to infer the payload facts that provider selection needs. Planned Load/Store -calls reuse those facts during provider analysis. This work runs before -Numba's main type-inference pass. - -After planning, Numba processes literals, reconstructs static single -assignment (SSA) form, and propagates literal values. ``MLIRTypeInference`` -then builds the type map and call signatures. Typed passes remove phi nodes, -inline overloads, and run the after-inference rewrites. MLIR lowering and -optimization follow. - -``GenericRewrites`` runs the compiler's ``before-inference`` rewrite -registry. ``NopythonRewrites`` runs its ``after-inference`` registry. -``CoopSinglePhaseRewrite`` is a helper inside the whole-function planner; -it has no separate registration in either registry. Its name does not imply -a second cooperative pass after typing. - -The next figure expands the handoffs between the compiler and the CUDA -tools. It shows successful bundling on a provider-cache miss. NVRTC compiles -the generated C++ providers. Numba-CUDA-MLIR compiles the surrounding Python -kernel and supplies both sets of link inputs to nvJitLink. - -.. tab-set:: - :class: coop-compiler-tabs - - .. tab-item:: Overview - - .. figure:: visualizations/compiler-swimlane.svg - :alt: Numba passes inlined IR to cuda.coop, which reuses planned payload facts and retained call analysis before sending providers to NVRTC. NVRTC returns LTO-IR and scratch layouts; Numba later links kernel and provider inputs with nvJitLink. - :width: 100% - - Provider LTO-IR returns before type inference. The final link happens - after Numba compiles the kernel. - - .. tab-item:: Python classes - - .. figure:: visualizations/compiler-swimlane-classes.svg - :alt: Class-level handoffs across Numba, cuda.coop, NVRTC, and nvJitLink. PayloadInference consumes planned Load/Store facts and retained matches supply Algorithm bundles. Invocable carries provider files into typing; MLIRLower emits symbolic calls before linking. - :width: 100% - :figclass: coop-compiler-diagram-desktop - - The rewrite binds an ``Invocable``. Linking resolves or inlines - the provider symbol. - - .. figure:: visualizations/compiler-swimlane-classes-mobile.svg - :alt: Class-level compilation handoffs in time order. PayloadInference consumes Load/Store plan facts; CoopSinglePhaseRewrite reuses matches and binds Invocable objects. Typing builds ExternFunction descriptors, MLIRLower emits calls, and nvJitLink returns a cubin. - :width: 100% - :figclass: coop-compiler-diagram-mobile - - The rewrite binds an ``Invocable``. Linking resolves or inlines - the provider symbol. - - Solid arrows show inputs or calls. Dashed arrows show returned results. - This path uses ``Algorithm`` providers. Direct C-ABI providers and extra - support images use other paths. - -The figure shows the default LTO path. Debug settings or extra PTX inputs -can change how Numba compiles the kernel's link input. The provider bundle -still contains LTO-IR. See :ref:`coop-numba-compilation-reuse` for cache -controls and artifact lifetime. - -Carry payload facts into provider analysis -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -For the tile copy, group planning has already established that ``items`` -holds two ``int32`` values per thread. The ``GroupLoadStoreSemantics`` in -``plan.call.operation`` records that dtype and item count. The private -provider call carries its ``GroupLoweringPlan`` in reserved compile-time -metadata, which is removed before the provider factory is called. - -``_validate_and_split_args()`` passes this plan to ``PayloadInference``. -The Load/Store hook normalizes the plan's dtype to a Numba type and merges -both facts through ``infer_kwarg()``. Existing factory arguments must agree; -the payload constructor's known extent must agree too. The hook records -the dtype through constructor and alias state so ``apply()`` can allocate -the local array. It does not repeat the memory/payload dtype walks or the -Store payload-write scan that established the plan. - -Common-API numeric validation and runtime-control checks still run. -Scalar Store boxing and scratch allocation remain part of emission. -Direct private provider calls without a supported Load/Store plan use full -operand inference. Other operation families retain their own inference -hooks. The group and rewrite resolvers handle different IR forms, including -loops and emitted payload markers, so sharing plan facts does not replace -them with one resolver. - -These facts belong to the current function IR and compilation attempt. -Provider analysis resolves the original operands and constructors; the -completed matches retain assignment identity across block visits. A launch -retry prepares a fresh rewriter. Provider artifact caches have a separate -lifetime from these call facts. - -Collect the providers before rewriting calls -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -``prepare_calls_and_storage()`` examines the whole function before any block is -rewritten. A consumer in a later block can determine an earlier payload's -dtype or add a requirement to a shared scratch descriptor. Keeping the -original calls and definitions available lets the compiler collect these -facts across blocks and inlined helpers. Provider preparation then supplies -the C++ size and alignment of each scratch type. - -#. ``prepare_calls_and_storage()`` calls ``_collect_function_calls()`` to scan - constructors first, then analyze each provider call once with - ``_analyze_provider_call()``. Each ``_RewriteMatch`` records the resolved - factory arguments, runtime operands, and lowering metadata. Planned - Load/Store calls consume the payload facts described above during this - analysis. -#. ``_collect_temp_storage_uses()`` consumes those matches and validates - storage ownership and descriptor uses before provider compilation. -#. ``_prepare_ltoir_bundle_for_matches()`` removes duplicate specializations. - It calls their factories inside ``collect_specializations()``. In this - context, a factory records an ``Algorithm`` instead of building an - ``Invocable`` immediately. -#. ``prepare_ltoir_bundle()`` combines compatible algorithms into one C++ - source unit. It shares the includes and emits each distinct provider - body and its ABI wrappers. -#. Bundle preparation calls ``compile_with_layouts()`` to get the provider - LTO-IR and scratch layouts. On a cache miss, one NVRTC compilation supplies both. On a cache - hit, the cache supplies the saved image and its matching layouts. -#. The algorithms retain a shared LTO-IR file. Their invocables use that - file as a link input. ``_compute_func_temp_storage_requirements()`` uses - the validated storage uses and provider layouts to collect sizes and - alignments before ``apply()`` emits arrays and scratch pointers. - -After preparation, ``_func_matches`` retains the completed call records by -original assignment identity. Block matching reuses them; an assignment -replaced by ``apply()`` no longer matches on a later visit. - -For the tile-copy example, the Load and Store providers can share one -NVRTC compilation. Each call keeps its own signature and runtime operands. -Bundling requires compatible specializations and at least two distinct -providers. A single provider uses its individual compilation path. If the -optional bundle cannot be prepared, the rewrite can materialize providers -individually. Extra support images and kernel compilation remain separate. -Thus, one provider bundle does not guarantee one NVRTC call for every kernel -compilation. - -Rewrite calls before the device address is known -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -The rewrite needs the provider's symbol name and calling convention. -It does not need the provider's final device address. ``Algorithm`` uses -the same qualified name for the generated wrapper and its call descriptor. - -#. ``prepare_calls_and_storage()`` prepares the provider bundle and creates - invocables for - calls that need scratch storage. Their LTO-IR and scratch layouts are - already available. ``apply()`` reuses these invocables or creates others. -#. ``apply()`` binds an ``Invocable`` with ``ir.Global`` and emits an - ``ir.Expr.call``. The global contains a Python object for the compiler, - not a device function pointer. The object retains the provider link files. -#. During typing, ``Invocable._numba_type_`` builds local overload templates. - ``Algorithm.codegen_method()`` creates an ``ExternFunction`` with the - symbol name, signature, ABI, and link files. The compiler uses this - description; it does not execute ``Invocable.__call__``. -#. ``MLIRLower`` declares the external symbol, emits a symbolic ``func.call``, - and collects the link files. Kernel compilation produces the other - LTO-IR input. nvJitLink resolves the symbol across these inputs and can - inline the provider into the kernel. -#. ``CodeLibrary.get_cufunc()`` loads the linked cubin and obtains the kernel's - ``CUfunction`` handle. This path does not request a separate provider handle. - -Follow the function rewrite lifecycle -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -``_CallRewriting._rewrite_calls()`` owns preparation, storage staging, block -replacement, and final cleanup. It calls ``prepare_calls_and_storage()`` and -``begin_rewrite()`` before visiting blocks. For each block with matches, it -installs the result of ``apply()``. Only after all replacements are installed -does it call ``finish_rewrite()``. These are steps within the existing -whole-function planner, which still runs before type inference. - -.. list-table:: Main responsibilities in ``CoopSinglePhaseRewrite`` - :header-rows: 1 - :widths: 22 38 40 - - * - Phase - - Entry points - - Purpose - * - Prepare the function - - ``prepare_calls_and_storage()`` - - Analyze each provider call, reusing planned Load/Store payload facts. - Validate storage uses, prepare the provider bundle, then collect - scratch requirements before changing statements. - * - Stage shared storage - - ``begin_rewrite()``: ``_stage_temp_storage_backing()`` - - Check shared-memory conflicts and insert the backing allocation in - the entry block, so every scratch view can use it. - * - Match one block - - ``match()``: ``_prepare_block()``, ``_match_assignment()`` - - Select prepared call records by assignment identity. Record this - block's constructors and payload extents. - * - Prepare replacements - - ``apply()``: ``_prepare_call_invocables()`` - - Reuse or materialize each call's provider and bind its specialization - before emitting the block's statements. - * - Emit replacements - - ``_emit_thread_data_array()``, ``_emit_temp_storage()``, - ``_emit_provider_call()`` - - Replace statements in source order with payload arrays, scratch - views, provider calls, result handling, and storage-reuse barriers. - Refresh the typing context and return the replacement block. - * - Clean unused bindings - - ``finish_rewrite()`` - - Remove unused compile-time arguments and constructor references. - Check uses across the completed function so another block's - unreplaced call cannot lose a binding it still needs. - -``ThreadData`` holds each thread's payload. Its dtype and extent determine -the local array emitted by ``apply()``. A standalone payload can obtain its -dtype from typed indexed writes even when no primitive consumes it; this -last inference step still runs during array emission. ``TempStorage`` is an -opaque scratch descriptor whose ownership and reuse policy are checked -against all consumers. Its concrete size, alignment, and slices depend on -the compiled providers, so backing allocation follows provider preparation. - -Putting shared backing in the entry block makes it dominate every scratch -view: execution reaches the allocation before any use, regardless of the -order in which the compiler visits blocks. ``apply()`` builds only its -replacement block and records cleanup candidates. Original definitions stay -available during emission; ``finish_rewrite()`` retires eligible bindings -after checking final uses. The planner registry then repairs IR analysis. - -Preparation can compile providers and update caches without changing block -statements. An incomplete launch-dependent preparation leaves the IR intact -for a fresh rewriter. Individual provider materialization can still occur in -``apply()``, before that block's payload allocations and calls are emitted. - -To trace these stages, start with :github:`the planner -` and -:github:`the block rewrite -`. -``prepare_calls_and_storage()`` coordinates :github:`storage validation and layout -` with -:github:`bundle preparation -`. -:github:`Provider types and source generation -` connect those helpers to -:github:`the NVRTC interface -`. -The :ref:`cuda.coop.debugger_walkthrough` provides breakpoints for the same -flow. The remaining passes belong to the installed compiler's -``mlir_compiler.py``, ``mlir_lowering.py``, and ``mlir_optimization.py``. - -.. _the-portable-core-and-the-numba-backend: -.. _coop-implementation-families: - -The shared core and implementation families -------------------------------------------- - -The common API is exposed through ``cuda.coop`` and implemented in -``_core/api/``. The private ``_core/`` package also contains shared -implementation used by the backends. The package name describes that -implementation layer; the user-facing API is called the common API. - -A :term:`family` groups related primitives and their implementation. The -Load/Store family has shared declarations in ``_core/api/load_store.py`` -and ``load_store.pyi``, semantic descriptions in -``_core/group/load_store.py``, and Numba-specific entry points in -``numba_mlir/_group/_load_store.py`` and ``_load_store.pyi``. Compiler -analysis and lowering have their own Load/Store modules. A family spans -common operations and qualified extensions. - -The shared core describes what a primitive means and which C++ implementation -can perform it. It does not import Numba or invoke a compiler. The Numba -backend reads compiler IR and types, then converts the core's plan into -code that Numba-CUDA-MLIR can compile. - -For example, the Load/Store family normalizes the group, algorithm, dtype, -item count, and optional arguments into a ``GroupPrimitiveCall``. -``plan_group_primitive()`` resolves that call to a ``GroupLoweringPlan``. -The plan records more than the selected CUB class: - -.. list-table:: - :header-rows: 1 - :widths: 25 75 - - * - Plan information - - What the backend needs to know - * - Topology and participation - - Which threads form a group, how many group instances exist in a - block, and which participants must reach the call. - * - Temporary storage - - Who owns scratch, how many instances are needed, and which layout - requirements must be obtained from the compiled provider. - * - Synchronization - - Which barrier is required before another call can reuse scratch. - * - Implementation - - The CUB specialization, along with its source library - and header. - -A supported plan must contain its execution requirements. The shared -``_build_execution_requirements()`` helper in -``_core/group/_execution_requirements.py`` returns a -``GroupExecutionRequirements`` record. Its named ``topology``, -``participation``, ``synchronization``, and ``temp_storage`` fields hold -``GroupTopologyRequirements``, ``ParticipationRequirements``, -``SynchronizationRequirements``, and ``TempStorageRequirements`` respectively. -Family planners copy these fields into the corresponding ``GroupLoweringPlan`` -fields. They describe what execution requires; constructing the records -does not prove uniform participation or insert device barriers. - -An unsupported -combination carries a reason that the backend reports before provider -compilation. For instance, having a ``ThreadGroup`` descriptor for a scope -does not imply that every primitive supports that scope. - -``Algorithm`` describes a CUB template specialization and its method -parameters without compiler types. Parameter descriptors such as ``Value``, -``Pointer``, and ``Array`` describe the method's arguments. For example, -``Array`` records a parameter's type and extent; ``ThreadData`` constructs -the payload passed to that parameter. ``NumbaMlirCoreAdapter`` maps the -core types and parameter descriptors to the Numba backend's representation. - - -Construct ``Algorithm(..., template_arguments={...})`` with the template -arguments and auxiliary dependency values. Construction validates and -freezes those bindings, so the resulting algorithm is ready for backend -materialization. - -The Block and Warp Load/Store specialization factories return this -``Algorithm`` directly, including its bound arguments, method parameters, -and operation metadata. - -Payloads, layouts, and results ------------------------------- - -``ThreadData`` becomes a fixed local array in Numba-CUDA-MLIR. Its extent -must be known at compile time. The compiler may keep its elements in -registers; indexing, address-taking, and register pressure determine the -final placement. - -Both common and qualified ``ThreadData`` constructors accept an optional -``alignment`` keyword. It specifies a minimum power-of-two alignment in -bytes when the compiler materializes payload storage. It does not assert -alignment of the source or destination arrays passed to Load and Store. - -The dtype can also come from the surrounding operation. In this kernel -fragment, ``items_per_thread`` is a kernel argument and the Load's source -establishes the dtype: - -.. code-block:: python - - items = coop.ThreadData(items_per_thread) - coop.load(block, source, items) - coop.store(block, destination, items) - -The planner propagates the dtype from ``source`` to ``items``. A later Store -can then use ``items`` even though its constructor did not specify a dtype. -Load fills ``items`` in place and returns ``None``. - -Layout describes which logical tile elements each thread owns. A striped -Load gives thread ``t`` elements ``t + i * block_size``. A blocked -Load gives it ``items_per_thread * t + i`` for each local index ``i``. -The ``transpose`` algorithm uses striped memory transactions internally -and fills the payload in blocked order; -``striped`` exposes the striped payload to the caller. The caller must -choose Load and Store algorithms that agree on that arrangement. - -Load and Store return ``None``. Load fills the supplied output in place. -Store follows CUB: transpose algorithms may rearrange the input payload -in place. Invalid Load slots are unspecified unless a default is supplied. - -The qualified namespace accepts additional compiler-specific values, such -as local-array payloads where supported. Type support is still checked by -each primitive. An ABI helper for aggregate values does not imply that -public payload operations accept arbitrary structures. The current -common payload APIs require their supported numeric dtypes. - -Inferring scalar types across loops -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -The group planner needs a scalar's dtype before ordinary Numba type inference -runs. Following assignments works directly for an array element or an explicit -scalar cast. A loop can introduce a cycle: the next iteration's input depends -on a result computed from that same input. Stopping at the cycle would lose -the known dtype of the value that entered the loop. - -For example, this fragment carries a value from a ``float32`` source array -through arithmetic that preserves its type: - -.. code-block:: python - - value = source[cuda.threadIdx.x] - for _ in range(repetitions): - value = value + np.float32(1) - coop.store(coop.this_block(), destination, value) - -``GroupPlanningContext.dtype`` first tries its ordinary recursive analysis. -If that cannot resolve the value, ``_loop_dtype`` propagates candidate types -from known definitions until they stop changing. A loop backedge means an -assignment that supplies the next iteration's value; it may temporarily use -the candidate for that variable. The planner then repeats the analysis with -strict checks on every reaching definition before accepting the candidate. - -That final check is essential. A known initial value does not establish the -type of a later opaque helper result. Such a path remains unresolved, and -conflicting known types are rejected. Candidate facts are discarded after -each query, including failed queries, so they cannot supply a type to a -later unrelated operation. This analysis supports type-preserving loops; -general type promotion remains the compiler's responsibility. The separate -:ref:`ThreadData inference guidance ` still -applies to payload construction and unsupported helper producers. - -Shared memory and reuse ------------------------ - -The direct tile copy needs no CUB scratch. Direct, striped, and vectorize -Load/Store providers have ``StorageABI.NONE`` and do not add a scratch -pointer or a reuse barrier. - -Changing the Load to ``algorithm="transpose"`` introduces shared-memory -communication. CUB defines a ``TempStorage`` type whose size and alignment -depend on the specialization. ``cuda.coop`` obtains both during the NVRTC -compilation that builds the provider's LTO-IR. It appends a variable template -whose arguments are ``sizeof(TempStorage)`` and ``alignof(TempStorage)`` and -registers its address with ``nvrtcAddNameExpression``. After compilation, -``nvrtcGetLoweredName`` returns a symbol that encodes the two integer values. -The decoder checks the expected symbol, encoding, and layout before handing -the size and alignment to the host planner. - -Both individual providers and shared bundles cache the ordered layouts with -the provider image. A memory or disk cache hit therefore supplies the image -and scratch requirements together. Storage-free providers skip the queries. -The provider's LTO-IR is retained for the final kernel link; layout discovery -requires no separate PTX link or GPU kernel launch. See -``numba_mlir/_compiler/_layout.py`` for the probe and decoder, and -``numba_mlir/_compiler/_nvrtc.py`` for compilation and cache integration. - -Storage-bearing CUB providers use ``StorageABI.LEADING_POINTER``. The -rewrite supplies a pointer to the appropriate slice of the kernel's -shared-memory allocation. This is the first *provider operand*; the -external ABI's return-value slot is separate. - -The planner considers uses across the function so compatible sequential -calls can reuse storage. It also keeps independent group instances -separate. A transpose in each physical warp needs a distinct scratch -slice for each warp; logical warps need slices and synchronization masks -for their smaller groups. Equal byte counts do not make storage from -different participation domains interchangeable. - -After a block call that uses compiler-managed scratch or a descriptor with -``auto_sync=True``, the rewrite emits a block reuse barrier. For supported -physical and logical Warp calls, it emits ``syncwarp`` with the participating -group's mask. CUB's synchronization inside a primitive does not generally -establish that a later primitive can immediately overwrite the same scratch. -Automatic synchronization adds a trailing barrier after each storage-consuming -call. It does not establish that arbitrary user control flow is safe: callers -must still ensure that all group members reach the primitive and barrier. - -A block operation can expose that reuse choice through ``TempStorage``: - -.. code-block:: python - - # Inside a kernel with an items_per_thread argument and one full tile. - scratch = coop.TempStorage(auto_sync=True) - items = coop.ThreadData(items_per_thread) - coop.load(block, source, items, - algorithm="transpose", temp_storage=scratch) - coop.store(block, destination, items, - algorithm="transpose", temp_storage=scratch) - -This shared descriptor requests automatic reuse barriers. Its -size and alignment can be inferred from its uses. An explicit capacity -must satisfy the compiled provider's requirements. - -Only ``size_in_bytes`` may be positional in ``TempStorage``; the other -options are keyword-only. An explicit ``alignment`` requests a minimum, -which the planner can strengthen to meet the requirements of its uses. - -``sharing="exclusive"`` allocates separate slices for distinct uses. Shared -and exclusive descriptors both default to ``auto_sync=False``, leaving -synchronization to the caller. Layout and synchronization are independent. A -loop that reaches the same call site again still needs safe reuse, including -with an exclusive descriptor. For a block primitive, put the required block -barrier where every thread reaches it before the next use. The planner -conservatively rejects collapsing multiple manually synchronized -constructors into one descriptor. This is a validation limit, not proof that -each rejected program races. Planner and rewrite contracts are cross-checked -before emission so parser disagreement cannot silently remove a reuse -barrier. - -Group planning and storage rewriting share two policies in -``_compiler/_operations.py``: ``expected_storage_reuse_barrier()`` derives -the barrier from storage ownership and ``auto_sync``, while -``provider_synchronization_matches()`` checks the provider's declaration -against that plan. The emitter still checks which topology it can implement -and compares the original ``TempStorage`` constructor with the plan before -emitting a barrier. Sharing the policies preserves these independent checks. - -The planner can switch its backing allocation to dynamic shared memory -when the required size exceeds the static allocation limit, subject to -the device's opt-in limit. It reports the required launch bytes through -Numba-CUDA-MLIR's compiler metadata. The user-facing call does not need a -manually maintained byte count. -The launcher treats that requirement as a minimum, not an allocation added to -user-supplied dynamic bytes. Dynamic backing accepts alignment up to 16 bytes. - -Until a released compiler with the shared-memory allocation fix is qualified, -the rewrite rejects user dynamic or runtime-sized shared allocations alongside -cooperative backing, and user static shared allocations when cooperative -backing becomes dynamic. It inspects user allocations after helper inlining, -including aliases and implicit oversized cooperative scratch. Static/static -combinations remain valid, and storage-free operations introduce no conflict. -Diagnostics identify both allocations and suggest keeping them static within -the device limit, moving the user buffer to global memory, or using separate -kernels. Passing coexistence tests against a development compiler alone does -not remove the compatibility guard. - -These controls are operation-specific. Warp Load/Store uses -compiler-owned storage and rejects an explicit ``TempStorage``. - - -.. _coop-numba-compilation-reuse: - -Activation and compilation reuse --------------------------------- - -Importing ``cuda.coop`` after ``numba_cuda_mlir`` activates the Numba backend -hooks. An isolated common API import does not load optional compilers. -Register explicitly to make initialization independent of import order: - -.. code-block:: python - - from cuda import coop - - coop.register("numba-cuda-mlir") - -This host-side call imports the selected backend and activates its hooks. -It is safe to repeat. Importing ``cuda.coop.numba_mlir as numba_coop`` also -activates the hooks and exposes the backend namespace. Every install includes -the same DSL integration modules. An extra only adds dependency requirements -from ``pyproject.toml``; it does not change the wheel or register hooks in a -running process. - -``_compiler/_activation.py`` checks the runtime and compiler compatibility, -imports the planner, and registers ``CoopWholeFunctionPlanner`` as its final -step. There is one compiler hook to register, so activation needs no -registry snapshots or rollback. ``_compiler/_numba_mlir_compat.py`` checks -the installed compiler version when the Numba backend is activated, before -importing its compiler integration. The Numba extras declare -``numba-cuda-mlir>=0.5.0,<0.6`` for package installers; the activation check -also covers applications that install bare ``cuda-coop`` and manage compiler -dependencies separately. An unused Numba installation has no effect on other -backends. Once the version check succeeds, integration modules import the -required compiler APIs directly. Missing launch metadata and other -compilation errors retain their specific diagnostics. A new compiler series -needs its integration checked before the supported range changes. The -compatibility module documents the reasons for this division in detail. - -The rewrite can collect compatible CUB specializations and compile them -in a single NVRTC source bundle. This reduces repeated header parsing and -compilation. Each call still has its own signature, result handling, and -storage contract. - -There are caches at several stages. Compiled Python operators are reused -within the process. Provider compilation can use a persistent cache when -``CUDA_COOP_ENABLE_CACHE`` is enabled before backend import. Numba-CUDA-MLIR -separately owns the compiled kernel's reuse and lifetime. The provider cache -uses ``XDG_CACHE_HOME/cccl`` on POSIX, falling back to ``~/.cache/cccl``; -on Windows it uses ``LOCALAPPDATA\cccl``, falling back to -``~\AppData\Local\cccl``. Unset, empty, or relative base directories use -the fallback. These settings are read at backend cache import. - -A provider cache key must identify the code being compiled: the operation, -dtype, shape, static arguments, wrapper ABI, target architecture, compiler -options, and selected headers and compiler libraries. Runtime array -addresses and runtime tail counts are operands, not specialization values. -The header and toolkit identity matters because the same CUB template -arguments can produce different code or scratch layouts with different -headers. - -The ``cuda-coop`` wheel bundles matching CCCL headers. A source checkout -uses its checkout's headers, and ``CUDA_COOP_CCCL_ROOT`` can select an -explicit root. See :doc:`configuration` for header selection, provider-cache -settings, and generated-source dumps. NVRTC, its builtins, nvJitLink, and CUDA -headers must also resolve coherently. ``_headers/`` and -``_compiler/_nvrtc.py`` handle that selection. An invocable retains its -temporary LTO-IR files for linking; -those files are compilation inputs, not loaded kernel handles. - -.. _cuda.coop.debugger_walkthrough: - -Debugger Walkthrough --------------------- - -Follow the tile copy through a Python debugger to see how the compiler -recognizes a primitive, chooses its CUB specialization, and connects the -generated device code to the kernel. These breakpoints stop in the host -Python code doing the compilation. GPU threads execute the compiled kernel; -the Python debugger cannot stop inside that device execution. - -An example to debug -^^^^^^^^^^^^^^^^^^^ - -Open ``docs/python/coop/debugger_walkthrough.py`` in your checkout, or -:download:`download the example `. It needs no -command-line arguments: - -.. literalinclude:: debugger_walkthrough.py - :language: python - :start-at: # Import the kernel DSL - -By default, all 128 threads copy two integers each. ``main`` accepts an -``items_per_thread`` argument and forwards it to both kernel launches. -The second launch uses the same argument types, item count, and block -shape so you can observe compilation reuse. -Both launches check the result. Warnings about a small grid and host-array -copies are expected for this small example. - -Configure VS Code -^^^^^^^^^^^^^^^^^ - -Use a checkout containing the Numba-CUDA-MLIR stack described above. Open -that CCCL checkout as the VS Code folder, and use **Python: Select -Interpreter** to choose an environment with the Numba-CUDA-MLIR dependencies -installed as described in the :doc:`installation instructions <../coop>`. -The Python and Python Debugger extensions must be installed in the -environment where VS Code runs the program, including the remote side -when using Remote SSH. - -Add this configuration to ``.vscode/launch.json`` in the checkout. If that -file already exists, add the configuration to its ``configurations`` list: - -.. code-block:: json - - { - "version": "0.2.0", - "configurations": [ - { - "name": "cuda.coop: Debug active Python file", - "type": "debugpy", - "request": "launch", - "program": "${file}", - "console": "integratedTerminal", - "cwd": "${workspaceFolder}", - "justMyCode": false, - "env": { - "PYTHONPATH": "${workspaceFolder}/python/cuda_coop", - "CUDA_COOP_CCCL_ROOT": "${workspaceFolder}", - "CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION": "0", - "CUDA_COOP_ENABLE_CACHE": "0", - "CUDA_COOP_SOURCE_DUMP_DIR": "${workspaceFolder}/build/coop-debug-sources" - } - } - ] - } - -``justMyCode: false`` lets you step into the compiler and library modules. -``PYTHONPATH`` and ``CUDA_COOP_CCCL_ROOT`` select this checkout's Python -sources and C++ headers. Automatic registration is enabled, and the -provider disk cache is disabled so a fresh debug session reaches NVRTC. -On a machine with several GPUs, add ``CUDA_VISIBLE_DEVICES`` to ``env`` -to select the device you want to use. - -With the example file active, set your **first breakpoint** on -``from cuda import coop``, using the gutter or **F9**. Select -``cuda.coop: Debug active Python file`` in **Run and Debug**, then press -**F5**. Keep the example active when starting: ``${file}`` means the -currently selected editor file. - -VS Code also has a **Python Debugger: Debug Python File** editor action. -Use the named launch configuration for this walkthrough so the source -paths and library-stepping setting above apply. See the -`VS Code Python debugging documentation -`_ for those controls. - -At a breakpoint, **F10** steps over a statement, **F11** steps into a call, -**Shift+F11** steps out, and **F5** continues to the next breakpoint. -Expressions in **Debug Console** use the selected **Call Stack** frame. -The instructions below name functions and statements so you can find the -breakpoints even as line numbers change. Put breakpoints on the executable -statement, rather than the ``def`` line or its decorator. - -Initial import and registration -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -The first stop is just before ``cuda.coop`` imports. Numba-CUDA-MLIR has -already loaded because of the preceding import. Under -``python/cuda_coop/cuda/coop/``, open ``__init__.py`` and put a breakpoint -on ``_auto_register_known_dsls()``. Continue to it and step into the call. - -In ``_core/_auto_registration.py``, follow the loop to -``candidate.activate()``. Inspect ``candidate.runtime_module`` and its -membership in ``sys.modules``: the runtime name is ``numba_cuda_mlir``. -This is why the import order matters. The automatic path recognizes a -runtime the application has already imported. - -Before continuing, set a breakpoint on ``_require_runtime()`` inside -``_initialize_runtime_hooks()`` in -``numba_mlir/_compiler/_activation.py``. At this stop, the call stack -connects the root import to the qualified backend import and then to hook -registration. Step over the compatibility check and the import of -``_planner``. Inspect ``planner_module.CoopWholeFunctionPlanner`` before -stepping over ``register_planner(...)``. This final call registers -the single compiler hook that resolves and rewrites cooperative calls. - -Disable these import breakpoints. Set a breakpoint on the first -``copy_tile[1, 128](source, destination, items_per_thread)`` in the example -and continue. At this point, check ``coop.__file__`` in Debug Console. It -should point inside the checkout you opened. The ``@cuda.jit`` decorator has -made a dispatcher; this first launch will trigger compilation for its -arguments. - -Fast-forward to the Numba hooks -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Paths from here through the provider steps are relative to -``python/cuda_coop/cuda/coop/numba_mlir/``. - -In ``_compiler/_planner.py``, find ``CoopWholeFunctionPlanner.run()`` and -set a breakpoint on ``groups_changed = self._resolve_groups()``. Continue -from the example's launch line. You have crossed from application code -into a hook Numba calls while compiling it. Inspect the Call Stack to see -that caller. - -In ``_compiler/_group_planner.py``, find -``_GroupPlanning._resolve_groups()`` and set a breakpoint on -``launch_config = require_launch_config(planner.state)``. Continue from the -planner's entry point to inspect the first phase. - -Step over the assignment, then evaluate: - -.. code-block:: python - - self.state.func_ir.func_id.func_qualname - launch_config - self.state.func_ir.dump() - -The function name is ``copy_tile``. The launch configuration contains -``block: (128, 1, 1)`` and ``grid: (1, 1, 1)``. The IR dump shows Numba's -intermediate representation of the Python function, including the -``this_block``, ``ThreadData``, ``load``, and ``store`` calls. The dump -prints to the debuggee's output; its return value is ``None``. - -The block dimensions came from ``[1, 128]`` at the host launch site. They -give ``this_block()`` a concrete group shape for C++ specialization. You -are inspecting compiler values and descriptors here; ``items`` has not -become a particular GPU thread's two integers. - -The planner can also run for generated helper functions. The launch-config -breakpoint comes after the no-group-markers guard, which avoids many -uninteresting stops. This configured launch obtains its launch facts -directly; it does not require a failed first compilation to discover them. - -.. _from-a-collective-call-to-a-plan: - -From a primitive call to a plan -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Set the next breakpoint in ``_compiler/_group_load_store.py``, inside -``_LoadStorePlanning._lower_load_store()``, on -``planned_operation = self._plan_provider_operation(plan)``. Continue. -The preceding call has produced a ``GroupLoweringPlan`` for Load. -Inspect these expressions: - -.. code-block:: python - - operation - plan.call.operation - plan.target - dict(plan.implementation.template_arguments) - plan.participation - plan.temp_storage - plan.synchronization - -Expect ``operation == "load"``, target ``CUB_BLOCK``, ``T`` equal to -``int32``, ``BLOCK_DIM_X`` equal to 128, and ``ITEMS_PER_THREAD`` equal to -2. ``ALGORITHM`` selects ``::cub::BLOCK_LOAD_DIRECT``. The participation -requirements specify the complete block. Direct Load needs no shared scratch, -and its storage-reuse barrier is ``NONE``. - -In the same function, advance to ``statements.extend(...)`` near the end. -Compare ``factory_kwargs`` with ``runtime_args``. The block shape, dtype, -item count, and algorithm are compile-time choices in ``factory_kwargs``. -``runtime_args`` holds IR variables for ``source`` and ``items``. Those -become operands of the generated device call. - -Continue to the same stops for Store. Its template selects -``BLOCK_STORE_DIRECT``, and its runtime operands are ``destination`` and -``items``. Disable the Load/Store breakpoints after inspecting both calls. -The public calls now have private provider calls carrying those choices. - -Generate and compile the C++ providers -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Set a breakpoint on ``rewrite = CoopSinglePhaseRewrite(...)`` in -``_CallRewriting._rewrite_calls()`` in ``_compiler/_rewrite.py``. Continue -and use ``self.state.func_ir.dump()`` again. Compare it with the earlier -dump: group resolution has introduced private factories and constants -for the resolved operations. - -Before continuing, set a breakpoint in ``_compiler/_rewrite_load_store.py``, -inside ``_infer_load_store_payload()``, on -``semantics = plan.call.operation``. -Continue and inspect ``inference.lowering_plan``, ``semantics`` after that -assignment, and ``payload_specification``. For the Load call, the plan's -operation records ``int32`` and two items per thread. - -Step through the ``infer_kwarg()`` calls to see the plan and factory inputs -checked for agreement, followed by constructor dtype recording. This branch -returns before the operand dtype walks used by unplanned provider calls. -Disable this breakpoint after inspecting Load. - -Before continuing from the Load hook, set a breakpoint in -``_analyze_provider_call()`` on ``return _RewriteMatch(...)``. Inspect -``factory_kwargs``, ``runtime_args``, and ``lowering_plan`` for Load and -Store. This analysis runs once per call -during function preparation; later block matching reuses the records. - -Before continuing, set a breakpoint in ``_types.py``, inside -``prepare_ltoir_bundle()``, on the call to -``nvrtc.compile_with_layouts(...)``. -At that stop, inspect ``src`` and: - -.. code-block:: python - - [algo.struct_name for algo in algorithms] - -This example produces both the BlockLoad and BlockStore specializations -in one C++ source unit. Find ``BLOCK_LOAD_DIRECT``, ``BLOCK_STORE_DIRECT``, -and the ``extern "C"`` ABI wrappers in ``src``. Their template arguments -should agree with the plan you just inspected. To trace an individual -provider's source generation on another run, stop in -``Algorithm._source_code()`` in the same file. - -The Call Stack at the bundle stop also explains its timing: -``_CallRewriting._rewrite_calls()`` calls ``prepare_calls_and_storage()``, which -analyzes calls and validates storage uses before preparing the bundle. -After compilation supplies the layouts, it collects scratch requirements -and saves the matches for block emission. This still occurs before -``apply()`` replaces the calls and before Numba's main type-inference pass. - -Set a breakpoint in ``_compiler/_nvrtc.py``, inside ``compile_impl()``, -on ``err, prog = nvrtc.nvrtcCreateProgram(...)``. Continue and inspect -``cpp``, ``cc``, ``rdc``, ``code``, and ``compiler_options``. ``cpp`` is the -source you just saw; ``cc`` identifies your device's target architecture. -``rdc`` is true and ``code`` is ``"lto"``. The options include C++17, -the selected include directories, and ``-dlto``. - -Step over ``nvrtcCompileProgram`` and the subsequent error check. NVRTC -has compiled the C++ provider code. The ``nvrtcGetLTOIR`` calls retrieve -the bytes for linking it into the Python kernel. The launch configuration -also saves the generated ``.cu`` source in ``build/coop-debug-sources``. -Open that file for a more convenient view of the complete source. - -Materialize the payload and device calls -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Set a breakpoint on ``if self._has_temp_storage_requirements():`` in -``begin_rewrite()`` in ``_compiler/_rewrite.py``. Continue to it and inspect -``self._func_matches`` and ``self._thread_data_specifications``. Function -preparation is complete; the original call assignments are still in the IR. -The direct Load/Store algorithms need no shared backing, so this stop -proceeds to block matching without allocating scratch. The optional -transpose example below takes the backing-allocation branch here. - -Set a breakpoint in the ``_prepare_call_invocables()`` method of -``CoopSinglePhaseRewrite`` in -``_compiler/_rewrite.py`` on -``invocable, _ = self._materialize_invocable(match)``. Continue to it, then -step over the statement to obtain the callable provider for this match. -Inspect: - -.. code-block:: python - - match.op_name - match.factory_kwargs - match.runtime_args - invocable.files - invocable.storage_abi - invocable.temp_storage_bytes - -Load and Store each have an ``Invocable``. Their ``files`` refer to the -shared provider LTO-IR file from the previous step. The direct algorithms -have storage ABI ``NONE`` and zero temporary-storage bytes. - -After seeing both matches, disable that breakpoint and stop on -``return new_block`` at the end of ``apply()``. Evaluate -``new_block.dump()``. Find the local-array construction for ``items`` and -the calls through globals holding the ``Invocable`` objects. The local -array has extent 2 and dtype ``int32``. Payload construction and primitive -calls have been replaced. Constructor aliases and compile-time arguments -can still be present because cleanup waits for every block replacement. - -Before continuing, set a breakpoint on -``self._remove_unused_factory_arguments()`` in ``finish_rewrite()``. At that -stop, ``self._func_ir.dump()`` shows all installed replacement blocks. Step -over both cleanup calls and compare the dump: eligible unused factory -arguments have been removed and unused payload constructor bindings have -become ``None``. This is the last rewrite step before the planner registry -repairs definitions and other IR analysis. - -This is the point where the compiler's representation has concrete -per-thread storage and device calls in place of the public primitive -syntax. Later compilation decides whether that local array's values can -live in registers. - -Hand the call and link inputs to Numba -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Before continuing from ``finish_rewrite()``, put a breakpoint in -``Algorithm.codegen_method()`` in ``_types.py`` on -``extern_fn = ExternFunction(...)``. At this stop, inspect ``abi_name``, -``abi_input_types``, ``arg_transforms``, and ``link_files``. - -``abi_name`` matches an ABI wrapper in the generated C++. Pointer -arguments have a ``"ptr"`` transform; scalar arguments, when present, -use ``"value"``. ``link_files`` supplies the compiled provider to -Numba-CUDA-MLIR. Each specialization can expose several supported call -signatures, so this breakpoint may fire more than once per primitive. - -``ExternFunction`` gives the compiler a device symbol, a signature, and -the files needed to resolve it. Its default Numba ABI includes the -status return and return-value slot shown earlier in this guide. The -Python ``Invocable.__call__`` body is a marker that rejects host calls; -kernel compilation consumes its compiler type and overload instead. - -For an optional look across the dependency boundary, open the installed -``numba_cuda_mlir/mlir_lowering.py`` from the selected interpreter. Find -``lower_call_external_function()`` and stop on its call to -``self._link_external_function(fn_value)``. Inspect ``fn_value.name``, -``fn_value.sig``, ``fn_value.abi``, and ``fn_value.link``. Here the kernel -compiler receives the external function and its provider link inputs. -Step through the remainder to see it construct the MLIR call. This file -belongs to Numba-CUDA-MLIR, so use its installed source path rather than -looking for it under CCCL. - -Finish the launch and observe reuse -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Disable the compiler breakpoints and set a breakpoint on the second -``copy_tile[1, 128](source, destination, items_per_thread)`` in the example. -Continue. The first launch has completed and its assertion has passed. -Inspect ``source[:8]`` and ``destination[:8]`` in the example's ``main`` -frame: ``source`` starts at 0, while ``destination`` contains the ``-1`` -values written immediately before this second launch. - -Re-enable the group-planner and NVRTC breakpoints, then continue. This -launch uses the existing kernel specialization, so it should finish -without those compilation stops. Both verification messages should print. - -Start a new debug session to repeat the whole tour. Disabling the provider -disk cache does not disable in-process provider or kernel reuse. Editing -source while paused also does not replace the function already loaded -into that process. - -A second pass: shared scratch and synchronization -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Replace just the kernel in the example with this version. Before restarting, -set the two additional storage breakpoints listed below: - -.. code-block:: python - - @cuda.jit - def copy_tile(source, destination, items_per_thread): - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - scratch = coop.TempStorage(auto_sync=True) - coop.load( - block, source, items, algorithm="transpose", temp_storage=scratch - ) - coop.store( - block, destination, items, algorithm="transpose", temp_storage=scratch - ) - -The expected result stays the same. At the Load/Store plan breakpoint, -inspect ``plan.temp_storage`` and ``plan.synchronization`` again. The -transpose algorithms exchange data through shared memory and require a -block barrier before that scratch can be reused. - -At the invocable breakpoint, the storage ABI is now ``LEADING_POINTER``, and -the temporary-storage size is nonzero. The compiled provider supplies its -size and alignment; avoid hard-coding the numbers from one toolkit. - -The additional stops are in ``_compiler/_rewrite_storage.py``. Set both -before starting this pass. ``begin_rewrite()`` stages the allocation after -function preparation and before any block is matched or emitted. - -* In ``_emit_temp_storage_backing()``, stop on - ``alloc_size = 0 if plan.uses_dynamic_smem else int(plan.total_size)``. - Inspect ``plan.total_size``, ``plan.max_alignment``, and - ``plan.uses_dynamic_smem``. The rewrite emits shared storage satisfying - the provider requirements in the entry block. Both uses of ``scratch`` - refer to this planned allocation, which precedes every scratch view. -* In ``_emit_temp_storage_auto_sync()``, stop on ``sync_args = []``. - Inspect ``synchronization_scope`` and ``sync_attr``. For this block - example they select a block barrier, emitted as ``syncthreads``. - -Compare ``new_block.dump()`` at the end of ``apply()`` with the direct -version. It now includes the shared allocation, scratch arguments, and -synchronization in addition to the per-thread ``items`` array. -``TempStorage`` describes shared workspace for the primitive; -``ThreadData`` holds each thread's payload. Seeing both in the rewritten -IR makes their different lifetimes and uses concrete. - -If a breakpoint does not stop -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -* Check ``coop.__file__`` at the first launch and confirm the breakpoint - belongs to that checkout. An older installed wheel or another worktree - can otherwise supply the running code. -* Confirm you started the named launch configuration with - ``justMyCode: false``. A hollow breakpoint can remain pending until its - module loads; check it again after the import. -* If registration is skipped, restart with the imports in the example's - order and ``CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION=0``. Read any - automatic-registration warning for an incompatible compiler dependency. -* If NVRTC is skipped, use a fresh process with - ``CUDA_COOP_ENABLE_CACHE=0``. Break in ``_nvrtc.compile()`` to see the - request before the compilation cache, and check which launch you are on. - -Working on a primitive ----------------------- - -For a new operation, start with the C++ overload and a concrete kernel -that should use it. Work out the payload arrangement, participation, -output ownership, and scratch lifetime before writing its public wrapper. -These choices determine which signatures the backend can implement. - -The existing Load/Store family shows the usual path: - -#. Add the shared signature and type declarations in ``_core/api/`` when - the operation belongs in the common API. Put compiler-specific - extensions in the qualified namespace. -#. Describe the operation and C++ overload in the core family. Its group - planner selects a supported implementation and returns complete result, - topology, storage, and synchronization requirements. -#. Add the Numba family binding under ``numba_mlir/_compiler/`` and its - provider lowering under ``numba_mlir/_lowering/``. - ``register_group_primitive()`` connects group-call planning; - ``register_rewrite_operation()`` describes the provider rewrite's - arguments and family-specific analysis. -#. Check that the existing materialization and storage code can consume - that description. Change the shared rewrite only when the new operation - needs a compiler behavior that the existing requirements cannot express. - -Tests live under ``python/cuda_coop/tests/``. Core contract tests check -normalization and plan selection without a compiler. Backend unit tests -check argument binding, inference, rewrites, and diagnostics. Compile -tests use real NVRTC and nvJitLink with devices hidden; their fixtures -provide an explicit target. Runtime tests check the resulting kernels. - - -Use tests that exercise the part you changed. A result-ownership change -needs a check of the operation's documented mutation behavior. A storage -change needs repeated calls and multiple independent groups. A callable ABI -change needs a real link and a runtime result. A mocked compiler test cannot -establish that the generated wrapper and operator agree on their ABI. - -For example, from the CCCL root in an environment with the Numba backend -and test dependencies installed: - -.. code-block:: console - - CUDA_VISIBLE_DEVICES="" PYTHONPATH="$PWD/python/cuda_coop" python -m pytest \ - python/cuda_coop/tests/test_group_load_store.py \ - python/cuda_coop/tests/backends/numba_mlir/unit/test_group_lowering_plan.py - - CUDA_VISIBLE_DEVICES="" PYTHONPATH="$PWD/python/cuda_coop" python -m pytest \ - python/cuda_coop/tests/backends/numba_mlir/compile/test_block_load_store_compile.py - -The source path prevents an older installed ``cuda-coop`` wheel from -silently supplying the module under test. The compile tests' fixed target -is a test fixture; it does not imply that every public kernel compilation -path is available without a GPU or configured launch. - -To inspect generated C++, set ``CUDA_COOP_SOURCE_DUMP_DIR`` to a -directory before compiling the kernel. Files are named -``cuda_coop__.cu``, allowing backends to share the directory. -Use the dump to check template arguments, wrapper signatures, and scratch -metadata. Inspect the final kernel's PTX or SASS separately for inlining, -barriers, and register behavior. - -.. _coop-numba-validation: - -Numba-CUDA-MLIR validation scope --------------------------------- - -The package supports Python 3.10 through 3.14. The CI matrix configures these -parts of that range: - -.. list-table:: - :header-rows: 1 - :widths: 36 64 - - * - Environment - - Automated checks - * - Linux x86-64, Python 3.14, CUDA 13 - - Installed-wheel compilation with GPUs hidden and L4 runtime tests - in pull requests - * - Linux x86-64, Python 3.14, CUDA 12 - - Installed-wheel compilation and GPU runtime tests in the nightly matrix - * - Linux x86-64, Python 3.10 and 3.14 - - Common API host contracts and wheel packaging - * - Windows x86-64, Python 3.10 and 3.14 - - Universal-wheel build, base import, and bundled-header checks - -The Windows checks do not compile or launch Numba-CUDA-MLIR kernels. Other -Python versions and platform combinations need separate backend runtime -qualification. Dependency bounds allow releases in the supported series; -they do not mean that every patch release in that series has been tested. -These checks cover Block and Warp Load/Store. They do not qualify operations -on thread, cluster, or grid groups. - -Source map ----------- - -Paths below are relative to ``python/cuda_coop/cuda/coop/``: - -.. list-table:: - :header-rows: 1 - :widths: 47 53 - - * - Path - - Responsibility - * - ``_core/api/`` and the adjacent ``.pyi`` files - - Common signatures, descriptors, and argument rules. - * - ``_core/group/`` - - Group resolution, primitive semantics, and lowering requirements. - * - ``_core/block/`` and ``_core/warp/`` - - CUB algorithm specializations and generated support code. - * - ``numba_mlir/_compiler/_planner.py`` - - Sequence group resolution and call rewriting in one compiler hook. - * - ``numba_mlir/_compiler/_group_planner.py`` and ``_group_*.py`` - - Recover group calls from compiler IR and invoke family planners. - * - ``numba_mlir/_compiler/_rewrite.py`` and ``_rewrite_*.py`` - - Materialize payloads, invocables, result handling, and storage. - * - ``numba_mlir/_lowering/`` - - Translate core specializations and generate CUB providers. - * - ``numba_mlir/_types.py`` - - Device ABIs, Python operator compilation, source generation, and - invocable link inputs. - * - ``_headers/`` and ``numba_mlir/_compiler/_nvrtc.py`` - - Header and toolkit selection, compile identity, and NVRTC invocation. - -Runnable kernels are in ``python/cuda_coop/examples/numba_mlir/``. The -installed ``.pyi`` files and :doc:`API reference <../coop_api>` describe -the public signatures; private provider factories are implementation -details. diff --git a/docs/python/coop/faqs.rst b/docs/python/coop/faqs.rst deleted file mode 100644 index 855d1f183ef8..000000000000 --- a/docs/python/coop/faqs.rst +++ /dev/null @@ -1,189 +0,0 @@ -.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -.. _coop-faqs: - -FAQs -==== - -.. _coop-faq-namespaces: - -Why are there both ``cuda.coop`` and ``cuda.coop.numba_mlir``? --------------------------------------------------------------- - -``cuda.coop`` provides the common API for cooperative operations. A kernel -compiler's backend implements those calls. Start with this namespace when -its groups, ``ThreadData`` payloads, and operations cover your needs: - -.. code-block:: python - - from cuda import coop - - coop.register("numba-cuda-mlir") - -``cuda.coop.numba_mlir`` exposes that backend's API, including extensions -specific to Numba-CUDA-MLIR. Use it for features such as fixed-size Numba -local-array payloads: - -.. code-block:: python - - import cuda.coop.numba_mlir as numba_coop - -Importing this namespace also registers the backend. Both namespaces can -appear in one kernel, and shared operations follow the same contracts. -See the :ref:`API comparison ` for the -operation-specific differences. - -Numba-CUDA-MLIR is the first backend; CUTLASS support is planned. The common -API gives libraries a compiler-independent way to express cooperative -operations. Kernel launch syntax and other DSL-specific code still need -adaptation when moving to another compiler. - -.. _coop-faq-numba-only: - -I only use Numba-CUDA-MLIR. Can I import its namespace as ``coop``? -------------------------------------------------------------------- - -Yes. This is supported and registers the backend: - -.. code-block:: python - - from numba_cuda_mlir import cuda - - import cuda.coop.numba_mlir as coop - -You can use common operations and backend extensions through that one name. -The documentation uses ``numba_coop`` when showing backend calls alongside -common calls, so readers can see which API an example needs. - -Keep the alias on a dotted import. Bare ``import cuda.coop.numba_mlir`` -assigns the top-level package to ``cuda`` in that scope, replacing the name -previously imported from ``numba_cuda_mlir``. - -.. _coop-faq-thread-data-dtype: - -When does ``ThreadData`` need an explicit element type? -------------------------------------------------------- - -Start with an inferred payload: - -.. code-block:: python - - # Inside a kernel; source and items_per_thread are kernel arguments. - items = coop.ThreadData(items_per_thread) - coop.load(coop.this_block(), source, items) - -Load supplies the source element type. In Numba-CUDA-MLIR, supported indexed -assignments and a Store destination can also establish the type. Other -cooperative producers define their output types where their contracts say so. -Inference follows the backend's supported operations and assignments. - -Numba's current planner cannot infer a payload type solely from the results of -a non-inlined device helper when no other supported operation supplies type -context. Cast the assigned result to the intended scalar type, such as -``numpy.int32``, or provide that context through a supported operation such as -Store. This is a limit of the current planner; the helper's result may already -have a type that the later compiler phases can determine. - -An explicit signature on the scalar helper does not supply that missing -context to this early planner. - -Helpers containing cooperative operations follow the separate -:ref:`device-helper inlining rules `. Those rules -also cover passing groups and ``ThreadData`` payloads between helpers and -the calling kernel. - -Use typed values when the computation needs a particular width or precision. -The optional ``dtype`` parameter supplies element-type information when the -surrounding program cannot establish it. It does not initialize the payload, -resolve conflicting typed values, or enable unsupported types. Initialize -every item before reading it. - -.. _coop-faq-exclusive-storage: - -Why use ``sharing="exclusive"`` instead of omitting storage? ------------------------------------------------------------- - -Omitting ``temp_storage`` lets the compiler choose the scratch layout and -insert reuse barriers. It may reuse scratch across compatible calls; omission -does not guarantee a separate slice for each call site. - -``TempStorage(sharing="exclusive")`` gives distinct call sites separate -slices while retaining the descriptor's capacity, alignment, and -synchronization controls. Two calls with separate slices need no barrier -solely to reuse each other's scratch. This can save that synchronization when -``auto_sync=False``, at the cost of more shared memory. Barriers required by -the algorithm or by application data dependencies still apply. - -Repeated execution of one call site, including a loop, reuses its slice. -Synchronize before that reuse or set ``auto_sync=True`` to request trailing -reuse barriers. ``sharing`` controls layout independently of ``auto_sync``. -The backend accounts for cooperative scratch with either explicit or omitted -storage; ``exclusive`` is a choice about which calls may share its bytes. - -.. _coop-faq-temp-storage: - -Why or when would I provide my own ``TempStorage``? ---------------------------------------------------- - -Usually, you can omit it. The compiler allocates the scratch required by -an operation and inserts a barrier to make reuse safe. Direct, striped, and -vectorized Load/Store need no shared scratch. - -An explicit descriptor lets you choose which supported block operations -share an allocation. For example, a transpose Load and Store can -share scratch while the values stay in each thread's payload: - -.. code-block:: python - - # Inside a kernel; source, destination, and items_per_thread are arguments. - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - scratch = coop.TempStorage(auto_sync=True) - coop.load( - block, source, items, algorithm="transpose", temp_storage=scratch - ) - coop.store( - block, destination, items, algorithm="transpose", temp_storage=scratch - ) - -The compiler sizes and aligns the shared region for both calls. It inserts -a barrier after each use. Construct the descriptor inside the kernel, and -keep application values in ``ThreadData`` or your own arrays. - -A descriptor also lets you request capacity or alignment, or choose separate -slices with ``sharing="exclusive"``. Explicit descriptors default to -``auto_sync=False``, so the kernel must provide reuse barriers, including across -loop iterations. The example requests ``auto_sync=True`` to insert those -barriers automatically. See :ref:`exclusive scratch slices -` for the tradeoff between memory and reuse -synchronization. - -Explicit descriptors control scratch for block transpose-family Load/Store. -Warp operations reject explicit descriptors. The compiler allocates any -scratch required by CUB operations. See :ref:`temporary storage -` for the complete contract and shared-memory restrictions. - -.. _coop-faq-installed-extra: - -Can ``cuda.coop`` tell which extra I installed? ------------------------------------------------ - -Not reliably. Package metadata lists the extras a distribution offers and -the dependencies associated with them. It does not provide a portable record -of which extra was requested when installing. The same dependencies may also -have been installed separately or by another package. - -``pip install cuda-coop`` and -``pip install "cuda-coop[numba-cuda-mlir-cu13]"`` install the same wheel, -including ``cuda.coop.numba_mlir`` and every other shipped DSL integration. -The base install declares no Python package dependencies. The extra only -adds the requirements in ``pyproject.toml`` that install the supported -Numba-CUDA-MLIR stack for CUDA 13. If those dependencies are already present -at supported versions, either command gives you the same usable API. -Registration selects which compiler hooks to activate in the running process. - -Use ``coop.register("numba-cuda-mlir")`` to state that intent explicitly. -It works regardless of import order, is safe to repeat, and also accepts -``"numba_cuda_mlir"``. The backend dependencies must already be installed. -See :ref:`backend registration `. diff --git a/docs/python/coop/glossary.rst b/docs/python/coop/glossary.rst deleted file mode 100644 index 486c8ef116d2..000000000000 --- a/docs/python/coop/glossary.rst +++ /dev/null @@ -1,156 +0,0 @@ -.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -.. _coop-glossary: - -Glossary -======== - -Terms ------ - -.. glossary:: - :sorted: - - backend - The integration that lets a kernel compiler recognize ``cuda.coop`` - calls and generate CUDA code for them. Numba-CUDA-MLIR is the first - supported backend. See :ref:`registration `. - - blocked - A layout in which each thread owns consecutive elements of a - :term:`tile`. With ``K`` items per thread, thread ``t`` owns tile - positions ``t * K`` through ``t * K + K - 1``. See - :ref:`blocked versus striped `. - - primitive - .. raw:: html - - - - A cooperative operation provided by ``cuda.coop``, such as ``load`` - or ``store``. Each primitive defines which threads - participate, how it uses their data, and where its results are - available. See - :ref:`participation and synchronization `. - - family - A group of related :term:`primitives ` that share semantics - and implementation. The Load/Store family includes ``load`` and - ``store``. Families organize implementation modules; a - :term:`thread group` describes the threads executing a primitive. - See :ref:`implementation families `. - - payload - The values contributed or received by one thread. ``ThreadData(K)`` - describes a fixed-size payload of ``K`` items. Some operations also - accept scalars or backend-specific arrays. See - :ref:`thread data `. - - phi - A compiler intermediate representation (IR) operation that selects a - value according to the control-flow path taken into a block. For - example, after an ``if`` statement assigns different values to a - variable, a phi selects the value from the branch that ran. At a loop - header, it can select the initial value or the value from the previous - iteration. Comments refer to these alternatives as *phi inputs*. - - common API - .. raw:: html - - - - The backend-independent API exposed through ``from cuda import coop``. - It describes operations on thread groups, values, and storage. - Support for particular operations and argument types depends on the - backend. Qualified APIs provide backend-specific extensions. See - :ref:`choosing an API `. - - qualified API - A backend's namespace, such as ``cuda.coop.numba_mlir``. It provides - the common operations and compiler-specific extensions. See - :ref:`namespace choices `. - - striped - A layout in which consecutive threads own consecutive tile elements - at each per-thread item position. For ``G`` threads, thread ``t`` owns - tile positions ``t``, ``t + G``, ``t + 2 * G``, and so on. See - :ref:`blocked versus striped `. - - temporary storage - Shared-memory scratch used internally by a primitive. The compiler - allocates it automatically when needed. An explicit ``TempStorage`` - descriptor can control allocation sharing and synchronization for - supported operations. Keep application data in payloads or arrays. - See :ref:`the temporary-storage FAQ `. - - thread group - The threads participating in a primitive, such as a block, a physical - warp, or a logical group within a warp. A thread's rank identifies its - position within that group. See :ref:`groups `. - - tile - The sequence of values processed by one group in an operation. With - ``G`` threads and ``K`` items per thread, a full tile has ``G * K`` - values. ``valid_items`` in Load and Store selects a prefix of that - sequence. It counts tile elements, not elements per thread. - -.. _coop-glossary-layouts: - -Blocked versus striped ----------------------- - -Both layouts distribute the same tile across threads. They differ in which -thread owns each position. For four illustrative threads and two items per -thread, a tile of positions ``0`` through ``7`` is distributed as follows: - -.. list-table:: - :header-rows: 1 - - * - 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 slot ``i``: - -.. list-table:: - :header-rows: 1 - - * - Layout - - Tile position owned by ``items[i]`` in thread ``t`` - - Owner of tile position ``p`` - * - Blocked - - ``t * K + i`` - - Thread ``p // K``, slot ``p % K`` - * - Striped - - ``t + i * G`` - - Thread ``p % G``, slot ``p // G`` - -The four-thread table illustrates ownership, not a supported physical CUDA -warp size. A physical warp has 32 threads. - -Layout describes the values a caller sees. The memory-access algorithm may -use a different intermediate layout. For example, ``load(..., -algorithm="transpose")`` reads striped values, then rearranges them into -blocked payloads. ``algorithm="striped"`` leaves the payloads striped. - -Store expects the layout selected by its algorithm. Combining a striped -Load with a direct Store without conversion changes the output order. -``ThreadData`` does not carry a layout tag that corrects this mismatch. - -Follow the values in the :doc:`Load ` and :doc:`Store -` visualizations. The :ref:`programming guide -` explains these layouts. diff --git a/docs/python/coop/programming_guide.rst b/docs/python/coop/programming_guide.rst deleted file mode 100644 index f53b34a61cda..000000000000 --- a/docs/python/coop/programming_guide.rst +++ /dev/null @@ -1,609 +0,0 @@ -.. Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -.. -.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -.. _cuda.coop.programming_guide: - -``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 :doc:`visualizations ` show how values move through -these operations. Each explorer includes an example kernel and lets you -step through the algorithm. The :doc:`glossary ` defines terms and -layouts; the :doc:`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 -:doc:`installation instructions <../coop>` describe the matching -``cuda-coop`` extra; the current backend requires -``numba-cuda-mlir>=0.5.0,<0.6``. Check the -:ref:`validation scope ` for environment coverage. -Keep a compiled kernel in its original CUDA context; see the -:ref:`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. - -.. _coop-numba-requirements: - -Requirements ------------- - -Use Python 3.10 through 3.14 and ``numba-cuda-mlir>=0.5.0,<0.6``. The -:doc:`installation instructions <../coop>` describe the ``cuda-coop`` extra -matching your CUDA major version. Running the examples requires a -CUDA-capable GPU. The developer guide's -:ref:`validation scope ` describes the environments -covered by CI. Before reusing compiled kernels, check the -:ref:`device and context-lifetime rules ` and -:ref:`disk-cache behavior `. - -A first kernel: copying a tile --------------------------------- - -The :func:`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: - -.. code-block:: python - - import numpy as np - from numba_cuda_mlir import cuda, types - from cuda import coop - -.. _coop-programming-api-choice: - -Choosing the common or qualified API ------------------------------------- - -The common API is imported with: - -.. code-block:: python - - 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: - -.. code-block:: python - - 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: - -.. list-table:: - :header-rows: 1 - :widths: 24 34 42 - - * - 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: - -.. code-block:: python - - 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 :ref:`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 ----------------------------- - -.. _coop-numba-context-lifetime: - -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. - -.. _coop-numba-kernel-cache: - -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 -:ref:`compiler activation and reuse ` for its -configuration and lifetime. - -.. _coop-thread-groups: - -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. - - -.. _coop-participation: - -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. - - -.. _coop-thread-data: - -``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. - -:class:`~cuda.coop.ThreadDataLike` names the common payload interface used -in API signatures. It describes the item count, dtype, and indexed reads and -writes. Use :func:`~cuda.coop.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: - -.. code-block:: python - - 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. - -.. _coop-data-layouts: - -Blocked and striped order -^^^^^^^^^^^^^^^^^^^^^^^^^ - -The glossary defines :term:`blocked` and :term:`striped` ownership and -compares their :ref:`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. - -.. list-table:: - :header-rows: 1 - - * - 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 :doc:`Load ` and -:doc:`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 :ref:`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 -^^^^^^^^^^^^^^^^^ - -.. list-table:: - :header-rows: 1 - :widths: 25 20 35 20 - - * - 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: - -.. literalinclude:: ../../../python/cuda_coop/tests/backends/numba_mlir/runtime/test_programming_guide_examples.py - :language: python - :name: coop-pg-warp-copy - :start-after: # coop-pg-warp-copy-begin - :end-before: # coop-pg-warp-copy-end - :dedent: 4 - -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. - -.. _coop-numba-device-helpers: - -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``: - -.. literalinclude:: ../../../python/cuda_coop/tests/backends/numba_mlir/runtime/test_device_helpers_example.py - :language: python - :name: coop-pg-device-helpers - :start-after: # coop-pg-device-helpers-begin - :end-before: # coop-pg-device-helpers-end - :dedent: 4 - -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 -:ref:`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. - -.. _coop-temp-storage: - -.. _tempstorage-scratch-used-during-a-collective: - -``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. - -:class:`~cuda.coop.TempStorageLike` names the common interface for explicit -scratch descriptors in API signatures. Construct one inside the kernel with -:func:`~cuda.coop.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. - -.. list-table:: - :header-rows: 1 - - * - 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 -:ref:`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 :doc:`API reference -<../coop_api>` records exact signatures; the :doc:`overview <../coop>` -collects operation restrictions and configuration. For generated-source -diagnostics and the compiler integration, see the -:doc:`Developer Overview `. diff --git a/docs/python/coop/source_dumps/direct.cpp.txt b/docs/python/coop/source_dumps/direct.cpp.txt deleted file mode 100644 index eb8c279a9f85..000000000000 --- a/docs/python/coop/source_dumps/direct.cpp.txt +++ /dev/null @@ -1,33 +0,0 @@ -// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -// -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -// NVRTC capture excerpt: generated identifiers shortened and whitespace formatted. -// The archive contains the original translation unit and the identifier mapping. -// excerpt-begin -#include - -#include - -using block_load_t = - cub::BlockLoad<::cuda::std::int32_t, 128, 2, - ::cub::BLOCK_LOAD_DIRECT, 1, 1>; - -extern "C" __device__ void -load_impl(::cuda::std::int32_t *src, - ::cuda::std::int32_t *dst) { - block_load_t().Load( - src, - *reinterpret_cast<::cuda::std::int32_t (*)[2]>( - dst)); -} - -extern "C" __device__ int -load_abi(void *__ret, void *src, void *dst) { - ::cuda::std::int32_t *cast_src = - reinterpret_cast<::cuda::std::int32_t *>(src); - ::cuda::std::int32_t *cast_dst = - reinterpret_cast<::cuda::std::int32_t *>(dst); - load_impl(cast_src, cast_dst); - return 0; -} diff --git a/docs/python/coop/source_dumps/transpose.cpp.txt b/docs/python/coop/source_dumps/transpose.cpp.txt deleted file mode 100644 index 4d081a848782..000000000000 --- a/docs/python/coop/source_dumps/transpose.cpp.txt +++ /dev/null @@ -1,48 +0,0 @@ -// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -// -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -// NVRTC capture excerpt: generated identifiers shortened and whitespace formatted. -// The archive contains the original translation unit and the identifier mapping. -// excerpt-begin -#include - -#include - -using block_load_t = - cub::BlockLoad<::cuda::std::int32_t, 128, 2, - ::cub::BLOCK_LOAD_TRANSPOSE, 1, 1>; - -using load_storage_t = - typename block_load_t::TempStorage; - -__device__ constexpr unsigned load_storage_bytes = - sizeof(load_storage_t); - -__device__ constexpr unsigned load_storage_alignment = - alignof(load_storage_t); - -extern "C" __device__ void -load_impl(load_storage_t *temp_storage, - ::cuda::std::int32_t *src, - ::cuda::std::int32_t *dst) { - block_load_t(*temp_storage) - .Load(src, - *reinterpret_cast< - ::cuda::std::int32_t (*)[2]>(dst)); -} - -extern "C" __device__ int load_abi(void *__ret, - void *temp_storage, - void *src, - void *dst) { - load_storage_t *cast_temp_storage = - reinterpret_cast( - temp_storage); - ::cuda::std::int32_t *cast_src = - reinterpret_cast<::cuda::std::int32_t *>(src); - ::cuda::std::int32_t *cast_dst = - reinterpret_cast<::cuda::std::int32_t *>(dst); - load_impl(cast_temp_storage, cast_src, cast_dst); - return 0; -} diff --git a/docs/python/coop/visualizations/compiler-phases-classes-mobile.svg b/docs/python/coop/visualizations/compiler-phases-classes-mobile.svg deleted file mode 100644 index bb4583c30279..000000000000 --- a/docs/python/coop/visualizations/compiler-phases-classes-mobile.svg +++ /dev/null @@ -1,261 +0,0 @@ - - - - Compiler phases — Python classes - Seven compiler phases in order. Cooperative planning remains before main type inference. Group planning computes Load/Store payload dtype and item count. PayloadInference reuses the carried GroupLoweringPlan, checks explicit inputs and constructor extent, and records constructor state. prepare_calls_and_storage analyzes all original calls, retains _RewriteMatch records, validates descriptor ownership, prepares providers, and collects scratch requirements from their compiled layouts. Later consumers can supply facts needed by earlier constructors. begin_rewrite stages shared backing in the entry block so it dominates every scratch view. match selects each block; apply reuses or materializes invocables and emits statements in source order. ThreadData holds each thread's payload; TempStorage describes cooperative scratch. Standalone payload-write dtype inference remains in apply. finish_rewrite cleans unused bindings after all replacement blocks are installed, preserving original definitions throughout emission. These steps belong to one planner, not separate compiler passes. Facts belong to the current IR and compilation attempt. Provider artifact caches and pipeline placement are unchanged. - - - PYTHON CLASSES · NUMBA-CUDA-MLIR - Compiler phases - Read - - from - - top - - to - - bottom. - - Cooperative - - planning - runs - - before - - main - - type - - inference. - - - 01 - Build untyped IR - - Numba - TranslateByteCode - → IRProcessing - MLIRCompiler - - defines - - the - - pipeline - - in - a - - PassManager - . - Produces - - FunctionIR - - with - - ir.Block - - and - - ir.Expr - - objects. - - - - - 02 - Early rewrites - - Numba - GenericRewrites - NumbaCudaMlirInlineInlinables - Run - - before-inference - - rewrites, - - then - - inline - - supported - - device - - helpers. - - - - - 03 - Cooperative planning - - cuda.coop - NUMBA COMPILER PASS - PostInlineWholeFunctionPlanners - REGISTERED PLANNER - CoopWholeFunctionPlanner - .run() - repairs IR between phases - - 1 RESOLVE GROUPS - _GroupPlanning - ._resolve_groups() - Plans Load/Store dtype and count. - Carries a GroupLoweringPlan. - 2 REWRITE CALLS - _CallRewriting - ._rewrite_calls() - PayloadInference reuses facts; - checks inputs; records ctor state. - - REWRITER - CoopSinglePhaseRewrite - prepare_calls_and_storage():all calls - Provider layouts set scratch needs. - begin_rewrite(): entry backing - match() / apply(): emit blocks - finish_rewrite(): final cleanup - INHERITED HELPERS / CARRIED DATA - _StorageRewrite - - — - - requirements - - and - - scratch - - allocation - _InvocableRewrite - - — - - providers - - and - - Invocable - - objects - _RewriteMatch — retained call - and GroupLoweringPlan - - Cooperative - - planning - - complete - - - - 04 - Finish untyped passes - - Numba - FindLiterallyCalls - NumbaCudaMlirLiteralUnroll - ReconstructSSA - LiteralPropagationSubPipelinePass - - - - - 05 - Infer types - - Numba - MLIRTypeInference - Build - - the - - type - - map - - and - - call - - signatures. - Invocable._numba_type_ - - describes - - the - - provider - - call. - - - - - 06 - Typed rewrites - - Numba - PreLowerStripPhis - → InlineOverloads - NopythonRewrites - Run - - after-inference - - rewrites. - - - - - 07 - Lower, link, load - - Numba - MLIRBackend - → MLIRLower - Lower - - symbolic - - calls - - and - - collect - - link - - inputs. - optimize() - - compiles - - the - - kernel; - - nvJitLink - - links - - it. - CodeLibrary.get_cufunc() - - loads - the - - kernel - - handle. - - diff --git a/docs/python/coop/visualizations/compiler-phases-classes.svg b/docs/python/coop/visualizations/compiler-phases-classes.svg deleted file mode 100644 index 76567fe08347..000000000000 --- a/docs/python/coop/visualizations/compiler-phases-classes.svg +++ /dev/null @@ -1,257 +0,0 @@ - - - - Compiler phases — Python classes - Seven compiler phases in order. Cooperative planning remains before main type inference. Group planning computes Load/Store payload dtype and item count. PayloadInference reuses the carried GroupLoweringPlan, checks explicit inputs and constructor extent, and records constructor state. prepare_calls_and_storage analyzes all original calls, retains _RewriteMatch records, validates descriptor ownership, prepares providers, and collects scratch requirements from their compiled layouts. Later consumers can supply facts needed by earlier constructors. begin_rewrite stages shared backing in the entry block so it dominates every scratch view. match selects each block; apply reuses or materializes invocables and emits statements in source order. ThreadData holds each thread's payload; TempStorage describes cooperative scratch. Standalone payload-write dtype inference remains in apply. finish_rewrite cleans unused bindings after all replacement blocks are installed, preserving original definitions throughout emission. These steps belong to one planner, not separate compiler passes. Facts belong to the current IR and compilation attempt. Provider artifact caches and pipeline placement are unchanged. - - - PYTHON CLASSES · NUMBA-CUDA-MLIR - Compiler phases - Read - - from - - top - - to - - bottom. - - Cooperative - - planning - - runs - - before - - main - - type - - inference. - - - 01 - Build untyped IR - - Numba - TranslateByteCode → IRProcessing - MLIRCompiler - - defines - - the - - pipeline - - in - - a - - PassManager - . - Produces - - FunctionIR - - with - - ir.Block - - and - - ir.Expr - - objects. - - - - - 02 - Early rewrites - - Numba - GenericRewrites - NumbaCudaMlirInlineInlinables - Run - - before-inference - - rewrites, - - then - - inline - - supported - - device - - helpers. - - - - - 03 - Cooperative planning - - cuda.coop - NUMBA COMPILER PASS - PostInlineWholeFunctionPlanners - REGISTERED PLANNER - CoopWholeFunctionPlanner.run() - Repairs changed IR between phases. - - 1 RESOLVE GROUPS - _GroupPlanning._resolve_groups() - _GroupCallPlanner: Load/Store dtype/count. - GroupLoweringPlan → private calls - 2 REWRITE CALLS - _CallRewriting._rewrite_calls() - PayloadInference reuses plan facts; - checks inputs; records constructor state. - - REWRITER - CoopSinglePhaseRewrite - prepare_calls_and_storage() - All calls; compiled scratch layouts. - begin_rewrite() - Entry backing before any block. - match() / apply(): emit blocks; finish_rewrite(): final cleanup - INHERITED HELPERS / CARRIED DATA - _StorageRewrite - - — - - requirements - - and - - scratch - - allocation - _InvocableRewrite - - — - - providers - - and - - Invocable - - objects - _RewriteMatch — retained call and GroupLoweringPlan - - Cooperative - - planning - - complete - - - - 04 - Finish untyped passes - - Numba - FindLiterallyCalls - NumbaCudaMlirLiteralUnroll - ReconstructSSA - LiteralPropagationSubPipelinePass - - - - - 05 - Infer types - - Numba - MLIRTypeInference - Build - - the - - type - - map - - and - - call - - signatures. - Invocable._numba_type_ - - describes - - the - - provider - - call. - - - - - 06 - Typed rewrites - - Numba - PreLowerStripPhis → InlineOverloads - NopythonRewrites - Run - - after-inference - - rewrites. - - - - - 07 - Lower, link, load - - Numba - MLIRBackend → MLIRLower - Lower - - symbolic - - calls - - and - - collect - - link - - inputs. - optimize() - - compiles - - the - - kernel; - - nvJitLink - - links - - it. - CodeLibrary.get_cufunc() - - loads - - the - - kernel - - handle. - - diff --git a/docs/python/coop/visualizations/compiler-phases.svg b/docs/python/coop/visualizations/compiler-phases.svg deleted file mode 100644 index 93372e7c497d..000000000000 --- a/docs/python/coop/visualizations/compiler-phases.svg +++ /dev/null @@ -1,64 +0,0 @@ - - - - Numba-CUDA-MLIR compilation phases - Seven compiler phases in order. Cooperative planning remains before main type inference. Group planning computes Load/Store payload dtype and item count. PayloadInference reuses the carried GroupLoweringPlan, checks explicit inputs and constructor extent, and records constructor state. prepare_calls_and_storage analyzes all original calls, retains _RewriteMatch records, validates descriptor ownership, prepares providers, and collects scratch requirements from their compiled layouts. Later consumers can supply facts needed by earlier constructors. begin_rewrite stages shared backing in the entry block so it dominates every scratch view. match selects each block; apply reuses or materializes invocables and emits statements in source order. ThreadData holds each thread's payload; TempStorage describes cooperative scratch. Standalone payload-write dtype inference remains in apply. finish_rewrite cleans unused bindings after all replacement blocks are installed, preserving original definitions throughout emission. These steps belong to one planner, not separate compiler passes. Facts belong to the current IR and compilation attempt. Provider artifact caches and pipeline placement are unchanged. - - - COMPILATION PHASES - One compiler attempt; default LTO path for a kernel with cooperative calls. - - 1 - Build untyped IR - Numba-CUDA-MLIR - Translate Python bytecode. - Form blocks, expressions, and call sites. - - - 2 - Run early rewrites - Numba-CUDA-MLIR - Fold constants and prune dead branches. - Run before-inference rewrites. Inline device helpers. - - - 3 - Plan cooperative calls - cuda.coop - Plan groups; reuse payload facts; prepare providers. - Stage entry backing; match and replace each block. - Clean unused bindings after every block is installed. - - - 4 - Finish untyped passes - Numba-CUDA-MLIR - Process literals and supported unrolling. - Reconstruct SSA and propagate literals. - - - 5 - Infer types - Numba-CUDA-MLIR - MLIRTypeInference assigns expression types - and call signatures to the rewritten kernel. - - - 6 - Run typed rewrites - Numba-CUDA-MLIR - Strip phi nodes and inline overloads. - Run after-inference rewrites. - - - 7 - Lower, link, and launch - Numba-CUDA-MLIR + CUDA tools - Lower through MLIR. Compile kernel LTO-IR. - nvJitLink combines it with provider LTO-IR. - Load the cubin and launch the kernel. - Phase 3 completes before the main type-inference pass. The next figure expands its handoffs. - diff --git a/docs/python/coop/visualizations/compiler-swimlane-classes-mobile.svg b/docs/python/coop/visualizations/compiler-swimlane-classes-mobile.svg deleted file mode 100644 index 81fd1011c6fc..000000000000 --- a/docs/python/coop/visualizations/compiler-swimlane-classes-mobile.svg +++ /dev/null @@ -1,354 +0,0 @@ - - - - Compilation handoffs — Python classes - Ordered handoffs between Numba, cuda.coop, NVRTC and nvJitLink. Successful provider bundling on a cache miss and the default kernel LTO path are shown. Group planning computes Load/Store payload dtype and item count. PayloadInference reuses the carried GroupLoweringPlan, checks explicit inputs and constructor extent, and records constructor state. prepare_calls_and_storage analyzes all original calls, retains _RewriteMatch records, validates descriptor ownership, prepares providers, and collects scratch requirements from their compiled layouts. Later consumers can supply facts needed by earlier constructors. begin_rewrite stages shared backing in the entry block so it dominates every scratch view. match selects each block; apply reuses or materializes invocables and emits statements in source order. ThreadData holds each thread's payload; TempStorage describes cooperative scratch. Standalone payload-write dtype inference remains in apply. finish_rewrite cleans unused bindings after all replacement blocks are installed, preserving original definitions throughout emission. These steps belong to one planner, not separate compiler passes. Facts belong to the current IR and compilation attempt. Provider artifact caches and pipeline placement are unchanged. - - - PYTHON CLASSES · NUMBA-CUDA-MLIR - Compilation handoffs - A - - provider-cache - - miss, - - successful - - bundling, - - and - - the - - default - - kernel - - LTO - - path. - 1 Prepare providers - Before - - main - - type - - inference - - - 01 - Invoke the planner - - Numba - PostInlineWholeFunctionPlanners - - Inlined, - - untyped - - IR - - - - 02 - Resolve and prepare - - cuda.coop - CoopWholeFunctionPlanner - .run() - CoopSinglePhaseRewrite - .prepare_calls_and_storage() - Plan Load/Store dtype and count. - PayloadInference reuses facts. - Check inputs; record ctor state. - Retain each _RewriteMatch. - - During prepare_calls_and_storage() - - - - 03 - Build the provider bundle - - cuda.coop - Algorithm - collect_specializations() - prepare_ltoir_bundle() - Collect - - providers - - and - - emit - - named - - C++ - wrappers. - - C++ - - source - - + - - layout - - queries - - - - 04 - Compile providers - - NVRTC - nvrtcCompileProgram() - nvrtcGetLTOIR() - - Provider - - LTO-IR - - + - - scratch - - layouts - - - - 05 - Keep the bundle alive - - cuda.coop - _SharedTempFile - Own - - the - - shared - - bundle - - file - - retained - - by - Invocable - - objects. - - - - - 06 - Rewrite cooperative calls - - cuda.coop - CoopSinglePhaseRewrite - begin_rewrite() stages backing. - Entry storage reaches every view. - match() / apply(): emit blocks - finish_rewrite(): final cleanup - Arrays, calls, and reuse barriers. - - Rewritten - - IR; - - invocables - - retain - - link - - files - 2 Type and lower - Providers - - are - - ready; - - the - - kernel - - is - - still - - being - - compiled. - - - - 07 - Infer callable types - - Numba - MLIRTypeInference - - Request - - callable - - type - - - - 08 - Describe the provider call - - cuda.coop - Invocable._numba_type_ - Algorithm.codegen_method() - Build - - an - - ExternFunction - : - - symbol - - name, - - signature, - - ABI, - - and - - link - - files. - - Callable - - descriptor - - - - 09 - Lower the kernel - - Numba - MLIRLower - After - - typed - - rewrites, - - emit - - func.call @symbol - - and - - collect - - link - inputs. - - - - - 10 - Compile the kernel - - Numba - optimize() - Compile - - kernel - - LTO-IR. - - Kernel - - LTO-IR - - + - - provider - - link - - inputs - 3 Link and load - - - - 11 - Link device code - - nvJitLink - Linker - → nvJitLink - Resolve - - or - - inline - - provider - - symbols. - - Linked - - cubin - - - - 12 - Load the kernel - - Numba - CodeLibrary.get_cufunc() - Load - - the - - cubin; - - look - - up - - the - - kernel - - CUfunction - - handle. - - Solid - - arrows: - - calls - - or - - inputs. - - Dashed - - arrows: - - returned - - results. - diff --git a/docs/python/coop/visualizations/compiler-swimlane-classes.svg b/docs/python/coop/visualizations/compiler-swimlane-classes.svg deleted file mode 100644 index 3fbd5ff84624..000000000000 --- a/docs/python/coop/visualizations/compiler-swimlane-classes.svg +++ /dev/null @@ -1,354 +0,0 @@ - - - - Compilation handoffs — Python classes - Ordered handoffs between Numba, cuda.coop, NVRTC and nvJitLink. Successful provider bundling on a cache miss and the default kernel LTO path are shown. Group planning computes Load/Store payload dtype and item count. PayloadInference reuses the carried GroupLoweringPlan, checks explicit inputs and constructor extent, and records constructor state. prepare_calls_and_storage analyzes all original calls, retains _RewriteMatch records, validates descriptor ownership, prepares providers, and collects scratch requirements from their compiled layouts. Later consumers can supply facts needed by earlier constructors. begin_rewrite stages shared backing in the entry block so it dominates every scratch view. match selects each block; apply reuses or materializes invocables and emits statements in source order. ThreadData holds each thread's payload; TempStorage describes cooperative scratch. Standalone payload-write dtype inference remains in apply. finish_rewrite cleans unused bindings after all replacement blocks are installed, preserving original definitions throughout emission. These steps belong to one planner, not separate compiler passes. Facts belong to the current IR and compilation attempt. Provider artifact caches and pipeline placement are unchanged. - - - PYTHON CLASSES · NUMBA-CUDA-MLIR - Compilation handoffs - A - - provider-cache - - miss, - - successful - - bundling, - - and - - the - - default - - kernel - - LTO - - path. - 1 Prepare providers - Before - - main - - type - - inference - - - 01 - Invoke the planner - - Numba - PostInlineWholeFunctionPlanners - - Inlined, - - untyped - - IR - - - - 02 - Resolve and prepare - - cuda.coop - CoopWholeFunctionPlanner.run() - CoopSinglePhaseRewrite - .prepare_calls_and_storage() - Plan Load/Store dtype and count. - PayloadInference reuses facts; - check inputs; retain _RewriteMatch. - - During prepare_calls_and_storage() - - - - 03 - Build the provider bundle - - cuda.coop - Algorithm - collect_specializations() - prepare_ltoir_bundle() - Collect - - providers - - and - - emit - - named - - C++ - - wrappers. - - C++ - - source - - + - - layout - - queries - - - - 04 - Compile providers - - NVRTC - nvrtcCompileProgram() - nvrtcGetLTOIR() - - Provider - - LTO-IR - - + - - scratch - - layouts - - - - 05 - Keep the bundle alive - - cuda.coop - _SharedTempFile - Own - - the - - shared - - bundle - - file - - retained - - by - - Invocable - - objects. - - - - - 06 - Rewrite cooperative calls - - cuda.coop - CoopSinglePhaseRewrite - begin_rewrite() → match() / apply() - → finish_rewrite() - Stage entry backing before block visits. - Emit arrays, calls, and reuse barriers. - Clean bindings after all blocks are installed. - - Rewritten - - IR; - - invocables - - retain - - link - - files - 2 Type and lower - Providers - - are - - ready; - - the - - kernel - - is - - still - - being - - compiled. - - - - 07 - Infer callable types - - Numba - MLIRTypeInference - - Request - - callable - - type - - - - 08 - Describe the provider call - - cuda.coop - Invocable._numba_type_ - Algorithm.codegen_method() - Build - - an - - ExternFunction - : - - symbol - - name, - - signature, - - ABI, - - and - - link - - files. - - Callable - - descriptor - - - - 09 - Lower the kernel - - Numba - MLIRLower - After - - typed - - rewrites, - - emit - - func.call @symbol - - and - - collect - - link - - inputs. - - - - - 10 - Compile the kernel - - Numba - optimize() - Compile - - kernel - - LTO-IR. - - Kernel - - LTO-IR - - + - - provider - - link - - inputs - 3 Link and load - - - - 11 - Link device code - - nvJitLink - Linker → nvJitLink - Resolve - - or - - inline - - provider - - symbols. - - Linked - - cubin - - - - 12 - Load the kernel - - Numba - CodeLibrary.get_cufunc() - Load - - the - - cubin; - - look - - up - - the - - kernel - - CUfunction - - handle. - - Solid - - arrows: - - calls - - or - - inputs. - - Dashed - - arrows: - - returned - - results. - diff --git a/docs/python/coop/visualizations/compiler-swimlane.svg b/docs/python/coop/visualizations/compiler-swimlane.svg deleted file mode 100644 index 08d6b90875d3..000000000000 --- a/docs/python/coop/visualizations/compiler-swimlane.svg +++ /dev/null @@ -1,89 +0,0 @@ - - - - Compilation swimlanes: Numba-CUDA-MLIR, cuda.coop, NVRTC, and nvJitLink - Ordered handoffs between Numba, cuda.coop, NVRTC and nvJitLink. Successful provider bundling on a cache miss and the default kernel LTO path are shown. Group planning computes Load/Store payload dtype and item count. PayloadInference reuses the carried GroupLoweringPlan, checks explicit inputs and constructor extent, and records constructor state. prepare_calls_and_storage analyzes all original calls, retains _RewriteMatch records, validates descriptor ownership, prepares providers, and collects scratch requirements from their compiled layouts. Later consumers can supply facts needed by earlier constructors. begin_rewrite stages shared backing in the entry block so it dominates every scratch view. match selects each block; apply reuses or materializes invocables and emits statements in source order. ThreadData holds each thread's payload; TempStorage describes cooperative scratch. Standalone payload-write dtype inference remains in apply. finish_rewrite cleans unused bindings after all replacement blocks are installed, preserving original definitions throughout emission. These steps belong to one planner, not separate compiler passes. Facts belong to the current IR and compilation attempt. Provider artifact caches and pipeline placement are unchanged. - - - COMPILATION HANDOFFS - Time runs down. Shown: successful provider bundle, cache miss, default kernel LTO path. - - - Numba-CUDA-MLIR - - - cuda.coop - - - NVRTC - - - nvJitLink - - Build untyped IR. - Run early rewrites. - Inline device helpers. - - Inlined, untyped IR - - Resolve groups; plan - Load/Store dtype and count. - - - prepare_calls_and_storage(): - Reuse facts; retain calls. - Bundle providers as C++. - - - Compile C++ - with -dlto. - Read scratch layouts. - - Provider LTO-IR - + scratch sizes and alignments - - Keep shared link file. - Use retained call matches. - Read scratch requirements. - - - begin_rewrite(): backing - match()/apply(): emit - finish_rewrite(): cleanup - - Rewritten IR + link inputs - - Finish untyped passes. - Infer types. - Run typed rewrites. - - - Lower to MLIR. - Collect ExternFunction - symbols and link inputs. - - - Optimize and compile - kernel LTO-IR. - - Kernel LTO-IR + provider LTO-IR - (plus any extra link inputs) - - Resolve device symbols. - Optimize across inputs. - Produce cubin. - - Linked cubin - - Load with CUDA driver. - Launch the kernel. - Reuse compiled kernel. - Numba infers types here. - Providers are already - compiled. - Solid arrows: requests and inputs. Dashed arrows: returned results. - A provider-cache hit supplies the saved LTO-IR and layouts without an NVRTC call. - diff --git a/docs/python/coop/visualizations/index.rst b/docs/python/coop/visualizations/index.rst deleted file mode 100644 index 0c83169d7ad7..000000000000 --- a/docs/python/coop/visualizations/index.rst +++ /dev/null @@ -1,18 +0,0 @@ -.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -.. _coop-visualizations: - -Visualizations -============== - -Explore how ``cuda.coop`` Load and Store move values between memory and -per-thread payloads. Change the settings, step through the stages, and select -a value to follow its ownership. These diagrams show data movement; their -timing and geometry do not predict GPU performance. - -.. toctree:: - :maxdepth: 1 - - load - store diff --git a/docs/python/coop/visualizations/load-transpose.svg b/docs/python/coop/visualizations/load-transpose.svg deleted file mode 100644 index 5b6ae5f33fd7..000000000000 --- a/docs/python/coop/visualizations/load-transpose.svg +++ /dev/null @@ -1,80 +0,0 @@ - - - Transpose load: striped reads become blocked ownership - Eight illustrative threads read two values each. Thread zero reads positions zero and eight. Shared scratch exchanges values so thread zero finishes with positions zero and one, thread one with two and three, and so on. - - STRIPED MEMORY READS - 8 teaching threads / 2 items each - Thread 0 - - 0 8 - - Thread 1 - - 1 9 - - Thread 2 - - 2 10 - - Thread 3 - - 3 11 - - Thread 4 - - 4 12 - - Thread 5 - - 5 13 - - Thread 6 - - 6 14 - - Thread 7 - - 7 15 - - - SHARED SCRATCH - Rearrange by value index into consecutive pairs - - - 0 1 - Thread 0 - - - 2 3 - Thread 1 - - - 4 5 - Thread 2 - - - 6 7 - Thread 3 - - - 8 9 - Thread 4 - - - 10 11 - Thread 5 - - - 12 13 - Thread 6 - - - 14 15 - Thread 7 - BLOCKED OUTPUT - Each thread owns two consecutive values. - diff --git a/docs/python/coop/visualizations/load.rst b/docs/python/coop/visualizations/load.rst deleted file mode 100644 index 016c3d2e192b..000000000000 --- a/docs/python/coop/visualizations/load.rst +++ /dev/null @@ -1,84 +0,0 @@ -.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -.. _coop-visualization-load: - -Load -==== - -:func:`cuda.coop.load` fills each thread's ``ThreadData`` from a contiguous -tile in memory. The algorithm determines which thread reads each value and -whether a shared-memory exchange rearranges it afterward. - -The explorer shows eight illustrative threads. For the warp-transpose -algorithms, it groups them into two four-lane teaching warps. **CUDA warps -have 32 lanes**; executable blocks using these algorithms must contain a -multiple of 32 threads. All illustrated loads read a full tile. - -.. coop-visualization:: load - - .. only:: html - - .. figure:: load-transpose.svg - :alt: Transpose load: threads read striped values, exchange through shared scratch, and finish with consecutive values per thread. - :width: 100% - - Transpose load with eight illustrative threads and two items per thread. - - Each column below belongs to one thread. Transpose first reads in striped - order, then exchanges values to give each thread a consecutive pair. - - .. code-block:: text - - Thread 0 1 2 3 4 5 6 7 - Memory reads 0,8 1,9 2,10 3,11 4,12 5,13 6,14 7,15 - After exchange 0,1 2,3 4,5 6,7 8,9 10,11 12,13 14,15 - -Reading the stages ------------------- - -See :ref:`blocked versus striped ` for a compact -ownership table and the index formulas used here. - -The memory row numbers values by their position in the tile. The next row -shows the threads that read them. For algorithms with an exchange, shared -scratch connects those readers to the final owners. Select a value to -inspect its source index and destination thread slot. - -``direct`` and ``vectorize`` produce blocked ownership: thread ``t`` receives -``K`` consecutive values starting at ``t * K``, where ``K`` is the number of -items per thread. Vectorization can combine a thread's adjacent values into -a wider access when the pointer, alignment, type, and item count permit it. -The diagram does not guarantee a vector instruction. - -``striped`` gives thread ``t`` the values at ``t + j * T`` for item slot -``j`` and group size ``T``. Neighboring threads read neighboring addresses -at each slot, and the output stays striped. ``transpose`` starts with those -same reads and exchanges the values into blocked ownership. - -``warp_transpose`` reads a separate striped tile within each physical warp -and exchanges it into blocked ownership. ``warp_transpose_timesliced`` uses -the same memory reads, then lets warps take turns through a shared scratch -region sized for one warp. The animation separates those exchange rounds; -it does not serialize the global-memory loads. - -Using Load in a kernel ----------------------- - -This fragment uses the common API inside a Numba-CUDA-MLIR kernel that -accepts ``items_per_thread``. Import ``cuda`` from ``numba_cuda_mlir``, -``numpy as np``, and ``cuda.coop as coop``. Launch with 128 threads and -provide at least ``128 * items_per_thread`` source elements for each block. - -.. code-block:: python - - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - offset = cuda.blockIdx.x * 128 * items_per_thread - coop.load(block, source, items, algorithm="transpose", offset=offset) - # Each thread owns items_per_thread consecutive values. Load returns None. - -For a partial final tile, supply ``valid_items`` and, when needed, -``oob_default``. See :func:`cuda.coop.load` for the full parameter contract -and a complete example, and the :doc:`../programming_guide` for backend -activation and thread groups. diff --git a/docs/python/coop/visualizations/store-transpose.svg b/docs/python/coop/visualizations/store-transpose.svg deleted file mode 100644 index e5c27f071d11..000000000000 --- a/docs/python/coop/visualizations/store-transpose.svg +++ /dev/null @@ -1,86 +0,0 @@ - - - -Transpose store: blocked inputs become striped writer registers -Eight illustrative threads own two values each. Blocked pairs are zero and one, two and three, and so on. Striped pairs are zero and eight, one and nine, and so on. Shared scratch connects the two layouts. - - -BLOCKED INPUT PAYLOAD -8 teaching threads / 2 items each -T0 - -0 1 - -T1 - -2 3 - -T2 - -4 5 - -T3 - -6 7 - -T4 - -8 9 - -T5 - -10 11 - -T6 - -12 13 - -T7 - -14 15 - -T0 - -0 8 - -T1 - -1 9 - -T2 - -2 10 - -T3 - -3 11 - -T4 - -4 12 - -T5 - -5 13 - -T6 - -6 14 - -T7 - -7 15 - - -SHARED SCRATCH -Logical positions connect source threads to destination slots; padding omitted. -STRIPED WRITER REGISTERS - - -Memory: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 -Transpose Store may rearrange the input payload. Copy values before Store if needed later. - - diff --git a/docs/python/coop/visualizations/store.rst b/docs/python/coop/visualizations/store.rst deleted file mode 100644 index e2ab21ff54f4..000000000000 --- a/docs/python/coop/visualizations/store.rst +++ /dev/null @@ -1,90 +0,0 @@ -.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -.. SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -.. _coop-visualization-store: - -Store -===== - -:func:`cuda.coop.store` writes each thread's values into a contiguous memory -tile. The algorithm determines which input ownership it expects and whether -a shared-memory exchange rearranges the values before the writes. - -The explorer shows eight illustrative threads and full-tile stores. The -warp-transpose variants use two four-lane teaching warps; **CUDA warps have -32 lanes**. Executable blocks using those variants must contain a multiple -of 32 threads. - -.. coop-visualization:: store - - .. only:: html - - .. figure:: store-transpose.svg - :alt: Transpose store exchanges blocked pairs into striped writer registers, then stores the values in consecutive memory locations. - :width: 100% - - Transpose store with eight illustrative threads and two items per thread. - - Each thread begins with a consecutive pair. The exchange redistributes - those values so adjacent threads can write adjacent memory locations. - - .. code-block:: text - - Thread 0 1 2 3 4 5 6 7 - Input 0,1 2,3 4,5 6,7 8,9 10,11 12,13 14,15 - Writer values 0,8 1,9 2,10 3,11 4,12 5,13 6,14 7,15 - Memory 0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 - -Reading the stages ------------------- - -See :ref:`blocked versus striped ` for a compact -ownership table and the index formulas used here. - -The first row shows the input payload. Transpose Store algorithms may -rearrange that payload in place, as in CUB; copy values before Store if they -are needed later. For algorithms with an exchange, values pass through shared -scratch before reaching the writer-register row. The final row shows memory -positions. Select a value to follow its input thread, writer thread, and -destination address. - -``direct`` and ``vectorize`` consume blocked input: thread ``t`` provides -``K`` consecutive values beginning at logical position ``t * K``. -Vectorization can combine adjacent values into wider stores when the -pointer, alignment, type, and item count allow it. The illustration does -not guarantee a particular instruction or memory-transaction count. - -``striped`` consumes striped input: thread ``t`` writes slot ``j`` to -``t + j * T``, where ``T`` is the group size. The input must already have -that ownership. ``transpose`` instead accepts blocked input and exchanges -it into striped writer registers before performing the same writes. - -``warp_transpose`` performs the corresponding exchange within each physical -warp. ``warp_transpose_timesliced`` reuses one warp's shared scratch across -serialized exchange rounds. The writer registers then feed warp-striped -stores; the diagram does not impose a serial global-memory store schedule. - -Using Store in a kernel ------------------------ - -This fragment uses the common API inside a Numba-CUDA-MLIR kernel that -accepts ``items_per_thread``. Import ``cuda`` from ``numba_cuda_mlir``, -``numpy as np``, and ``cuda.coop as coop``. Launch with 128 threads and -provide at least ``128 * items_per_thread`` source and destination elements -for each block. - -.. code-block:: python - - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - offset = cuda.blockIdx.x * 128 * items_per_thread - coop.load(block, source, items, algorithm="direct", offset=offset) - coop.store(block, destination, items, algorithm="transpose", offset=offset) - # The block's values are written in their original logical order. - # Store writes memory and returns None; items may now be rearranged. - -For a partial final tile, pass ``valid_items`` to limit the stored prefix; -memory beyond that prefix remains untouched. Supply a matching valid count -when loading the input tile. See :func:`cuda.coop.store` for the parameter -contract and the :doc:`../programming_guide` for backend activation and -thread groups. diff --git a/docs/python/coop_api.rst b/docs/python/coop_api.rst deleted file mode 100644 index 660ed095dbc7..000000000000 --- a/docs/python/coop_api.rst +++ /dev/null @@ -1,116 +0,0 @@ -.. _cuda_coop-module: - -``cuda.coop`` API Reference -=========================== - -.. warning:: - ``cuda.coop`` is an experimental API and is subject to change. - -.. _coop-portable-api: -.. _portable-api: -.. _coop-common-api: - -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. - -.. currentmodule:: cuda.coop - -Backend registration -^^^^^^^^^^^^^^^^^^^^ - -See :ref:`registering a backend `. - -.. autofunction:: register - -Thread groups -^^^^^^^^^^^^^ - -See :ref:`thread groups ` and -:ref:`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. - -.. autofunction:: this_thread -.. autofunction:: this_warp -.. autofunction:: this_block -.. autofunction:: this_cluster -.. autofunction:: this_grid - -.. autoclass:: ThreadGroup - :no-members: - :no-special-members: - - .. automethod:: group_by - -.. autoclass:: ThreadHierarchy - :no-members: - :no-special-members: - -.. py:class:: Hierarchy - - Alias for :class:`ThreadHierarchy`. - -Payloads and temporary storage -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -``ThreadData`` constructs the per-thread payload supplied to a primitive; -thread-group descriptors describe the participating threads. - -.. autofunction:: ThreadData - -.. autoclass:: ThreadDataLike - :no-members: - :no-special-members: - -.. autofunction:: TempStorage - -.. autoclass:: TempStorageLike - :no-members: - :no-special-members: - -Memory operations -^^^^^^^^^^^^^^^^^ - -.. autofunction:: load -.. autofunction:: store - - -.. _coop-numba-extensions: - -Numba-CUDA-MLIR-qualified API ------------------------------ - -.. py:module:: cuda.coop.numba_mlir - -Use this module for the extensions below. Shared parameters and behavior -follow the :ref:`Common API `. - -.. code-block:: python - - 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. - -.. currentmodule:: cuda.coop.numba_mlir - -Payloads and temporary storage -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -.. autofunction:: ThreadData - -.. autoclass:: TempStorage - :no-members: - :no-special-members: - -Memory operations -^^^^^^^^^^^^^^^^^ - -.. autofunction:: load -.. autofunction:: store diff --git a/docs/python/index.rst b/docs/python/index.rst index ee91e15b1e1d..0e590d351966 100644 --- a/docs/python/index.rst +++ b/docs/python/index.rst @@ -11,9 +11,6 @@ abstractions for CUDA Python developers. * :doc:`cuda.compute ` — Composable device-level primitives for building custom parallel algorithms, without writing CUDA kernels directly. -* :doc:`cuda.coop ` — Cooperative Block and Warp Load and Store - primitives for Python GPU kernel compilers. - * :doc:`cuda.stf._experimental ` — Sequential Task Flow for CUDA: define logical data and tasks with read/write annotations; STF orchestrates execution and data movement. @@ -37,7 +34,6 @@ Who is this for? setup compute/index - coop stf resources api_reference diff --git a/python/cuda_coop/AGENTS.md b/python/cuda_coop/AGENTS.md deleted file mode 100644 index 7cf738471d93..000000000000 --- a/python/cuda_coop/AGENTS.md +++ /dev/null @@ -1,81 +0,0 @@ -# `cuda.coop` agent guidance - -Apply this guidance to package code, tests, documentation, API docstrings, and -runnable examples. - -# Design Tenets - -As a guiding principle, the design tenets behind `cuda.coop` are -**performant, consistent, delightful**: - -- **Performant:** target negligible overhead over the underlying primitives - through compile-time planning, inlining, and LTO; avoid hidden work. -- **Consistent:** the same operation should look and behave the same across - DSLs, including argument order and result ownership. -- **Delightful:** Pythonic APIs, useful defaults, and inference should make - common cases work out of the box, with clear errors for unsupported uses. - -# Conventions - -Write the API/project name as `cuda.coop` and the distribution as `cuda-coop`. - -Use these import conventions in each module: - -- Import the common API with `from cuda import coop` and use - `coop.` in DSL kernels. -- If the module does not need the common API, import its qualified API as - `coop`: `import cuda.coop.numba_mlir as coop` or - `import cuda.coop.cutlass as coop`. Use `coop.` in kernels. -- When the module needs both common and qualified APIs, keep - `from cuda import coop` for the common API. Use - `import cuda.coop.numba_mlir as numba_coop` and/or - `import cuda.coop.cutlass as cutlass_coop` for the qualified APIs, and use - their respective aliases in kernels. - -Regarding `coop.ThreadData` (including qualified APIs and import aliases): - -- Prefer an `items_per_thread` kernel parameter for the main per-thread - payload size, so callers can reuse the kernel with different item counts. - Carry it through host launchers and device helpers as needed. The value - must be known during compilation: Numba-CUDA-MLIR specializes the kernel - argument, and CUTLASS kernel and launcher parameters use - `items_per_thread: cutlass.Constexpr`. -- When the argument is a variable named `items_per_thread`, prefer - `coop.ThreadData(items_per_thread)`. Otherwise, use the explicit - `items_per_thread=` kwarg, including for integer literals and expressions. -- Keep fixed-size auxiliary payloads at the sizes required by their - operations; they need not match the kernel's main `items_per_thread`. -- Omit the `dtype=` kwarg and positional dtype arguments from examples. -- Infer the element type from supported producers or typed assignments, - preserving intended widths. If inference fails, investigate or revise the - example; do not silently add `dtype=`. -- Document the actual constructor signature, including its supported `dtype=` - parameter. These example conventions do not change the accepted arguments. -- Tests specifically checking `dtype=`, positional arguments, or invalid inputs - may keep those calls. Do not use those tests as public examples. - -Treat `cuda.coop` as the common API contract. Qualified APIs such as -`cuda.coop.numba_mlir` and `cuda.coop.cutlass` must be otherwise identical -supersets built on that contract: the same names, argument order, defaults, -behavior, and result ownership for shared operations. Add only the extensions -needed by each DSL. - -Ensure new primitives have thorough docstrings and are captured in the API -reference docs; see existing precedence in code base. - -Preserve CUB behavior for all primitives and don't invent new or stronger -semantics than CUB provides. - -Tests should be geared toward testing the end-user primitives, not internal -helper functions or supporting glue--that infrastructure is exercised and -implicitly tested when doing primitive-level tests. - -Parametrize kernel tests over `items_per_thread` where the operation supports -multiple item counts; prefer `pytest.mark.parametrize("items_per_thread", [1, 4])` -as the default coverage. Derive input sizes, launch arguments, and expected -results from that value. Preserve operation-specific constraints and focused -coverage for literal constructor arguments and invalid inputs. - -Update the source of `literalinclude` snippets and run affected examples when -changing them. Apply changes at the first owning PR and preserve them when doing -PR restacking. diff --git a/python/cuda_coop/CMakeLists.txt b/python/cuda_coop/CMakeLists.txt deleted file mode 100644 index 3c5610b44a03..000000000000 --- a/python/cuda_coop/CMakeLists.txt +++ /dev/null @@ -1,178 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -cmake_minimum_required(VERSION 3.21...3.31 FATAL_ERROR) - -project(cuda_coop DESCRIPTION "Python package cuda-coop" LANGUAGES NONE) - -option( - CUDA_COOP_INSTALL_HEADER_BUNDLE - "Install cuda-coop's private CCCL header bundle" - ON -) -option( - CUDA_COOP_ALLOW_DIRTY_HEADER_BUNDLE - "Allow locally modified CCCL headers with unknown bundle provenance" - OFF -) -set( - CUDA_COOP_CCCL_SOURCE_REVISION - "" - CACHE STRING - "CCCL source revision recorded in the cuda-coop header bundle" -) -if (CUDA_COOP_INSTALL_HEADER_BUNDLE) - # cuda-coop owns its bundled CCCL headers. Follow the repository's install - # rules so CUB, Thrust, and libcu++ stay on one source revision. - set(CCCL_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../..") - set( - _cuda_coop_cccl_install_rules - "${CCCL_SOURCE_DIR}/cmake/CCCLInstallRules.cmake" - ) - if (NOT EXISTS "${_cuda_coop_cccl_install_rules}") - message( - FATAL_ERROR - "CUDA_COOP_INSTALL_HEADER_BUNDLE=ON requires a full CCCL source tree, " - "but CCCL install rules were not found at " - "${_cuda_coop_cccl_install_rules}." - ) - endif() - - set(CCCL_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}") - set(CCCL_TOPLEVEL_PROJECT OFF) - set(CCCL_ENABLE_CUDAX OFF) - set(CCCL_ENABLE_INSTALL_RULES ON) - set(libcudacxx_ENABLE_INSTALL_RULES ON) - set(CUB_ENABLE_INSTALL_RULES ON) - set(Thrust_ENABLE_INSTALL_RULES ON) - set(cudax_ENABLE_INSTALL_RULES OFF) - set(CMAKE_INSTALL_INCLUDEDIR "cuda/coop/_headers/include" CACHE PATH "" FORCE) - set(CMAKE_INSTALL_LIBDIR "cuda/coop/_headers/lib" CACHE PATH "" FORCE) - include("${_cuda_coop_cccl_install_rules}") - - # Git searches parent directories for repository metadata. Require the CCCL - # source directory itself to be the worktree root so an unpacked archive - # nested in an unrelated checkout cannot inherit that checkout's revision. - set(_cuda_coop_source_is_worktree_root OFF) - execute_process( - COMMAND git -C "${CCCL_SOURCE_DIR}" rev-parse --show-prefix - RESULT_VARIABLE _cuda_coop_source_prefix_result - OUTPUT_VARIABLE _cuda_coop_source_prefix - ERROR_QUIET - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - if ( - "${_cuda_coop_source_prefix_result}" STREQUAL "0" - AND "${_cuda_coop_source_prefix}" STREQUAL "" - ) - set(_cuda_coop_source_is_worktree_root ON) - endif() - - set(_cuda_coop_header_bundle_unverified OFF) - if (_cuda_coop_source_is_worktree_root) - # Match the headers and CMake package sources selected by the CCCL install - # rules and the package configuration controlling the bundle. Include - # ignored files because directory install rules do not consult Git ignores. - set( - _cuda_coop_bundle_input_pathspecs - ":(glob)cub/cub/**/*.cuh" - ":(glob)thrust/thrust/**/*.h" - ":(glob)thrust/thrust/**/*.inl" - ":(glob)libcudacxx/include/cuda/**" - ":(glob)libcudacxx/include/nv/**" - ":(exclude,glob)libcudacxx/include/cuda/**/CMakeLists.txt" - ":(exclude,glob)libcudacxx/include/nv/**/CMakeLists.txt" - ":(glob)lib/cmake/cccl/**" - ":(glob)lib/cmake/cub/**" - ":(glob)lib/cmake/libcudacxx/**" - ":(glob)lib/cmake/thrust/**" - "cmake/CCCLInstallRules.cmake" - "cmake/install/cccl.cmake" - "cmake/install/cub.cmake" - "cmake/install/libcudacxx.cmake" - "cmake/install/thrust.cmake" - "python/cuda_coop/CMakeLists.txt" - ) - execute_process( - COMMAND - git -C "${CCCL_SOURCE_DIR}" status --porcelain=v1 --untracked-files=all - --ignored=matching -- ${_cuda_coop_bundle_input_pathspecs} - RESULT_VARIABLE _cuda_coop_header_status_result - OUTPUT_VARIABLE _cuda_coop_header_status - ERROR_QUIET - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - if (NOT "${_cuda_coop_header_status_result}" STREQUAL "0") - set(_cuda_coop_header_bundle_unverified ON) - set( - _cuda_coop_header_bundle_diagnostic - "Git could not verify the CCCL header bundle inputs" - ) - elseif (NOT "${_cuda_coop_header_status}" STREQUAL "") - set(_cuda_coop_header_bundle_unverified ON) - set( - _cuda_coop_header_bundle_diagnostic - "The CCCL header bundle inputs contain local changes" - ) - endif() - endif() - - if (_cuda_coop_header_bundle_unverified) - if (NOT CUDA_COOP_ALLOW_DIRTY_HEADER_BUNDLE) - message( - FATAL_ERROR - "${_cuda_coop_header_bundle_diagnostic}. Commit or remove those changes, " - "or explicitly configure CUDA_COOP_ALLOW_DIRTY_HEADER_BUNDLE=ON to " - "build with unknown provenance." - ) - endif() - set(_cuda_coop_source_revision "unknown") - message( - WARNING - "${_cuda_coop_header_bundle_diagnostic}; recording an unknown CCCL " - "source revision" - ) - elseif (CUDA_COOP_CCCL_SOURCE_REVISION) - set(_cuda_coop_source_revision "${CUDA_COOP_CCCL_SOURCE_REVISION}") - elseif (_cuda_coop_source_is_worktree_root) - execute_process( - COMMAND git -C "${CCCL_SOURCE_DIR}" rev-parse HEAD - RESULT_VARIABLE _cuda_coop_source_revision_result - OUTPUT_VARIABLE _cuda_coop_source_revision - ERROR_QUIET - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - else() - set(_cuda_coop_source_revision_result 1) - endif() - if ( - DEFINED _cuda_coop_source_revision_result - AND NOT "${_cuda_coop_source_revision_result}" STREQUAL "0" - ) - set(_cuda_coop_source_revision "unknown") - message( - WARNING - "Git metadata is unavailable; recording an unknown CCCL source revision" - ) - endif() - if (NOT _cuda_coop_source_revision MATCHES "^[0-9A-Za-z._+-]+$") - message( - FATAL_ERROR - "CUDA_COOP_CCCL_SOURCE_REVISION must be a revision token; got " - "${_cuda_coop_source_revision}" - ) - endif() - set( - _cuda_coop_bundle_provenance - "${CMAKE_CURRENT_BINARY_DIR}/cccl-bundle-provenance.json" - ) - file( - WRITE "${_cuda_coop_bundle_provenance}" - "{\n \"cccl_source_commit\": \"${_cuda_coop_source_revision}\"\n}\n" - ) - install( - FILES "${_cuda_coop_bundle_provenance}" - DESTINATION cuda/coop/_headers - ) -endif() diff --git a/python/cuda_coop/LICENSE b/python/cuda_coop/LICENSE deleted file mode 100644 index fd50777c3050..000000000000 --- a/python/cuda_coop/LICENSE +++ /dev/null @@ -1 +0,0 @@ -../../LICENSE diff --git a/python/cuda_coop/README.md b/python/cuda_coop/README.md deleted file mode 100644 index e8d879dfc7c1..000000000000 --- a/python/cuda_coop/README.md +++ /dev/null @@ -1,364 +0,0 @@ -# `cuda.coop` - -`cuda.coop` provides cooperative primitives for CUDA thread groups in Python -kernel DSLs. The first backend targets Numba-CUDA-MLIR. The table below -lists the available operations. - -The distribution is a universal Python wheel containing a coherent bundle of -CUB, Thrust, and libcu++ headers. Installed-wheel compilation uses that -bundle by default. Development from a CCCL source checkout uses the matching -checkout headers, and `CUDA_COOP_CCCL_ROOT` can select a different source -checkout or `cuda-coop` header bundle. None of these modes substitutes CUB -headers from the active CUDA Toolkit. - -## Installation - -Install `cuda-coop` without adding Python package dependencies: - -```bash -python -m pip install cuda-coop -``` - -The wheel includes the common API, every shipped DSL integration (including -`cuda.coop.numba_mlir`), type declarations, and bundled CCCL headers. The base -install declares no Python package dependencies. You can import `cuda.coop` -without a compiler or GPU; using an integration requires its backend -dependencies to be installed. Numba-CUDA-MLIR is the first supported backend; -CUTLASS support is planned. - -For Numba-CUDA-MLIR, choose the extra matching the CUDA Toolkit major version: - -```bash -python -m pip install "cuda-coop[numba-cuda-mlir-cu13]" -# Use numba-cuda-mlir-cu12 with CUDA 12. -``` - -Both commands install the same `cuda-coop` wheel with the same DSL -integrations. The extra only adds the dependency requirements declared in -`pyproject.toml` so pip installs the supported Numba-CUDA-MLIR stack for the -selected CUDA major version. - -Python 3.10 through 3.14 is supported. The current backend integration -requires `numba-cuda-mlir>=0.5.0,<0.6`. - -Backend compiler and runtime CI is configured for Linux x86-64 with Python -3.14: CUDA 13 in pull requests and CUDA 12 in the nightly matrix. Linux host -contracts cover Python 3.10 and 3.14. Windows checks build and import the -universal wheel and verify its headers; they do not execute the compiler -backend. Other combinations need separate runtime qualification. See the -[validation scope](https://nvidia.github.io/cccl/unstable/python/coop/developer_overview.html#coop-numba-validation) -for tested platforms and coverage. - -With Numba-CUDA-MLIR 0.5.0 through 0.5.4, keep a compiled kernel's dispatcher -and configured launch callables in their original CUDA context. Reuse on -another device or after context teardown is not qualified because cached -architecture or launch state can belong to the original context. The upstream -[context-isolation fix](https://github.com/NVIDIA/numba-cuda-mlir/pull/314) -must be released and qualified before relying on that reuse. - -## Backend registration and imports - -Register the backend on the host before compiling kernels: - -```python -from cuda import coop - -coop.register("numba-cuda-mlir") - -from numba_cuda_mlir import cuda -``` - -`register` loads the backend and installs its compiler hooks regardless of -import order. It also accepts `"numba_cuda_mlir"`. Repeated calls are safe and -return `None`. It requires installed backend dependencies and does not install -packages. Installing an extra is separate from registering a compiler in the -running process; installed package metadata does not reliably record which -extra was requested. - -Importing `cuda.coop` after `numba_cuda_mlir` also registers the backend -automatically. A standalone `cuda.coop` import does not load optional -compilers. Explicit registration works when a dependency or earlier notebook -cell already imported `cuda.coop`. - -The common API exposed by `cuda.coop` describes operations independently of -a particular compiler backend. Its public entry points are implemented in -`cuda/coop/_core/api/`; the private `_core` package also contains shared -implementation used by the backends. Each backend implements its supported -operations, so a common spelling does not guarantee support in every compiler. - -To use Numba-specific features, import the qualified backend namespace; this -also registers it: - -```python -import cuda.coop.numba_mlir as numba_coop -``` - -Use `numba_coop` alongside common calls. In a program using only this -backend, `import cuda.coop.numba_mlir as coop` is also supported. Keep the -alias: a bare `import cuda.coop.numba_mlir` binds `cuda` to the top-level -Python package and can replace the name used for Numba's `cuda.jit`. - -Shared operations retain the common signatures, string selectors, and -inference rules. The backend namespace adds Numba local-array payloads and -memory namespaces. The common and qualified namespaces both accept -`ThreadData(items_per_thread, alignment=None)`. Use a compile-time positive -power of two in bytes to request minimum payload storage alignment, or omit -it to let the compiler choose. This does not assert alignment of Load or -Store arrays. - -Pass `items_per_thread` as a kernel argument. Numba-CUDA-MLIR specializes -the kernel for its value. The payload count stays fixed during execution. - -The [FAQs](https://nvidia.github.io/cccl/unstable/python/coop/faqs.html) -explain namespace choices and temporary storage. The -[Glossary](https://nvidia.github.io/cccl/unstable/python/coop/glossary.html) -explains terms and concepts, including blocked and striped layouts. - -## Primitive families - -| Family | Entry points | -| --- | --- | -| Memory operations | `load`, `store` | - -Each operation documents its supported groups and result ownership in the -[API reference](https://nvidia.github.io/cccl/unstable/python/coop_api.html). The -[visualizations](https://nvidia.github.io/cccl/unstable/python/coop/visualizations/index.html) -explain these contracts with interactive diagrams and tested kernel -examples. - - -## Configuration - -Runtime configuration is controlled by these environment variables: - -| Variable | Effect | -| --- | --- | -| `CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION` | A truthy value disables automatic backend activation during `cuda.coop` import. Explicit qualified-backend import still works. | -| `CUDA_COOP_CCCL_ROOT` | Selects a CCCL source checkout or a `cuda-coop` header bundle. An invalid configured root is an error; resolution does not fall back to another CCCL source. | -| `CUDA_COOP_ENABLE_CACHE` | A truthy value enables the persistent compiler cache. The value is read when the backend cache module is imported. | -| `XDG_CACHE_HOME` | On Linux and other POSIX systems, sets the cache base directory; entries are stored in `/cccl`. Unset, empty, or relative values fall back to `~/.cache/cccl`. Read when the backend cache module is imported. | -| `LOCALAPPDATA` | On Windows, sets the cache base directory; entries are stored in `\cccl`. Unset, empty, or relative values fall back to `~\AppData\Local\cccl`. Read when the backend cache module is imported. | -| `CUDA_COOP_SOURCE_DUMP_DIR` | Writes generated CUDA source as `cuda_coop__.cu` files. Set before compiling; Numba provider cache hits also dump source. Unset or empty disables dumping. | -| `CUDA_PATH` | Supplies `/include` as a CUDA header candidate if `cuda-pathfinder` does not resolve one. | -| `CUDA_HOME` | Supplies `/include` after `CUDA_PATH` under the same fallback rule. | -| `CUDA_ROOT` | Supplies `/include` after `CUDA_HOME` under the same fallback rule. | - -On POSIX, `/usr/local/cuda/include` is tried last if those mechanisms do -not resolve CUDA headers. - -The build recognizes these CMake cache variables: - -| Variable | Default | Effect | -| --- | --- | --- | -| `CUDA_COOP_INSTALL_HEADER_BUNDLE` | `ON` | Installs the private CCCL header and CMake-package bundle into the wheel. | -| `CUDA_COOP_ALLOW_DIRTY_HEADER_BUNDLE` | `OFF` | Allows a Git-worktree bundle when selected inputs are changed or `git status` cannot verify them, and records its source revision as `unknown`. | -| `CUDA_COOP_CCCL_SOURCE_REVISION` | empty | Supplies the revision token recorded instead of deriving it from Git. A dirty or unverifiable Git worktree still records `unknown`. | - -For the two Boolean runtime switches, values are case-insensitive; `0`, -`false`, `no`, `off`, and the empty string are false. - -## Block and Warp Load and Store - -The common `cuda.coop` entry points and the qualified -`cuda.coop.numba_mlir` entry points have matching signatures. The following -kernel-body example clamps a grid tile tail, where `source`, `destination`, -`count`, and `items_per_thread` are kernel arguments: - -```python -from numba_cuda_mlir import cuda, types - -from cuda import coop - -block = coop.this_block() -items = coop.ThreadData(items_per_thread) -tile_items = cuda.blockDim.x * items_per_thread -tile_offset = cuda.blockIdx.x * tile_items -valid_items = count - tile_offset -if valid_items < 0: - valid_items = 0 -elif valid_items > tile_items: - valid_items = tile_items -coop.load( - block, - source, - items, - algorithm="direct", - valid_items=valid_items, - oob_default=0, - offset=tile_offset, -) -coop.store( - block, - destination, - items, - algorithm="direct", - valid_items=valid_items, - offset=tile_offset, -) -``` - -`load` fills the caller's output in place and returns `None`. `valid_items` -counts items across the selected group tile, while `offset` is a nonnegative -element offset. Runtime offsets are caller-validated. Source and destination -arrays must be one-dimensional and contiguous. Without `oob_default`, -invalid Load slots have unspecified values, even if initialized before Load. -Every supplied runtime control (`valid_items`, `oob_default`, and `offset`) -must be uniform within its selected group; different groups may use -different values. - -Runtime `valid_items` and `offset` accept signed integer types through 64 bits -and unsigned integer types through 32 bits. Boolean, floating-point, and -`uint64` runtime values are rejected. A runtime `oob_default` is already typed -by the compiler and must exactly match the Load payload dtype. Ordinary Python -integer and floating-point literals are converted contextually and checked -against that dtype before provider generation. - -> **`valid_items` must satisfy -> `0 <= valid_items <= group_size * items_per_thread`.** Static values outside -> that range are rejected while planning. Runtime values are checked rather -> than saturated; do not rely on CUB's oversized-count behavior. An invalid -> runtime value executes a deterministic device trap before narrowing to CUB's -> integer parameter, and that trap poisons the current CUDA context. Clamp -> grid-stride and tail counts as above. Run intentional failure probes in -> disposable processes. For a Warp-group call, subtract that group's tile -> origin from a block-wide remainder and clamp the result to -> `[0, group_size * items_per_thread]`. - -Store payloads must have exactly the destination dtype. Numba-CUDA-MLIR may -promote integer arithmetic even when its operands are 32-bit, so explicitly -cast computed values before storing them: - -```python -value = types.int32(source[cuda.threadIdx.x] + 1) -coop.store(block, destination, value, algorithm="direct") -``` - -Both common and qualified entry points use the same lowercase string -algorithm vocabulary: `direct`, `striped`, `vectorize`, `transpose`, -`warp_transpose`, and `warp_transpose_timesliced`. All six are executable. -`striped` exposes a striped per-thread payload; the other Load algorithms -expose blocked payloads. Store consumes the matching arrangement. As in CUB, -transpose Store algorithms may rearrange the input payload in place. Copy -values before Store if they are needed later. The two warp-transpose modes -require a block size divisible by 32. - -Algorithm selectors are normalized to lowercase underscore-delimited strings. -Enum and integer selectors, including `0`, are rejected. - -Warp Load and Store accept `this_warp()` and support `direct`, `striped`, -`vectorize`, and `transpose`. Partition a physical warp into consecutive -logical groups with `this_warp().group_by(width)`, where `width` is 1, 2, 4, -8, 16, or 32. The enclosing block must contain a multiple of 32 threads and -must not have an incomplete final physical warp. Every member of a -participating group must reach the collective; complete sibling logical -groups may diverge. `direct` and `vectorize` expose blocked payloads, -`striped` exposes a striped payload, and `transpose` uses striped memory -transactions while exposing a blocked payload. - -Each Warp group addresses a distinct tile. The compiler advances the memory -base by `group_index * (group_size * items_per_thread)` and then applies the -caller's element `offset`. The offset must be uniform within each -participating group; different groups may use different offsets. The group -index is the x-major linear thread rank divided by the selected group size. -For a multi-block traversal, include the block's global tile origin in the -caller offset; the compiler-provided origin distinguishes the physical or -logical Warp groups within that block and must not be added again. Runtime -offsets must also leave enough signed 64-bit range for the last group origin -in the block; static offsets are checked during planning. `valid_items` is -relative to each group's own tile, not the entire block, and must be uniform -within that group. - -`ThreadGroup` objects are descriptor-only in this release. Runtime query, -membership, and synchronization methods such as `rank`, `count`, `rank_as`, -`count_as`, `sync`, `sync_aligned`, and `is_member` are not exposed. - -## Temporary storage - -Block Load and Store accept an optional caller descriptor: - -```python -storage = coop.TempStorage( - size_in_bytes=None, - alignment=None, - auto_sync=False, - sharing="shared", -) -coop.load(block, source, items, algorithm="transpose", temp_storage=storage) -``` - -For block, physical Warp, and logical Warp calls, `direct`, `striped`, and -`vectorize` are storage-free: they default-construct CUB primitives without -shared-memory allocation, pointer arguments, or barriers. For block calls, -an explicit descriptor is validated but does not change their code -generation. Construct `TempStorage` inside the kernel; module-global storage -descriptors cannot be resolved. A descriptor may be passed to a device -helper that Numba-CUDA-MLIR inlines into the kernel, which is the default. - -The three block transpose algorithms use CUB temporary storage. Without a -descriptor, the compiler allocates the specialization's exact storage and -inserts a block reuse barrier. A caller descriptor selects shared or -exclusive slices and may request capacity and alignment. Both explicit and -omitted storage participate in the shared-memory plan and launch accounting. -The provider determines the required byte count and alignment. - -A descriptor's `sharing` selects only the slice layout: `"shared"` overlaps -every call that passes the same descriptor on one region, while -`"exclusive"` gives each call site its own slice. A call site inside a loop -reuses its slice under either layout, so `auto_sync` is independent of -`sharing` and defaults to `False` for both. - -Exclusive slices use 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 must be synchronized. Omitting storage -lets the compiler choose the layout and insert reuse barriers; it does not -guarantee a separate slice per call site. - -A descriptor names one region. Distinct descriptors and compiler-owned storage -never alias each other. With `auto_sync=True`, the compiler appends -`cuda.syncthreads()` immediately after every block call that consumes the -storage, including the last one. Warp calls do not accept a descriptor. The -compiler never inserts a barrier before a call. The trailing barrier orders -only reuse of the temporary storage. Arrange explicit barriers for the -kernel's own shared-memory traffic. With `auto_sync=False`, the caller issues -`cuda.syncthreads()` between consecutive uses, and a call site inside a loop -counts as a reuse on every iteration. Compiler-owned storage always -synchronizes. - -All descriptors and compiler-owned requirements of a kernel share one -shared-memory backing. Up to 48 KiB, the compiler uses static shared memory -without querying the device. Larger requests trigger a query of the device's -default and opt-in limits. Backing above the default limit uses dynamic shared -memory, and the launch reserves its exact byte count within the opt-in limit. -Supported Numba-CUDA-MLIR releases do not separate static and dynamic shared -allocations reliably. A kernel using cooperative temporary storage must not -also declare a zero-sized or runtime-sized `cuda.shared.array`. When -cooperative backing becomes dynamic, user static shared arrays are also -unsupported. Keep both user arrays and cooperative backing static, or move the -user data out of shared memory. Storage-free operations do not add this -restriction. - -With `auto_sync=False`, a descriptor must originate from exactly one -constructor site. Selecting between multiple manual-sync constructors is -unsupported: the compiler cannot prove that caller barriers protect the -merged region, even when a particular program supplies sufficient barriers. - -Cooperative calls in device helpers must be inlined into the kernel; use -`@cuda.jit(device=True, inline="always")` when selecting the helper's -policy explicitly. Standalone collective helpers and collectives inside -standalone callbacks are unsupported. `literal_unroll` values cannot determine -cooperative payload extents, group dimensions, selectors, or descriptor -constructor arguments. Use separate calls with explicit constants, or an -ordinary loop with one fixed cooperative shape. An unrelated `literal_unroll` -loop does not add this restriction. - -Warp `transpose` uses compiler-owned storage with one disjoint slice per -physical or logical group and inserts `syncwarp` with the exact group mask. -Explicit `TempStorage` is rejected by both the common and qualified APIs for -every Warp Load and Store algorithm, including the storage-free modes. - -These APIs are compile-time kernel constructs. Calling them outside a -compatible compiler context reports a structured context error. - -See the -[CCCL documentation](https://nvidia.github.io/cccl/unstable/python/coop.html) -for the complete signatures. diff --git a/python/cuda_coop/cuda/coop/__init__.py b/python/cuda_coop/cuda/coop/__init__.py deleted file mode 100644 index c3390e93a118..000000000000 --- a/python/cuda_coop/cuda/coop/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Expose cooperative operations through a common CUDA Python interface. - -Compiler adapters translate these calls into device code. The root import -probes supported compiler runtimes already loaded by the application and can -activate their adapters. Use :func:`register` or a qualified backend import -when importing the common interface before the compiler runtime. -""" - -import importlib.metadata -from pkgutil import extend_path - -from ._registration import register - -__path__ = extend_path(__path__, __name__) - -from ._core import api as _common_api -from ._core._auto_registration import _auto_register_known_dsls - -globals().update( - {name: getattr(_common_api, name) for name in _common_api.__all__} -) - - -def _package_version() -> str: - """Return the installed version or ``0+unknown`` without metadata.""" - - try: - return importlib.metadata.version("cuda-coop") - except importlib.metadata.PackageNotFoundError: - return "0+unknown" - - -__version__ = _package_version() - -__all__ = ["__version__", "register"] -__all__.extend(_common_api.__all__) - - -def __dir__() -> list[str]: - """List public names for interactive completion.""" - - return sorted(__all__) - - -_auto_register_known_dsls() -del _auto_register_known_dsls diff --git a/python/cuda_coop/cuda/coop/__init__.pyi b/python/cuda_coop/cuda/coop/__init__.pyi deleted file mode 100644 index 2f177254c91c..000000000000 --- a/python/cuda_coop/cuda/coop/__init__.pyi +++ /dev/null @@ -1,44 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from typing import Literal - -from ._core.api.load_store import load, store -from ._core.api.temp_storage import TempStorage, TempStorageLike -from ._core.api.thread_data import ThreadData, ThreadDataLike -from ._core.api.thread_group import ( - Hierarchy, - ThreadGroup, - ThreadHierarchy, - this_block, - this_cluster, - this_grid, - this_thread, - this_warp, -) - -__version__: str - -def register( - backend: Literal["numba-cuda-mlir", "numba_cuda_mlir"], -) -> None: ... - -__all__ = [ - "Hierarchy", - "TempStorage", - "TempStorageLike", - "ThreadData", - "ThreadDataLike", - "ThreadGroup", - "ThreadHierarchy", - "__version__", - "load", - "register", - "store", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -] diff --git a/python/cuda_coop/cuda/coop/_core/__init__.py b/python/cuda_coop/cuda/coop/_core/__init__.py deleted file mode 100644 index e6131abfe07c..000000000000 --- a/python/cuda_coop/cuda/coop/_core/__init__.py +++ /dev/null @@ -1,191 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Cooperative call descriptions used by compiler adapters. - -Importing this package requires only the Python standard library. -""" - -from ._adapter import CoreBackendAdapter, lower_method_parameters -from ._algorithm import Algorithm, TypeDefinition -from ._bindings import ArgumentBinding, BindingKind, binding, i32_parameter -from ._symbols import semantic_token -from ._types import ( - FLOAT32, - FLOAT64, - INT8, - INT16, - INT32, - INT64, - UINT8, - UINT16, - UINT32, - UINT64, - ArgumentKind, - Array, - BuiltinDType, - Constant, - CxxFunction, - Dependency, - ParameterClassification, - ParameterRole, - Pointer, - PointerOffset, - Reference, - RuntimeValue, - SubstitutionFailure, - TemplateParameter, - TempStorageParameter, - Value, - classify_parameter, -) -from .group import ( - ArgumentPrecondition, - GroupExecutionRequirements, - GroupLoadStoreAlgorithm, - GroupLoadStoreKind, - GroupLoadStoreSemantics, - GroupLoweringPlan, - GroupLoweringTarget, - GroupOperationSemantics, - GroupPrimitiveCall, - GroupTopologyRequirements, - ImplementationProvenance, - ParticipationRequirements, - PreconditionEnforcement, - ResultVisibility, - StorageOwnership, - SynchronizationRequirements, - SynchronizationScope, - TempStorageRequirements, - ThreadGroupLaunchResolution, - UnsupportedReason, - UnsupportedReasonCode, - make_group_primitive_call, - plan_group_primitive, - resolve_thread_group, -) -from .launch import ( - Dim3, - LaunchFactConflict, - LaunchFactOrigin, - LaunchFacts, - merge_launch_facts, -) -from .thread_group import ( - COMPLETE_WARP_GROUP_KINDS, - MAPPED_GROUP_KINDS, - PHYSICAL_GROUP_KINDS, - THREAD_GROUP_KINDS, - THREAD_LEVELS, - CoopCompilerContextRequiredError, - GroupByMapping, - Hierarchy, - ThreadGroup, - ThreadHierarchy, - cpp_level_expr, - make_thread_group, - normalize_thread_dim, - normalize_thread_group_kind, - normalize_thread_level, - render_group_decl, - render_group_decl_lines, - render_hierarchy_decl, - this_block, - this_cluster, - this_grid, - this_thread, - this_warp, -) - -__all__ = [ - "COMPLETE_WARP_GROUP_KINDS", - "FLOAT32", - "FLOAT64", - "INT8", - "INT16", - "INT32", - "INT64", - "MAPPED_GROUP_KINDS", - "PHYSICAL_GROUP_KINDS", - "THREAD_GROUP_KINDS", - "THREAD_LEVELS", - "UINT8", - "UINT16", - "UINT32", - "UINT64", - "Algorithm", - "ArgumentBinding", - "ArgumentKind", - "ArgumentPrecondition", - "Array", - "BindingKind", - "BuiltinDType", - "Constant", - "CoopCompilerContextRequiredError", - "CoreBackendAdapter", - "CxxFunction", - "Dependency", - "Dim3", - "GroupByMapping", - "GroupExecutionRequirements", - "GroupLoadStoreAlgorithm", - "GroupLoadStoreKind", - "GroupLoadStoreSemantics", - "GroupLoweringPlan", - "GroupLoweringTarget", - "GroupOperationSemantics", - "GroupPrimitiveCall", - "GroupTopologyRequirements", - "Hierarchy", - "ImplementationProvenance", - "LaunchFactConflict", - "LaunchFactOrigin", - "LaunchFacts", - "ParameterClassification", - "ParameterRole", - "ParticipationRequirements", - "Pointer", - "PointerOffset", - "PreconditionEnforcement", - "Reference", - "ResultVisibility", - "RuntimeValue", - "StorageOwnership", - "SubstitutionFailure", - "SynchronizationRequirements", - "SynchronizationScope", - "TempStorageParameter", - "TempStorageRequirements", - "TemplateParameter", - "ThreadGroup", - "ThreadGroupLaunchResolution", - "ThreadHierarchy", - "TypeDefinition", - "UnsupportedReason", - "UnsupportedReasonCode", - "Value", - "binding", - "classify_parameter", - "cpp_level_expr", - "i32_parameter", - "lower_method_parameters", - "make_group_primitive_call", - "make_thread_group", - "merge_launch_facts", - "normalize_thread_dim", - "normalize_thread_group_kind", - "normalize_thread_level", - "plan_group_primitive", - "render_group_decl", - "render_group_decl_lines", - "render_hierarchy_decl", - "resolve_thread_group", - "semantic_token", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -] diff --git a/python/cuda_coop/cuda/coop/_core/_adapter.py b/python/cuda_coop/cuda/coop/_core/_adapter.py deleted file mode 100644 index 939a5abe0386..000000000000 --- a/python/cuda_coop/cuda/coop/_core/_adapter.py +++ /dev/null @@ -1,128 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Translate shared parameter descriptions into a backend's compiler objects. - -Core factories describe C++ calls without depending on a Python GPU compiler. -An adapter supplies the conversion rules for that compiler. The shared helper -preserves parameter order and can omit temporary storage for backends that -supply it another way. -""" - -from __future__ import annotations - -from typing import Any, Protocol - -from ._algorithm import Algorithm -from ._types import ( - TempStorageParameter, -) - - -class CoreBackendAdapter(Protocol): - """Define how a backend turns an ``Algorithm`` into a callable object. - - Lowering means converting a shared descriptor into the backend's type or - parameter representation. Materialization uses those representations to - build the backend's callable algorithm. The backend also owns tracing, - linking, caches, launch integration, and compiler hook registration. - """ - - def normalize_dtype(self, dtype: Any) -> Any: - """Convert a core dtype into the representation this backend uses. - - Built-in core tokens let a factory request types such as signed - 32-bit integers without importing the backend's type system. - """ - ... - - def cpp_type(self, dtype: Any) -> str: - """Spell a dtype as C++ source for generated declarations or calls.""" - ... - - def lower_parameter( - self, - parameter: Any, - *, - specialization: Algorithm, - ) -> Any: - """Convert one non-storage descriptor into a backend parameter. - - ``specialization`` supplies the bound values needed to resolve a - descriptor's ``Dependency`` entries. The returned object belongs to - the backend's signature representation. - """ - ... - - def lower_temp_storage( - self, - parameter: TempStorageParameter, - *, - specialization: Algorithm, - ) -> Any: - """Describe a temporary-storage argument in the backend's signature. - - The separate hook lets the backend apply its storage conventions. - ``specialization`` provides the primitive and its bound arguments. - """ - ... - - def materialize(self, specialization: Algorithm, **kwargs: Any) -> Any: - """Build a backend callable from the fully bound core description. - - ``kwargs`` carries options specific to that backend's construction - path. The protocol does not prescribe the returned object's type. - """ - ... - - -def lower_method_parameters( - adapter: CoreBackendAdapter, - specialization: Algorithm, - method: tuple[Any, ...], - *, - include_temp_storage: bool, -) -> tuple[Any, ...]: - """Convert one signature in order, optionally omitting temporary storage. - - Some backends supply storage outside the ordinary argument list. The - ``include_temp_storage`` flag lets them skip its descriptor while keeping - the remaining arguments in the factory's order. - - Parameters - ---------- - adapter : CoreBackendAdapter - Backend that converts each retained descriptor. - specialization : Algorithm - Bound primitive passed to each conversion hook. - method : tuple - Parameter descriptors for one method signature. - include_temp_storage : bool - If true, send storage descriptors to ``lower_temp_storage``. If - false, omit them. Send all other descriptors to ``lower_parameter``. - - Returns - ------- - tuple - Backend parameter objects in retained call order. - """ - - lowered = [] - for parameter in method: - if isinstance(parameter, TempStorageParameter): - if include_temp_storage: - lowered.append( - adapter.lower_temp_storage( - parameter, - specialization=specialization, - ) - ) - else: - lowered.append( - adapter.lower_parameter( - parameter, - specialization=specialization, - ) - ) - return tuple(lowered) diff --git a/python/cuda_coop/cuda/coop/_core/_algorithm.py b/python/cuda_coop/cuda/coop/_core/_algorithm.py deleted file mode 100644 index cbdd12958f6b..000000000000 --- a/python/cuda_coop/cuda/coop/_core/_algorithm.py +++ /dev/null @@ -1,269 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe a concrete primitive before a backend generates its callable code. - -Factories bind template arguments and describe the C++ method's parameters in -an ``Algorithm``. Backends translate that description into their own compiler -types and generated wrappers. The core record lets factories share this work -without importing a compiler. Its identity includes the bound values so cache -lookups can distinguish different specializations of the same primitive. -""" - -from __future__ import annotations - -from collections.abc import Iterable, Mapping -from dataclasses import dataclass, field -from types import MappingProxyType -from typing import Any - -from ._symbols import semantic_token -from ._types import ( - TemplateParameter, -) - - -def _freeze_methods( - methods: Iterable[Iterable[Any]], -) -> tuple[tuple[Any, ...], ...]: - """Copy method sequences into tuples and retain their descriptors.""" - - return tuple(tuple(method) for method in methods) - - -def _freeze_semantic_value(value: Any, active: set[int]) -> Any: - """Copy supported containers so later caller edits cannot change them. - - Mappings become read-only mappings, lists and tuples become tuples, sets - become frozen sets, and byte arrays become bytes. Other objects stay as - supplied. ``active`` holds container identities on the current recursion - path. A cycle raises ``ValueError``; repeated references are allowed. - """ - - if not isinstance(value, (Mapping, tuple, list, set, frozenset, bytearray)): - return value - - identity = id(value) - if identity in active: - raise ValueError( - "specialization values must not contain container cycles" - ) - active.add(identity) - try: - if isinstance(value, Mapping): - return MappingProxyType( - { - _freeze_semantic_value(key, active): _freeze_semantic_value( - item, active - ) - for key, item in value.items() - } - ) - if isinstance(value, (tuple, list)): - return tuple(_freeze_semantic_value(item, active) for item in value) - if isinstance(value, (set, frozenset)): - return frozenset( - _freeze_semantic_value(item, active) for item in value - ) - return bytes(value) - finally: - active.remove(identity) - - -def _freeze_mapping(mapping: Mapping[str, Any]) -> Mapping[str, Any]: - """Snapshot named values for an algorithm's cached identity.""" - - return MappingProxyType( - { - key: _freeze_semantic_value(value, set()) - for key, value in mapping.items() - } - ) - - -@dataclass(frozen=True) -class TypeDefinition: - """Supply a C++ declaration needed by the generated primitive wrapper. - - A factory can use this for a helper type referenced by its parameters or - template arguments. The backend emits ``code`` before the wrapper. - - Attributes - ---------- - name : str - Name that identifies this definition in the algorithm description. - code : str - Complete C++ source for the declaration, including any terminator. - """ - - name: str - code: str - - @property - def semantic_key(self) -> tuple[Any, ...]: - return self.name, semantic_token(self.code) - - -@dataclass(frozen=True, eq=False) -class Algorithm: - """Describe one C++ primitive with its template arguments already bound. - - Construction checks that every declared template name has a binding. - A backend then resolves parameter dependencies and builds the callable - wrapper. This record contains the information for that step; constructing - it does not compile or link code. - - Attributes - ---------- - struct_name : str - C++ primitive class name, such as ``BlockLoad``. - method_name : str - Method to call on that class, such as ``Load``. - c_name : str - Stem the backend uses to name generated wrappers. - includes : tuple of str - Header paths needed to compile the primitive and its parameters. - template_parameters : tuple of TemplateParameter - Template declarations in C++ argument order. - parameters : tuple of tuple - Candidate method signatures. Each inner tuple contains parameter - descriptors in call order, including any temporary-storage marker. - template_arguments : Mapping - Values for the declared template parameters and any extra named - dependencies used by parameter descriptors. - metadata : Mapping - Additional planning facts. These also contribute to record identity. - type_definitions : tuple of TypeDefinition - Helper declarations the backend emits before the wrapper. - fake_return : bool - If true, pass every output descriptor to the C++ method as an ordinary - argument. The wrapper ignores the method's return value. - output_by_reference : bool - Used when ``fake_return`` is false. The signature can then have at - most one output, which receives the result. ``True`` passes that output - to the method as an argument. ``False`` leaves it out of the method's - arguments and assigns the method's return value to it. - - Notes - ----- - Construction copies supported containers in ``template_arguments`` and - ``metadata`` into immutable forms, then computes the equality and hash - key. It retains other leaf objects as supplied. Parameter descriptors and - those leaf objects must remain suitable for the backend that reads them. - - Extra dependency names need not be C++ template parameters. For example, - a parameter's array extent can use a named value that the class template - itself does not accept. - - Return-handling flags describe the wrapper behavior for the adapter to - implement. The core stores these flags and includes them in identity; - wrapper generation belongs to the backend. - """ - - struct_name: str - method_name: str - c_name: str - includes: tuple[str, ...] - template_parameters: tuple[TemplateParameter, ...] - parameters: tuple[tuple[Any, ...], ...] - template_arguments: Mapping[str, Any] = field(kw_only=True) - metadata: Mapping[str, Any] = field(default_factory=dict, kw_only=True) - type_definitions: tuple[TypeDefinition, ...] = () - fake_return: bool = False - output_by_reference: bool = False - _semantic_key: tuple[Any, ...] = field(init=False, repr=False) - - def __post_init__(self) -> None: - object.__setattr__(self, "includes", tuple(self.includes)) - object.__setattr__( - self, "template_parameters", tuple(self.template_parameters) - ) - object.__setattr__(self, "parameters", _freeze_methods(self.parameters)) - object.__setattr__( - self, "type_definitions", tuple(self.type_definitions) - ) - - names = self.template_parameter_names - if len(set(names)) != len(names): - raise ValueError("template parameter names must be unique") - missing = [ - name for name in names if name not in self.template_arguments - ] - if missing: - joined = ", ".join(missing) - raise ValueError(f"Template argument(s) not provided: {joined}") - - object.__setattr__( - self, "template_arguments", _freeze_mapping(self.template_arguments) - ) - object.__setattr__(self, "metadata", _freeze_mapping(self.metadata)) - object.__setattr__( - self, - "_semantic_key", - ( - ( - self.struct_name, - self.method_name, - self.c_name, - self.includes, - semantic_token(self.template_parameters), - semantic_token(self.parameters), - semantic_token(self.type_definitions), - self.fake_return, - self.output_by_reference, - ), - semantic_token(self.template_arguments), - semantic_token(self.metadata), - ), - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, Algorithm): - return NotImplemented - return self.semantic_key == other.semantic_key - - def __hash__(self) -> int: - return hash(self.semantic_key) - - @property - def template_parameter_names(self) -> tuple[str, ...]: - return tuple(parameter.name for parameter in self.template_parameters) - - @property - def ordered_template_arguments(self) -> tuple[tuple[str, Any], ...]: - """Return template bindings in C++ order.""" - - return tuple( - (name, self.template_arguments[name]) - for name in self.template_parameter_names - ) - - @property - def ordered_specialization_arguments(self) -> tuple[tuple[str, Any], ...]: - """List template bindings first, then extra dependencies by name. - - Backends use this order to resolve a signature without depending on - the caller's mapping insertion order. - """ - - template_names = set(self.template_parameter_names) - auxiliary = sorted( - ( - (name, value) - for name, value in self.template_arguments.items() - if name not in template_names - ), - key=lambda item: item[0], - ) - return (*self.ordered_template_arguments, *auxiliary) - - @property - def semantic_key(self) -> tuple[Any, ...]: - """Return the cached identity of the description and bound values. - - Equality and hashing use this key so equivalent records can share a - provider or cache entry. Construction computes it once. - """ - - return self._semantic_key diff --git a/python/cuda_coop/cuda/coop/_core/_auto_registration.py b/python/cuda_coop/cuda/coop/_core/_auto_registration.py deleted file mode 100644 index d81c81445c1d..000000000000 --- a/python/cuda_coop/cuda/coop/_core/_auto_registration.py +++ /dev/null @@ -1,258 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Register compatible compiler runtimes already imported by the application. - -The ``cuda.coop`` root import calls this module's allowlisted registration -probe. A runtime must already be in ``sys.modules`` before its backend is -considered; merely installing an optional compiler does not cause the root -import to load it. This makes compiler-first imports convenient while keeping -the common API available without a compiler. Root-first applications can call -``cuda.coop.register("numba-cuda-mlir")`` or import the qualified backend. - -``CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION`` controls this probe. Unset, empty, -``"0"``, ``"false"``, ``"no"``, and ``"off"`` leave automatic registration -enabled; surrounding whitespace and letter case are ignored. Any other value, -including ``"1"``, disables it. The value is read on each probe, normally -during the first ``cuda.coop`` import. Changing it does not unregister an -active backend or trigger another probe. Explicit registration and qualified -backend imports remain available when automatic registration is disabled. - -An absent optional runtime is skipped silently. A detected runtime that fails -activation produces ``CudaCoopAutoRegistrationWarning`` and leaves the common -API import usable under normal warning handling. The probe removes newly -imported modules for the failed backend. The Numba backend registers its -planner only after its implementation modules have loaded successfully. -Warning filters may promote the warning to an exception. -""" - -from __future__ import annotations - -import importlib -import importlib.metadata -import os -import sys -import warnings -from collections.abc import Callable -from dataclasses import dataclass -from types import ModuleType - -_DISABLE_ENV = "CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION" -_FALSE_ENV_VALUES = frozenset({"", "0", "false", "no", "off"}) -_WARNING_PREFIX = "cuda.coop automatic DSL registration:" - -# Keep this an explicit allowlist so installing an unrelated compiler package -# cannot change the cuda.coop root import. -_AUTO_DSL_CANDIDATES = ("numba_mlir",) - - -class CudaCoopAutoRegistrationWarning(UserWarning): - """A detected optional DSL could not activate the common API.""" - - -class _BackendUnavailable(ImportError): - """The optional backend's top-level runtime is genuinely absent.""" - - -@dataclass(frozen=True) -class _Candidate: - """Describe one runtime that the root import is allowed to activate. - - The runtime module name controls detection. Distribution names and the - installation hint supply diagnostics without importing another compiler. - The activation callable performs the backend's own compatibility checks. - """ - - display_name: str - runtime_module: str - distributions: tuple[str, ...] - install_hint: str - activate: Callable[[], ModuleType] - - -def _auto_registration_disabled(value: str | None = None) -> bool: - """Return whether automatic probing is disabled by the environment.""" - - if value is None: - value = os.environ.get(_DISABLE_ENV) - if value is None: - return False - return value.strip().lower() not in _FALSE_ENV_VALUES - - -def _import_optional(module_name: str, *, top_level: str) -> ModuleType: - """Import an optional runtime without hiding dependency failures. - - Automatic registration may silently skip an absent optional runtime, but - a dependency failure inside an installed runtime should be reported. Use - the missing module recorded on ``ImportError`` to distinguish these cases; - catching every import failure as absence would hide incompatible installs. - - ``_BackendUnavailable`` is an internal control signal, caught by - ``_auto_register_known_dsls`` around the candidate activation call. That - catcher cleans up newly imported backend modules and skips the candidate - without a warning or user-facing exception. Other import failures reach - the probe's general exception handler and become an incompatibility - warning. The public ``register`` entry point and the qualified backend's - activation code report import failures through their own error path. - - Parameters - ---------- - module_name : str - Fully qualified module to import. - top_level : str - Runtime module name whose absence is an expected optional dependency. - Only an exact match with the exception's ``name`` counts as absent. - - Returns - ------- - ModuleType - Imported module. - - Raises - ------ - _BackendUnavailable - The import reports that ``top_level`` itself is missing. The automatic - registration probe consumes this private signal as a silent skip. - ImportError - Any other import failure, propagated unchanged for diagnostics. - """ - - try: - return importlib.import_module(module_name) - except ImportError as error: - if getattr(error, "name", None) == top_level: - raise _BackendUnavailable(top_level) from error - raise - - -def _activate_numba_mlir() -> ModuleType: - """Load the runtime, then let the Numba backend validate and activate.""" - - _import_optional("numba_cuda_mlir", top_level="numba_cuda_mlir") - return importlib.import_module("cuda.coop.numba_mlir") - - -_CANDIDATES = { - "numba_mlir": _Candidate( - display_name="Numba-CUDA-MLIR", - runtime_module="numba_cuda_mlir", - distributions=("numba-cuda-mlir",), - install_hint=( - "cuda-coop[numba-cuda-mlir-cu12] for CUDA 12 or " - "cuda-coop[numba-cuda-mlir-cu13] for CUDA 13" - ), - activate=_activate_numba_mlir, - ), -} - - -def _detected_version(candidate: _Candidate) -> str | None: - """Find a version for diagnostics without importing the runtime again. - - Prefer the loaded module's version, then installed distribution metadata. - An unknown version does not prevent the probe from reporting its error. - """ - - runtime = sys.modules.get(candidate.runtime_module) - version = getattr(runtime, "__version__", None) - if isinstance(version, str) and version: - return version - for distribution in candidate.distributions: - try: - return importlib.metadata.version(distribution) - except importlib.metadata.PackageNotFoundError: - continue - return None - - -def _remove_failed_backend_modules(prefix: str, before: frozenset[str]) -> None: - """Remove only backend modules added by the failed activation attempt. - - Keep modules that were present before probing so cleanup does not remove - imports owned by the application or an earlier successful activation. - """ - - for module_name in tuple(sys.modules): - if ( - module_name == prefix or module_name.startswith(f"{prefix}.") - ) and module_name not in before: - sys.modules.pop(module_name, None) - - -def _warn_incompatible(candidate: _Candidate, error: Exception) -> None: - """Explain an activation failure and how to select a compatible install. - - The warning normally lets the common API import finish. Application - warning filters can still turn it into an exception. - """ - - version = _detected_version(candidate) - detected = f" (detected version {version})" if version is not None else "" - reason = str(error).strip() or type(error).__name__ - missing = getattr(error, "name", None) - if isinstance(missing, str) and missing and missing not in reason: - reason = f"dependency {missing!r} failed to import: {reason}" - warnings.warn( - f"{_WARNING_PREFIX} {candidate.display_name}{detected} was detected " - f"but was not enabled because {reason}. The cuda.coop root import " - "continued and other DSL backends were unaffected. " - f"Install a compatible {candidate.install_hint}. " - f"Set {_DISABLE_ENV}=1 to disable automatic DSL probing.", - CudaCoopAutoRegistrationWarning, - stacklevel=2, - ) - - -def _auto_register_known_dsls() -> tuple[str, ...]: - """Activate allowlisted runtimes already imported by the application. - - The root package import must remain usable without loading an optional - compiler or CUDA bindings. Inspect ``sys.modules`` first: installing a - runtime is insufficient to activate it. Compiler-first imports get this - automatic activation. Root-first callers can activate the backend with - ``cuda.coop.register("numba-cuda-mlir")``. - - Respect ``CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION`` and reuse qualified - backends already in ``sys.modules``. For a new attempt, snapshot loaded - modules so failure cleanup removes only newly imported backend modules. - An absent optional runtime is skipped silently. Other activation failures - produce ``CudaCoopAutoRegistrationWarning``, and probing continues. - - Returns - ------- - tuple of str - Internal backend names successfully activated or already loaded, in - candidate order. Empty when probing is disabled or no candidate - qualifies. Does not include missing or unsuccessfully activated DSLs. - """ - - if _auto_registration_disabled(): - return () - - registered = [] - for name in _AUTO_DSL_CANDIDATES: - candidate = _CANDIDATES[name] - if candidate.runtime_module not in sys.modules: - continue - package_prefix = f"cuda.coop.{name}" - if package_prefix in sys.modules: - registered.append(name) - continue - before = frozenset(sys.modules) - try: - candidate.activate() - except _BackendUnavailable: - _remove_failed_backend_modules(package_prefix, before) - except Exception as error: # noqa: BLE001 - optional activation must not break the root import. - _remove_failed_backend_modules(package_prefix, before) - _warn_incompatible(candidate, error) - else: - registered.append(name) - return tuple(registered) - - -__all__ = [ - "_auto_register_known_dsls", -] diff --git a/python/cuda_coop/cuda/coop/_core/_bindings.py b/python/cuda_coop/cuda/coop/_core/_bindings.py deleted file mode 100644 index e2b30dc48464..000000000000 --- a/python/cuda_coop/cuda/coop/_core/_bindings.py +++ /dev/null @@ -1,434 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Carry scalar option choices from a frontend into a primitive factory. - -A value alone cannot tell a factory which call to build. An omitted count can -select an unguarded overload; a known count can be embedded in C++ source; a -runtime count needs a device-call argument. ``ArgumentBinding`` keeps that -choice separate from the backend expression that supplies a runtime value. - -The helpers here validate scalar representations and build parameter records. -Operation factories apply further rules, such as keeping a count within one -tile. In these helpers, i32 and i64 mean signed 32-bit and 64-bit integers. -""" - -from __future__ import annotations - -import math -from dataclasses import dataclass -from enum import Enum -from numbers import Integral, Real -from typing import Any - -from ._types import INT32, ArgumentKind, CxxFunction, RuntimeValue, Value - -_I32_MIN = -(1 << 31) -_I32_MAX = (1 << 31) - 1 -_I64_MIN = -(1 << 63) -_I64_MAX = (1 << 63) - 1 -_U64_MAX = (1 << 64) - 1 - - -class BindingKind(str, Enum): - """How a scalar option is supplied to a generated primitive call. - - A factory needs more than the option's value: it must know whether to - select an overload without that option, embed a compile-time constant, - or accept an argument on each device call. These choices affect the - generated wrapper signature and whether one specialization can be reused. - - For example, an omitted ``valid_items`` can select an unguarded load; - a static count of 12 produces a guarded load with 12 embedded in its - wrapper; a runtime count produces a guarded load whose count is supplied - on each invocation. A static zero is therefore different from omission. - - Attributes - ---------- - OMITTED - The caller did not supply the option. The consuming factory chooses - an overload or supplies its own default; no payload is stored. - STATIC - The value is known while building the specialization and can be - embedded in generated code. The binding retains that value. - RUNTIME - The value is supplied when the device code executes. The binding - records that requirement without retaining a value or compiler IR. - - Notes - ----- - ``ArgumentKind`` classifies parameters after a signature is selected and - has only static and runtime cases. ``BindingKind`` also represents the - earlier decision to omit an option entirely. - """ - - OMITTED = "omitted" - STATIC = "static" - RUNTIME = "runtime" - - -@dataclass(frozen=True, eq=False) -class ArgumentBinding: - """Record how a scalar option is supplied and retain its static value. - - Frontends use this record to pass argument decisions into shared - planning code without carrying backend-specific runtime expressions. - Factories then select overloads and turn the bindings into parameter - descriptors: for example, ``i32_parameter`` creates an embedded - ``CxxFunction`` for a static count or a runtime ``Value`` descriptor. - Lowering connects runtime descriptors to the actual device operands. - - Prefer ``omitted()``, ``static(value)``, and ``runtime()`` when building - a record explicitly, or ``binding`` to classify a frontend value. - A static value contributes to specialization identity; runtime bindings - carry no payload, allowing calls with different runtime values to share - the same description. - - Attributes - ---------- - kind : BindingKind - Whether the option is omitted, known statically, or runtime-provided. - value : Any, optional - Payload for ``STATIC``; must be ``None`` for the other two modes. - Consuming helpers and factories validate payload types, integer - widths, and operation-specific bounds. - - Notes - ----- - Equality and hashing use ``semantic_key`` rather than Python numeric - equality, so static ``True`` and ``1``, or ``0.0`` and ``-0.0``, remain - distinct requests. Integer normalization can give equivalent integral - types the same representation before they enter a specialization key. - """ - - kind: BindingKind - value: Any = None - - def __post_init__(self) -> None: - if self.kind is not BindingKind.STATIC and self.value is not None: - raise ValueError("only static argument bindings may carry a value") - - @classmethod - def omitted(cls) -> ArgumentBinding: - """Leave the option's overload or default selection to the factory.""" - - return cls(BindingKind.OMITTED) - - @classmethod - def static(cls, value: Any) -> ArgumentBinding: - """Retain a compile-time value for later validation and embedding.""" - - return cls(BindingKind.STATIC, value) - - @classmethod - def runtime(cls) -> ArgumentBinding: - """Request an argument without storing its runtime value.""" - - return cls(BindingKind.RUNTIME) - - @property - def semantic_key(self) -> tuple[str, ...]: - """Return the binding identity used by equality and hashing. - - The kind alone identifies omitted and runtime bindings. - Static bindings also include the payload type's module and qualified - name and the payload's ``repr``. This distinguishes values that - Python considers numerically equal but may generate different code. - The key reflects the stored representation; consumers normalize - values first when equivalent input types should share a key. - """ - - if self.kind is not BindingKind.STATIC: - return (self.kind.value,) - value_type = type(self.value) - return ( - self.kind.value, - value_type.__module__, - value_type.__qualname__, - repr(self.value), - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, ArgumentBinding): - return NotImplemented - return self.semantic_key == other.semantic_key - - def __hash__(self) -> int: - return hash(self.semantic_key) - - @property - def argument_kind(self) -> ArgumentKind | None: - """Classify a supplied parameter, or return ``None`` for omission. - - Planners use this when describing static and runtime call arguments. - A factory may separately replace an omitted option with a default. - """ - - if self.kind is BindingKind.OMITTED: - return None - if self.kind is BindingKind.STATIC: - return ArgumentKind.STATIC - return ArgumentKind.RUNTIME - - -def binding(value: Any, *, omitted: Any = None) -> ArgumentBinding: - """Classify an option as omitted, static, or runtime-provided. - - Frontends mark runtime expressions with ``RuntimeValue`` before calling - this helper. Every other supplied object is treated as a static payload; - its type alone does not imply that it is a device-time expression. - In particular, booleans are static values here. The boolean flags accepted - by some factories to select overloads are a separate convention. - - Parameters - ---------- - value : Any - Frontend option value or a ``RuntimeValue`` marker. An existing - ``ArgumentBinding`` is not passed through; use it directly instead - of classifying it again. - omitted : Any, optional - Sentinel for an absent option, defaulting to ``None``. Compared by - object identity before checking for ``RuntimeValue``. - - Returns - ------- - ArgumentBinding - Omitted when ``value is omitted``, runtime for a ``RuntimeValue``, - and static with the original payload otherwise. A runtime marker's - name is discarded; parameter builders supply their own names. - """ - - if value is omitted: - return ArgumentBinding.omitted() - if isinstance(value, RuntimeValue): - return ArgumentBinding.runtime() - return ArgumentBinding.static(value) - - -def i32_parameter( - option: ArgumentBinding, - *, - name: str, - omitted_value: int | None = None, -) -> Value | CxxFunction | None: - """Turn a count-like option into a signed-i32 parameter descriptor. - - A runtime ``Value`` becomes an operand of the generated wrapper. A - ``CxxFunction`` embeds the constant expression in the underlying CUB - call without adding a runtime operand. An omitted binding can remove - the parameter or substitute a factory-provided constant default. - - Parameters - ---------- - option : ArgumentBinding - Binding that determines whether and how the parameter is supplied. - name : str - Parameter name, also used to identify the option in validation errors. - omitted_value : int or None, optional - Constant to embed when ``option`` is omitted. With the default - ``None``, omission produces no descriptor. Ignored for static and - runtime bindings. - - Returns - ------- - Value or CxxFunction or None - Runtime i32 descriptor, embedded i32 constant, or no parameter. - This constructs metadata; wrapper code generation happens later. - - Raises - ------ - TypeError - A used static value or omitted default is a boolean or non-integer. - ValueError - A used static value or omitted default does not fit signed i32. - The consuming factory or planner checks operation-specific bounds, - such as a tile's item count. - """ - - if option.kind is BindingKind.OMITTED: - if omitted_value is None: - return None - value = _normalize_i32(omitted_value, name=name, source="omitted") - return CxxFunction(str(value), INT32, name=name) - if option.kind is BindingKind.RUNTIME: - return Value(INT32, name=name) - value = _normalize_i32(option.value, name=name, source="static") - return CxxFunction(str(value), INT32, name=name) - - -def _normalize_i32(value: Any, *, name: str, source: str) -> int: - """Convert an integral value to Python ``int`` within signed-i32 bounds. - - Reject booleans and non-integral values with ``TypeError`` and overflow - with ``ValueError``. ``source`` and ``name`` identify the offending - binding in diagnostics, such as ``static valid_items``. - """ - - if isinstance(value, bool) or not isinstance(value, Integral): - raise TypeError(f"{source} {name} must be an integer") - normalized = int(value) - if not _I32_MIN <= normalized <= _I32_MAX: - raise ValueError(f"{source} {name} must fit a signed 32-bit integer") - return normalized - - -def normalize_i32_binding( - option: ArgumentBinding, - *, - name: str, -) -> ArgumentBinding: - """Validate a static i32 binding and normalize its identity. - - Equivalent integral values, such as Python ``1`` and NumPy ``int32(1)``, - become the same Python ``int`` payload. This avoids separate semantic - keys for constants that generate the same i32 argument. - - Parameters - ---------- - option : ArgumentBinding - Binding to normalize; omitted and runtime bindings pass through. - name : str - Option name included in validation errors. - - Returns - ------- - ArgumentBinding - A new static binding with a Python ``int`` payload, or the original - non-static binding. This helper accepts negative values within the - signed-i32 range. Factories check operation-specific bounds. - - Raises - ------ - TypeError - The static payload is not integral or is a boolean. - ValueError - The static payload is outside ``[-2**31, 2**31 - 1]``. - """ - - if option.kind is not BindingKind.STATIC: - return option - return ArgumentBinding.static( - _normalize_i32(option.value, name=name, source="static") - ) - - -def _normalize_i64(value: Any, *, name: str, source: str) -> int: - """Convert an integral value to Python ``int`` within signed-i64 bounds. - - Apply the same type checks and diagnostic labels as ``_normalize_i32``. - Accept negative values here; pointer-offset factories apply their own - nonnegative bounds after this representation check. - """ - - if isinstance(value, bool) or not isinstance(value, Integral): - raise TypeError(f"{source} {name} must be an integer") - normalized = int(value) - if not _I64_MIN <= normalized <= _I64_MAX: - raise ValueError(f"{source} {name} must fit a signed 64-bit integer") - return normalized - - -def normalize_i64_binding( - option: ArgumentBinding, - *, - name: str, -) -> ArgumentBinding: - """Validate a static i64 binding and normalize its identity. - - Pointer-offset planning uses this to give equal integral offsets the - same Python ``int`` payload and semantic key. Omitted and runtime - bindings already describe their supply mode without a value to normalize. - - Parameters - ---------- - option : ArgumentBinding - Binding to normalize; omitted and runtime bindings pass through. - name : str - Option name included in validation errors. - - Returns - ------- - ArgumentBinding - A new static binding with a Python ``int`` payload, or the original - non-static binding. Factories and group planning check nonnegative - offsets and combined tile-origin bounds separately. - - Raises - ------ - TypeError - The static payload is not integral or is a boolean. - ValueError - The static payload is outside ``[-2**63, 2**63 - 1]``. - """ - - if option.kind is not BindingKind.STATIC: - return option - return ArgumentBinding.static( - _normalize_i64(option.value, name=name, source="static") - ) - - -def cxx_scalar_literal(value: Any, *, name: str) -> str: - """Render a static scalar as an expression for a generated CUB call. - - Factories use this for embedded values such as a load's ``oob_default``. - The result is source text for a ``CxxFunction`` descriptor, whose dtype - is supplied separately by the factory. - - Parameters - ---------- - value : Any - Boolean, integral, or real scalar, or an object whose ``value`` - attribute contains one. Integral values may span signed-i64 minimum - through unsigned-i64 maximum. Real values must remain finite after - conversion to Python ``float``. - name : str - Option name included in validation errors. - - Returns - ------- - str - C++ boolean literal, integer expression, or floating-point literal. - Integers above signed-i64 maximum receive a ``ULL`` suffix; signed - i64 minimum uses an expression that avoids an overflowing positive - literal. Conversion into the eventual CUB element type is left to - later compilation. - - Raises - ------ - TypeError - The unwrapped value is not a supported numeric scalar. - ValueError - An integer is outside ``[-2**63, 2**64 - 1]`` or a real is nonfinite. - """ - - scalar = getattr(value, "value", value) - if isinstance(scalar, bool): - return "true" if scalar else "false" - if isinstance(scalar, Integral): - normalized = int(scalar) - if not _I64_MIN <= normalized <= _U64_MAX: - raise ValueError(f"static {name} must fit a 64-bit integer") - if normalized == _I64_MIN: - return "(-9223372036854775807LL - 1LL)" - if normalized > _I64_MAX: - return f"{normalized}ULL" - return str(normalized) - if isinstance(scalar, Real): - normalized = float(scalar) - if not math.isfinite(normalized): - raise ValueError(f"static {name} must be finite") - return repr(normalized) - raise TypeError(f"static {name} must be a numeric scalar") - - -__all__ = [ - "ArgumentBinding", - "BindingKind", - "binding", - "cxx_scalar_literal", - "i32_parameter", - "normalize_i32_binding", - "normalize_i64_binding", -] diff --git a/python/cuda_coop/cuda/coop/_core/_source_dump.py b/python/cuda_coop/cuda/coop/_core/_source_dump.py deleted file mode 100644 index 9bae3fdb5e58..000000000000 --- a/python/cuda_coop/cuda/coop/_core/_source_dump.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Save generated CUDA source when source inspection is enabled. - -Backends call this before cache lookup so cached compilations remain -inspectable. The backend, source, and compile identity determine the -filename. An atomic write keeps concurrent calls from leaving a partial -source file. -""" - -from __future__ import annotations - -import hashlib -import os -import tempfile -from pathlib import Path - - -def dump_source( - source: str, - *, - backend: str, - identity: tuple[object, ...] = (), -) -> Path | None: - """Write generated CUDA source to a file named by its content and backend. - - ``CUDA_COOP_SOURCE_DUMP_DIR`` selects the directory and is read on every - call; unset or empty disables dumping. Other values are directory names, - including strings such as ``"0"`` or ``"false"``. Expand ``~`` and resolve - relative paths against the current working directory when the call runs. - Environment changes affect subsequent dumps without reimporting a backend. - - Call before compiler or cache lookup to capture sources on cache hits too. - The filename combines the backend tag with a hash of the identity - representation and source bytes, allowing different targets and generated - variants to coexist. An existing file is reused without rewriting it; - changing or unsetting the directory does not remove previous dumps. - - Publish through a flushed and synced temporary file in the same directory - followed by atomic replacement, so concurrent dumps leave complete files. - Filesystem failures propagate to the caller. - - Parameters - ---------- - source : str - Complete generated CUDA source, written as UTF-8 without modification. - backend : str - Backend label embedded directly in the filename, such as - ``"numba_mlir"``. Callers supply a filename-safe label. - identity : tuple of object, optional - Extra compile identity, such as target architecture and output kind. - Its ``repr`` contributes to the digest. Use reproducible - representations so repeated compilations select the same path. - - Returns - ------- - pathlib.Path or None - Absolute path of the existing or newly written source file, or - ``None`` when dumping is disabled. - - Raises - ------ - OSError - The destination directory or source file cannot be created or written. - """ - - dump_dir = os.environ.get("CUDA_COOP_SOURCE_DUMP_DIR") - if not dump_dir: - return None - - contents = source.encode("utf-8") - digest = hashlib.sha256() - digest.update(repr(identity).encode("utf-8", errors="surrogateescape")) - digest.update(b"\0") - digest.update(contents) - root = Path(dump_dir).expanduser().resolve() - root.mkdir(mode=0o700, parents=True, exist_ok=True) - destination = root / f"cuda_coop_{backend}_{digest.hexdigest()}.cu" - if destination.is_file(): - return destination - - fd, temporary = tempfile.mkstemp(dir=root, prefix=f".{destination.name}.") - try: - with os.fdopen(fd, "wb") as source_file: - source_file.write(contents) - source_file.flush() - os.fsync(source_file.fileno()) - os.replace(temporary, destination) - finally: - Path(temporary).unlink(missing_ok=True) - return destination diff --git a/python/cuda_coop/cuda/coop/_core/_symbols.py b/python/cuda_coop/cuda/coop/_core/_symbols.py deleted file mode 100644 index 0d31003d17e5..000000000000 --- a/python/cuda_coop/cuda/coop/_core/_symbols.py +++ /dev/null @@ -1,292 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe specialization values for equality and cache lookup. - -Python equality can merge values that must generate different code, such as -``True`` and ``1`` or positive and negative floating-point zero. These helpers -preserve those distinctions while turning containers and records into nested -tokens. They also describe object cycles without using process-local addresses -as the identity of a back-reference. -""" - -from __future__ import annotations - -import dataclasses -import math -import re -import struct -from collections import defaultdict -from collections.abc import Mapping -from enum import Enum -from functools import partial -from types import BuiltinFunctionType, FunctionType, MethodType, ModuleType -from typing import Any - -_ADDRESS_IN_REPR = re.compile(r"(?<= at )0x[0-9a-fA-F]+") - - -def _defined_module_name(value: Any) -> str | None: - """Use a module name only when the object supplies it as text.""" - - module_name = getattr(value, "__module__", None) - return module_name if isinstance(module_name, str) else None - - -@dataclasses.dataclass -class _TokenState: - """Track recursion and reuse completed work within one token query. - - ``active`` maps object identities to their depth on the current path. - ``completed`` retains each reusable token and its source object; retaining - the object prevents Python from reusing its identity during this query. - ``cycle_hits`` counts back-references so a parent can tell whether its - token depends on the current recursion path. - """ - - active: dict[int, int] = dataclasses.field(default_factory=dict) - completed: dict[tuple[str, int], tuple[Any, Any]] = dataclasses.field( - default_factory=dict - ) - cycle_hits: int = 0 - - -def _object_state_token( - value: Any, - state: _TokenState, -) -> tuple[Any, ...] | None: - """Describe stored instance attributes, including inherited private slots. - - Python mangles private slot names using the defining class. Resolve those - names before reading values, and skip slots that have not been assigned. - Return ``None`` when these attributes expose no stored state. - """ - - object_state: list[tuple[str, Any]] = [] - attributes = getattr(value, "__dict__", None) - if attributes: - object_state.append(("__dict__", _semantic_token(attributes, state))) - - seen_slots: set[str] = set() - for cls in type(value).__mro__: - slots = cls.__dict__.get("__slots__", ()) - if isinstance(slots, str): - slots = (slots,) - for slot in slots: - storage_name = slot - if slot.startswith("__") and not slot.endswith("__"): - class_name = cls.__name__.lstrip("_") - if class_name: - storage_name = f"_{class_name}{slot}" - if storage_name in {"__dict__", "__weakref__"} or ( - storage_name in seen_slots - ): - continue - seen_slots.add(storage_name) - try: - slot_value = getattr(value, storage_name) - except AttributeError: - continue - token = ( - ("self",) - if slot_value is value - else _semantic_token(slot_value, state) - ) - object_state.append((storage_name, token)) - - return tuple(object_state) if object_state else None - - -def _cycle_token( - value: Any, back_reference_depth: int -) -> tuple[str, str, str, int]: - """Identify a cycle by object kind and distance to its active ancestor.""" - - return ( - "cycle", - getattr(value, "__module__", type(value).__module__), - getattr(value, "__qualname__", type(value).__qualname__), - back_reference_depth, - ) - - -def _container_kind(value: Any) -> tuple[str | None, str]: - """Keep a container subclass distinct from its built-in counterpart.""" - - return _defined_module_name(type(value)), type(value).__qualname__ - - -def _container_state_token( - value: Any, - state: _TokenState, -) -> tuple[tuple[str, Any], ...] | None: - """Capture container state that equal elements alone cannot describe. - - This includes instance attributes and a ``defaultdict``'s factory, which - determines the value supplied for a missing key. - """ - - container_state = list(_object_state_token(value, state) or ()) - if isinstance(value, defaultdict): - default_factory = value.default_factory - default_factory_token = ( - ("self",) - if default_factory is value - else _semantic_token(default_factory, state) - ) - container_state.append(("default_factory", default_factory_token)) - return tuple(container_state) if container_state else None - - -def _semantic_token(value: Any, state: _TokenState) -> Any: - """Describe a value while sharing cycle detection and memoized results. - - Scalar cases preserve type-sensitive distinctions before recursion starts. - Container tokens include their type, stored state, and elements. Sort - mapping and set tokens so insertion order does not affect their identity. - Sequences retain their element order. - """ - - if isinstance(value, Enum): - return type(value).__module__, type(value).__qualname__, value.value - if isinstance(value, bool): - return "bool", value - if isinstance(value, float): - if math.isnan(value): - return "float", "nan", struct.pack(">d", value).hex() - return "float", value.hex() - if value is None or isinstance(value, (int, str, bytes)): - return value - if isinstance(value, ModuleType): - return "module", value.__name__ - if isinstance(value, type): - return "type", _defined_module_name(value), value.__qualname__ - - value_id = id(value) - if value_id in state.active: - state.cycle_hits += 1 - return _cycle_token(value, len(state.active) - state.active[value_id]) - - memo_key = ("value", value_id) - completed = state.completed.get(memo_key) - if completed is not None: - return completed[1] - - state.active[value_id] = len(state.active) - cycle_hits_before = state.cycle_hits - try: - if isinstance(value, Mapping): - items_snapshot = tuple(value.items()) - items = [ - (_semantic_token(key, state), _semantic_token(item, state)) - for key, item in items_snapshot - ] - token = ( - "mapping", - _container_kind(value), - _container_state_token(value, state), - tuple(sorted(items, key=lambda item: repr(item[0]))), - ) - elif isinstance(value, (tuple, list)): - token = ( - "sequence", - _container_kind(value), - _container_state_token(value, state), - tuple(_semantic_token(item, state) for item in value), - ) - elif isinstance(value, (set, frozenset)): - token = ( - "set", - _container_kind(value), - _container_state_token(value, state), - tuple( - sorted( - (_semantic_token(item, state) for item in value), - key=repr, - ) - ), - ) - elif dataclasses.is_dataclass(value) and not isinstance(value, type): - token = ( - type(value).__module__, - type(value).__qualname__, - tuple( - ( - field.name, - _semantic_token(getattr(value, field.name), state), - ) - for field in dataclasses.fields(value) - ), - ) - elif isinstance( - value, - ( - BuiltinFunctionType, - FunctionType, - MethodType, - partial, - staticmethod, - classmethod, - ), - ): - raise TypeError( - "callable values are not supported in specialization keys" - ) - else: - attributes = _object_state_token(value, state) - if attributes is not None: - token = ( - type(value).__module__, - type(value).__qualname__, - attributes, - ) - else: - stable_repr = _ADDRESS_IN_REPR.sub("0x?", repr(value)) - token = ( - type(value).__module__, - type(value).__qualname__, - stable_repr, - ) - finally: - del state.active[value_id] - - if state.cycle_hits == cycle_hits_before: - # A cycle token contains path-relative depth, so reuse only results - # that did not encounter a back-reference while visiting this value. - state.completed[memo_key] = (value, token) - return token - - -def semantic_token(value: Any) -> Any: - """Describe a specialization value for comparison and hashing. - - Algorithms use this token to compare bound values and planning metadata. - A token records container types as well as contents, preserves floating - zero signs and NaN bit patterns, and handles cycles in object state. - Each call starts a fresh traversal so earlier queries cannot supply stale - descriptions of mutable objects. - - Parameters - ---------- - value : object - Scalar, container, dataclass, type, module, or state-bearing object - used in a specialization. Objects with no stored state use their - type and ``repr``, with ordinary memory-address text removed. - - Returns - ------- - object - A scalar or nested tuple description. Supported specialization values - produce hashable tokens; this function does not compute a hash digest. - - Raises - ------ - TypeError - A visited value is a Python function, built-in function, bound method, - ``functools.partial``, ``staticmethod``, or ``classmethod``. Backends - must supply their own compiled-callback identity instead of using a - function object as a core specialization value. - """ - - return _semantic_token(value, _TokenState()) diff --git a/python/cuda_coop/cuda/coop/_core/_types.py b/python/cuda_coop/cuda/coop/_core/_types.py deleted file mode 100644 index d447dc1e7852..000000000000 --- a/python/cuda_coop/cuda/coop/_core/_types.py +++ /dev/null @@ -1,404 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe primitive arguments without importing a backend type system. - -A factory uses parameter descriptors such as ``Value``, ``Pointer``, and -``Array`` to describe a C++ method's signature. These records carry types, -directions, and dependencies; they do not hold the device values passed to a -call. For example, ``Array`` describes an array parameter's type and extent, -while ``ThreadData`` constructs the per-thread payload supplied to it. An -adapter turns the records into its compiler's parameter objects. - -Argument kind and parameter role answer different questions. Kind says whether -the generated code embeds an argument or receives it at runtime. Role says -whether the primitive reads it, writes it, or uses it as temporary storage. -Return-value handling is a separate adapter choice. The ``is_return`` field -controls that choice where it is available. -""" - -from __future__ import annotations - -from collections.abc import Mapping -from dataclasses import dataclass -from enum import Enum -from numbers import Integral -from typing import Any - -_I64_MIN = -(1 << 63) -_I64_MAX = (1 << 63) - 1 - - -class ArgumentKind(str, Enum): - """Distinguish embedded method arguments from runtime operands. - - ``STATIC`` arguments appear in generated C++ source. ``RUNTIME`` arguments - need a value when the device call executes. This classification describes - a selected signature; ``BindingKind`` also covers omitted options while a - factory chooses that signature. - """ - - STATIC = "static" - RUNTIME = "runtime" - - -class ParameterRole(str, Enum): - """Describe how a primitive uses an argument, apart from its C++ type. - - ``INPUT`` supplies a value, ``OUTPUT`` receives a result, and ``INOUT`` - does both. ``CONSTANT`` describes an embedded argument. ``TEMP_STORAGE`` - identifies scratch memory that the implementation needs during the call. - An output role does not by itself make the argument a Python return value. - """ - - INPUT = "input" - OUTPUT = "output" - INOUT = "inout" - CONSTANT = "constant" - TEMP_STORAGE = "temp_storage" - - -@dataclass(frozen=True) -class ParameterClassification: - """Keep an argument's name, supply kind, and role in a group call plan. - - Planners use this smaller record when they need argument roles without - the full type or array extent from a parameter descriptor. ``name`` may - be ``None`` for an unnamed parameter. - """ - - name: str | None - kind: ArgumentKind - role: ParameterRole - - -@dataclass(frozen=True) -class BuiltinDType: - """Name a scalar type without choosing a compiler's representation. - - Adapters map ``name`` to a supported backend dtype. For example, ``INT32`` - requests a signed 32-bit integer. The record itself does not validate - arbitrary names or choose their size and alignment. - """ - - name: str - - -INT8 = BuiltinDType("int8") -UINT8 = BuiltinDType("uint8") -INT16 = BuiltinDType("int16") -UINT16 = BuiltinDType("uint16") -INT32 = BuiltinDType("int32") -UINT32 = BuiltinDType("uint32") -INT64 = BuiltinDType("int64") -UINT64 = BuiltinDType("uint64") -FLOAT32 = BuiltinDType("float32") -FLOAT64 = BuiltinDType("float64") - - -class SubstitutionFailure(ValueError): - """Report that a named dependency cannot supply a signature's value. - - A missing or ``None`` binding can make one candidate overload unavailable. - A backend can reject that candidate while considering other signatures. - """ - - -@dataclass(frozen=True) -class TemplateParameter: - """Name one C++ template parameter in the declaration's argument order. - - ``Algorithm`` uses ``name`` to find the bound value and to distinguish - template arguments from extra dependencies used only by its parameters. - """ - - name: str - - -@dataclass(frozen=True) -class Dependency: - """Defer a parameter's type or extent to a named algorithm binding. - - For example, ``Array(Dependency("T"), Dependency("ITEMS_PER_THREAD"))`` - describes a signature that a factory can reuse for different element - types and array sizes. ``resolve`` looks up each name in the algorithm's - ``template_arguments`` mapping. - """ - - name: str - - def resolve(self, template_arguments: Mapping[str, Any]) -> Any: - """Read the named binding, rejecting a missing or ``None`` value. - - Raise ``SubstitutionFailure`` when this dependency cannot form part - of the selected signature. Leave value-type checks to the adapter. - """ - - if self.name not in template_arguments: - raise SubstitutionFailure( - f"Template argument {self.name} not provided" - ) - value = template_arguments[self.name] - if value is None: - raise SubstitutionFailure(f"Template argument {self.name} is None") - return value - - -@dataclass(frozen=True) -class Constant: - """Supply a fixed value through the same interface as ``Dependency``. - - A parameter can use this record for a type or extent that does not depend - on the algorithm's bindings. Resolution returns ``value`` unchanged. - """ - - value: Any - - def resolve(self, _template_arguments: Mapping[str, Any]) -> Any: - return self.value - - -@dataclass(frozen=True) -class RuntimeValue: - """Tell a factory that a scalar comes from the eventual device call. - - Frontends use this marker in place of a compiler expression when building - shared semantics. ``name`` is an optional label; the ``binding`` helper - discards it and records only that a runtime value is required. - """ - - name: str | None = None - - -class _RuntimeParameter: - """Share runtime classification and direction rules among descriptors. - - An inout flag takes precedence over an output flag. A descriptor with - neither flag has the input role. Subclasses can override these rules. - """ - - @property - def argument_kind(self) -> ArgumentKind: - return ArgumentKind.RUNTIME - - @property - def role(self) -> ParameterRole: - if getattr(self, "is_inout", False): - return ParameterRole.INOUT - if getattr(self, "is_output", False): - return ParameterRole.OUTPUT - return ParameterRole.INPUT - - -class _StaticParameter: - """Mark an embedded argument that needs no runtime operand.""" - - @property - def argument_kind(self) -> ArgumentKind: - return ArgumentKind.STATIC - - -@dataclass(frozen=True) -class Value(_RuntimeParameter): - """Describe a scalar argument and how the primitive uses it. - - ``Pointer``, ``Reference``, and ``Array`` use the same name and direction - fields. The direction describes the operation; ``is_return`` lets an - adapter represent an output as a result or keep it as a caller argument. - - Attributes - ---------- - dtype : object - Scalar type understood by the consuming adapter. - name : str or None - Argument name for planning and generated signatures. - is_output : bool - Mark a value that the primitive writes. - is_inout : bool - Mark a value that the primitive reads and writes. This takes - precedence over ``is_output`` when classifying its role. - is_return : bool or None - Select return-value handling independently of the output role. - ``True`` asks the adapter to expose this argument as a result. - ``None`` leaves the choice to the adapter; ``False`` keeps an output - as an argument. - """ - - dtype: Any - name: str | None = None - is_output: bool = False - is_inout: bool = False - is_return: bool | None = None - - -@dataclass(frozen=True) -class PointerOffset(Value): - """Scalar element offset applied to an earlier pointer argument. - - ``pointer_arg_index`` is the zero-based method-argument position after the - leading :class:`TempStorageParameter` is removed. It must identify an - earlier pointer in the same method signature; backends enforce that - relationship when lowering the signature. - - ``static_value`` embeds a fixed element offset. With ``None``, the - wrapper receives the offset at runtime. In either case, the offset - adjusts the pointer passed to the primitive instead of adding a C++ - method argument. Its unit is elements, not bytes. - - Construction checks the index and normalizes a static offset to a Python - integer within signed 64-bit bounds. Operation factories impose any - further bounds, such as requiring a nonnegative offset. - """ - - pointer_arg_index: int = 0 - static_value: int | None = None - - def __post_init__(self) -> None: - if ( - not isinstance(self.pointer_arg_index, int) - or isinstance(self.pointer_arg_index, bool) - or self.pointer_arg_index < 0 - ): - raise ValueError("pointer_arg_index must be a non-negative integer") - if self.static_value is not None and ( - not isinstance(self.static_value, Integral) - or isinstance(self.static_value, bool) - ): - raise TypeError("static pointer offset must be an integer") - if self.static_value is not None: - normalized = int(self.static_value) - if not _I64_MIN <= normalized <= _I64_MAX: - raise ValueError( - "static pointer offset must fit a signed 64-bit integer" - ) - object.__setattr__(self, "static_value", normalized) - - @property - def argument_kind(self) -> ArgumentKind: - if self.static_value is not None: - return ArgumentKind.STATIC - return ArgumentKind.RUNTIME - - @property - def role(self) -> ParameterRole: - if self.static_value is not None: - return ParameterRole.CONSTANT - return ParameterRole.INPUT - - -@dataclass(frozen=True) -class Pointer(_RuntimeParameter): - """Describe an address passed to a primitive or its generated wrapper. - - The common dtype, name, and direction fields follow ``Value``. Extra - fields record information that an adapter can use when building the call. - - Attributes - ---------- - is_array_pointer : bool - Identify a pointer to an array payload rather than a single value. - restrict : bool - Record the factory's non-aliasing intent for adapters that express - this property in generated code. - deref_on_call : bool - Pass the pointed-to value to the C++ method. The wrapper still - receives the address, which lets a method reference caller storage. - """ - - dtype: Any - name: str | None = None - is_output: bool = False - is_inout: bool = False - is_return: bool | None = None - is_array_pointer: bool = False - restrict: bool = False - deref_on_call: bool = False - - -@dataclass(frozen=True) -class Reference(_RuntimeParameter): - """Describe a C++ reference parameter that can access the caller's value. - - The fields follow ``Value``. The adapter chooses how its call interface - supplies the address needed by the generated C++ reference. - """ - - dtype: Any - name: str | None = None - is_output: bool = False - is_inout: bool = False - is_return: bool | None = None - - -@dataclass(frozen=True) -class Array(_RuntimeParameter): - """Describe an array parameter with a count known during compilation. - - ``dtype`` gives the element type and ``size`` gives the element count. - Either can be a ``Dependency`` to resolve from the algorithm's bindings. - In a Load/Store signature, this count is the items per thread, not the tile - size. The name and direction fields follow ``Value``. - """ - - dtype: Any - size: Any - name: str | None = None - is_output: bool = False - is_inout: bool = False - is_return: bool | None = None - - -@dataclass(frozen=True) -class TempStorageParameter(_RuntimeParameter): - """Mark scratch storage that the backend supplies to the primitive. - - The default dtype describes bytes. Allocation size, alignment, and - ownership come from storage planning. An adapter can retain this marker - as an argument or omit it when it supplies storage through another path. - """ - - dtype: Any = UINT8 - name: str | None = "temp_storage" - is_output: bool = False - is_array_pointer: bool = True - - @property - def role(self) -> ParameterRole: - return ParameterRole.TEMP_STORAGE - - -@dataclass(frozen=True) -class CxxFunction(_StaticParameter): - """Embed a C++ expression as a method argument without a runtime operand. - - ``cpp`` can name a functor or contain a scalar expression such as a fixed - item count. ``dtype`` describes that argument to the adapter. ``name`` - identifies the parameter in the shared signature. The factory supplies - source text here; this record does not parse or validate it. - """ - - cpp: str - dtype: Any - name: str | None = None - - @property - def role(self) -> ParameterRole: - return ParameterRole.CONSTANT - - -def classify_parameter(parameter: Any) -> ParameterClassification: - """Extract the argument facts needed by a group call plan. - - Read ``argument_kind`` and ``role`` from the descriptor, with an optional - ``name``. Raise ``TypeError`` if either required attribute is absent. - This preserves the descriptor's classification without resolving its - dtype or other dependencies. - """ - - try: - kind = parameter.argument_kind - role = parameter.role - except AttributeError as exc: - raise TypeError(f"Unsupported core parameter {parameter!r}") from exc - return ParameterClassification(getattr(parameter, "name", None), kind, role) diff --git a/python/cuda_coop/cuda/coop/_core/api/__init__.py b/python/cuda_coop/cuda/coop/_core/api/__init__.py deleted file mode 100644 index 8c5d75cc54f0..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/__init__.py +++ /dev/null @@ -1,55 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Collect the common kernel API and identify its constructor functions. - -The root package exports these names. Compiler adapters use their identities -to recognize cooperative operations without changing the public call syntax. -""" - -from .load_store import load, store -from .temp_storage import TempStorage, TempStorageLike -from .thread_data import ThreadData, ThreadDataLike -from .thread_group import ( - Hierarchy, - ThreadGroup, - ThreadHierarchy, - this_block, - this_cluster, - this_grid, - this_thread, - this_warp, -) - -# Constructors and group factories bypass the operation registration decorator. -# The Numba rewrite requires this tag and the exported function identity, so -# an unrelated callable with the same name is not treated as a constructor. -for _member_name in ( - "TempStorage", - "ThreadData", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -): - globals()[_member_name].__cuda_coop_backend_member__ = _member_name -del _member_name - -__all__ = [ - "Hierarchy", - "TempStorage", - "TempStorageLike", - "ThreadData", - "ThreadDataLike", - "ThreadGroup", - "ThreadHierarchy", - "load", - "store", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -] diff --git a/python/cuda_coop/cuda/coop/_core/api/__init__.pyi b/python/cuda_coop/cuda/coop/_core/api/__init__.pyi deleted file mode 100644 index 1a84626d5aa5..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/__init__.pyi +++ /dev/null @@ -1,34 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from .load_store import load, store -from .temp_storage import TempStorage, TempStorageLike -from .thread_data import ThreadData, ThreadDataLike -from .thread_group import ( - Hierarchy, - ThreadGroup, - ThreadHierarchy, - this_block, - this_cluster, - this_grid, - this_thread, - this_warp, -) - -__all__ = [ - "Hierarchy", - "TempStorage", - "TempStorageLike", - "ThreadData", - "ThreadDataLike", - "ThreadGroup", - "ThreadHierarchy", - "load", - "store", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -] diff --git a/python/cuda_coop/cuda/coop/_core/api/_dispatch.py b/python/cuda_coop/cuda/coop/_core/api/_dispatch.py deleted file mode 100644 index c189d50e9efe..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/_dispatch.py +++ /dev/null @@ -1,166 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Let compiler adapters recognize public cooperative calls. - -Importing an operation records its function object and supported group kinds. -Lookup by function identity also recognizes imported aliases. The original -callable stays in place, so its signature and documentation remain available. -""" - -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass -from typing import TypeVar - -from ..thread_group import ThreadGroup - -_CallableT = TypeVar("_CallableT", bound=Callable[..., object]) - - -@dataclass(frozen=True) -class _CommonGroupOperation: - """Store the operation name and allowed groups for one public callable.""" - - name: str - group_kinds: tuple[str, ...] - function: Callable[..., object] - - -_COMMON_GROUP_OPERATIONS_BY_NAME: dict[str, _CommonGroupOperation] = {} -_COMMON_GROUP_OPERATIONS_BY_FUNCTION: dict[ - Callable[..., object], _CommonGroupOperation -] = {} - - -def _common_group_operation( - name: str, - *, - group_kinds: tuple[str, ...], -) -> Callable[[_CallableT], _CallableT]: - """Record a public operation and the groups it accepts. - - The decorator runs when the API module is imported and records the - function object without wrapping it. Compiler adapters use that object - to recognize calls, including imported aliases. An unrelated function - with the same name does not acquire this registration. - """ - - if not name or not group_kinds: - raise ValueError( - "common group operations require a name and group kinds" - ) - - def decorate(function: _CallableT) -> _CallableT: - """Register one function in both tables and reject conflicts.""" - - registration = _CommonGroupOperation(name, tuple(group_kinds), function) - existing = _COMMON_GROUP_OPERATIONS_BY_NAME.get(name) - if existing is not None and existing != registration: - raise RuntimeError( - f"common group operation {name!r} is already registered" - ) - existing_function = _COMMON_GROUP_OPERATIONS_BY_FUNCTION.get(function) - if existing_function is not None and existing_function != registration: - raise RuntimeError( - f"common group marker {function!r} is already registered" - ) - _COMMON_GROUP_OPERATIONS_BY_NAME[name] = registration - _COMMON_GROUP_OPERATIONS_BY_FUNCTION[function] = registration - function.__cuda_coop_backend_member__ = name - return function - - return decorate - - -def _common_group_operation_name(function: object) -> str | None: - """Return the registered name used to recognize a common API call. - - The planner passes a resolved callable, including an imported alias. - Return ``None`` when its identity has no operation registration. - """ - - registration = _COMMON_GROUP_OPERATIONS_BY_FUNCTION.get(function) - return None if registration is None else registration.name - - -class UnsupportedCoopBackendOperationError(NotImplementedError): - """Report that an operation is unavailable in the selected API or backend. - - Keep the module, operation, and stable reason code available to callers - that need to classify the failure without parsing its display message. - """ - - def __init__(self, backend_module: str, operation: str) -> None: - self.backend_module = backend_module - self.operation = operation - self.reason_code = "cuda-coop-backend-operation-unavailable" - super().__init__( - f"cuda.coop.{operation} is not implemented by {backend_module!r}" - ) - - -def _common_group_name(kind: str) -> str: - """Return the common API spelling for one internal group kind.""" - - return "physical_warp" if kind == "warp" else kind - - -def _validate_common_operation_group( - operation: str, - group: ThreadGroup, -) -> None: - """Reject groups outside a registered operation's common API contract. - - Compiler adapters call this after reconstructing a symbolic group - descriptor and before lowering a recognized common operation. The registry - defines the common set of group kinds; a backend-qualified operation can - support more kinds without broadening that common contract. - - The descriptor may still have unresolved launch dimensions. This check - validates only its type and group kind, not launch dimensions, primitive - participation, or backend implementation. - - Parameters - ---------- - operation : str - Common operation name used to look up its registered group kinds and - identify it in diagnostics. - group : ThreadGroup - Symbolic group descriptor whose ``kind`` must be supported by the - registered operation. It is not modified. - - Raises - ------ - TypeError - ``group`` is not a ``ThreadGroup``. - UnsupportedCoopBackendOperationError - ``operation`` has no common operation registration. - NotImplementedError - The registered common operation does not support this group kind. - """ - - if not isinstance(group, ThreadGroup): - raise TypeError(f"cuda.coop.{operation} group must be a ThreadGroup") - registration = _COMMON_GROUP_OPERATIONS_BY_NAME.get(operation) - if registration is None: - raise UnsupportedCoopBackendOperationError("cuda.coop", operation) - supported = registration.group_kinds - if group.kind in supported: - return - group_name = _common_group_name(group.kind) - supported_names = ", ".join(map(_common_group_name, supported)) - raise NotImplementedError( - f"cuda.coop.{operation} does not support group kind {group_name!r} in " - f"the common API; supported group kinds: {supported_names}; use a " - "backend-qualified import for backend-specific group support" - ) - - -__all__ = [ - "_common_group_operation", - "_common_group_operation_name", - "_validate_common_operation_group", -] diff --git a/python/cuda_coop/cuda/coop/_core/api/_payload.py b/python/cuda_coop/cuda/coop/_core/api/_payload.py deleted file mode 100644 index 9f688a3a0380..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/_payload.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Define the payload and scratch interfaces used by common API annotations. - -Each compiler backend supplies concrete values with these attributes and -operations. The backend must also recognize those values during compilation. -This lets the common API describe inputs without importing compiler types. -""" - -from __future__ import annotations - -import operator -from typing import Protocol, SupportsIndex, TypeVar, runtime_checkable - -_ItemT = TypeVar("_ItemT") - - -@runtime_checkable -class _ReadableThreadDataLike(Protocol[_ItemT]): - """Describe a fixed number of readable values owned by one thread. - - ``items_per_thread`` is the fixed extent. ``dtype`` can remain unknown - until a supported producer establishes the element type. The compiler - backend supplies the storage and indexed access. - """ - - items_per_thread: int - dtype: object | None - - def __len__(self) -> int: ... - - def __getitem__(self, index: int) -> _ItemT: ... - - -@runtime_checkable -class ThreadDataLike(_ReadableThreadDataLike[_ItemT], Protocol[_ItemT]): - """Mutable fixed-size per-thread payload understood by supported backends. - - See :ref:`per-thread payloads ` for construction, item - access, and dtype requirements. - """ - - def __setitem__(self, index: int, value: _ItemT) -> None: ... - - -@runtime_checkable -class TempStorageLike(Protocol): - """Explicit scratch descriptor understood by supported backends. - - See :ref:`temporary storage ` for construction, - allocation sharing, and synchronization. - """ - - size_in_bytes: int | None - alignment: int | None - auto_sync: bool - sharing: str - - -def _normalize_alignment(alignment: SupportsIndex | None) -> int | None: - """Normalize a payload or scratch alignment request to bytes. - - Preserve ``None`` so the compiler can choose the alignment. Explicit - requests must be positive powers of two. Accept integer-like values - through ``__index__``, but reject booleans as accidental requests. - """ - - if alignment is None: - return None - if isinstance(alignment, bool): - raise TypeError("alignment must be an integer or None") - try: - alignment = operator.index(alignment) - except TypeError as exc: - raise TypeError("alignment must be an integer or None") from exc - if alignment <= 0: - raise ValueError("alignment must be a positive integer") - if alignment & (alignment - 1): - raise ValueError("alignment must be a power of 2") - return alignment - - -__all__ = [ - "TempStorageLike", - "ThreadDataLike", - "_ReadableThreadDataLike", - "_normalize_alignment", -] diff --git a/python/cuda_coop/cuda/coop/_core/api/load_store.py b/python/cuda_coop/cuda/coop/_core/api/load_store.py deleted file mode 100644 index 9d5b36aed2ac..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/load_store.py +++ /dev/null @@ -1,224 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Define the shared Load and Store call signatures and memory contracts. - -The decorators register each function's identity and supported group kinds. -Compiler adapters recognize these calls and generate the selected memory -operation. The Python bodies reject calls outside a supported GPU kernel. -""" - -from __future__ import annotations - -from cuda.coop._typing import ( - BlockLoadStoreAlgorithm, - CommonThreadDataLike, - IntegerValue, - ValidItems, - WarpLoadStoreAlgorithm, - _CommonNumericT, -) - -from ..thread_group import CoopCompilerContextRequiredError -from ._dispatch import ( - _common_group_operation, -) -from ._payload import ( - TempStorageLike, - ThreadDataLike, -) -from .thread_group import BlockGroup, WarpGroup - - -@_common_group_operation( - "load", - group_kinds=("block", "warp", "threads_within_warp"), -) -def 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 :ref:`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 :ref:`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 - :ref:`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 - ------- - None - The call populates ``output`` in place. - - See Also - -------- - :cpp:class:`cub::BlockLoad`, :cpp:class:`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. - - .. literalinclude:: - ../../python/cuda_coop/tests/backends/numba_mlir/runtime/test_load_example.py - :language: python - :start-after: # example-begin - :end-before: # example-end - :dedent: 4 - - The qualified import activates the Numba-CUDA-MLIR backend even if another - module imported ``cuda.coop`` first. Use the qualified - ``cuda.coop.`` API for backend-specific behavior. - """ - - raise CoopCompilerContextRequiredError( - "cuda.coop.load must be called from a supported GPU kernel." - ) - - -@_common_group_operation( - "store", - group_kinds=("block", "warp", "threads_within_warp"), -) -def 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 :ref:`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 :ref:`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 :ref:`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 - :func:`cuda.coop.load`. - temp_storage : cuda.coop.TempStorageLike, optional - :ref:`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 - ------- - None - Writes to ``destination`` and may rearrange the input payload. - - See Also - -------- - :cpp:class:`cub::BlockStore`, :cpp:class:`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. - - .. literalinclude:: - ../../python/cuda_coop/tests/backends/numba_mlir/runtime/test_store_example.py - :language: python - :start-after: # example-begin - :end-before: # example-end - :dedent: 4 - - See :ref:`participation and synchronization ` for - control-flow requirements at primitive calls. - """ - - raise CoopCompilerContextRequiredError( - "cuda.coop.store must be called from a supported GPU kernel." - ) - - -__all__ = ["load", "store"] diff --git a/python/cuda_coop/cuda/coop/_core/api/load_store.pyi b/python/cuda_coop/cuda/coop/_core/api/load_store.pyi deleted file mode 100644 index 8970544b47d8..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/load_store.pyi +++ /dev/null @@ -1,110 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from typing import overload - -from typing_extensions import TypeVar - -from cuda.coop._typing import ( - BlockLoadStoreAlgorithm, - CommonNumericScalar, - CommonThreadDataLike, - IntegerValue, - TempStorageLike, - ThreadDataLike, - ValidItems, - WarpLoadStoreAlgorithm, -) - -from .thread_group import BlockGroup, WarpGroup - -_CommonNumericT = TypeVar("_CommonNumericT", bound=CommonNumericScalar) - -@overload -def load( - group: BlockGroup, - source: object, - output: ThreadDataLike[_CommonNumericT], - /, - *, - algorithm: BlockLoadStoreAlgorithm = "direct", - valid_items: ValidItems | None = None, - oob_default: None = None, - offset: IntegerValue | None = None, - temp_storage: TempStorageLike | None = None, -) -> None: - """Populate ``output`` in place with one cooperative block tile.""" - -@overload -def load( - group: BlockGroup, - source: object, - output: ThreadDataLike[_CommonNumericT], - /, - *, - algorithm: BlockLoadStoreAlgorithm = "direct", - valid_items: ValidItems, - oob_default: _CommonNumericT | float, - offset: IntegerValue | None = None, - temp_storage: TempStorageLike | None = None, -) -> None: - """Populate a partial block tile and fill invalid items.""" - -@overload -def load( - group: WarpGroup, - source: object, - output: ThreadDataLike[_CommonNumericT], - /, - *, - algorithm: WarpLoadStoreAlgorithm = "direct", - valid_items: ValidItems | None = None, - oob_default: None = None, - offset: IntegerValue | None = None, - temp_storage: None = None, -) -> None: - """Populate ``output`` in place with one physical or logical warp tile.""" - -@overload -def load( - group: WarpGroup, - source: object, - output: ThreadDataLike[_CommonNumericT], - /, - *, - algorithm: WarpLoadStoreAlgorithm = "direct", - valid_items: ValidItems, - oob_default: _CommonNumericT | float, - offset: IntegerValue | None = None, - temp_storage: None = None, -) -> None: - """Load a partial warp tile and fill invalid output items.""" - -@overload -def store( - group: BlockGroup, - destination: object, - value: _CommonNumericT | CommonThreadDataLike[_CommonNumericT], - /, - *, - algorithm: BlockLoadStoreAlgorithm = "direct", - valid_items: ValidItems | None = None, - offset: IntegerValue | None = None, - temp_storage: TempStorageLike | None = None, -) -> None: - """Store one scalar or per-thread payload cooperatively across a block.""" - -@overload -def store( - group: WarpGroup, - destination: object, - value: _CommonNumericT | CommonThreadDataLike[_CommonNumericT], - /, - *, - algorithm: WarpLoadStoreAlgorithm = "direct", - valid_items: ValidItems | None = None, - offset: IntegerValue | None = None, - temp_storage: None = None, -) -> None: - """Store one scalar or per-thread payload across a Warp group.""" diff --git a/python/cuda_coop/cuda/coop/_core/api/temp_storage.py b/python/cuda_coop/cuda/coop/_core/api/temp_storage.py deleted file mode 100644 index aa617ae086bc..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/temp_storage.py +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe explicit scratch storage for compiler-supported block operations. - -The constructor marks a request inside a GPU kernel. A compiler that supports -explicit scratch collects the descriptor's uses to choose shared-memory space -and any requested reuse barriers. The Python body raises an error outside a -supported kernel; it never allocates storage. -""" - -from __future__ import annotations - -from ..thread_group import CoopCompilerContextRequiredError -from ._payload import TempStorageLike - - -def 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 :ref:`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 - ------- - cuda.coop.TempStorageLike - Compiler-recognized scratch descriptor. The storage contents are - opaque; keep application data in :func:`cuda.coop.ThreadData` or - application-owned arrays. - - 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. - - .. literalinclude:: - ../../python/cuda_coop/tests/backends/numba_mlir/runtime/test_storage_examples.py - :language: python - :start-after: # temp-storage-example-begin - :end-before: # temp-storage-example-end - :dedent: 4 - """ - - raise CoopCompilerContextRequiredError( - "cuda.coop.TempStorage must be called from a supported GPU kernel." - ) - - -__all__ = ["TempStorage", "TempStorageLike"] diff --git a/python/cuda_coop/cuda/coop/_core/api/temp_storage.pyi b/python/cuda_coop/cuda/coop/_core/api/temp_storage.pyi deleted file mode 100644 index d9016f1e7d01..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/temp_storage.pyi +++ /dev/null @@ -1,16 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from cuda.coop._typing import TempStorageLike, TempStorageSharing - -__all__ = ["TempStorage", "TempStorageLike"] - -def TempStorage( - size_in_bytes: int | None = None, - *, - alignment: int | None = None, - auto_sync: bool | None = False, - sharing: TempStorageSharing = "shared", -) -> TempStorageLike: - """Construct the selected backend's explicit scratch descriptor.""" diff --git a/python/cuda_coop/cuda/coop/_core/api/thread_data.py b/python/cuda_coop/cuda/coop/_core/api/thread_data.py deleted file mode 100644 index 5bdce48f2de2..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/thread_data.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Define the common constructor for per-thread values in a GPU kernel. - -A compiler replaces this call with storage for the calling thread. The Python -body rejects host execution because it cannot provide that device storage. -""" - -from __future__ import annotations - -from typing import Any - -from ..thread_group import CoopCompilerContextRequiredError -from ._payload import ThreadDataLike - - -def 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 :ref:`per-thread payloads - ` for their relationship to a group tile and the - :ref:`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 :func:`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 - ------- - cuda.coop.ThreadDataLike - Writable per-thread payload with indexed reads and writes. Its - contents are uninitialized; write every item before reading it. - Construction does not synchronize threads. - - Examples - -------- - Construct a payload from the kernel's item-count argument, fill its - items with squared indices, and store the resulting blocked tiles: - - .. literalinclude:: - ../../python/cuda_coop/tests/backends/numba_mlir/runtime/test_storage_examples.py - :language: python - :start-after: # thread-data-example-begin - :end-before: # thread-data-example-end - :dedent: 4 - """ - - raise CoopCompilerContextRequiredError( - "cuda.coop.ThreadData must be called from a supported GPU kernel." - ) - - -__all__ = ["ThreadData", "ThreadDataLike"] diff --git a/python/cuda_coop/cuda/coop/_core/api/thread_data.pyi b/python/cuda_coop/cuda/coop/_core/api/thread_data.pyi deleted file mode 100644 index df75cdf391d6..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/thread_data.pyi +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from typing import Any, overload - -from typing_extensions import TypeVar - -from cuda.coop._typing import CommonNumericScalar, ThreadDataLike - -_CommonNumericT = TypeVar("_CommonNumericT", bound=CommonNumericScalar) - -__all__ = ["ThreadData", "ThreadDataLike"] - -@overload -def ThreadData( - items_per_thread: int, - dtype: type[_CommonNumericT], - *, - alignment: int | None = None, -) -> ThreadDataLike[_CommonNumericT]: - """Construct a payload; builtin int and float mean 32-bit dtypes.""" - -@overload -def ThreadData( - items_per_thread: int, - dtype: object | None = None, - *, - alignment: int | None = None, -) -> ThreadDataLike[Any]: - """Omit dtype or use an ``Any``-typed external compiler dtype token.""" diff --git a/python/cuda_coop/cuda/coop/_core/api/thread_group.py b/python/cuda_coop/cuda/coop/_core/api/thread_group.py deleted file mode 100644 index 74f9ad1d8274..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/thread_group.py +++ /dev/null @@ -1,135 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Expose symbolic thread groups through the common API. - -These factories describe which threads will cooperate. A compiler resolves the -description against the kernel launch before it selects an implementation. -Constructing a group does not execute an operation or synchronize its threads. -""" - -from __future__ import annotations - -from ..thread_group import ( - Hierarchy, - ThreadGroup, - ThreadHierarchy, - this_block, - this_cluster, - this_grid, - this_thread, - this_warp, -) - -_core_this_block = this_block -_core_this_cluster = this_cluster -_core_this_grid = this_grid -_core_this_thread = this_thread -_core_this_warp = this_warp - -# These names support explicit imports used by the adjacent typing stubs. They -# are overload helpers, not additional root-package exports. -MemoryGroup = ThreadGroup -ReductionGroup = ThreadGroup -BlockGroup = ThreadGroup -WarpGroup = ThreadGroup - - -def this_thread() -> ThreadGroup: - """Describe the calling thread as a one-thread group. - - Returns - ------- - cuda.coop.ThreadGroup - A symbolic descriptor. Load and Store do not accept one-thread groups. - See :ref:`thread groups `. - """ - - return _core_this_thread() - - -def this_warp() -> ThreadGroup: - """Describe the calling thread's physical warp. - - Returns - ------- - cuda.coop.ThreadGroup - A group of 32 consecutive threads in the block's linear thread - order. Use :meth:`~cuda.coop.ThreadGroup.group_by` to form smaller - logical warps. - - 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 :ref:`participation requirements `. - """ - - return _core_this_warp() - - -def this_block() -> ThreadGroup: - """Describe all threads in the calling thread's CUDA block. - - Returns - ------- - cuda.coop.ThreadGroup - 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. - - Notes - ----- - The factory takes no size argument and does not synchronize the block. - See :ref:`thread groups ` and the primitive's - :ref:`participation requirements `. - """ - - return _core_this_block() - - -def this_cluster() -> ThreadGroup: - """Describe the thread-block cluster containing the calling thread. - - Returns - ------- - cuda.coop.ThreadGroup - A symbolic descriptor. Cluster groups are not Load or Store targets. - - Notes - ----- - Construction does not create a cluster or enable cluster scheduling. - See :ref:`thread groups `. - """ - - return _core_this_cluster() - - -def this_grid() -> ThreadGroup: - """Describe all threads in the current kernel grid. - - Returns - ------- - cuda.coop.ThreadGroup - A symbolic descriptor. Grid groups are not Load or Store targets. - - Notes - ----- - Construction does not request a cooperative launch or synchronize - the grid. See :ref:`thread groups `. - """ - - return _core_this_grid() - - -__all__ = [ - "Hierarchy", - "ThreadGroup", - "ThreadHierarchy", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -] diff --git a/python/cuda_coop/cuda/coop/_core/api/thread_group.pyi b/python/cuda_coop/cuda/coop/_core/api/thread_group.pyi deleted file mode 100644 index 60369e22d7d2..000000000000 --- a/python/cuda_coop/cuda/coop/_core/api/thread_group.pyi +++ /dev/null @@ -1,110 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from typing import Generic, Literal, TypeAlias, overload - -from typing_extensions import Self, TypeVar - -from cuda.coop._typing import ThreadGroupKind - -_GroupKindT_co = TypeVar( - "_GroupKindT_co", - bound=ThreadGroupKind, - covariant=True, - default=ThreadGroupKind, -) - -class ThreadHierarchy: - """CUDA hierarchy descriptor resolved from the current kernel launch.""" - - block_dim: tuple[int, int, int] | None - grid_dim: tuple[int, int, int] | None - cluster_dim: tuple[int, int, int] | None - implicit: bool - - def __init__(self) -> None: - """Describe the compiler's current launch hierarchy.""" - - @classmethod - def current(cls) -> Self: - """Describe the compiler's current launch hierarchy.""" - - @property - def is_static(self) -> bool: - """Return whether the hierarchy carries static dimensions.""" - - @property - def block_thread_count(self) -> int | None: - """Return the statically known number of threads per block, if any.""" - -Hierarchy: TypeAlias = ThreadHierarchy - -class ThreadGroup(Generic[_GroupKindT_co]): - """Common compiler-facing contract for one CUDA thread group.""" - - hierarchy: ThreadHierarchy - block_dim: tuple[int, int, int] | None - parent: ThreadGroup | None - source: str - - @property - def kind(self) -> _GroupKindT_co: - """Return the kind carried by this group descriptor.""" - - @property - def static_size(self) -> int | None: - """Return the statically known group extent, if any.""" - - @property - def group_thread_count(self) -> int: - """Return the group extent, requiring it to be statically known.""" - - @property - def is_current(self) -> bool: - """Return whether this descriptor refers to the current hierarchy.""" - - @property - def is_static(self) -> bool: - """Return whether dimensions needed by this group are static.""" - - @overload - def group_by( - self: ThreadGroup[Literal["warp"]], - count: int, - *, - exhaustive: bool = True, - ) -> ThreadGroup[Literal["threads_within_warp"]]: - """Partition a physical warp into groups of threads.""" - @overload - def group_by( - self: ThreadGroup[Literal["block"]], - count: int, - *, - exhaustive: bool = True, - ) -> ThreadGroup[Literal["warps_within_block"]]: - """Partition a block into groups of physical warps.""" - -MemoryGroup: TypeAlias = ThreadGroup[ - Literal["warp", "threads_within_warp", "block"] -] -ReductionGroup: TypeAlias = ThreadGroup[ - Literal["thread", "warp", "threads_within_warp", "block", "cluster"] -] -BlockGroup: TypeAlias = ThreadGroup[Literal["block"]] -WarpGroup: TypeAlias = ThreadGroup[Literal["warp", "threads_within_warp"]] - -def this_thread() -> ThreadGroup[Literal["thread"]]: - """Describe the current thread.""" - -def this_warp() -> ThreadGroup[Literal["warp"]]: - """Describe the current physical warp.""" - -def this_block() -> ThreadGroup[Literal["block"]]: - """Describe the current CUDA thread block.""" - -def this_cluster() -> ThreadGroup[Literal["cluster"]]: - """Describe the current thread-block cluster.""" - -def this_grid() -> ThreadGroup[Literal["grid"]]: - """Describe the grid; common operations do not support this group.""" diff --git a/python/cuda_coop/cuda/coop/_core/block/__init__.py b/python/cuda_coop/cuda/coop/_core/block/__init__.py deleted file mode 100644 index f2c6f40eef01..000000000000 --- a/python/cuda_coop/cuda/coop/_core/block/__init__.py +++ /dev/null @@ -1,34 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from .._bindings import ArgumentBinding, BindingKind, binding -from ._common import normalize_block_dim, normalize_positive_int -from .load_store import ( - BlockLoadAlgorithm, - BlockLoadStoreAlgorithm, - BlockLoadStoreKind, - BlockLoadStoreSemantics, - BlockStoreAlgorithm, - make_block_load_specialization, - make_block_load_store_semantics, - make_block_load_store_specialization, - make_block_store_specialization, -) - -__all__ = [ - "ArgumentBinding", - "BindingKind", - "BlockLoadAlgorithm", - "BlockLoadStoreAlgorithm", - "BlockLoadStoreKind", - "BlockLoadStoreSemantics", - "BlockStoreAlgorithm", - "binding", - "make_block_load_specialization", - "make_block_load_store_semantics", - "make_block_load_store_specialization", - "make_block_store_specialization", - "normalize_block_dim", - "normalize_positive_int", -] diff --git a/python/cuda_coop/cuda/coop/_core/block/_common.py b/python/cuda_coop/cuda/coop/_core/block/_common.py deleted file mode 100644 index 48c834668d82..000000000000 --- a/python/cuda_coop/cuda/coop/_core/block/_common.py +++ /dev/null @@ -1,57 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Normalize integer dimensions before CUB block templates are specialized. - -Template parameters need concrete integers. These helpers accept integral -scalar types, reject booleans, and produce ordinary Python integers for the -shared specialization descriptions. -""" - -from __future__ import annotations - -from numbers import Integral -from typing import Any - - -def normalize_positive_int(name: str, value: Any) -> int: - """Return a positive integral option without accepting boolean values.""" - - if not isinstance(value, Integral) or isinstance(value, bool) or value < 1: - raise ValueError(f"{name} must be a positive integer") - return int(value) - - -def normalize_block_dim(value: Any) -> tuple[int, int, int]: - """Require three positive dimensions and return them as Python integers. - - CUB specialization receives an already resolved ``(x, y, z)`` shape. - This check therefore requires all three entries instead of supplying - missing dimensions as a kernel launch interface might do. - """ - - try: - dimensions = tuple(value) - except TypeError as exc: - raise ValueError( - "block_dim must contain three positive dimensions" - ) from exc - if len(dimensions) != 3: - raise ValueError("block_dim must contain three positive dimensions") - try: - x, y, z = ( - normalize_positive_int("block_dim", dimension) - for dimension in dimensions - ) - except ValueError as exc: - raise ValueError( - "block_dim must contain three positive dimensions" - ) from exc - return x, y, z - - -__all__ = [ - "normalize_block_dim", - "normalize_positive_int", -] diff --git a/python/cuda_coop/cuda/coop/_core/block/load_store.py b/python/cuda_coop/cuda/coop/_core/block/load_store.py deleted file mode 100644 index d6e7723c2222..000000000000 --- a/python/cuda_coop/cuda/coop/_core/block/load_store.py +++ /dev/null @@ -1,607 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe CUB BlockLoad and BlockStore wrappers for compiler adapters. - -The semantics builder selects full-tile, guarded, and pointer-offset -signatures. The specialization builder then binds the block shape and CUB -template arguments. These descriptions let each backend generate wrappers for -the same operation without importing compiler types into the shared model. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from enum import Enum -from numbers import Integral -from typing import Any - -from .._algorithm import Algorithm -from .._bindings import ( - ArgumentBinding, - BindingKind, - cxx_scalar_literal, - i32_parameter, - normalize_i32_binding, - normalize_i64_binding, -) -from .._symbols import semantic_token -from .._types import ( - INT64, - Array, - CxxFunction, - Dependency, - Pointer, - PointerOffset, - TemplateParameter, - TempStorageParameter, - Value, -) -from ._common import normalize_block_dim, normalize_positive_int - - -class BlockLoadStoreKind(str, Enum): - """Select CUB's ``BlockLoad.Load`` or ``BlockStore.Store`` operation.""" - - LOAD = "load" - STORE = "store" - - -class BlockLoadStoreAlgorithm(str, Enum): - """CUB memory-access strategies and per-thread item arrangements. - - ``DIRECT`` and ``VECTORIZE`` use blocked items: thread ``t`` owns tile - positions ``t * items_per_thread + i``. ``STRIPED`` uses positions - ``t + i * block_threads`` instead. ``VECTORIZE`` requests vectorized - accesses where CUB's type, alignment, and item-count requirements allow. - - The transpose variants also present blocked items to the caller, using - shared memory to reorder coalesced accesses. ``TRANSPOSE`` exchanges a - block-striped arrangement; ``WARP_TRANSPOSE`` exchanges a warp-striped - arrangement. ``WARP_TRANSPOSE_TIMESLICED`` reuses one warp's exchange - storage across warps to reduce shared memory use. Both warp-transpose - variants require the block's thread count to be a multiple of 32. - - ``BlockLoadAlgorithm`` and ``BlockStoreAlgorithm`` alias this enum; the - operation kind selects the corresponding CUB load or store constant. - """ - - DIRECT = "direct" - STRIPED = "striped" - VECTORIZE = "vectorize" - TRANSPOSE = "transpose" - WARP_TRANSPOSE = "warp_transpose" - WARP_TRANSPOSE_TIMESLICED = "warp_transpose_timesliced" - - -BlockLoadAlgorithm = BlockLoadStoreAlgorithm -BlockStoreAlgorithm = BlockLoadStoreAlgorithm - - -_LOAD_ALGORITHM_CPP = { - BlockLoadAlgorithm.DIRECT: "::cub::BLOCK_LOAD_DIRECT", - BlockLoadAlgorithm.STRIPED: "::cub::BLOCK_LOAD_STRIPED", - BlockLoadAlgorithm.VECTORIZE: "::cub::BLOCK_LOAD_VECTORIZE", - BlockLoadAlgorithm.TRANSPOSE: "::cub::BLOCK_LOAD_TRANSPOSE", - BlockLoadAlgorithm.WARP_TRANSPOSE: "::cub::BLOCK_LOAD_WARP_TRANSPOSE", - BlockLoadAlgorithm.WARP_TRANSPOSE_TIMESLICED: ( - "::cub::BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED" - ), -} -_STORE_ALGORITHM_CPP = { - BlockStoreAlgorithm.DIRECT: "::cub::BLOCK_STORE_DIRECT", - BlockStoreAlgorithm.STRIPED: "::cub::BLOCK_STORE_STRIPED", - BlockStoreAlgorithm.VECTORIZE: "::cub::BLOCK_STORE_VECTORIZE", - BlockStoreAlgorithm.TRANSPOSE: "::cub::BLOCK_STORE_TRANSPOSE", - BlockStoreAlgorithm.WARP_TRANSPOSE: "::cub::BLOCK_STORE_WARP_TRANSPOSE", - BlockStoreAlgorithm.WARP_TRANSPOSE_TIMESLICED: ( - "::cub::BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED" - ), -} -_T = Dependency("T") -_ITEMS_PER_THREAD = Dependency("ITEMS_PER_THREAD") -_TEMPLATE_PARAMETERS = ( - TemplateParameter("T"), - TemplateParameter("BLOCK_DIM_X"), - TemplateParameter("ITEMS_PER_THREAD"), - TemplateParameter("ALGORITHM"), - TemplateParameter("BLOCK_DIM_Y"), - TemplateParameter("BLOCK_DIM_Z"), -) - - -def _algorithm_cpp_map( - kind: BlockLoadStoreKind, -) -> dict[BlockLoadStoreAlgorithm, str]: - return ( - _LOAD_ALGORITHM_CPP - if kind is BlockLoadStoreKind.LOAD - else _STORE_ALGORITHM_CPP - ) - - -def _normalize_algorithm( - kind: BlockLoadStoreKind, - algorithm: str | BlockLoadStoreAlgorithm, -) -> BlockLoadStoreAlgorithm: - """Accept an enum, a short name, or the matching CUB operation's token.""" - - mapping = _algorithm_cpp_map(kind) - if isinstance(algorithm, BlockLoadStoreAlgorithm): - return algorithm - if isinstance(algorithm, str): - for candidate, cpp in mapping.items(): - if algorithm in {candidate.value, cpp}: - return candidate - raise ValueError( - f"unsupported Block{kind.value.title()} algorithm {algorithm!r}" - ) - - -def _base_parameters(kind: BlockLoadStoreKind) -> list[Any]: - """Describe scratch, the global pointer, and the per-thread item array. - - CUB takes the global pointer first for both operations: ``src`` for a - load and ``dst`` for a store. Outputs are written through arguments; - neither operation returns the item array or pointer. - """ - - if kind is BlockLoadStoreKind.LOAD: - return [ - TempStorageParameter(), - Pointer(_T, name="src", is_array_pointer=True, restrict=True), - Array( - _T, - _ITEMS_PER_THREAD, - name="dst", - is_output=True, - is_return=False, - ), - ] - return [ - TempStorageParameter(), - Pointer( - _T, - name="dst", - is_output=True, - is_return=False, - is_array_pointer=True, - restrict=True, - ), - Array(_T, _ITEMS_PER_THREAD, name="src"), - ] - - -def _normalize_optional_binding( - value: bool | ArgumentBinding, - *, - name: str, -) -> ArgumentBinding: - """Use booleans to select signatures and keep explicit value bindings.""" - - if isinstance(value, ArgumentBinding): - return value - if not isinstance(value, bool): - raise TypeError(f"{name} must be a bool or ArgumentBinding") - return ArgumentBinding.runtime() if value else ArgumentBinding.omitted() - - -def _with_pointer_offset( - parameters: list[Any], - offset: ArgumentBinding, -) -> tuple[Any, ...]: - """Append an element offset targeting the first non-scratch argument. - - The offset adjusts the global pointer during wrapper generation; it is - not an extra argument to CUB's ``Load`` or ``Store`` method. A static - binding embeds the offset, while a runtime binding adds an i64 input. - """ - - static_value = offset.value if offset.kind is BindingKind.STATIC else None - return ( - *parameters, - PointerOffset( - INT64, - name="offset", - pointer_arg_index=0, - static_value=static_value, - ), - ) - - -@dataclass(frozen=True) -class BlockLoadStoreSemantics: - """Normalized BlockLoad/BlockStore options and wrapper overloads. - - Built by ``make_block_load_store_semantics`` before block dimensions - are known. ``make_block_load_store_specialization`` adds those - dimensions and validates constraints that depend on the tile size. - - Attributes - ---------- - kind : BlockLoadStoreKind - Operation to perform and CUB method to invoke. - dtype : Any - Element type to bind to CUB's ``T`` template parameter. - algorithm : BlockLoadStoreAlgorithm - Memory-access strategy and per-thread item arrangement. - items_per_thread : int - Positive number of elements in each thread's item array. - valid_items, oob_default, pointer_offset : ArgumentBinding - Omitted, static, or runtime scalar arguments. A static value is - embedded in the wrapper; a runtime binding describes an input - without retaining its eventual value. - has_full_tile : bool - Whether an overload without a valid-item count is included. This - describes available signatures, not the size of a particular call. - parameters : tuple[tuple[Any, ...], ...] - Ordered wrapper parameter descriptors, one tuple per overload. - Each starts with scratch storage, the global pointer, and the - per-thread array, followed by any count, default, and offset. - Static descriptors do not consume runtime arguments; ``T`` and - ``ITEMS_PER_THREAD`` dependencies are bound during specialization. - """ - - kind: BlockLoadStoreKind - dtype: Any - algorithm: BlockLoadStoreAlgorithm - items_per_thread: int - valid_items: ArgumentBinding - oob_default: ArgumentBinding - has_full_tile: bool - pointer_offset: ArgumentBinding - parameters: tuple[tuple[Any, ...], ...] - - @property - def has_valid_items(self) -> bool: - return self.valid_items.kind is not BindingKind.OMITTED - - @property - def has_oob_default(self) -> bool: - return self.oob_default.kind is not BindingKind.OMITTED - - @property - def has_pointer_offset(self) -> bool: - return self.pointer_offset.kind is not BindingKind.OMITTED - - @property - def method_name(self) -> str: - return self.kind.value.title() - - @property - def algorithm_cpp(self) -> str: - return _algorithm_cpp_map(self.kind)[self.algorithm] - - @property - def semantic_key(self) -> tuple[Any, ...]: - """Identify the normalized options and overloads before dimensions.""" - - return ( - f"block_{self.kind.value}", - semantic_token(self.dtype), - self.algorithm.value, - self.items_per_thread, - self.valid_items.semantic_key, - self.oob_default.semantic_key, - self.has_full_tile, - self.pointer_offset.semantic_key, - semantic_token(self.parameters), - ) - - -def make_block_load_store_semantics( - *, - kind: str | BlockLoadStoreKind, - dtype: Any, - items_per_thread: int, - algorithm: str | BlockLoadStoreAlgorithm, - valid_items: bool | ArgumentBinding = False, - oob_default: bool | ArgumentBinding = False, - include_full_tile: bool = False, - include_pointer_offset: bool | ArgumentBinding = False, -) -> BlockLoadStoreSemantics: - """Normalize BlockLoad/BlockStore options and describe its overloads. - - Assemble parameter descriptors for CUB wrappers before a block shape - is available. For scalar options, ``False`` omits the argument and - ``True`` requests a runtime argument. Use ``ArgumentBinding.static`` - to embed a value, including a literal boolean ``oob_default``. - - Parameters - ---------- - kind : str or BlockLoadStoreKind - ``"load"`` or ``"store"``; selects the CUB class and method. - dtype : Any - Element type for the global pointer, per-thread array, and optional - load default. Retained for later binding to CUB's ``T``. - items_per_thread : int - Positive integral number of elements per thread; booleans are - rejected. - algorithm : str or BlockLoadStoreAlgorithm - Enum member, its lowercase value, or the corresponding fully - qualified CUB token, such as ``"::cub::BLOCK_LOAD_TRANSPOSE"``. - A CUB token must match ``kind``. - valid_items : bool or ArgumentBinding, optional - Number of valid elements in the whole block tile, measured from - the possibly offset global pointer. A present binding selects a - guarded overload. Static counts must fit signed i32; validation - against the block tile size happens during specialization. - oob_default : bool or ArgumentBinding, optional - Value assigned to invalid load positions. Requires a load with - ``valid_items`` present. Static defaults must be representable as - finite C++ scalar literals. Without a default, invalid load values - are unspecified; guarded stores leave out-of-range memory alone. - include_full_tile : bool, optional - Add an unguarded overload alongside the guarded overload. Requires - ``valid_items`` to be present. When ``valid_items`` is omitted, - the unguarded overload is already included with the default - ``include_full_tile=False``. - include_pointer_offset : bool or ArgumentBinding, optional - Offset in elements applied to the global pointer. ``False`` omits - it; ``True`` includes both unoffset and runtime-offset overloads. - An explicit static or runtime binding includes only offset - overloads; an omitted binding includes only unoffset overloads. - Static offsets must be nonnegative and fit signed i64. - - Returns - ------- - BlockLoadStoreSemantics - Canonical options and ordered parameter descriptors. Runtime - values, pointer validity, and memory extents are checked or - established by the caller and later planning/lowering stages. - - Raises - ------ - TypeError - A scalar option is neither a boolean nor an ``ArgumentBinding``, - a static count or offset is not an integer or is a boolean, or a - static default is not a numeric scalar. - ValueError - The kind, algorithm, or item count is invalid; a static scalar is - out of range or nonfinite; a static offset is negative; or the - requested default/full-tile overload lacks a guarded load/store - signature. A default is also rejected for stores. - - See Also - -------- - make_block_load_store_specialization - Bind block dimensions and validate the tile-dependent constraints. - """ - - pointer_offset_overload_cohort = isinstance(include_pointer_offset, bool) - kind = BlockLoadStoreKind(kind) - items_per_thread = normalize_positive_int( - "items_per_thread", items_per_thread - ) - algorithm = _normalize_algorithm(kind, algorithm) - valid_items = _normalize_optional_binding(valid_items, name="valid_items") - valid_items = normalize_i32_binding(valid_items, name="valid_items") - oob_default = _normalize_optional_binding(oob_default, name="oob_default") - pointer_offset = _normalize_optional_binding( - include_pointer_offset, - name="include_pointer_offset", - ) - pointer_offset = normalize_i64_binding( - pointer_offset, name="pointer offset" - ) - if ( - pointer_offset.kind is BindingKind.STATIC - and int(pointer_offset.value) < 0 - ): - raise ValueError("static pointer offset must be nonnegative") - if ( - kind is BlockLoadStoreKind.STORE - and oob_default.kind is not BindingKind.OMITTED - ): - raise ValueError("oob_default is only valid for BlockLoad") - if ( - oob_default.kind is not BindingKind.OMITTED - and valid_items.kind is BindingKind.OMITTED - ): - raise ValueError("oob_default requires a valid_items signature") - if include_full_tile and valid_items.kind is BindingKind.OMITTED: - raise ValueError("include_full_tile requires a valid_items signature") - - base = _base_parameters(kind) - has_full_tile = valid_items.kind is BindingKind.OMITTED or include_full_tile - methods: list[tuple[Any, ...]] = [] - if has_full_tile and ( - pointer_offset.kind is BindingKind.OMITTED - or pointer_offset_overload_cohort - ): - methods.append(tuple(base)) - if valid_items.kind is not BindingKind.OMITTED: - num_valid_items = i32_parameter(valid_items, name="num_valid_items") - assert num_valid_items is not None - partial = [*base, num_valid_items] - if oob_default.kind is BindingKind.RUNTIME: - partial.append(Value(dtype, name="oob_default")) - elif oob_default.kind is BindingKind.STATIC: - partial.append( - CxxFunction( - cxx_scalar_literal(oob_default.value, name="oob_default"), - dtype, - name="oob_default", - ) - ) - if ( - pointer_offset.kind is BindingKind.OMITTED - or pointer_offset_overload_cohort - ): - methods.append(tuple(partial)) - if pointer_offset.kind is not BindingKind.OMITTED: - methods.append(_with_pointer_offset(partial, pointer_offset)) - if pointer_offset.kind is not BindingKind.OMITTED and has_full_tile: - methods.append(_with_pointer_offset(base, pointer_offset)) - - return BlockLoadStoreSemantics( - kind=kind, - dtype=dtype, - algorithm=algorithm, - items_per_thread=items_per_thread, - valid_items=valid_items, - oob_default=oob_default, - has_full_tile=has_full_tile, - pointer_offset=pointer_offset, - parameters=tuple(methods), - ) - - -def make_block_load_store_specialization( - *, - kind: str | BlockLoadStoreKind, - dtype: Any, - block_dim: tuple[int, int, int], - items_per_thread: int, - algorithm: str | BlockLoadStoreAlgorithm, - valid_items: bool | ArgumentBinding = False, - oob_default: bool | ArgumentBinding = False, - include_full_tile: bool = False, - include_pointer_offset: bool | ArgumentBinding = False, -) -> Algorithm: - """Bind a BlockLoad/BlockStore operation to a concrete block shape. - - Normalize the options with ``make_block_load_store_semantics``, then - construct ``Algorithm`` with all CUB template arguments bound. The - returned description is ready for materialization; it does not compile - device code, allocate scratch storage, or insert synchronization. - - Parameters - ---------- - block_dim : tuple[int, int, int] - Exact ``(x, y, z)`` launch dimensions, each a positive integer. - Booleans are rejected. The block tile contains - ``x * y * z * items_per_thread`` elements. ``warp_transpose`` and - ``warp_transpose_timesliced`` require ``x * y * z`` divisible by 32. - - Other Parameters - ---------------- - kind : str or BlockLoadStoreKind - ``"load"`` or ``"store"``. - dtype : Any - Element type bound to CUB's ``T``. - items_per_thread : int - Positive integral number of elements per thread, excluding booleans. - algorithm : str or BlockLoadStoreAlgorithm - CUB strategy, accepted in the forms documented by - ``make_block_load_store_semantics``. - valid_items : bool or ArgumentBinding, optional - Guarded-call count binding. Static counts must additionally lie - between zero and the block tile size, inclusive. Runtime counts - must satisfy the same bounds at execution time. - oob_default : bool or ArgumentBinding, optional - Optional invalid-item default for guarded loads. - include_full_tile : bool, optional - Include an unguarded overload alongside guarded overloads. - include_pointer_offset : bool or ArgumentBinding, optional - Select global-pointer offset overloads. Boolean and explicit-binding - selection follow ``make_block_load_store_semantics``. - - Returns - ------- - Algorithm - Bound CUB template arguments, wrapper parameters, and operation - metadata, ready for backend materialization. - - Raises - ------ - TypeError - An argument binding fails the semantics builder's type checks. - ValueError - An option fails the semantics builder's validation, the block - shape is invalid, a static count exceeds the tile bounds, or a - warp-transpose algorithm is selected for an incomplete warp. - """ - - block_dim = normalize_block_dim(block_dim) - semantics = make_block_load_store_semantics( - kind=kind, - dtype=dtype, - items_per_thread=items_per_thread, - algorithm=algorithm, - valid_items=valid_items, - oob_default=oob_default, - include_full_tile=include_full_tile, - include_pointer_offset=include_pointer_offset, - ) - if semantics.valid_items.kind is BindingKind.STATIC: - value = semantics.valid_items.value - if isinstance(value, bool) or not isinstance(value, Integral): - raise TypeError("static valid_items must be an integer") - value = int(value) - tile_items = ( - semantics.items_per_thread - * block_dim[0] - * block_dim[1] - * block_dim[2] - ) - if not 0 <= value <= tile_items: - raise ValueError( - "static valid_items must be between zero and the block tile " - f"size ({tile_items})" - ) - block_threads = block_dim[0] * block_dim[1] * block_dim[2] - if ( - semantics.algorithm - in { - BlockLoadStoreAlgorithm.WARP_TRANSPOSE, - BlockLoadStoreAlgorithm.WARP_TRANSPOSE_TIMESLICED, - } - and block_threads % 32 != 0 - ): - raise ValueError( - f"Block{semantics.kind.value.title()} algorithm " - f"{semantics.algorithm.value!r} " - "requires a block size that is a multiple of 32" - ) - title = semantics.kind.value.title() - return Algorithm( - struct_name=f"Block{title}", - method_name=title, - c_name=f"block_{semantics.kind.value}", - includes=(f"cub/block/block_{semantics.kind.value}.cuh",), - template_parameters=_TEMPLATE_PARAMETERS, - parameters=semantics.parameters, - template_arguments={ - "T": dtype, - "BLOCK_DIM_X": block_dim[0], - "ITEMS_PER_THREAD": semantics.items_per_thread, - "ALGORITHM": semantics.algorithm_cpp, - "BLOCK_DIM_Y": block_dim[1], - "BLOCK_DIM_Z": block_dim[2], - }, - metadata={ - "scope": "block", - "primitive": semantics.kind.value, - "algorithm": semantics.algorithm.value, - "valid_items": semantics.has_valid_items, - "oob_default": semantics.has_oob_default, - "full_tile": semantics.has_full_tile, - "pointer_offset": semantics.has_pointer_offset, - }, - ) - - -def make_block_load_specialization(**kwargs: Any) -> Algorithm: - """Call ``make_block_load_store_specialization`` with ``kind="load"``. - - Accepts the same keyword arguments except ``kind`` and returns the - resulting ``Algorithm``. - """ - - return make_block_load_store_specialization( - kind=BlockLoadStoreKind.LOAD, **kwargs - ) - - -def make_block_store_specialization(**kwargs: Any) -> Algorithm: - """Call ``make_block_load_store_specialization`` with ``kind="store"``. - - Accepts the same keyword arguments except ``kind`` and returns the - resulting ``Algorithm``. ``oob_default`` must be omitted for a store. - """ - - return make_block_load_store_specialization( - kind=BlockLoadStoreKind.STORE, **kwargs - ) diff --git a/python/cuda_coop/cuda/coop/_core/dtype_policy.py b/python/cuda_coop/cuda/coop/_core/dtype_policy.py deleted file mode 100644 index af5eb528487e..000000000000 --- a/python/cuda_coop/cuda/coop/_core/dtype_policy.py +++ /dev/null @@ -1,98 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Keep the common API's numeric type choices consistent across backends. - -A backend first converts its dtype object to a standard name such as -``int32``. These checks then enforce the common operation's supported set. -Names keep this policy independent of each compiler's type objects. -""" - -from __future__ import annotations - -_COMMON_NUMERIC_DTYPE_NAMES = ( - "int8", - "uint8", - "int16", - "uint16", - "int32", - "uint32", - "int64", - "uint64", - "float32", - "float64", -) - - -def _validate_common_dtype_name( - dtype_name: str, - *, - operation: str, - parameter: str | None, - supported_dtype_names: tuple[str, ...], -) -> str: - """Check that a dtype name is allowed; name the operation in any error. - - ``parameter`` adds argument-specific context when supplied. Return the - accepted name unchanged so callers can use it after validation. - - Parameters - ---------- - dtype_name : str - Canonical type name already normalized by the backend. - operation : str - Common operation name used to identify the failing call. - parameter : str or None - Public argument name for a parameter-specific diagnostic; - ``None`` describes the operation's dtype support generally. - supported_dtype_names : tuple of str - Canonical names accepted by this particular common operation - or operand. - """ - - if dtype_name not in supported_dtype_names: - supported = ", ".join(supported_dtype_names) - subject = "dtypes" if parameter is None else f"{parameter} dtypes" - raise TypeError( - f"cuda.coop.{operation} supports {subject} {supported} through the " - "common API; " - f"use a backend-qualified import for backend-specific {subject}" - ) - return dtype_name - - -def validate_common_numeric_dtype_name( - dtype_name: str, - *, - operation: str, - parameter: str | None = None, -) -> str: - """Require one of the common API's integer or floating-point type names. - - Backend adapters call this after translating compiler-specific type - objects to names. Checking the same names in one place keeps the - common API's numeric contract consistent even when backends have - different type representations. - - The backend must normalize aliases before calling this function. Its - qualified API may support additional types outside this common set. - - Parameters - ---------- - dtype_name : str - Canonical numeric type name, such as ``int32`` or - ``float64``; aliases must already be resolved. - operation : str - Common operation name included in any rejection. - parameter : str or None, optional - Operand name included in the diagnostic when only one - argument's type is being checked. - """ - - return _validate_common_dtype_name( - dtype_name, - operation=operation, - parameter=parameter, - supported_dtype_names=_COMMON_NUMERIC_DTYPE_NAMES, - ) diff --git a/python/cuda_coop/cuda/coop/_core/group/__init__.py b/python/cuda_coop/cuda/coop/_core/group/__init__.py deleted file mode 100644 index a25bc771aaa1..000000000000 --- a/python/cuda_coop/cuda/coop/_core/group/__init__.py +++ /dev/null @@ -1,61 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from ._dispatch import ( - GroupOperationSemantics, - make_group_primitive_call, - plan_group_primitive, -) -from ._model import ( - ArgumentPrecondition, - GroupExecutionRequirements, - GroupLoweringPlan, - GroupLoweringTarget, - GroupPrimitiveCall, - GroupTopologyRequirements, - ImplementationProvenance, - ParticipationRequirements, - PreconditionEnforcement, - ResultVisibility, - StorageOwnership, - SynchronizationRequirements, - SynchronizationScope, - TempStorageRequirements, - ThreadGroupLaunchResolution, - UnsupportedReason, - UnsupportedReasonCode, -) -from ._resolution import resolve_thread_group -from .load_store import ( - GroupLoadStoreAlgorithm, - GroupLoadStoreKind, - GroupLoadStoreSemantics, -) - -__all__ = [ - "ArgumentPrecondition", - "GroupExecutionRequirements", - "GroupLoadStoreAlgorithm", - "GroupLoadStoreKind", - "GroupLoadStoreSemantics", - "GroupLoweringPlan", - "GroupLoweringTarget", - "GroupOperationSemantics", - "GroupPrimitiveCall", - "GroupTopologyRequirements", - "ImplementationProvenance", - "ParticipationRequirements", - "PreconditionEnforcement", - "ResultVisibility", - "StorageOwnership", - "SynchronizationRequirements", - "SynchronizationScope", - "TempStorageRequirements", - "ThreadGroupLaunchResolution", - "UnsupportedReason", - "UnsupportedReasonCode", - "make_group_primitive_call", - "plan_group_primitive", - "resolve_thread_group", -] diff --git a/python/cuda_coop/cuda/coop/_core/group/_dispatch.py b/python/cuda_coop/cuda/coop/_core/group/_dispatch.py deleted file mode 100644 index 4245ee8af103..000000000000 --- a/python/cuda_coop/cuda/coop/_core/group/_dispatch.py +++ /dev/null @@ -1,279 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Select the shared planner for an operation's semantics type. - -Each operation family registers how to classify its parameters and how -to build a lowering plan. Dispatch uses the exact semantics type, so a -new family can supply its rules without adding operation-specific -branches here. - -Before calling the family planner, check its accepted group kinds, -required launch capabilities, and resolved group shape. The result is a -plan or a structured unsupported reason. This stage does not compile -code or execute a kernel. -""" - -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass - -from .._types import ParameterClassification -from ..launch import LaunchFacts -from ..thread_group import THREAD_GROUP_KINDS, ThreadGroup -from ._execution_requirements import _unsupported -from ._model import ( - GroupLoweringPlan, - GroupOperationSemantics, - GroupPrimitiveCall, - UnsupportedReasonCode, -) -from ._resolution import _resolve_group - - -@dataclass(frozen=True) -class _GroupOperationFamily: - """Keep the planning rules shared by one operation semantics type. - - The registry uses this record to route a ``GroupPrimitiveCall`` to its - family planner. The callback chooses the implementation and its - requirements after dispatch has checked the group and launch facts. - - Parameters - ---------- - classifications : callable - Map an operation's semantics to its parameter classifications. - These classifications separate specialization choices, runtime - arguments, and storage requirements in the call description. - planner : callable - Accept the call, resolved group, launch facts, and operation - semantics. Return a supported or unsupported ``GroupLoweringPlan``. - group_kinds : frozenset of str - Group kinds this family can consider. Membership permits planning; - the family can still reject an unsupported algorithm or shape. - unsupported_group_message : str - Diagnostic to use when a call has a group kind outside that set. - """ - - classifications: Callable[ - [GroupOperationSemantics], tuple[ParameterClassification, ...] - ] - planner: Callable[ - [GroupPrimitiveCall, ThreadGroup, LaunchFacts, GroupOperationSemantics], - GroupLoweringPlan, - ] - group_kinds: frozenset[str] - unsupported_group_message: str - - def __post_init__(self) -> None: - """Reject incomplete registrations before dispatch uses them.""" - - if not callable(self.classifications): - raise TypeError("classifications must be callable") - if not callable(self.planner): - raise TypeError("planner must be callable") - object.__setattr__(self, "group_kinds", frozenset(self.group_kinds)) - if not self.group_kinds: - raise ValueError("group_kinds must not be empty") - invalid_group_kinds = { - kind - for kind in self.group_kinds - if not isinstance(kind, str) or kind not in THREAD_GROUP_KINDS - } - if invalid_group_kinds: - names = ", ".join( - sorted(repr(kind) for kind in invalid_group_kinds) - ) - raise ValueError( - f"group_kinds contains unsupported values: {names}" - ) - if ( - not isinstance(self.unsupported_group_message, str) - or not self.unsupported_group_message.strip() - ): - raise ValueError( - "unsupported_group_message must be a non-empty string" - ) - - -_GROUP_OPERATION_FAMILIES: dict[type, _GroupOperationFamily] = {} - - -def _register_group_operation_family( - semantics_type: type, - *, - classifications: Callable[ - [GroupOperationSemantics], tuple[ParameterClassification, ...] - ], - planner: Callable[ - [GroupPrimitiveCall, ThreadGroup, LaunchFacts, GroupOperationSemantics], - GroupLoweringPlan, - ], - group_kinds: frozenset[str], - unsupported_group_message: str, -) -> None: - """Install the callbacks for one exact operation semantics type. - - A family module calls this when it loads. Repeating the same registration - is harmless; replacing it with different rules raises ``RuntimeError``. - This prevents import order from changing how an existing call is planned. - See ``_GroupOperationFamily`` for callbacks and group-kind checks. - """ - - if not isinstance(semantics_type, type): - raise TypeError("semantics_type must be a type") - registration = _GroupOperationFamily( - classifications=classifications, - planner=planner, - group_kinds=frozenset(group_kinds), - unsupported_group_message=unsupported_group_message, - ) - existing = _GROUP_OPERATION_FAMILIES.get(semantics_type) - if existing is not None and existing != registration: - raise RuntimeError( - f"group operation semantics {semantics_type!r} " - "are already registered" - ) - _GROUP_OPERATION_FAMILIES[semantics_type] = registration - - -def _group_operation_family(operation: object) -> _GroupOperationFamily | None: - """Find the exact semantics type; ignore rules for its base classes.""" - - return _GROUP_OPERATION_FAMILIES.get(type(operation)) - - -def _is_group_operation(operation: object) -> bool: - return _group_operation_family(operation) is not None - - -def _call_classifications( - operation: GroupOperationSemantics, -) -> tuple[ParameterClassification, ...]: - """Ask the registered family which roles its operation parameters have. - - Reject unregistered semantics before constructing a call identity. - """ - - family = _group_operation_family(operation) - if family is None: - raise TypeError("unsupported GroupPrimitiveCall operation") - return family.classifications(operation) - - -def make_group_primitive_call( - group: ThreadGroup, - operation: GroupOperationSemantics, -) -> GroupPrimitiveCall: - """Pair a group request with the operation semantics to plan. - - ``GroupPrimitiveCall`` validates the descriptor and registered semantics. - Launch-dependent resolution happens later in ``plan_group_primitive``. - """ - - return GroupPrimitiveCall(group=group, operation=operation) - - -def plan_group_primitive( - call: GroupPrimitiveCall, - launch: LaunchFacts, -) -> GroupLoweringPlan: - """Resolve a group call and ask its family to choose an implementation. - - Check the family's group-kind limit first. Multi-block cluster operations - need verified cluster-launch support, and grid operations need verified - cooperative-launch support. These checks protect implementations whose - threads must participate together; a shape alone does not establish the - required launch capability. - - Then resolve the group dimensions and pass the call to its registered - planner. Family-specific checks choose a supported implementation or - return a reason that explains why the request cannot be lowered. - - Parameters - ---------- - call : GroupPrimitiveCall - Group and operation semantics to plan. - launch : LaunchFacts - Exact launch dimensions and capability evidence. - - Returns - ------- - GroupLoweringPlan - Selected implementation and execution requirements, or an - unsupported plan with the failed requirement. The caller decides - when to turn an unsupported result into a user-facing error. - - Raises - ------ - TypeError - An argument has the wrong type or the semantics type is unregistered. - ValueError - Group dimensions contradict the launch, or an operation-specific - argument is invalid. - """ - - if not isinstance(call, GroupPrimitiveCall): - raise TypeError("call must be a GroupPrimitiveCall") - if not isinstance(launch, LaunchFacts): - raise TypeError("launch must be LaunchFacts") - family = _group_operation_family(call.operation) - if family is None: - raise TypeError("unsupported GroupPrimitiveCall operation") - if call.group.kind not in family.group_kinds: - return _unsupported( - call, - call.group, - UnsupportedReasonCode.GROUP_KIND, - family.unsupported_group_message, - ) - cluster_dim = launch.exact_cluster_dim - uses_multi_block_cluster = cluster_dim is not None and cluster_dim != ( - 1, - 1, - 1, - ) - if ( - call.group.kind in {"cluster", "grid"} - and uses_multi_block_cluster - and ( - launch.cluster_launch is not True - or not launch.is_verified("cluster_launch") - ) - ): - return _unsupported( - call, - call.group, - UnsupportedReasonCode.LAUNCH_CAPABILITY, - "multi-block cluster lowering requires verified cluster launch " - f"capability; observed {launch.cluster_launch!r} with verified=" - f"{launch.is_verified('cluster_launch')!r}", - ) - if call.group.kind == "grid" and ( - launch.cooperative_launch is not True - or not launch.is_verified("cooperative_launch") - ): - return _unsupported( - call, - call.group, - UnsupportedReasonCode.LAUNCH_CAPABILITY, - "grid group lowering requires verified cooperative launch " - f"capability; observed {launch.cooperative_launch!r} with " - f"verified={launch.is_verified('cooperative_launch')!r}", - ) - resolved, failure = _resolve_group(call, launch) - if failure is not None: - return failure - return family.planner(call, resolved, launch, call.operation) - - -__all__ = [ - "GroupOperationSemantics", - "_call_classifications", - "_is_group_operation", - "_register_group_operation_family", - "make_group_primitive_call", - "plan_group_primitive", -] diff --git a/python/cuda_coop/cuda/coop/_core/group/_execution_requirements.py b/python/cuda_coop/cuda/coop/_core/group/_execution_requirements.py deleted file mode 100644 index 12d26f24fad7..000000000000 --- a/python/cuda_coop/cuda/coop/_core/group/_execution_requirements.py +++ /dev/null @@ -1,340 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe how a resolved group must participate and use scratch storage. - -Operation planners call these helpers after choosing a group shape. The -result tells a backend how many group instances need storage, which -threads belong to each instance, and where storage reuse needs a barrier. -For example, several logical warps in one block need separate storage -instances even when they use the same primitive. - -These records state execution requirements. They do not prove that a -kernel reaches a call uniformly or insert synchronization themselves. -Backend lowering must honor the plan when it creates the device call. -""" - -from __future__ import annotations - -from ..launch import LaunchFacts -from ..thread_group import ThreadGroup -from ._model import ( - ArgumentPrecondition, - GroupExecutionRequirements, - GroupLoweringPlan, - GroupLoweringTarget, - GroupPrimitiveCall, - GroupTopologyRequirements, - ParticipationRequirements, - StorageOwnership, - SynchronizationRequirements, - SynchronizationScope, - TempStorageRequirements, - UnsupportedReason, - UnsupportedReasonCode, -) - - -def _unsupported( - call: GroupPrimitiveCall, - resolved_group: ThreadGroup, - code: UnsupportedReasonCode, - message: str, -) -> GroupLoweringPlan: - """Return a plan containing the call and the failed requirement. - - Leave implementation and execution requirements unset. Callers can inspect - or report the failure without mistaking it for a usable device plan. - """ - - return GroupLoweringPlan( - target=GroupLoweringTarget.UNSUPPORTED, - call=call, - resolved_group=resolved_group, - implementation=None, - topology=None, - participation=None, - result=None, - synchronization=None, - temp_storage=None, - provenance=None, - unsupported=UnsupportedReason(code=code, message=message), - ) - - -def _group_topology( - resolved_group: ThreadGroup, - launch: LaunchFacts, -) -> GroupTopologyRequirements: - """Describe group instances and each thread's rank within an instance. - - This calculation is shared by primitive families. For a warp-based group, - linear thread rank selects the instance by division and the rank within - it by remainder. Whole blocks, clusters, and grids use their enclosing - execution scope. - - Parameters - ---------- - resolved_group : ThreadGroup - Group with a known thread count. Mapped groups must fit the block. - launch : LaunchFacts - Exact block size needed to count thread or warp-based instances. - - Returns - ------- - GroupTopologyRequirements - Group width, instance count, rank expressions, and execution scope. - Expressions describe the rank calculation for backend lowering. - - Raises - ------ - ValueError - The group size is unknown, required block dimensions are missing, - or a warp-based group does not divide the block. - """ - - group_size = resolved_group.static_size - if group_size is None: - raise ValueError("group topology requires a static group size") - - block_threads = launch.exact_block_threads - kind = resolved_group.kind - if kind == "thread": - if block_threads is None: - raise ValueError("thread topology requires exact block dimensions") - instances = block_threads - index = "linear_thread_rank" - thread_rank = "0" - execution_scope = SynchronizationScope.NONE - elif kind in {"warp", "threads_within_warp"}: - if block_threads is None: - raise ValueError("warp topology requires exact block dimensions") - if block_threads % group_size != 0: - raise ValueError("group width must divide the enclosing block size") - instances = block_threads // group_size - index = f"linear_thread_rank / {group_size}" - thread_rank = f"linear_thread_rank % {group_size}" - execution_scope = SynchronizationScope.WARP - elif kind == "warps_within_block": - if block_threads is None: - raise ValueError( - "mapped block topology requires exact block dimensions" - ) - if block_threads % group_size != 0: - raise ValueError("group width must divide the enclosing block size") - instances = block_threads // group_size - index = f"linear_thread_rank / {group_size}" - thread_rank = f"linear_thread_rank % {group_size}" - execution_scope = ( - SynchronizationScope.WARP - if group_size == 32 - else SynchronizationScope.GROUP - ) - elif kind == "block": - instances = 1 - index = "cta" - thread_rank = "linear_thread_rank" - execution_scope = SynchronizationScope.BLOCK - elif kind == "cluster": - instances = 1 - index = "cluster" - thread_rank = "group.thread_rank()" - execution_scope = SynchronizationScope.GROUP - elif kind == "grid": - instances = 1 - index = "grid" - thread_rank = "group.thread_rank()" - execution_scope = SynchronizationScope.GROUP - else: - instances = resolved_group.groups_per_parent or 1 - index = "group.rank(parent)" - thread_rank = "group.thread_rank()" - execution_scope = SynchronizationScope.GROUP - - return GroupTopologyRequirements( - group_kind=kind, - logical_width=group_size, - instances=instances, - instance_index=index, - thread_rank=thread_rank, - execution_scope=execution_scope, - ) - - -def _build_execution_requirements( - resolved_group: ThreadGroup, - launch: LaunchFacts, - *, - storage_ownership: StorageOwnership, - cpp_type: str | None, - storage_sharing: str | None = None, - requested_size_in_bytes: int | None = None, - requested_alignment: int | None = None, - auto_sync: bool | None = None, - uniform_arguments: tuple[str, ...] = (), - valid_member_selection: str | None = None, - argument_preconditions: tuple[ArgumentPrecondition, ...] = (), -) -> GroupExecutionRequirements: - """Build participation, synchronization, and storage requirements. - - A resolved group's topology determines both the storage instance index - and the scope of a reuse barrier. Storage-free implementations need no - reuse barrier. For implementations with storage, ``auto_sync`` defaults - to true; a false value leaves storage reuse synchronization to the caller. - - The participation record requires contiguous, aligned members to enter - together. Operation planners add argument-specific requirements, such as - which counts must be uniform. These are requirements for later lowering - and execution, not runtime checks performed by this helper. - - Parameters - ---------- - resolved_group : ThreadGroup - Group with a known size, already accepted by the operation planner. - launch : LaunchFacts - Exact launch dimensions used to describe participation and instances. - storage_ownership : StorageOwnership - Whether the primitive uses no storage, storage supplied by its - implementation, or storage supplied by the caller. - cpp_type : str or None - C++ storage type selected by the implementation, when applicable. - storage_sharing : str, optional - Sharing policy to carry into the storage requirements. - requested_size_in_bytes : int, optional - Caller-requested storage capacity to check against the actual layout. - requested_alignment : int, optional - Caller-requested storage alignment to check against the actual layout. - auto_sync : bool, optional - Whether lowering should arrange a barrier before storage reuse. - Omission enables it for implementations that use storage. - uniform_arguments : tuple of str, optional - Argument names whose values must agree across participating threads. - valid_member_selection : str, optional - Description of which data elements the operation uses. For example, - a guarded load can use only the first ``valid_items`` tile elements. - argument_preconditions : tuple of ArgumentPrecondition, optional - Additional bounds or other requirements on operation arguments. - - Returns - ------- - GroupExecutionRequirements - Topology, participation, synchronization, and temporary-storage - requirements. Caller-owned storage requires an exact - layout check; this helper does not calculate the compiled layout. - """ - - group_size = resolved_group.static_size - assert group_size is not None - topology = _group_topology(resolved_group, launch) - storage_ownership = StorageOwnership(storage_ownership) - storage_free = storage_ownership is StorageOwnership.NONE - if auto_sync is None: - auto_sync = not storage_free - barrier = ( - SynchronizationScope.NONE - if storage_free or not auto_sync - else topology.execution_scope - ) - return GroupExecutionRequirements( - topology=topology, - participation=ParticipationRequirements( - group_kind=resolved_group.kind, - exact_group_size=group_size, - exact_block_dim=launch.exact_block_dim, - complete_membership=resolved_group.complete_membership is not False, - contiguous=True, - aligned=True, - converged_entry=True, - complete_parent_partition=( - resolved_group.kind == "warp" - or resolved_group.complete_membership is True - ), - uniform_arguments=uniform_arguments, - valid_member_selection=valid_member_selection, - argument_preconditions=argument_preconditions, - ), - synchronization=SynchronizationRequirements( - converged_entry=True, - storage_reuse_barrier=barrier, - ), - temp_storage=TempStorageRequirements( - ownership=storage_ownership, - address_space=None if storage_free else "shared", - cpp_type=cpp_type, - instances=( - None - if storage_ownership is StorageOwnership.NONE - else topology.instances - ), - instance_index=( - None - if storage_ownership is StorageOwnership.NONE - else topology.instance_index - ), - exact_layout_required=storage_ownership is StorageOwnership.CALLER, - sharing=storage_sharing, - requested_size_in_bytes=requested_size_in_bytes, - requested_alignment=requested_alignment, - auto_sync=auto_sync, - ), - ) - - -def _cub_warp_width(group: ThreadGroup) -> int: - """Select a width that CUB's warp primitives can represent. - - Physical warps use 32 threads. A logical warp must have a power-of-two - width from 1 through 32 and divide its physical warp. Raise ``ValueError`` - for a different group kind or width. Block participation is checked by - group resolution before a primitive uses this width. - """ - - if group.kind == "warp": - return 32 - if group.kind != "threads_within_warp": - raise ValueError("CUB warp primitives require a warp-based group") - width = group.static_size - if ( - not isinstance(width, int) - or isinstance(width, bool) - or width < 1 - or width > 32 - or width & (width - 1) - or 32 % width != 0 - ): - raise ValueError( - "CUB-backed logical-warp operations require a power-of-two group " - "width in [1, 32] that divides the 32-thread physical warp; " - f"got {width!r}" - ) - return width - - -def _unsupported_cub_warp_width( - call: GroupPrimitiveCall, - resolved: ThreadGroup, -) -> tuple[int | None, GroupLoweringPlan | None]: - """Return the CUB width or an unsupported plan for an invalid group. - - Translate the width check's ``ValueError`` into the same structured - failure that other primitive-planning checks return. - """ - - try: - return _cub_warp_width(resolved), None - except ValueError as exc: - return None, _unsupported( - call, - resolved, - UnsupportedReasonCode.GROUP_KIND, - str(exc), - ) - - -__all__ = [ - "_build_execution_requirements", - "_group_topology", - "_unsupported", - "_unsupported_cub_warp_width", -] diff --git a/python/cuda_coop/cuda/coop/_core/group/_model.py b/python/cuda_coop/cuda/coop/_core/group/_model.py deleted file mode 100644 index 3b43dbe422bc..000000000000 --- a/python/cuda_coop/cuda/coop/_core/group/_model.py +++ /dev/null @@ -1,947 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe a cooperative call and the requirements for implementing it. - -A frontend first records an operation and its requested thread group in a -``GroupPrimitiveCall``. Launch resolution establishes the group's concrete -membership, then the operation's planner chooses an implementation and returns -a ``GroupLoweringPlan``. That plan carries the participation, synchronization, -and scratch-storage requirements needed to preserve the operation's semantics. -For example, a warp load using a transpose needs a separate scratch instance -for each participating warp and synchronization before that storage is reused. - -Backends use these records to select a provider, compile specialized code, and -arrange storage and barriers. The records describe those requirements; -constructing them does not emit device code or check runtime participation. -Unsupported requests retain the original call and a structured reason so -callers can inspect a planning result before requiring executable support. - -Implementation provenance identifies the native library entry point selected -by planning. Together with the specialized implementation and execution -requirements, it contributes to the plan's artifact identity. The logical -request and a particular implementation have separate keys. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from enum import Enum -from typing import Any, Protocol - -from .._algorithm import Algorithm -from .._types import ParameterClassification -from ..launch import Dim3, LaunchFacts -from ..thread_group import MAPPED_GROUP_KINDS, ThreadGroup - - -class GroupLoweringTarget(str, Enum): - """Implementation family chosen by planning, or an unsupported outcome. - - CUB block and warp targets describe the cooperative primitive's scope. - A backend still selects and validates a provider for that target. - """ - - CUB_BLOCK = "cub_block" - CUB_WARP = "cub_warp" - UNSUPPORTED = "unsupported" - - -class ResultVisibility(str, Enum): - """Which group members have meaningful results from the operation. - - ``ALL_MEMBERS`` describes a group result available to every member; - ``GROUP_ROOT`` restricts it to the root; ``PER_MEMBER`` describes each - member's own result, such as its items from a cooperative load. - """ - - ALL_MEMBERS = "all_members" - GROUP_ROOT = "group_root" - PER_MEMBER = "per_member" - - -class PreconditionEnforcement(str, Enum): - """Whether planning checked a bound or the caller must satisfy it. - - ``PLANNER_VALIDATED`` records a check against a known static value. - ``CALLER`` records a precondition on a runtime value; it does not request - insertion of a runtime guard. - """ - - PLANNER_VALIDATED = "planner_validated" - CALLER = "caller" - - -class StorageOwnership(str, Enum): - """Who supplies the primitive's temporary storage. - - ``NONE`` selects a storage-free implementation. ``IMPLEMENTATION`` lets - lowering arrange scratch internally; ``CALLER`` requires the supplied - storage binding and its layout requests to be honored. - """ - - NONE = "none" - IMPLEMENTATION = "implementation" - CALLER = "caller" - - -class SynchronizationScope(str, Enum): - """Members covered by execution or a storage-reuse synchronization. - - ``NONE`` requires no barrier, while ``WARP`` and ``BLOCK`` identify - hardware scopes. ``GROUP`` refers to the resolved cooperative group; - a backend must support that group's synchronization before lowering it. - """ - - NONE = "none" - WARP = "warp" - BLOCK = "block" - GROUP = "group" - - -@dataclass(frozen=True) -class GroupTopologyRequirements: - """Describe the group instances and ranks needed by backend lowering. - - Topology supplies the indexing rules for per-group scratch and data - tiles. For logical warps of width 16 in a 128-thread block, for example, - there are eight instances; dividing the linear thread rank by 16 selects - the instance, and taking its remainder gives the rank within that group. - The expression strings name canonical formulas understood by backends. - They are not arbitrary expressions to evaluate as Python or C++. - - Attributes - ---------- - group_kind : str - Kind of the resolved group, such as ``"block"`` or ``"warp"``. - logical_width : int - Positive number of threads in each group instance. - instances : int - Positive number of group instances in the enclosing execution - context; for block and warp lowering, this is per thread block. - instance_index : str - Symbolic rule selecting the current group instance, such as - ``"cta"`` or ``"linear_thread_rank / 16"``. - execution_scope : SynchronizationScope - Scope within which the primitive's participating threads execute. - thread_rank : str, optional - Symbolic rule for the thread's rank within its group. Defaults to - the block's linear thread rank. - """ - - group_kind: str - logical_width: int - instances: int - instance_index: str - execution_scope: SynchronizationScope - thread_rank: str = field(default="linear_thread_rank", kw_only=True) - - def __post_init__(self) -> None: - if not self.group_kind: - raise ValueError("group topology kind must not be empty") - for name in ("logical_width", "instances"): - value = getattr(self, name) - if ( - not isinstance(value, int) - or isinstance(value, bool) - or value < 1 - ): - raise ValueError( - f"group topology {name} must be a positive integer" - ) - for name in ("instance_index", "thread_rank"): - value = getattr(self, name) - if not isinstance(value, str) or not value: - label = name.replace("_", " ") - raise ValueError( - f"group topology {label} must be a non-empty string" - ) - object.__setattr__( - self, - "execution_scope", - SynchronizationScope(self.execution_scope), - ) - - -class UnsupportedReasonCode(str, Enum): - """Stable categories for unsupported planning and resolution outcomes. - - These distinguish missing launch facts, incomplete membership, and - unsupported group or operation forms without matching diagnostic text. - """ - - MISSING_EXACT_BLOCK_DIM = "missing_exact_block_dim" - PARTIAL_PHYSICAL_WARP = "partial_physical_warp" - GROUP_KIND = "group_kind" - OPERAND_FORM = "operand_form" - CUB_BROADCAST = "cub_broadcast" - OPERATION_VARIANT = "operation_variant" - LAUNCH_CAPABILITY = "launch_capability" - - -class GroupOperationSemantics(Protocol): - """Operation facts required by the common group-call model. - - A primitive-family record supplies its semantic identity, result - visibility, and whether it returns a value. Its concrete type also has - to be registered with the dispatcher, which provides that family's - argument classification and planning functions. Satisfying this protocol - alone does not register a new operation. - """ - - @property - def semantic_key(self) -> tuple[Any, ...]: ... - - @property - def result_visibility(self) -> ResultVisibility: ... - - @property - def returns_value(self) -> bool: ... - - -def _requested_result_visibility( - operation: GroupOperationSemantics, -) -> ResultVisibility: - return ResultVisibility(operation.result_visibility) - - -def _group_key(group: ThreadGroup) -> tuple[Any, ...]: - """Identify a group using the hierarchy dimensions relevant to its kind. - - Physical warps share the same logical identity across enclosing block - shapes. Artifact identity separately retains the exact block dimensions - needed for lowering and scratch layout. - """ - - hierarchy = group.hierarchy - assert hierarchy is not None - if group.kind == "warp": - return "warp", "physical", 32 - if group.kind == "block": - return "block", hierarchy.block_dim - if group.kind == "cluster": - return "cluster", hierarchy.block_dim, hierarchy.cluster_dim - if group.kind == "grid": - return ( - "grid", - hierarchy.block_dim, - hierarchy.cluster_dim, - hierarchy.grid_dim, - ) - if group.kind in MAPPED_GROUP_KINDS: - return group.semantic_key - return (group.kind,) - - -@dataclass(frozen=True, eq=False) -class GroupPrimitiveCall: - """A requested group operation before selecting its implementation. - - Frontends provide a compile-time group description and a registered - operation record. The dispatcher derives the argument classifications; - the call itself contains no runtime operands or compiled provider. - ``plan_group_primitive`` combines it with launch facts to resolve the - group and select an implementation. - - Attributes - ---------- - group : ThreadGroup - Requested participating threads. Launch-dependent dimensions may - still need to be resolved. - operation : GroupOperationSemantics - Primitive-family options, including static bindings and requested - result behavior. - argument_classifications : tuple of ParameterClassification - Names, static/runtime kinds, and roles derived by the registered - operation family. Computed during construction. - - Notes - ----- - Equality and hashing use the requested group, the operation's semantic - key, and its requested result visibility. Backend provider selection and - executable-artifact details enter the later lowering plan. - """ - - group: ThreadGroup - operation: GroupOperationSemantics - argument_classifications: tuple[ParameterClassification, ...] = field( - init=False - ) - - def __post_init__(self) -> None: - if not isinstance(self.group, ThreadGroup): - raise TypeError("GroupPrimitiveCall group must be a ThreadGroup") - from ._dispatch import _call_classifications, _is_group_operation - - if not _is_group_operation(self.operation): - raise TypeError("unsupported GroupPrimitiveCall operation") - object.__setattr__( - self, - "argument_classifications", - _call_classifications(self.operation), - ) - - @property - def semantic_key(self) -> tuple[Any, ...]: - return ( - _group_key(self.group), - self.operation.semantic_key, - _requested_result_visibility(self.operation).value, - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, GroupPrimitiveCall): - return NotImplemented - return self.semantic_key == other.semantic_key - - def __hash__(self) -> int: - return hash(self.semantic_key) - - -@dataclass(frozen=True) -class ArgumentPrecondition: - """An inclusive scalar bound and who is responsible for checking it. - - Planners use these records for controls such as a tile's ``valid_items`` - count or a pointer offset. Construction checks the bounds' shape and - ordering; it does not receive or validate the argument's actual value. - - Attributes - ---------- - name : str - Nonempty name of the argument constrained by this record. - minimum, maximum : int or None - Inclusive bounds. ``None`` leaves that end unconstrained here. - enforcement : PreconditionEnforcement - Whether planning already checked a static value or a runtime - caller must satisfy the bound. - """ - - name: str - minimum: int | None - maximum: int | None - enforcement: PreconditionEnforcement - - def __post_init__(self) -> None: - if not self.name: - raise ValueError("argument precondition name must not be empty") - object.__setattr__( - self, - "enforcement", - PreconditionEnforcement(self.enforcement), - ) - for bound_name, bound in ( - ("minimum", self.minimum), - ("maximum", self.maximum), - ): - if bound is not None and ( - not isinstance(bound, int) or isinstance(bound, bool) - ): - raise TypeError(f"{bound_name} must be an integer or None") - if ( - self.minimum is not None - and self.maximum is not None - and self.minimum > self.maximum - ): - raise ValueError("argument precondition minimum exceeds maximum") - - -@dataclass(frozen=True) -class ParticipationRequirements: - """Membership and argument requirements for executing a group primitive. - - Backends use these facts to check that their provider and launch model - match the selected group. Runtime requirements, such as converged entry - and uniform arguments, remain obligations of the generated call; this - record does not inspect running threads or insert validation code. - - Attributes - ---------- - group_kind : str - Kind of the resolved participating group. - exact_group_size : int - Number of participating threads in one group instance. - exact_block_dim : Dim3 or None - Exact enclosing launch shape when known, including all three axes. - complete_membership : bool - Whether the group has its full required membership. - contiguous, aligned : bool - Whether membership is contiguous in thread rank and aligned to - the group's partition boundaries. - converged_entry : bool - Whether all participating threads must enter the primitive together. - complete_parent_partition : bool - Whether groups form complete partitions of the relevant parent. - uniform_arguments : tuple of str, optional - Argument names whose values must agree across participating members. - valid_member_selection : str or None, optional - Description of how a guarded operation selects valid data, such as - ``"first valid_items tile elements"``. A partial data tile still - requires the group's participating threads. - argument_preconditions : tuple of ArgumentPrecondition, optional - Static or caller-enforced scalar bounds. Names must be unique. - """ - - group_kind: str - exact_group_size: int - exact_block_dim: Dim3 | None - complete_membership: bool - contiguous: bool - aligned: bool - converged_entry: bool - complete_parent_partition: bool - uniform_arguments: tuple[str, ...] = () - valid_member_selection: str | None = None - argument_preconditions: tuple[ArgumentPrecondition, ...] = () - - def __post_init__(self) -> None: - object.__setattr__( - self, "uniform_arguments", tuple(self.uniform_arguments) - ) - object.__setattr__( - self, - "argument_preconditions", - tuple(self.argument_preconditions), - ) - if any( - not isinstance(precondition, ArgumentPrecondition) - for precondition in self.argument_preconditions - ): - raise TypeError( - "argument_preconditions must contain " - "ArgumentPrecondition records" - ) - names = [ - precondition.name for precondition in self.argument_preconditions - ] - if len(names) != len(set(names)): - raise ValueError("argument precondition names must be unique") - - -@dataclass(frozen=True) -class SynchronizationRequirements: - """Entry and scratch-reuse synchronization required by the plan. - - Attributes - ---------- - converged_entry : bool - Whether participating threads must reach the operation together. - Must agree with the plan's participation requirements. - storage_reuse_barrier : SynchronizationScope - Scope of synchronization required before scratch is reused, or - ``NONE`` when no automatic reuse barrier is requested. This does not - describe every synchronization performed inside the native primitive. - """ - - converged_entry: bool - storage_reuse_barrier: SynchronizationScope - - -@dataclass(frozen=True) -class TempStorageRequirements: - """Describe scratch ownership, per-group layout, and reuse policy. - - A storage-bearing primitive may need one independent scratch instance - per participating group. Planning records that topology and any caller - requests; backend materialization discovers the concrete layout and - lowering arranges the storage binding. Constructing this record does not - allocate memory or establish that a requested capacity is sufficient. - - Attributes - ---------- - ownership : StorageOwnership - Storage-free, internally arranged, or caller-supplied storage. - address_space : str or None - Required memory space, such as ``"shared"`` for CUB scratch. - cpp_type : str or None - Native scratch type when already named. ``None`` can defer its - resolution until the specialized implementation is materialized. - instances : int or None - Positive scratch-instance count for a storage-bearing operation. - instance_index : str or None - Symbolic rule selecting the current group's scratch instance. - exact_layout_required : bool - Whether the supplied storage must be checked against the concrete - implementation layout, as required for caller-owned storage. - sharing : {"shared", "exclusive"} or None, optional - Caller storage-sharing policy: reusable with other compatible - consumers or reserved exclusively. ``None`` for implementation-owned - storage; distinct from the ``"shared"`` address space. - requested_size_in_bytes : int or None, optional - Positive caller-requested capacity, or no explicit capacity request. - requested_alignment : int or None, optional - Positive caller-requested byte alignment, or no explicit request. - auto_sync : bool, optional - Whether lowering should arrange scratch-reuse synchronization. - Defaults to ``True``; storage-free operations must set it to ``False``. - - Notes - ----- - Storage-free requirements carry no layout, sharing, or size requests. - Implementation-owned storage carries no caller sharing, size, or - alignment requests. Caller-owned storage must select a sharing policy. - """ - - ownership: StorageOwnership - address_space: str | None - cpp_type: str | None - instances: int | None - instance_index: str | None - exact_layout_required: bool - sharing: str | None = None - requested_size_in_bytes: int | None = None - requested_alignment: int | None = None - auto_sync: bool = True - - def __post_init__(self) -> None: - """Check that layout and reuse requests agree with storage ownership. - - A storage-free operation cannot carry layout requests. An internally - managed allocation needs group-instance information, while - caller-owned storage also needs a sharing policy. Concrete capacity is - checked later. - """ - - object.__setattr__(self, "ownership", StorageOwnership(self.ownership)) - if self.sharing not in {None, "shared", "exclusive"}: - raise ValueError( - "temporary storage sharing must be shared or exclusive" - ) - if not isinstance(self.auto_sync, bool): - raise TypeError("auto_sync must be a bool") - if self.ownership is StorageOwnership.NONE: - if any( - value is not None - for value in ( - self.address_space, - self.cpp_type, - self.instances, - self.instance_index, - self.sharing, - self.requested_size_in_bytes, - self.requested_alignment, - ) - ): - raise ValueError( - "storage-free requirements cannot carry storage layout" - ) - if self.exact_layout_required: - raise ValueError( - "storage-free requirements cannot require an exact layout" - ) - if self.auto_sync: - raise ValueError( - "storage-free requirements cannot request automatic sync" - ) - else: - if ( - not isinstance(self.instances, int) - or isinstance(self.instances, bool) - or self.instances < 1 - ): - raise ValueError( - "storage-bearing operations require " - "a positive instance count" - ) - if ( - not isinstance(self.instance_index, str) - or not self.instance_index - ): - raise ValueError( - "storage-bearing operations require " - "a non-empty instance index" - ) - if self.ownership is StorageOwnership.IMPLEMENTATION: - if self.sharing is not None: - raise ValueError( - "implementation-owned storage has no sharing mode" - ) - if self.requested_size_in_bytes is not None: - raise ValueError( - "implementation-owned storage has no requested size" - ) - if self.requested_alignment is not None: - raise ValueError( - "implementation-owned storage has no requested alignment" - ) - elif self.ownership is StorageOwnership.CALLER and self.sharing is None: - raise ValueError("caller-owned storage requires a sharing mode") - for name in ("requested_size_in_bytes", "requested_alignment"): - value = getattr(self, name) - if value is not None and ( - not isinstance(value, int) - or isinstance(value, bool) - or value <= 0 - ): - raise ValueError(f"{name} must be a positive integer or None") - - -@dataclass(frozen=True) -class GroupExecutionRequirements: - """Group topology, participation, synchronization, and scratch requirements. - - Operation planners attach these records to the selected implementation's - lowering plan so backends can arrange execution and temporary storage. - """ - - topology: GroupTopologyRequirements - participation: ParticipationRequirements - synchronization: SynchronizationRequirements - temp_storage: TempStorageRequirements - - -@dataclass(frozen=True) -class ImplementationProvenance: - """Identify the native library entry point chosen for a group operation. - - A logical request such as a group load does not by itself identify the - native primitive that will execute it. The planner records that choice - here so a backend can recognize the implementation and route it to a - supported provider. For example, a CUB block load records ``"CUB"``, - ``"cub/block/block_load.cuh"``, ``"cub::BlockLoad"``, and ``"Load"``. - A backend can use this tuple to select its load or store route, then check - the provider against the plan's requirements. - - The same native entry point can serve many element types, tile sizes, - and algorithm choices. Those details belong to the plan's specialized - ``implementation`` and its metadata. Provenance supplies the entry-point - identity; both contribute to ``GroupLoweringPlan.artifact_key`` so the - identity retains the chosen native implementation as well as its - specialization and execution requirements. - - Attributes - ---------- - library : str - Underlying native library, such as ``"CUB"``. The backend compiler - or Python provider factory is selected separately. - header : str - Native header declaring the chosen primitive. - cpp_class : str - Qualified C++ class or namespace containing the operation, such as - ``"cub::BlockLoad"``. - method : str - Native operation name invoked by the implementation, such as - ``"Load"``. - - Notes - ----- - This record identifies an implementation choice. It does not record a - library version, compiled binary, or proof that a backend supports the - choice. A backend can reject an unrecognized tuple; planning failures - are represented separately by ``UnsupportedReason``. - """ - - library: str - header: str - cpp_class: str - method: str - - @property - def semantic_key(self) -> tuple[str, str, str, str]: - """Return the library/header/implementation/method identity.""" - - return self.library, self.header, self.cpp_class, self.method - - -@dataclass(frozen=True) -class UnsupportedReason: - """A machine-readable failure category with a human explanation. - - ``code`` identifies the unsupported case; ``message`` supplies details - for diagnostics and ``require_supported`` exceptions. Messages are - excluded from equality and hashing, allowing wording to change without - changing the failure's identity. - """ - - code: UnsupportedReasonCode - message: str = field(compare=False, hash=False) - - -@dataclass(frozen=True) -class ThreadGroupLaunchResolution: - """A thread group reconciled with launch facts, or a resolution failure. - - Resolution establishes dimensions and membership needed by subsequent - operation planning. It does not choose a primitive implementation. - - Attributes - ---------- - group : ThreadGroup - Resolved group on success, or the group retained with failure context. - unsupported : UnsupportedReason or None, optional - Why the launch facts cannot support this group, or ``None`` when - resolution succeeded. - """ - - group: ThreadGroup - unsupported: UnsupportedReason | None = None - - def __post_init__(self) -> None: - if not isinstance(self.group, ThreadGroup): - raise TypeError( - "ThreadGroupLaunchResolution group must be a ThreadGroup" - ) - if self.unsupported is not None and not isinstance( - self.unsupported, UnsupportedReason - ): - raise TypeError( - "ThreadGroupLaunchResolution unsupported must be an " - "UnsupportedReason" - ) - - def require_supported(self) -> ThreadGroup: - """Return the resolved group or raise its ``NotImplementedError``.""" - - if self.unsupported is not None: - raise NotImplementedError(self.unsupported.message) - return self.group - - -@dataclass(frozen=True, eq=False) -class GroupLoweringPlan: - """An implementation choice and its execution requirements. - - Operation planners construct this after resolving the requested group - against launch facts. A supported plan combines a specialized native - implementation with topology, participation, synchronization, storage, - and provenance. Backends use that information to select compatible - providers and produce code; a supported planning outcome alone does not - mean that compilation has occurred or every backend can lower the plan. - - Construction checks consistency between the resolved group and its - topology, participation, and convergence requirements. Unsupported plans - instead retain a typed reason and can be inspected before - ``require_supported`` turns that reason into an exception. - - Attributes - ---------- - target : GroupLoweringTarget - Chosen implementation family, or ``UNSUPPORTED``. - call : GroupPrimitiveCall - Original request, retained alongside the resolved group. - resolved_group : ThreadGroup - Group after applying launch facts, or the group retained on failure. - implementation : Algorithm or None - Specialized primitive description; ``None`` for an unsupported plan. - topology : GroupTopologyRequirements or None - Group instances and rank rules used for indexing and execution. - participation : ParticipationRequirements or None - Required membership, launch shape, uniformity, and scalar bounds. - result : None - Load/Store writes its results through the supplied item array or - memory pointer and has no separate returned value. - synchronization : SynchronizationRequirements or None - Converged-entry and scratch-reuse requirements. - temp_storage : TempStorageRequirements or None - Scratch ownership, layout requests, and automatic reuse policy. - provenance : ImplementationProvenance or None - Native library entry point used for routing and artifact identity. - unsupported : UnsupportedReason or None, optional - Required for ``UNSUPPORTED`` and absent for a supported plan. - - Notes - ----- - Supported plans specify all lowering requirements. Their equality and - hashing use ``artifact_key``. Unsupported plans have no artifact key; - they compare by logical semantics and reason code, excluding diagnostic - wording. ``semantic_key`` alone is not enough to identify an executable - implementation because it omits provider and storage choices. - """ - - target: GroupLoweringTarget - call: GroupPrimitiveCall - resolved_group: ThreadGroup - implementation: Algorithm | None - topology: GroupTopologyRequirements | None - participation: ParticipationRequirements | None - result: None - synchronization: SynchronizationRequirements | None - temp_storage: TempStorageRequirements | None - provenance: ImplementationProvenance | None - unsupported: UnsupportedReason | None = None - - def __post_init__(self) -> None: - """Require a complete implementation plan or an unsupported reason. - - For a supported plan, compare the execution requirements with the - resolved group. This catches inconsistent group sizes, scratch - indexing, launch dimensions, and convergence requirements before a - backend consumes them. - """ - - is_unsupported = self.target is GroupLoweringTarget.UNSUPPORTED - if is_unsupported != (self.unsupported is not None): - raise ValueError("unsupported plans require exactly one reason") - if not is_unsupported and ( - self.implementation is None - or self.topology is None - or self.participation is None - or self.synchronization is None - or self.temp_storage is None - or self.provenance is None - ): - raise ValueError( - "supported plans require complete lowering requirements" - ) - if not is_unsupported: - assert self.topology is not None - assert self.participation is not None - resolved_kind = self.resolved_group.kind - if ( - self.topology.group_kind != resolved_kind - or self.participation.group_kind != resolved_kind - ): - raise ValueError( - "supported plan group requirements must match " - "the resolved group kind" - ) - resolved_size = self.resolved_group.static_size - if resolved_size is None or ( - self.topology.logical_width != resolved_size - or self.participation.exact_group_size != resolved_size - ): - raise ValueError( - "supported plan group widths must match " - "the resolved group size" - ) - resolved_block_dim = self.resolved_group.block_dim - if ( - resolved_block_dim is not None - and self.participation.exact_block_dim != resolved_block_dim - ): - raise ValueError( - "supported plan block dimensions must match " - "the resolved group" - ) - from ._execution_requirements import _group_topology - - expected_topology = _group_topology( - self.resolved_group, - LaunchFacts(exact_block_dim=self.participation.exact_block_dim), - ) - if self.topology != expected_topology: - raise ValueError( - "supported plan topology must match the resolved group" - ) - complete_membership = ( - self.resolved_group.complete_membership is not False - ) - complete_parent_partition = ( - resolved_kind == "warp" - or self.resolved_group.complete_membership is True - ) - if ( - self.participation.complete_membership - is not complete_membership - or self.participation.complete_parent_partition - is not complete_parent_partition - ): - raise ValueError( - "supported plan participation must match the resolved group" - ) - assert self.synchronization is not None - if ( - self.participation.converged_entry - is not self.synchronization.converged_entry - ): - raise ValueError( - "supported plan participation and synchronization " - "must agree on converged entry" - ) - - @property - def semantic_key(self) -> tuple[Any, ...]: - """Identify the resolved request without its provider choice. - - Physical warp requests can share this key across enclosing block - shapes. Use ``artifact_key`` when execution and storage details must - also distinguish plans. - """ - - return ( - _group_key(self.resolved_group), - self.call.operation.semantic_key, - None, - ) - - @property - def artifact_key(self) -> tuple[Any, ...] | None: - """Identify a supported implementation and its lowering requirements. - - Include exact block dimensions, the specialized implementation, - execution/storage requirements, and native provenance. Backends can use - this identity when reusing generated artifacts, together with their - compiler- and target-specific cache inputs. Unsupported plans return - ``None`` because they have no executable implementation to reuse. - """ - - if self.unsupported is not None: - return None - implementation_key = ( - None - if self.implementation is None - else self.implementation.semantic_key - ) - return ( - self.target.value, - _group_key(self.resolved_group), - self.resolved_group.hierarchy.block_dim, - self.topology, - self.call.operation.semantic_key, - implementation_key, - self.participation, - self.result, - self.synchronization, - self.temp_storage, - None if self.provenance is None else self.provenance.semantic_key, - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, GroupLoweringPlan): - return NotImplemented - return self._identity_key == other._identity_key - - def __hash__(self) -> int: - return hash(self._identity_key) - - @property - def _identity_key(self) -> tuple[Any, ...]: - # Diagnostic prose can evolve without changing an unsupported request; - # successful requests must retain every artifact-affecting requirement. - if self.artifact_key is not None: - return "artifact", self.artifact_key - assert self.unsupported is not None - return "unsupported", self.semantic_key, self.unsupported.code.value - - def require_supported(self) -> GroupLoweringPlan: - """Return a supported plan or report its reason in an exception. - - Unsupported plans raise ``NotImplementedError`` with their message. - """ - - if self.unsupported is not None: - raise NotImplementedError(self.unsupported.message) - return self - - -__all__ = [ - "ArgumentPrecondition", - "GroupExecutionRequirements", - "GroupLoweringPlan", - "GroupLoweringTarget", - "GroupOperationSemantics", - "GroupPrimitiveCall", - "GroupTopologyRequirements", - "ImplementationProvenance", - "ParticipationRequirements", - "PreconditionEnforcement", - "ResultVisibility", - "StorageOwnership", - "SynchronizationRequirements", - "SynchronizationScope", - "TempStorageRequirements", - "ThreadGroupLaunchResolution", - "UnsupportedReason", - "UnsupportedReasonCode", -] diff --git a/python/cuda_coop/cuda/coop/_core/group/_resolution.py b/python/cuda_coop/cuda/coop/_core/group/_resolution.py deleted file mode 100644 index 24549c34227b..000000000000 --- a/python/cuda_coop/cuda/coop/_core/group/_resolution.py +++ /dev/null @@ -1,304 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Resolve group descriptions using the kernel's launch facts. - -A descriptor can name the current block before its dimensions are known. -Resolution fills in the exact dimensions needed for that group or for a -query about an enclosing group. It also checks complete physical warps -and mapped-group partitions. - -Missing facts and unsupported shapes produce a structured reason that an -operation planner can report. Contradictory explicit dimensions raise an -error. Cluster and grid resolution also require verified cluster-launch state. -Operation dispatch adds capability checks specific to the operation, such as -cooperative-launch support for grid operations. -""" - -from __future__ import annotations - -from ..launch import LaunchFacts -from ..thread_group import ( - COMPLETE_WARP_GROUP_KINDS, - ThreadGroup, - ThreadHierarchy, - normalize_thread_level, -) -from ._execution_requirements import _unsupported -from ._model import ( - GroupLoweringPlan, - GroupPrimitiveCall, - ThreadGroupLaunchResolution, - UnsupportedReason, - UnsupportedReasonCode, -) - -_THREAD_LEVEL_ORDER = { - "thread": 0, - "warp": 1, - "block": 2, - "cluster": 3, - "grid": 4, -} - -_MAPPED_PARENT_LEVEL = { - "threads_within_warp": "warp", - "warps_within_block": "block", -} - - -def _resolution_failure( - group: ThreadGroup, - code: UnsupportedReasonCode, - message: str, -) -> ThreadGroupLaunchResolution: - """Keep the requested group with a reason that prevents its resolution.""" - - return ThreadGroupLaunchResolution( - group=group, - unsupported=UnsupportedReason(code=code, message=message), - ) - - -def resolve_thread_group( - group: ThreadGroup, - launch: LaunchFacts, - *, - through_level: str | None = None, -) -> ThreadGroupLaunchResolution: - """Fill a group's required hierarchy from exact launch facts. - - The group's own level determines how much hierarchy an operation needs. - ``through_level`` can request more, for example when a group query asks - for a count at an enclosing level. A mapped group needs the dimensions - of its physical parent. - - Exact dimensions remain distinct from upper bounds. Cluster and grid - resolution require verified cluster-launch state. A verified non-cluster - launch lets a missing cluster shape resolve to one block. Each grid axis's - block count must be divisible by the cluster's block count on that axis. - The resolved hierarchy stores the resulting cluster counts. - - Warp-based groups require complete 32-thread physical warps in the - block. A mapped group's count must fit its parent, and an exhaustive - mapping must divide the parent's unit count exactly. - - Parameters - ---------- - group : ThreadGroup - Descriptor to resolve. Any dimensions already present must agree - with the required launch dimensions. - launch : LaunchFacts - Exact dimensions and capability evidence for this compilation. - through_level : str, optional - Enclosing level needed by a group query. Primitive planners omit - this argument when the group's own level is sufficient. - - Returns - ------- - ThreadGroupLaunchResolution - The resolved group, or the original group with an unsupported - reason. Missing launch facts and unsupported partitions use this - result instead of raising an exception. - - Raises - ------ - TypeError - The group or launch argument has the wrong type. - ValueError - The requested level is invalid or an existing group dimension - contradicts the exact launch. - """ - - if not isinstance(group, ThreadGroup): - raise TypeError("group must be a ThreadGroup") - if not isinstance(launch, LaunchFacts): - raise TypeError("launch must be LaunchFacts") - - group_level = _MAPPED_PARENT_LEVEL.get(group.kind, group.kind) - if through_level is not None: - through_level = normalize_thread_level( - through_level, - scope="resolve_thread_group", - feature="through_level", - ) - required_level = max( - (group_level, through_level or group_level), - key=_THREAD_LEVEL_ORDER.__getitem__, - ) - needs_complete_warp = ( - group.kind in COMPLETE_WARP_GROUP_KINDS or through_level == "warp" - ) - if required_level == "thread": - return ThreadGroupLaunchResolution(group) - - exact_block_dim = launch.exact_block_dim - if exact_block_dim is None: - return _resolution_failure( - group, - UnsupportedReasonCode.MISSING_EXACT_BLOCK_DIM, - "group operation requires exact block dimensions; max_block_dim " - "is only an upper bound", - ) - assert group.hierarchy is not None - if ( - group.hierarchy.block_dim is not None - and group.hierarchy.block_dim != exact_block_dim - ): - raise ValueError( - f"group block dimensions {group.hierarchy.block_dim!r} do not " - f"match the exact kernel launch dimensions {exact_block_dim!r}", - ) - needs_cluster = ( - _THREAD_LEVEL_ORDER[required_level] >= _THREAD_LEVEL_ORDER["cluster"] - ) - exact_cluster_dim = launch.exact_cluster_dim if needs_cluster else None - if needs_cluster: - cluster_launch_verified = launch.is_verified("cluster_launch") - if exact_cluster_dim is None: - if ( - launch.cluster_launch is not False - or not cluster_launch_verified - ): - return _resolution_failure( - group, - UnsupportedReasonCode.LAUNCH_CAPABILITY, - "cluster and grid group operations require exact static " - "cluster dimensions, or a backend-verified " - "non-cluster launch", - ) - exact_cluster_dim = (1, 1, 1) - elif launch.cluster_launch is None or not cluster_launch_verified: - return _resolution_failure( - group, - UnsupportedReasonCode.LAUNCH_CAPABILITY, - "cluster and grid group operations require backend-verified " - "cluster launch state", - ) - elif ( - exact_cluster_dim != (1, 1, 1) and launch.cluster_launch is not True - ): - return _resolution_failure( - group, - UnsupportedReasonCode.LAUNCH_CAPABILITY, - "multi-block cluster operations require verified " - "cluster launch capability", - ) - - hierarchy_grid_dim = None - needs_grid = required_level == "grid" - if needs_grid: - exact_grid_dim = launch.exact_grid_dim - if exact_grid_dim is None: - return _resolution_failure( - group, - UnsupportedReasonCode.LAUNCH_CAPABILITY, - "grid group operations require exact static grid dimensions", - ) - assert exact_cluster_dim is not None - if any( - grid_extent % cluster_extent != 0 - for grid_extent, cluster_extent in zip( - exact_grid_dim, - exact_cluster_dim, - ) - ): - return _resolution_failure( - group, - UnsupportedReasonCode.LAUNCH_CAPABILITY, - "physical CTA grid dimensions must be divisible by " - "the cluster dimensions", - ) - hierarchy_grid_dim = tuple( - grid_extent // cluster_extent - for grid_extent, cluster_extent in zip( - exact_grid_dim, - exact_cluster_dim, - ) - ) - - resolved_hierarchy = ThreadHierarchy._resolved( - block_dim=exact_block_dim, - cluster_dim=(exact_cluster_dim if needs_cluster else None), - grid_dim=hierarchy_grid_dim, - ) - if ( - needs_cluster - and group.hierarchy.cluster_dim is not None - and (group.hierarchy.cluster_dim != resolved_hierarchy.cluster_dim) - ): - raise ValueError( - f"group cluster dimensions {group.hierarchy.cluster_dim!r} do not " - f"match exact launch dimensions {resolved_hierarchy.cluster_dim!r}" - ) - if ( - needs_grid - and group.hierarchy.grid_dim is not None - and (group.hierarchy.grid_dim != resolved_hierarchy.grid_dim) - ): - raise ValueError( - f"group grid dimensions {group.hierarchy.grid_dim!r} do not match " - f"exact hierarchy dimensions {resolved_hierarchy.grid_dim!r}" - ) - if needs_complete_warp and launch.exact_block_threads % 32 != 0: # type: ignore[operator] - return _resolution_failure( - group, - UnsupportedReasonCode.PARTIAL_PHYSICAL_WARP, - "physical-warp operation requires complete 32-thread warps and " - "every physical warp in the enclosing CTA to be complete; got " - f"{launch.exact_block_threads} block threads", - ) - if group.mapping is not None: - parent_units = ( - 32 - if group.kind == "threads_within_warp" - else launch.exact_block_threads // 32 # type: ignore[operator] - ) - if group.mapping.count > parent_units: - return _resolution_failure( - group, - UnsupportedReasonCode.GROUP_KIND, - "mapped group count exceeds the resolved parent unit count", - ) - if group.mapping.exhaustive and parent_units % group.mapping.count != 0: - return _resolution_failure( - group, - UnsupportedReasonCode.GROUP_KIND, - "exhaustive mapped group count must divide " - "the resolved parent unit count", - ) - resolved = group.with_hierarchy( - resolved_hierarchy, - source="launch_facts", - ) - return ThreadGroupLaunchResolution(resolved) - - -def _resolve_group( - call: GroupPrimitiveCall, - launch: LaunchFacts, -) -> tuple[ThreadGroup, GroupLoweringPlan | None]: - """Resolve a primitive's group and convert failure to an operation plan. - - Return ``(resolved_group, None)`` on success. Otherwise return the group - with an unsupported plan that retains the call and the resolver's reason. - Dispatch returns this plan directly, so a resolution failure has the same - form as an unsupported result from a family planner. - """ - - resolution = resolve_thread_group(call.group, launch) - if resolution.unsupported is None: - return resolution.group, None - return resolution.group, _unsupported( - call, - resolution.group, - resolution.unsupported.code, - resolution.unsupported.message, - ) - - -__all__ = [ - "_resolve_group", - "resolve_thread_group", -] diff --git a/python/cuda_coop/cuda/coop/_core/group/load_store.py b/python/cuda_coop/cuda/coop/_core/group/load_store.py deleted file mode 100644 index 574f07c796dd..000000000000 --- a/python/cuda_coop/cuda/coop/_core/group/load_store.py +++ /dev/null @@ -1,612 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Plan a group's memory transfer using a CUB Block or Warp implementation. - -The operation record describes the requested transfer without compiler values. -After group resolution, the planner selects a CUB specialization and records -participation, argument bounds, and scratch requirements. Compiler adapters -use that plan to generate the call and provide any required storage. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from enum import Enum -from numbers import Integral -from typing import Any - -from .._bindings import ( - ArgumentBinding, - BindingKind, - normalize_i32_binding, - normalize_i64_binding, -) -from .._symbols import semantic_token -from .._types import ArgumentKind, ParameterClassification, ParameterRole -from ..block.load_store import ( - BlockLoadStoreAlgorithm, - make_block_load_specialization, - make_block_store_specialization, -) -from ..launch import LaunchFacts -from ..thread_group import ThreadGroup -from ..warp.load_store import ( - WarpLoadStoreAlgorithm, - make_warp_load_specialization, - make_warp_store_specialization, -) -from ._dispatch import _register_group_operation_family -from ._execution_requirements import ( - _build_execution_requirements, - _unsupported, - _unsupported_cub_warp_width, -) -from ._model import ( - ArgumentPrecondition, - GroupLoweringPlan, - GroupLoweringTarget, - GroupPrimitiveCall, - ImplementationProvenance, - PreconditionEnforcement, - ResultVisibility, - StorageOwnership, - UnsupportedReasonCode, -) - - -class GroupLoadStoreKind(str, Enum): - """Distinguish a load into per-thread items from a store to memory.""" - - LOAD = "load" - STORE = "store" - - -class GroupLoadStoreAlgorithm(str, Enum): - """Name transfer algorithms before the planner selects Block or Warp CUB. - - The set includes both families' choices. The planner checks whether the - resolved group supports the requested choice, including warp width limits. - """ - - DIRECT = "direct" - STRIPED = "striped" - VECTORIZE = "vectorize" - TRANSPOSE = "transpose" - WARP_TRANSPOSE = "warp_transpose" - WARP_TRANSPOSE_TIMESLICED = "warp_transpose_timesliced" - - -_STORAGE_FREE_ALGORITHMS = frozenset( - { - GroupLoadStoreAlgorithm.DIRECT, - GroupLoadStoreAlgorithm.STRIPED, - GroupLoadStoreAlgorithm.VECTORIZE, - } -) - - -@dataclass(frozen=True, eq=False) -class GroupLoadStoreSemantics: - """Describe one transfer and the values available when planning it. - - ``dtype`` and ``items_per_thread`` determine each thread's payload. - Argument bindings distinguish omitted arguments, known constants, and - runtime inputs. Storage fields carry the caller's requests until the - selected algorithm's scratch needs are known. Direct, striped, and - vectorized transfers need no scratch, so their semantic keys exclude those - storage requests. - """ - - kind: GroupLoadStoreKind - dtype: Any - items_per_thread: int - algorithm: GroupLoadStoreAlgorithm = GroupLoadStoreAlgorithm.DIRECT - valid_items: ArgumentBinding = field( - default_factory=ArgumentBinding.omitted - ) - oob_default: ArgumentBinding = field( - default_factory=ArgumentBinding.omitted - ) - offset: ArgumentBinding = field(default_factory=ArgumentBinding.omitted) - storage_ownership: StorageOwnership = StorageOwnership.IMPLEMENTATION - storage_sharing: str | None = None - storage_size_in_bytes: int | None = None - storage_alignment: int | None = None - storage_auto_sync: bool = True - - def __post_init__(self) -> None: - """Normalize selectors and check argument and storage requests. - - Check integer representation here. Bounds that depend on group size, - such as ``valid_items``, are checked after the group is resolved. - """ - - object.__setattr__(self, "kind", GroupLoadStoreKind(self.kind)) - object.__setattr__( - self, - "algorithm", - GroupLoadStoreAlgorithm(self.algorithm), - ) - object.__setattr__( - self, - "storage_ownership", - StorageOwnership(self.storage_ownership), - ) - if ( - not isinstance(self.items_per_thread, int) - or isinstance(self.items_per_thread, bool) - or self.items_per_thread <= 0 - ): - raise ValueError("items_per_thread must be a positive integer") - for name in ("valid_items", "oob_default", "offset"): - if not isinstance(getattr(self, name), ArgumentBinding): - raise TypeError(f"{name} must be an ArgumentBinding") - object.__setattr__( - self, - "valid_items", - normalize_i32_binding(self.valid_items, name="valid_items"), - ) - object.__setattr__( - self, - "offset", - normalize_i64_binding(self.offset, name="offset"), - ) - if ( - self.offset.kind is BindingKind.STATIC - and int(self.offset.value) < 0 - ): - raise ValueError("static offset must be nonnegative") - if self.kind is GroupLoadStoreKind.STORE and ( - self.oob_default.kind is not BindingKind.OMITTED - ): - raise ValueError("oob_default is valid only for group load") - if ( - self.oob_default.kind is not BindingKind.OMITTED - and self.valid_items.kind is BindingKind.OMITTED - ): - raise ValueError("oob_default requires valid_items") - if self.storage_sharing not in {None, "shared", "exclusive"}: - raise ValueError("storage_sharing must be shared or exclusive") - for name in ("storage_size_in_bytes", "storage_alignment"): - value = getattr(self, name) - if value is not None and ( - not isinstance(value, int) - or isinstance(value, bool) - or value <= 0 - ): - raise ValueError(f"{name} must be a positive integer or None") - if not isinstance(self.storage_auto_sync, bool): - raise TypeError("storage_auto_sync must be a bool") - if self.storage_ownership is StorageOwnership.NONE: - if self.algorithm not in _STORAGE_FREE_ALGORITHMS: - raise ValueError( - "storage-free group Load/Store is valid only for direct, " - "striped, or vectorize algorithms" - ) - if any( - value is not None - for value in ( - self.storage_sharing, - self.storage_size_in_bytes, - self.storage_alignment, - ) - ): - raise ValueError( - "storage-free operations cannot carry storage layout" - ) - if self.storage_auto_sync: - raise ValueError( - "storage-free operations cannot request automatic sync" - ) - elif self.storage_ownership is StorageOwnership.IMPLEMENTATION: - if any( - value is not None - for value in ( - self.storage_sharing, - self.storage_size_in_bytes, - self.storage_alignment, - ) - ): - raise ValueError( - "implementation-owned storage cannot carry caller requests" - ) - elif self.storage_sharing is None: - raise ValueError("caller-owned storage requires storage_sharing") - - @property - def has_valid_items(self) -> bool: - return self.valid_items.kind is not BindingKind.OMITTED - - @property - def has_oob_default(self) -> bool: - return self.oob_default.kind is not BindingKind.OMITTED - - @property - def has_offset(self) -> bool: - return self.offset.kind is not BindingKind.OMITTED - - @property - def result_visibility(self) -> ResultVisibility: - return ResultVisibility.PER_MEMBER - - @property - def returns_value(self) -> bool: - """Report that the call writes through its output argument.""" - - return False - - @property - def semantic_key(self) -> tuple[Any, ...]: - """Identify the transfer, excluding unused scratch requests.""" - - common = ( - f"group_{self.kind.value}", - semantic_token(self.dtype), - self.items_per_thread, - self.algorithm.value, - self.valid_items.semantic_key, - self.oob_default.semantic_key, - self.offset.semantic_key, - ) - if self.algorithm in _STORAGE_FREE_ALGORITHMS: - return (*common, StorageOwnership.NONE.value) - return ( - *common, - self.storage_ownership.value, - self.storage_sharing, - self.storage_size_in_bytes, - self.storage_alignment, - self.storage_auto_sync, - ) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, GroupLoadStoreSemantics): - return NotImplemented - return self.semantic_key == other.semantic_key - - def __hash__(self) -> int: - return hash(self.semantic_key) - - -def _call_classifications( - operation: GroupLoadStoreSemantics, -) -> tuple[ParameterClassification, ...]: - """Describe each argument's role and when its value is available. - - Both transfers take a memory pointer and per-thread items, but load and - store reverse their input/output roles. Omitted optional arguments do not - enter the call. Record constants separately from runtime inputs. - """ - - classifications = [ - ParameterClassification( - "source" - if operation.kind is GroupLoadStoreKind.LOAD - else "destination", - ArgumentKind.RUNTIME, - ( - ParameterRole.INPUT - if operation.kind is GroupLoadStoreKind.LOAD - else ParameterRole.OUTPUT - ), - ) - ] - classifications.append( - ParameterClassification( - "output" if operation.kind is GroupLoadStoreKind.LOAD else "value", - ArgumentKind.RUNTIME, - ( - ParameterRole.OUTPUT - if operation.kind is GroupLoadStoreKind.LOAD - else ParameterRole.INPUT - ), - ) - ) - for name, binding in ( - ("valid_items", operation.valid_items), - ("oob_default", operation.oob_default), - ("offset", operation.offset), - ): - if binding.argument_kind is None: - continue - classifications.append( - ParameterClassification( - name, - binding.argument_kind, - ( - ParameterRole.CONSTANT - if binding.kind is BindingKind.STATIC - else ParameterRole.INPUT - ), - ) - ) - classifications.append( - ParameterClassification( - "algorithm", - ArgumentKind.STATIC, - ParameterRole.CONSTANT, - ) - ) - return tuple(classifications) - - -def _plan_load_store( - call: GroupPrimitiveCall, - resolved: ThreadGroup, - launch: LaunchFacts, - operation: GroupLoadStoreSemantics, -) -> GroupLoweringPlan: - """Choose a CUB Load/Store specialization for a resolved thread group. - - Called by ``plan_group_primitive`` after group resolution. Select the - block or warp implementation and describe its participation, scratch - storage, and storage-reuse synchronization requirements for later - compilation and lowering. Direct, striped, and vectorized algorithms - require no scratch storage or storage-reuse barrier; other algorithms - retain the requested storage ownership, layout constraints, sharing, and - automatic sync policy. - - Validate static ``valid_items`` against the group tile size and check that - pointer offsets fit signed 64-bit arithmetic. Runtime bounds remain caller - preconditions recorded in the plan. Each physical or logical warp handles - a consecutive tile, so its specialization always takes a runtime pointer - offset. The backend combines that tile's origin with the user offset - preserved in ``operation``; the plan limits the user offset accordingly. - - Parameters - ---------- - call : GroupPrimitiveCall - Original group call, retained in the plan and unsupported diagnostics. - resolved : ThreadGroup - Group resolved against ``launch``. Supported block, physical-warp, and - logical-warp groups must have a static size and complete membership. - launch : LaunchFacts - Launch facts used during group resolution, including exact block - dimensions for specialization and counting warp-group instances. - operation : GroupLoadStoreSemantics - Normalized load or store semantics from ``call.operation``, including - dtype, item count, algorithm, argument bindings, and storage requests. - - Returns - ------- - GroupLoweringPlan - CUB specialization, execution requirements, and implementation - provenance, or an ``UNSUPPORTED`` plan with a reason when the group - kind, warp width, or algorithm variant is unsupported. This builds - metadata; compilation and device-storage allocation happen during - later lowering. - - Raises - ------ - TypeError - Static ``valid_items`` is not an integer or is a boolean. - ValueError - Static ``valid_items`` is outside the group tile, a warp tile origin - or its sum with the user offset exceeds signed 64-bit range, or the - specialization builder rejects an argument binding. - AssertionError - A supported group lacks a static size or exact block dimensions are - missing from ``launch``; group resolution must establish these first. - """ - - if resolved.kind not in {"block", "warp", "threads_within_warp"}: - return _unsupported( - call, - resolved, - UnsupportedReasonCode.GROUP_KIND, - "cuda.coop Load/Store supports this_block(), complete physical " - "this_warp(), and power-of-two logical-warp groups", - ) - group_size = resolved.static_size - assert group_size is not None - tile_items = group_size * operation.items_per_thread - if operation.valid_items.kind is BindingKind.STATIC: - valid_items = operation.valid_items.value - if isinstance(valid_items, bool) or not isinstance( - valid_items, Integral - ): - raise TypeError("static valid_items must be an integer") - valid_items = int(valid_items) - if not 0 <= valid_items <= tile_items: - raise ValueError( - "static valid_items must be between zero and the group tile " - f"size ({tile_items})" - ) - assert launch.exact_block_dim is not None - block_threads = launch.exact_block_threads - assert block_threads is not None - maximum_user_offset = (1 << 63) - 1 - if resolved.kind == "block": - algorithm = BlockLoadStoreAlgorithm(operation.algorithm.value) - if ( - algorithm - in { - BlockLoadStoreAlgorithm.WARP_TRANSPOSE, - BlockLoadStoreAlgorithm.WARP_TRANSPOSE_TIMESLICED, - } - and block_threads % 32 != 0 - ): - return _unsupported( - call, - resolved, - UnsupportedReasonCode.OPERATION_VARIANT, - f"cub::Block{operation.kind.value.title()} algorithm " - f"{operation.algorithm.value!r} requires a block size that is " - "a multiple of 32", - ) - make_specialization = ( - make_block_load_specialization - if operation.kind is GroupLoadStoreKind.LOAD - else make_block_store_specialization - ) - specialization = make_specialization( - dtype=operation.dtype, - block_dim=launch.exact_block_dim, - items_per_thread=operation.items_per_thread, - algorithm=algorithm, - valid_items=operation.valid_items, - oob_default=operation.oob_default, - include_full_tile=False, - include_pointer_offset=operation.offset, - ) - target = GroupLoweringTarget.CUB_BLOCK - header = f"cub/block/block_{operation.kind.value}.cuh" - else: - warp_width, width_error = _unsupported_cub_warp_width(call, resolved) - if width_error is not None: - return width_error - assert warp_width is not None - if operation.storage_ownership is StorageOwnership.CALLER: - return _unsupported( - call, - resolved, - UnsupportedReasonCode.OPERATION_VARIANT, - f"cub::Warp{operation.kind.value.title()} does not accept " - "explicit temp_storage", - ) - try: - algorithm = WarpLoadStoreAlgorithm(operation.algorithm.value) - except ValueError: - return _unsupported( - call, - resolved, - UnsupportedReasonCode.OPERATION_VARIANT, - f"cub::Warp{operation.kind.value.title()} does not support " - f"algorithm {operation.algorithm.value!r}", - ) - make_specialization = ( - make_warp_load_specialization - if operation.kind is GroupLoadStoreKind.LOAD - else make_warp_store_specialization - ) - specialization = make_specialization( - dtype=operation.dtype, - items_per_thread=operation.items_per_thread, - algorithm=algorithm, - threads_in_warp=warp_width, - valid_items=operation.valid_items, - oob_default=operation.oob_default, - include_full_tile=False, - # Every physical or logical Warp group receives a consecutive tile. - # The backend combines this runtime ABI argument with the preserved - # user offset recorded on ``operation``. - include_pointer_offset=ArgumentBinding.runtime(), - ) - target = GroupLoweringTarget.CUB_WARP - header = f"cub/warp/warp_{operation.kind.value}.cuh" - group_instances = block_threads // warp_width - maximum_tile_origin = (group_instances - 1) * tile_items - if maximum_tile_origin > maximum_user_offset: - raise ValueError( - "warp-group tile origin must fit a signed 64-bit offset" - ) - maximum_user_offset -= maximum_tile_origin - if operation.offset.kind is BindingKind.STATIC and ( - int(operation.offset.value) > maximum_user_offset - ): - raise ValueError( - "static offset plus the warp-group tile origin must fit a " - "signed 64-bit integer" - ) - cpp_class = f"cub::{specialization.struct_name}" - storage_free = operation.algorithm in _STORAGE_FREE_ALGORITHMS - requirements = _build_execution_requirements( - resolved, - launch, - storage_ownership=( - StorageOwnership.NONE - if storage_free - else operation.storage_ownership - ), - cpp_type=None, - storage_sharing=None if storage_free else operation.storage_sharing, - requested_size_in_bytes=( - None if storage_free else operation.storage_size_in_bytes - ), - requested_alignment=( - None if storage_free else operation.storage_alignment - ), - auto_sync=False if storage_free else operation.storage_auto_sync, - uniform_arguments=( - *(("valid_items",) if operation.has_valid_items else ()), - *(("oob_default",) if operation.has_oob_default else ()), - *(("offset",) if operation.has_offset else ()), - ), - valid_member_selection=( - "first valid_items tile elements" - if operation.has_valid_items - else None - ), - argument_preconditions=( - *( - ( - ArgumentPrecondition( - name="valid_items", - minimum=0, - maximum=tile_items, - enforcement=( - PreconditionEnforcement.PLANNER_VALIDATED - if operation.valid_items.kind is BindingKind.STATIC - else PreconditionEnforcement.CALLER - ), - ), - ) - if operation.has_valid_items - else () - ), - *( - ( - ArgumentPrecondition( - name="offset", - minimum=0, - maximum=maximum_user_offset, - enforcement=( - PreconditionEnforcement.PLANNER_VALIDATED - if operation.offset.kind is BindingKind.STATIC - else PreconditionEnforcement.CALLER - ), - ), - ) - if operation.has_offset - else () - ), - ), - ) - return GroupLoweringPlan( - target=target, - call=call, - resolved_group=resolved, - implementation=specialization, - topology=requirements.topology, - participation=requirements.participation, - result=None, - synchronization=requirements.synchronization, - temp_storage=requirements.temp_storage, - provenance=ImplementationProvenance( - library="CUB", - header=header, - cpp_class=cpp_class, - method=specialization.method_name, - ), - ) - - -_register_group_operation_family( - GroupLoadStoreSemantics, - classifications=_call_classifications, - planner=_plan_load_store, - group_kinds=frozenset({"block", "warp", "threads_within_warp"}), - unsupported_group_message=( - "cuda.coop Load/Store supports this_block(), complete physical " - "this_warp(), and power-of-two logical-warp groups" - ), -) - - -__all__ = [ - "GroupLoadStoreAlgorithm", - "GroupLoadStoreKind", - "GroupLoadStoreSemantics", -] diff --git a/python/cuda_coop/cuda/coop/_core/launch.py b/python/cuda_coop/cuda/coop/_core/launch.py deleted file mode 100644 index 7ec2791929fc..000000000000 --- a/python/cuda_coop/cuda/coop/_core/launch.py +++ /dev/null @@ -1,369 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Record the launch information that group planning can use. - -A compiler can know an exact block shape, an upper bound on that shape, -or only part of a launch configuration. ``LaunchFacts`` keeps these -cases separate. A bound can reject an impossible launch, but it cannot -supply the exact thread count needed to choose a group implementation. - -Origins record where each fact came from and whether its producer -verified it. Group planning uses this evidence for capabilities such as -cluster and cooperative launches. Merging facts keeps compatible values -and reports contradictions before an operation is compiled. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field, replace -from typing import Any - -from .thread_group import normalize_thread_dim - -Dim3 = tuple[int, int, int] - -_LAUNCH_FACT_VALUE_FIELDS = frozenset( - { - "exact_block_dim", - "max_block_dim", - "exact_grid_dim", - "exact_cluster_dim", - "cooperative_launch", - "cluster_launch", - } -) - - -class LaunchFactConflict(ValueError): - """Two individually valid launch facts contradict one another.""" - - -def _normalize_optional_dim(value: Any, *, label: str) -> Dim3 | None: - if value is None: - return None - return normalize_thread_dim(value, scope="LaunchFacts", label=label) - - -@dataclass(frozen=True) -class LaunchFactOrigin: - """Record the source of a launch fact and its verification status. - - A backend uses this record to explain how it obtained a value, such as - an exact block shape from the configured kernel launch. ``LaunchFacts`` - retains the record for diagnostics and tracks verified fact names for - planning. The record itself does not inspect or verify a launch. - - Parameters - ---------- - fact : str - Name of the fact being described. A verified origin must name a - value field in the containing ``LaunchFacts`` record. - source : str - Nonempty name identifying the producer of the fact. - detail : str, optional - Extra diagnostic text about how the producer obtained the value. - verified : bool, optional - Whether the producer has verified this fact. The default is false. - A verified origin requires a value in the same ``LaunchFacts``. - """ - - fact: str - source: str - detail: str | None = None - verified: bool = False - - def __post_init__(self) -> None: - if not self.fact: - raise ValueError("LaunchFactOrigin fact cannot be empty") - if not self.source: - raise ValueError("LaunchFactOrigin source cannot be empty") - if not isinstance(self.verified, bool): - raise TypeError("LaunchFactOrigin verified must be a bool") - - -@dataclass(frozen=True, eq=False) -class LaunchFacts: - """Keep exact launch dimensions, bounds, and capability evidence. - - Group planners need exact dimensions to count participants and allocate - storage for each group. ``max_block_dim`` only limits a possible block; - it does not identify the block that will run. Capability values use - ``None`` for unknown, which differs from a known false value. - - Construction normalizes dimensions to three positive integers. It checks - that an exact block fits its bound and that each verified origin names a - fact with a value. It records the producer's verification claim; it does - not query the driver or inspect a kernel launch. - - Parameters - ---------- - exact_block_dim : int or sequence of int, optional - Exact thread dimensions of one block. - max_block_dim : int or sequence of int, optional - Upper bound for each block dimension. - exact_grid_dim : int or sequence of int, optional - Exact number of blocks along each grid axis. The group - resolver converts these to cluster counts when needed. - exact_cluster_dim : int or sequence of int, optional - Exact number of blocks along each axis of one cluster. - cooperative_launch : bool, optional - Whether the launch supports cooperative grid participation. - cluster_launch : bool, optional - Whether the launch uses clusters. - provenance : LaunchFactOrigin or tuple of LaunchFactOrigin, optional - Sources for these facts. Sources and diagnostic details do not - affect equality or hashing. Verified fact names do, because a - planner can require verified capability evidence. - - Notes - ----- - Two frontends that supply the same values and verify the same fact names - share an identity, even if they describe different ways of obtaining - those facts. - """ - - exact_block_dim: Dim3 | int | tuple[int, ...] | list[int] | None = None - max_block_dim: Dim3 | int | tuple[int, ...] | list[int] | None = None - exact_grid_dim: Dim3 | int | tuple[int, ...] | list[int] | None = None - exact_cluster_dim: Dim3 | int | tuple[int, ...] | list[int] | None = None - cooperative_launch: bool | None = None - cluster_launch: bool | None = None - provenance: tuple[LaunchFactOrigin, ...] | LaunchFactOrigin = field( - default=(), - compare=False, - hash=False, - ) - _verified_facts: tuple[str, ...] = field(init=False, repr=False) - - def __post_init__(self) -> None: - """Normalize dimensions and check the supplied evidence.""" - - for field_name, label in ( - ("exact_block_dim", "exact block"), - ("max_block_dim", "maximum block"), - ("exact_grid_dim", "exact grid"), - ("exact_cluster_dim", "exact cluster"), - ): - object.__setattr__( - self, - field_name, - _normalize_optional_dim(getattr(self, field_name), label=label), - ) - - for field_name in ("cooperative_launch", "cluster_launch"): - value = getattr(self, field_name) - if value is not None and not isinstance(value, bool): - raise TypeError( - f"LaunchFacts {field_name} must be bool or None" - ) - - provenance = self.provenance - if isinstance(provenance, LaunchFactOrigin): - provenance = (provenance,) - else: - provenance = tuple(provenance) - if any(not isinstance(item, LaunchFactOrigin) for item in provenance): - raise TypeError( - "LaunchFacts provenance entries must be " - "LaunchFactOrigin records" - ) - object.__setattr__(self, "provenance", provenance) - verified_facts = set() - for origin in provenance: - if not origin.verified: - continue - if origin.fact not in _LAUNCH_FACT_VALUE_FIELDS: - raise ValueError( - "verified LaunchFactOrigin fact must name a LaunchFacts " - f"value field; got {origin.fact!r}" - ) - if getattr(self, origin.fact) is None: - raise ValueError( - f"verified LaunchFactOrigin for {origin.fact!r} requires " - "the same LaunchFacts record to carry its value" - ) - verified_facts.add(origin.fact) - object.__setattr__( - self, - "_verified_facts", - tuple(sorted(verified_facts)), - ) - - exact = self.exact_block_dim - maximum = self.max_block_dim - if ( - exact is not None - and maximum is not None - and any(required > limit for required, limit in zip(exact, maximum)) - ): - raise ValueError( - "LaunchFacts exact_block_dim exceeds max_block_dim" - ) - - @property - def exact_block_threads(self) -> int | None: - """Return the exact block thread count, or ``None`` if it is unknown. - - An upper bound does not contribute to this count. - """ - - if self.exact_block_dim is None: - return None - x, y, z = self.exact_block_dim - return x * y * z - - @property - def semantic_key(self) -> tuple[Any, ...]: - """Identify the values and verified fact names that affect planning. - - Exclude source labels and diagnostic details so equivalent backend - facts can share a cache entry. - """ - - return ( - self.exact_block_dim, - self.max_block_dim, - self.exact_grid_dim, - self.exact_cluster_dim, - self.cooperative_launch, - self.cluster_launch, - self._verified_facts, - ) - - def is_verified(self, fact: str) -> bool: - """Return whether a producer marked this fact's value as verified. - - Construction checks that each verified name has a value in the record. - This query does not perform another verification. - """ - - return fact in self._verified_facts - - def __eq__(self, other: object) -> bool: - if not isinstance(other, LaunchFacts): - return NotImplemented - return self.semantic_key == other.semantic_key - - def __hash__(self) -> int: - return hash(self.semantic_key) - - -def _merge_exact_dimension( - facts: tuple[LaunchFacts, ...], - field_name: str, -) -> Dim3 | None: - """Keep one known exact shape; reject different known shapes.""" - - values = { - value - for fact in facts - if (value := getattr(fact, field_name)) is not None - } - if len(values) > 1: - raise LaunchFactConflict( - f"conflicting {field_name} launch facts: {sorted(values)!r}" - ) - return next(iter(values), None) - - -def _merge_max_block_dimension(facts: tuple[LaunchFacts, ...]) -> Dim3 | None: - """Intersect known bounds by taking the minimum along each axis.""" - - values = tuple( - fact.max_block_dim for fact in facts if fact.max_block_dim is not None - ) - if not values: - return None - return tuple(min(dimensions) for dimensions in zip(*values)) # type: ignore[return-value] - - -def _merge_capability( - facts: tuple[LaunchFacts, ...], - field_name: str, -) -> bool | None: - """Combine known capability values and reject a true/false conflict.""" - - values = { - value - for fact in facts - if (value := getattr(fact, field_name)) is not None - } - if len(values) > 1: - raise LaunchFactConflict(f"conflicting {field_name} launch facts") - return next(iter(values), None) - - -def merge_launch_facts(*facts: LaunchFacts) -> LaunchFacts: - """Combine compatible launch knowledge from several producers. - - Exact dimensions and known capability flags must agree. Block bounds - combine by taking the smallest limit on each axis. An empty input gives - an empty record. Bounds remain bounds; merging never uses one as an - exact launch shape. - - Keep distinct origin records for diagnostics. If a merged value differs - from the value an origin verified, clear that origin's verification - flag. This matters when the intersection of several block bounds - produces a new bound that no one producer supplied. - - Parameters - ---------- - *facts : LaunchFacts - Records to combine. Unknown values leave known values unchanged. - - Returns - ------- - LaunchFacts - Combined values, bounds, and origin records. - - Raises - ------ - TypeError - An input is not a ``LaunchFacts`` record. - LaunchFactConflict - Exact dimensions or capability flags disagree, or the merged exact - block exceeds the merged bound. - """ - - if any(not isinstance(fact, LaunchFacts) for fact in facts): - raise TypeError("merge_launch_facts expects LaunchFacts records") - facts = tuple(facts) - try: - merged_values = { - "exact_block_dim": _merge_exact_dimension(facts, "exact_block_dim"), - "max_block_dim": _merge_max_block_dimension(facts), - "exact_grid_dim": _merge_exact_dimension(facts, "exact_grid_dim"), - "exact_cluster_dim": _merge_exact_dimension( - facts, "exact_cluster_dim" - ), - "cooperative_launch": _merge_capability( - facts, "cooperative_launch" - ), - "cluster_launch": _merge_capability(facts, "cluster_launch"), - } - provenance = tuple( - dict.fromkeys( - replace(origin, verified=False) - if origin.verified - and getattr(fact, origin.fact) != merged_values[origin.fact] - else origin - for fact in facts - for origin in fact.provenance - ) - ) - return LaunchFacts(**merged_values, provenance=provenance) - except ValueError as exc: - if isinstance(exc, LaunchFactConflict): - raise - raise LaunchFactConflict(str(exc)) from exc - - -__all__ = [ - "Dim3", - "LaunchFactConflict", - "LaunchFactOrigin", - "LaunchFacts", - "merge_launch_facts", -] diff --git a/python/cuda_coop/cuda/coop/_core/thread_group.py b/python/cuda_coop/cuda/coop/_core/thread_group.py deleted file mode 100644 index 66b722e08bfe..000000000000 --- a/python/cuda_coop/cuda/coop/_core/thread_group.py +++ /dev/null @@ -1,1054 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe which threads participate in a cooperative operation. - -``ThreadHierarchy`` holds the launch dimensions known to a planner. -``ThreadGroup`` selects a physical group or partitions it into smaller groups. -Public ``this_*`` factories start with the current launch; a backend later -resolves the dimensions from ``LaunchFacts``. - -A descriptor is a Python planning value. Constructing one does not launch -a kernel, synchronize threads, or establish that a primitive supports -the requested group. Group resolution and operation planning check those -requirements before a backend uses the CUDAX declaration helpers here. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from functools import reduce -from operator import mul -from types import GenericAlias -from typing import Any, TypeVar - -# Hierarchy levels identify the coordinate spaces used by group descriptors. -THREAD_LEVELS = frozenset({"thread", "warp", "block", "cluster", "grid"}) -# Physical groups use CUDA's thread, warp, block, cluster, and grid levels. -PHYSICAL_GROUP_KINDS = frozenset({"thread", "warp", "block", "cluster", "grid"}) -MAPPED_GROUP_KINDS = frozenset({"threads_within_warp", "warps_within_block"}) -THREAD_GROUP_KINDS = PHYSICAL_GROUP_KINDS | MAPPED_GROUP_KINDS -COMPLETE_WARP_GROUP_KINDS = frozenset({"warp"}) | MAPPED_GROUP_KINDS -_ThreadGroupT = TypeVar("_ThreadGroupT", bound="ThreadGroup") -_CPP_LEVEL_EXPR = { - "thread": "::cuda::gpu_thread", - "warp": "::cuda::warp", - "block": "::cuda::block", - "cluster": "::cuda::cluster", - "grid": "::cuda::grid", -} - - -class CoopCompilerContextRequiredError(RuntimeError): - """Report a cooperative value used outside the compiler stage it needs.""" - - -def normalize_thread_dim( - value: Any, - *, - scope: str, - label: str, -) -> tuple[int, int, int]: - """Convert a positive launch shape to an ``(x, y, z)`` tuple. - - Accept an integer or a tuple/list of one to three integers. Pad missing - axes with one so equivalent spellings share the same descriptor and - cache identity. Booleans, empty shapes, and nonpositive dimensions are - invalid. ``scope`` and ``label`` identify the caller in diagnostics. - """ - - if isinstance(value, bool): - raise TypeError(f"{scope} {label} shape must be int-like") - if isinstance(value, int): - dims = (value,) - elif isinstance(value, (tuple, list)): - dims = tuple(value) - if not dims: - raise ValueError(f"{scope} {label} shape cannot be empty") - if len(dims) > 3: - raise ValueError( - f"{scope} {label} shape must have at most 3 dimensions" - ) - else: - raise TypeError(f"{scope} {label} shape must be an int or tuple/list") - - normalized = [] - for dim in dims: - if not isinstance(dim, int) or isinstance(dim, bool): - raise TypeError(f"{scope} {label} dimensions must be integers") - if dim <= 0: - raise ValueError(f"{scope} {label} dimensions must be positive") - normalized.append(dim) - - while len(normalized) < 3: - normalized.append(1) - return tuple(normalized) # type: ignore[return-value] - - -def normalize_thread_level(level: str, *, scope: str, feature: str) -> str: - """Validate a hierarchy level and accept ``gpu_thread`` as ``thread``. - - Use ``scope`` and ``feature`` to identify an invalid level in the error. - """ - - if level == "gpu_thread": - level = "thread" - if level not in THREAD_LEVELS: - names = ", ".join(sorted(THREAD_LEVELS)) - raise ValueError(f"{scope}.{feature} level must be one of: {names}") - return level - - -def normalize_thread_group_kind(kind: str, *, scope: str, feature: str) -> str: - """Validate a physical or mapped group kind. - - Accept ``gpu_thread`` as ``thread``. Group kinds include static partitions - as well as physical levels; ``scope`` and ``feature`` identify invalid - input in the diagnostic. - """ - - if kind == "gpu_thread": - kind = "thread" - if kind not in THREAD_GROUP_KINDS: - names = ", ".join(sorted(THREAD_GROUP_KINDS)) - raise ValueError( - f"{scope}.{feature} group kind must be one of: {names}" - ) - return kind - - -def _thread_count(block_dim: tuple[int, int, int] | None) -> int | None: - if block_dim is None: - return None - return reduce(mul, block_dim, 1) - - -def _dims_token(prefix: str, dims: tuple[int, int, int]) -> str: - x, y, z = dims - if y == 1 and z == 1: - return f"{prefix}{x}" - if z == 1: - return f"{prefix}{x}x{y}" - return f"{prefix}{x}x{y}x{z}" - - -@dataclass(frozen=True, init=False) -class 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. - - Attributes - ---------- - block_dim : tuple of int or None - Threads along each axis of one block, once known. - cluster_dim : tuple of int or None - Blocks along each axis of one cluster, once needed and known. - grid_dim : tuple of int or None - Clusters along each axis of the resolved grid. On a non-cluster - launch, each cluster contains one block. - implicit : bool - 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. - """ - - block_dim: tuple[int, int, int] | None - grid_dim: tuple[int, int, int] | None - cluster_dim: tuple[int, int, int] | None - implicit: bool - - def __init__(self) -> None: - object.__setattr__(self, "block_dim", None) - object.__setattr__(self, "grid_dim", None) - object.__setattr__(self, "cluster_dim", None) - object.__setattr__(self, "implicit", True) - - @classmethod - def _resolved( - cls, - *, - block_dim: int | tuple[int, ...] | list[int], - grid_dim: int | tuple[int, ...] | list[int] | None = None, - cluster_dim: int | tuple[int, ...] | list[int] | None = None, - ) -> ThreadHierarchy: - """Construct a hierarchy from dimensions supplied by a planner. - - Normalize each supplied shape and mark the result as explicit. The - caller must obtain and check the launch facts; this constructor cannot - verify an actual launch. ``grid_dim`` counts clusters, so the resolver - must first convert physical grid dimensions when needed. - """ - - hierarchy = object.__new__(cls) - object.__setattr__( - hierarchy, - "block_dim", - normalize_thread_dim( - block_dim, - scope="ThreadHierarchy", - label="block", - ), - ) - object.__setattr__( - hierarchy, - "grid_dim", - None - if grid_dim is None - else normalize_thread_dim( - grid_dim, - scope="ThreadHierarchy", - label="grid", - ), - ) - object.__setattr__( - hierarchy, - "cluster_dim", - None - if cluster_dim is None - else normalize_thread_dim( - cluster_dim, - scope="ThreadHierarchy", - label="cluster", - ), - ) - object.__setattr__(hierarchy, "implicit", False) - return hierarchy - - @classmethod - def current(cls) -> ThreadHierarchy: - """Describe the current launch before its dimensions are known.""" - - return cls() - - @property - def is_static(self) -> bool: - """Return whether a planner supplied explicit hierarchy dimensions. - - Higher levels can still be unknown. Use ``has_static_extents_for`` - when a particular group level needs them. - """ - - return not self.implicit - - @property - def block_thread_count(self) -> int | None: - """Return the block thread count, or ``None`` if it is unknown.""" - - return _thread_count(self.block_dim) # type: ignore[arg-type] - - @property - def symbol_suffix(self) -> str: - """Encode known hierarchy dimensions for generated symbol names. - - Use ``current`` for an unresolved launch and include each known outer - level before the block shape. - """ - - if self.implicit: - return "current" - parts: list[str] = [] - if self.grid_dim is not None: - parts.append(_dims_token("g", self.grid_dim)) # type: ignore[arg-type] - if self.cluster_dim is not None: - parts.append(_dims_token("c", self.cluster_dim)) # type: ignore[arg-type] - parts.append(self.block_dim_token) - return "_".join(parts) - - @property - def block_dim_token(self) -> str: - """Encode the block shape, or use ``current`` if it is unknown.""" - - if self.block_dim is None: - return "current" - return _dims_token("b", self.block_dim) - - @property - def semantic_key(self) -> tuple[Any, ...]: - """Identify the dimensions and their current-launch status.""" - - return self.block_dim, self.grid_dim, self.cluster_dim, self.implicit - - def has_static_extents_for(self, group_kind: str) -> bool: - """Check whether this hierarchy has the extents used by a group level. - - Thread groups have a known size even in an implicit hierarchy. Other - physical groups need the block shape and, for cluster or grid groups, - the corresponding outer shape. This checks available dimensions; - operation support and launch capabilities are separate checks. - """ - - group_kind = normalize_thread_level( - group_kind, - scope="ThreadHierarchy", - feature="has_static_extents_for", - ) - if self.implicit: - return group_kind == "thread" - if group_kind in {"thread", "warp", "block"}: - return self.block_dim is not None - if group_kind == "cluster": - return self.block_dim is not None and self.cluster_dim is not None - if group_kind == "grid": - return self.block_dim is not None and self.grid_dim is not None - return False - - -Hierarchy = ThreadHierarchy - - -@dataclass(frozen=True) -class GroupByMapping: - """Describe a static partition of a physical warp or block. - - ``ThreadGroup.group_by`` creates this record. A warp partitions into - sets of threads; a block partitions into sets of complete physical - warps. The record keeps the unit count and synchronization choice so - C++ generation and planning agree on each subgroup's membership. - - Parameters - ---------- - unit : str - Unit placed in each subgroup: ``thread`` or ``warp``. - parent : str - Physical parent level: ``warp`` or ``block``. ``ThreadGroup`` checks - that it matches the selected unit and mapped group kind. - count : int - Positive number of units in one subgroup, known during compilation. - exhaustive : bool - Whether the count must divide the parent's unit count exactly. - A non-exhaustive mapping can leave units outside complete groups. - synchronizer : str - Must be ``lane`` for thread units and ``barrier`` for warp units. - Construction rejects any other pairing. Code generation emits the - matching synchronizer. - """ - - unit: str - parent: str - count: int - exhaustive: bool - synchronizer: str - - def __post_init__(self) -> None: - """Check partition controls before attaching them to a group.""" - - if self.unit not in {"thread", "warp"}: - raise ValueError("GroupByMapping unit must be thread or warp") - if self.parent not in {"warp", "block"}: - raise ValueError("GroupByMapping parent must be warp or block") - if not isinstance(self.count, int) or isinstance(self.count, bool): - raise TypeError("GroupByMapping count must be a static integer") - if self.count <= 0: - raise ValueError("GroupByMapping count must be positive") - if not isinstance(self.exhaustive, bool): - raise TypeError("GroupByMapping exhaustive must be a bool") - expected_synchronizer = "lane" if self.unit == "thread" else "barrier" - if self.synchronizer != expected_synchronizer: - raise ValueError( - f"GroupByMapping {self.unit} unit requires " - f"{expected_synchronizer!r} synchronization" - ) - - @property - def semantic_key(self) -> tuple[Any, ...]: - """Identify partition size, coverage, and synchronization.""" - - return ( - self.unit, - self.parent, - self.count, - self.exhaustive, - self.synchronizer, - ) - - -def _validate_mapped_group_extent( - kind: str, - hierarchy: ThreadHierarchy, - mapping: GroupByMapping, -) -> None: - """Check a partition against the parent dimensions already known. - - Mapped groups require complete physical warps. Their count cannot exceed - the parent's units, and an exhaustive count must divide those units. - For a block with an unknown size, defer these extent checks to launch - resolution. A warp's 32-thread extent is already known. - """ - - block_threads = hierarchy.block_thread_count - if block_threads is not None and block_threads % 32 != 0: - raise ValueError( - "mapped group_by requires an enclosing block " - "composed of complete warps" - ) - if kind == "threads_within_warp": - parent_units = 32 - if mapping.count > parent_units: - raise ValueError("warp group_by count cannot exceed 32 threads") - else: - if block_threads is None: - return - parent_units = block_threads // 32 - if mapping.count > parent_units: - raise ValueError( - "block group_by count cannot exceed the parent warp count" - ) - if mapping.exhaustive and parent_units % mapping.count != 0: - raise ValueError( - "exhaustive ThreadGroup.group_by requires the count to divide " - "the parent unit count" - ) - - -@dataclass(frozen=True) -class ThreadGroup: - """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 :ref:`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. - """ - - __class_getitem__ = classmethod(GenericAlias) - - kind: str - hierarchy: ThreadHierarchy = field(default_factory=ThreadHierarchy.current) - parent: ThreadGroup | None = None - mapping: GroupByMapping | None = None - # Source labels do not affect structural identity or cache keys. - # Planners can still use them to preserve public API policy. - source: str = field(default="explicit", compare=False, hash=False) - - def __post_init__(self) -> None: - """Check that the kind, parent, and partition describe one group. - - Mapped groups must use the expected physical parent and the same - hierarchy. Physical groups carry no mapping metadata. Known dimensions - also let construction reject impossible partition sizes early. - """ - - kind = normalize_thread_group_kind( - self.kind, - scope="ThreadGroup", - feature="kind", - ) - object.__setattr__(self, "kind", kind) - - hierarchy = self.hierarchy - if not isinstance(hierarchy, ThreadHierarchy): - raise TypeError("ThreadGroup hierarchy must be a ThreadHierarchy") - - if kind in MAPPED_GROUP_KINDS: - if not isinstance(self.parent, ThreadGroup): - raise TypeError( - "mapped ThreadGroup requires a parent ThreadGroup" - ) - if not isinstance(self.mapping, GroupByMapping): - raise TypeError("mapped ThreadGroup requires GroupByMapping") - expected_parent = ( - "warp" if kind == "threads_within_warp" else "block" - ) - expected_unit = ( - "thread" if kind == "threads_within_warp" else "warp" - ) - if self.parent.kind != expected_parent: - raise ValueError( - f"{kind} requires a physical {expected_parent} parent" - ) - if ( - self.mapping.parent != expected_parent - or self.mapping.unit != expected_unit - ): - raise ValueError( - f"{kind} mapping does not match its group kind" - ) - if self.parent.hierarchy != hierarchy: - raise ValueError( - "mapped ThreadGroup hierarchy must match its parent" - ) - _validate_mapped_group_extent(kind, hierarchy, self.mapping) - elif self.parent is not None or self.mapping is not None: - raise ValueError( - "physical ThreadGroup cannot carry mapping metadata" - ) - - @property - def block_dim(self) -> tuple[int, int, int] | None: - """Return the resolved block dimensions, or ``None`` if unknown.""" - - return self.hierarchy.block_dim - - @property - def group_thread_count(self) -> int: - """Return the known number of threads, or reject an unresolved size. - - A physical warp always has an extent of 32. The operation planner - still checks whether all 32 lanes can participate. The descriptor does - not observe active lanes. Use ``static_size`` to return ``None`` - instead of raising ``ValueError`` when the size is unknown. - """ - - count = self.static_size - if count is None: - raise ValueError( - f"ThreadGroup.{self.kind} uses a runtime hierarchy with no " - "static group size" - ) - return count - - @property - def static_size(self) -> int | None: - """Return the known group thread count, or ``None`` if unresolved. - - Thread, warp, and mapped group sizes follow from their kind and - mapping. Block, cluster, and grid sizes need hierarchy dimensions. The - count describes membership, but does not prove that all threads in - the parent participate. - """ - - if self.kind == "thread": - return 1 - if self.kind == "warp": - return 32 - if self.kind == "threads_within_warp": - assert self.mapping is not None - return self.mapping.count - if self.kind == "warps_within_block": - assert self.mapping is not None - return self.mapping.count * 32 - hierarchy = self.hierarchy - assert hierarchy is not None - if self.kind == "block": - return hierarchy.block_thread_count - if self.kind == "cluster": - block_threads = hierarchy.block_thread_count - if block_threads is None or hierarchy.cluster_dim is None: - return None - cluster_blocks = _thread_count(hierarchy.cluster_dim) - assert cluster_blocks is not None - return block_threads * cluster_blocks - if self.kind == "grid": - block_threads = hierarchy.block_thread_count - if block_threads is None or hierarchy.grid_dim is None: - return None - grid_groups = _thread_count(hierarchy.grid_dim) - cluster_blocks = ( - 1 - if hierarchy.cluster_dim is None - else _thread_count(hierarchy.cluster_dim) - ) - assert grid_groups is not None - assert cluster_blocks is not None - return block_threads * cluster_blocks * grid_groups - return None - - @property - def is_current(self) -> bool: - """Return whether the hierarchy still needs launch resolution.""" - - return self.hierarchy.implicit # type: ignore[union-attr] - - @property - def is_static(self) -> bool: - """Return whether the hierarchy has the extents this group needs. - - Mapped groups use their physical parent's result. A known subgroup - size alone does not supply an unknown enclosing block shape. - """ - - if self.kind in MAPPED_GROUP_KINDS: - assert self.parent is not None - return self.parent.is_static - return self.hierarchy.has_static_extents_for(self.kind) # type: ignore[union-attr] - - @property - def parent_unit_count(self) -> int | None: - """Return the parent's known unit count for a mapped group. - - A warp parent has 32 thread units. A block parent has one unit per - complete physical warp. Return ``None`` for a physical group or an - unknown parent count. - """ - - if self.kind == "threads_within_warp": - return 32 - if self.kind == "warps_within_block": - hierarchy = self.hierarchy - assert hierarchy is not None - block_threads = hierarchy.block_thread_count - if block_threads is None or block_threads % 32 != 0: - return None - return block_threads // 32 - return None - - @property - def groups_per_parent(self) -> int | None: - """Count complete mapped subgroups within one physical parent. - - Exclude any remainder allowed by a non-exhaustive mapping. Return - ``None`` when the group is physical or its parent count is unknown. - """ - - if self.mapping is None: - return None - parent_units = self.parent_unit_count - if parent_units is None: - return None - return parent_units // self.mapping.count - - @property - def remainder_count(self) -> int | None: - """Count parent units left outside complete mapped subgroups. - - The unit is a thread for a warp parent and a warp for a block parent. - Return ``None`` for a physical group or an unknown parent count. - """ - - if self.mapping is None: - return None - parent_units = self.parent_unit_count - if parent_units is None: - return None - return parent_units % self.mapping.count - - @property - def complete_membership(self) -> bool | None: - """Report whether the mapping covers all units of its physical parent. - - Return ``None`` until a mapped parent's unit count is known. Physical - groups return true because no mapping excludes members. Execution can - still require complete warps and converged participation. - """ - - if self.mapping is None: - return True - remainder = self.remainder_count - if remainder is None: - return None - return remainder == 0 - - @property - def symbol_suffix(self) -> str: - """Encode the group kind, partition, and hierarchy for C++ symbols.""" - - if self.mapping is None: - return f"{self.kind}_{self.hierarchy.symbol_suffix}" # type: ignore[union-attr] - mode = "all" if self.mapping.exhaustive else "partial" - return ( - f"{self.kind}_{self.mapping.count}_{mode}_" - f"{self.hierarchy.symbol_suffix}" # type: ignore[union-attr] - ) - - @property - def block_dim_token(self) -> str: - return self.hierarchy.block_dim_token # type: ignore[union-attr] - - @property - def semantic_key(self) -> tuple[Any, ...]: - """Identify group shape and mapping without their source label. - - Include the parent for mapped groups. Planners can use ``source`` for - API policy, but that label does not change this structural identity. - """ - - if self.mapping is None: - return self.kind, self.hierarchy.semantic_key # type: ignore[union-attr] - assert self.parent is not None - return ( - self.kind, - self.parent.semantic_key, - self.mapping.semantic_key, - self.hierarchy.semantic_key, # type: ignore[union-attr] - ) - - # Keep runtime annotations dependency-free on Python 3.10. - def with_hierarchy( # noqa: PYI019 - self: _ThreadGroupT, - hierarchy: ThreadHierarchy, - *, - source: str = "resolved", - ) -> _ThreadGroupT: - """Copy the descriptor with a given hierarchy and source label. - - Keep its concrete backend type. For a mapped group, give its physical - parent the same hierarchy so both descriptions remain consistent. - Construction rechecks the mapping against the known dimensions. - """ - - if self.mapping is None: - return type(self)( - kind=self.kind, hierarchy=hierarchy, source=source - ) - assert self.parent is not None - return type(self)( - kind=self.kind, - hierarchy=hierarchy, - parent=self.parent.with_hierarchy(hierarchy, source=source), - mapping=self.mapping, - source=source, - ) - - # Keep runtime annotations dependency-free on Python 3.10. - def group_by( # noqa: PYI019 - self: _ThreadGroupT, - 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 - ------- - cuda.coop.ThreadGroup - 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. - - Examples - -------- - Descriptors can be inspected without compiling or launching a kernel: - - .. code-block:: python - - from cuda import coop - - group = coop.this_warp().group_by(8) - assert group.kind == "threads_within_warp" - assert group.static_size == 8 - """ - - if self.mapping is not None: - raise NotImplementedError( - "nested ThreadGroup.group_by is not supported" - ) - if not isinstance(count, int) or isinstance(count, bool): - raise TypeError( - "ThreadGroup.group_by count must be a static integer" - ) - if count <= 0: - raise ValueError("ThreadGroup.group_by count must be positive") - if not isinstance(exhaustive, bool): - raise TypeError("ThreadGroup.group_by exhaustive must be a bool") - - if self.kind == "warp": - kind = "threads_within_warp" - unit = "thread" - synchronizer = "lane" - elif self.kind == "block": - kind = "warps_within_block" - unit = "warp" - synchronizer = "barrier" - else: - raise NotImplementedError( - "ThreadGroup.group_by supports only physical warp " - "and block parents" - ) - - mapping = GroupByMapping( - unit=unit, - parent=self.kind, - count=count, - exhaustive=exhaustive, - synchronizer=synchronizer, - ) - return type(self)( - kind=kind, - hierarchy=self.hierarchy, - parent=self, - mapping=mapping, - source="group_by", - ) - - -def make_thread_group( - kind: str, - *, - group_type: type[_ThreadGroupT] = ThreadGroup, - scope: str = "cuda.coop", -) -> _ThreadGroupT: - """Create a physical group descriptor for the current launch. - - Backends can supply a ``group_type`` subclass while using the same kind - normalization and unresolved hierarchy. ``scope`` names the API in an - invalid-kind diagnostic. - """ - - kind = normalize_thread_level(kind, scope=scope, feature="ThreadGroup") - return group_type( - kind=kind, - hierarchy=ThreadHierarchy.current(), - source="current", - ) - - -def _cpp_dims_expr(level: str, dims: tuple[int, int, int]) -> str: - x, y, z = dims - if y == 1 and z == 1: - return f"::cuda::{level}_dims<{x}>()" - if z == 1: - return f"::cuda::{level}_dims<{x}, {y}>()" - return f"::cuda::{level}_dims<{x}, {y}, {z}>()" - - -def render_hierarchy_decl( - hierarchy: ThreadHierarchy, - *, - var_name: str = "hierarchy", - indent: str = " ", -) -> list[str]: - """Generate a CUDAX hierarchy declaration from known dimensions. - - An implicit hierarchy needs no declaration and returns an empty list. - For an explicit hierarchy, emit the known grid and cluster levels before - the required block level. ``var_name`` and ``indent`` control the - surrounding generated source. - """ - - if hierarchy.implicit: - return [] - exprs: list[str] = [] - if hierarchy.grid_dim is not None: - exprs.append(_cpp_dims_expr("grid", hierarchy.grid_dim)) - if hierarchy.cluster_dim is not None: - exprs.append(_cpp_dims_expr("cluster", hierarchy.cluster_dim)) - assert hierarchy.block_dim is not None - exprs.append(_cpp_dims_expr("block", hierarchy.block_dim)) - if len(exprs) == 1: - return [f"{indent}auto {var_name} = ::cuda::hierarchy{{{exprs[0]}}};"] - lines = [f"{indent}auto {var_name} = ::cuda::hierarchy{{"] - for idx, expr in enumerate(exprs): - comma = "," if idx < len(exprs) - 1 else "" - lines.append(f"{indent} {expr}{comma}") - lines.append(f"{indent}}};") - return lines - - -def render_group_decl( - group: ThreadGroup, - *, - var_name: str = "group", - hierarchy_var: str = "hierarchy", - indent: str = " ", -) -> str: - """Generate one CUDAX declaration for a physical group. - - Use the supplied ``hierarchy_var`` for an explicit hierarchy, or CUDAX's - implicit hierarchy for the current launch. The caller must declare an - explicit hierarchy first. Mapped groups need additional declarations; - reject them here and use ``render_group_decl_lines``. - """ - - if group.mapping is not None: - raise ValueError( - "render_group_decl supports physical groups only; use " - "render_group_decl_lines for mapped groups" - ) - assert group.hierarchy is not None - if group.hierarchy.implicit: - return ( - f"{indent}::cuda::experimental::coop::this_{group.kind} " - f"{var_name}{{::cuda::experimental::implicit_hierarchy()}};" - ) - return ( - f"{indent}::cuda::experimental::coop::this_{group.kind} " - f"{var_name}{{{hierarchy_var}}};" - ) - - -def render_group_decl_lines( - group: ThreadGroup, - *, - var_name: str = "group", - hierarchy_var: str = "hierarchy", - indent: str = " ", -) -> list[str]: - """Generate a CUDAX group and any declarations its mapping needs. - - Physical groups need one declaration. A thread subgroup also needs its - physical warp parent and a lane synchronizer. A subgroup of physical - warps needs its block parent and shared barrier storage, with one - barrier slot per complete subgroup. - - The mapped block's subgroup count must be known to size that storage; - otherwise raise ``ValueError``. These lines describe the group and its - storage. They do not establish the enclosing launch's capabilities or - prove that a primitive supports the group. - """ - - if group.mapping is None: - return [ - render_group_decl( - group, - var_name=var_name, - hierarchy_var=hierarchy_var, - indent=indent, - ) - ] - - assert group.parent is not None - mapping = group.mapping - exhaustive = "true" if mapping.exhaustive else "false" - mapping_args = str(mapping.count) - if not mapping.exhaustive: - mapping_args = ( - f"::cuda::experimental::coop::non_exhaustive, {mapping.count}" - ) - parent_name = f"{var_name}_parent" - lines = [ - render_group_decl( - group.parent, - var_name=parent_name, - hierarchy_var=hierarchy_var, - indent=indent, - ) - ] - if group.kind == "threads_within_warp": - lines.extend( - [ - ( - f"{indent}::cuda::experimental::coop::generic_group " - f"{var_name}{{" - ), - f"{indent} ::cuda::gpu_thread, {parent_name},", - ( - f"{indent} ::cuda::experimental::coop::group_by<" - f"{mapping.count}, {exhaustive}>{{{mapping_args}}}," - ), - ( - f"{indent} " - "::cuda::experimental::coop::lane_synchronizer{}};" - ), - ] - ) - return lines - - groups_per_parent = group.groups_per_parent - if groups_per_parent is None: - raise ValueError( - "mapped warp group requires a static parent group count" - ) - lines.extend( - [ - f"{indent}using {var_name}_barriers_type =", - ( - f"{indent} ::cuda::barrier<::cuda::thread_scope_block>" - f"[{groups_per_parent}];" - ), - f"{indent}__shared__ ::cuda::std::aligned_storage_t<", - f"{indent} sizeof({var_name}_barriers_type),", - ( - f"{indent} alignof({var_name}_barriers_type)> " - f"{var_name}_barriers_storage;" - ), - f"{indent}auto& {var_name}_barriers =", - f"{indent} reinterpret_cast<{var_name}_barriers_type&>(", - f"{indent} {var_name}_barriers_storage);", - f"{indent}::cuda::experimental::coop::generic_group {var_name}{{", - f"{indent} ::cuda::warp, {parent_name},", - ( - f"{indent} ::cuda::experimental::coop::group_by<" - f"{mapping.count}, {exhaustive}>{{{mapping_args}}}," - ), - ( - f"{indent} " - "::cuda::experimental::coop::barrier_synchronizer{" - f"{var_name}_barriers}}}};" - ), - ] - ) - return lines - - -def cpp_level_expr(level: str) -> str: - """Translate a hierarchy level to its CUDAX tag expression. - - For example, ``thread`` selects ``::cuda::gpu_thread``. Reuse level - normalization so the accepted Python spellings stay consistent. - """ - - return _CPP_LEVEL_EXPR[ - normalize_thread_level( - level, scope="cuda.coop", feature="cpp_level_expr" - ) - ] - - -def this_thread() -> ThreadGroup: - """Describe the calling thread as a one-thread group.""" - - return make_thread_group("thread") - - -def this_warp() -> ThreadGroup: - """Describe the calling thread's physical 32-thread warp.""" - - return make_thread_group("warp") - - -def this_block() -> ThreadGroup: - """Describe the calling thread's block in the current launch.""" - - return make_thread_group("block") - - -def this_cluster() -> ThreadGroup: - """Describe the calling thread's cluster for later launch validation.""" - - return make_thread_group("cluster") - - -def this_grid() -> ThreadGroup: - """Describe the launch grid for later cooperative-launch validation.""" - - return make_thread_group("grid") - - -__all__ = [ - "COMPLETE_WARP_GROUP_KINDS", - "MAPPED_GROUP_KINDS", - "PHYSICAL_GROUP_KINDS", - "THREAD_GROUP_KINDS", - "THREAD_LEVELS", - "CoopCompilerContextRequiredError", - "GroupByMapping", - "Hierarchy", - "ThreadGroup", - "ThreadHierarchy", - "cpp_level_expr", - "make_thread_group", - "normalize_thread_dim", - "normalize_thread_group_kind", - "normalize_thread_level", - "render_group_decl", - "render_group_decl_lines", - "render_hierarchy_decl", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -] diff --git a/python/cuda_coop/cuda/coop/_core/warp/__init__.py b/python/cuda_coop/cuda/coop/_core/warp/__init__.py deleted file mode 100644 index b33583bca976..000000000000 --- a/python/cuda_coop/cuda/coop/_core/warp/__init__.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from .load_store import ( - WarpLoadAlgorithm, - WarpLoadStoreAlgorithm, - WarpLoadStoreKind, - WarpLoadStoreSemantics, - WarpStoreAlgorithm, - make_warp_load_specialization, - make_warp_load_store_semantics, - make_warp_load_store_specialization, - make_warp_store_specialization, -) - -__all__ = [ - "WarpLoadAlgorithm", - "WarpLoadStoreAlgorithm", - "WarpLoadStoreKind", - "WarpLoadStoreSemantics", - "WarpStoreAlgorithm", - "make_warp_load_specialization", - "make_warp_load_store_semantics", - "make_warp_load_store_specialization", - "make_warp_store_specialization", -] diff --git a/python/cuda_coop/cuda/coop/_core/warp/load_store.py b/python/cuda_coop/cuda/coop/_core/warp/load_store.py deleted file mode 100644 index c8595fd59b2a..000000000000 --- a/python/cuda_coop/cuda/coop/_core/warp/load_store.py +++ /dev/null @@ -1,643 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe CUB WarpLoad and WarpStore wrappers for physical or logical warps. - -Builders validate each warp's tile and describe the selected call signatures. -The resulting specialization includes the logical width and pointer-offset -metadata. Group planning and backend lowering use that metadata to give each -warp its own consecutive tile within the block's input or output. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from enum import Enum -from numbers import Integral -from typing import Any - -from .._algorithm import Algorithm -from .._bindings import ( - ArgumentBinding, - BindingKind, - cxx_scalar_literal, - i32_parameter, - normalize_i32_binding, - normalize_i64_binding, -) -from .._symbols import semantic_token -from .._types import ( - INT64, - Array, - CxxFunction, - Dependency, - Pointer, - PointerOffset, - TemplateParameter, - TempStorageParameter, - Value, -) - -_MAX_WARP_THREADS = 32 -_SUPPORTED_LOGICAL_WARP_THREADS = frozenset({1, 2, 4, 8, 16, 32}) - - -class WarpLoadStoreKind(str, Enum): - """Direction of data movement between memory and per-thread items.""" - - LOAD = "load" - STORE = "store" - - -class WarpLoadStoreAlgorithm(str, Enum): - """CUB Warp Load/Store algorithms and their per-thread arrangements. - - ``DIRECT``, ``VECTORIZE``, and ``TRANSPOSE`` use a blocked arrangement: - lane ``r`` owns tile indices ``r * items_per_thread + i``. ``STRIPED`` - uses indices ``r + i * threads_in_warp`` instead. In both expressions, - ``i`` indexes the thread's items and ``r`` is its logical-warp lane. - - ``DIRECT`` accesses that blocked arrangement directly. ``VECTORIZE`` - requests vectorized access where CUB's type, size, and alignment - requirements permit it; guarded calls use direct access. ``TRANSPOSE`` - exchanges items within the logical warp so memory accesses are striped - while the per-thread arrangement remains blocked. Stores reverse the - corresponding load's data movement. - """ - - DIRECT = "direct" - STRIPED = "striped" - VECTORIZE = "vectorize" - TRANSPOSE = "transpose" - - -WarpLoadAlgorithm = WarpLoadStoreAlgorithm -WarpStoreAlgorithm = WarpLoadStoreAlgorithm - - -_LOAD_ALGORITHM_CPP = { - WarpLoadAlgorithm.DIRECT: "::cub::WARP_LOAD_DIRECT", - WarpLoadAlgorithm.STRIPED: "::cub::WARP_LOAD_STRIPED", - WarpLoadAlgorithm.VECTORIZE: "::cub::WARP_LOAD_VECTORIZE", - WarpLoadAlgorithm.TRANSPOSE: "::cub::WARP_LOAD_TRANSPOSE", -} -_STORE_ALGORITHM_CPP = { - WarpStoreAlgorithm.DIRECT: "::cub::WARP_STORE_DIRECT", - WarpStoreAlgorithm.STRIPED: "::cub::WARP_STORE_STRIPED", - WarpStoreAlgorithm.VECTORIZE: "::cub::WARP_STORE_VECTORIZE", - WarpStoreAlgorithm.TRANSPOSE: "::cub::WARP_STORE_TRANSPOSE", -} -_T = Dependency("T") -_ITEMS_PER_THREAD = Dependency("ITEMS_PER_THREAD") -_TEMPLATE_PARAMETERS = ( - TemplateParameter("T"), - TemplateParameter("ITEMS_PER_THREAD"), - TemplateParameter("ALGORITHM"), - TemplateParameter("LOGICAL_WARP_THREADS"), -) - - -def _algorithm_cpp_map( - kind: WarpLoadStoreKind, -) -> dict[WarpLoadStoreAlgorithm, str]: - return ( - _LOAD_ALGORITHM_CPP - if kind is WarpLoadStoreKind.LOAD - else _STORE_ALGORITHM_CPP - ) - - -def _normalize_algorithm( - kind: WarpLoadStoreKind, - algorithm: str | WarpLoadStoreAlgorithm, -) -> WarpLoadStoreAlgorithm: - """Accept an enum, a short name, or the matching CUB operation's token.""" - - mapping = _algorithm_cpp_map(kind) - if isinstance(algorithm, WarpLoadStoreAlgorithm): - return algorithm - if isinstance(algorithm, str): - for candidate, cpp in mapping.items(): - if algorithm in {candidate.value, cpp}: - return candidate - raise ValueError( - f"unsupported Warp{kind.value.title()} algorithm {algorithm!r}" - ) - - -def _normalize_items_per_thread(items_per_thread: Any) -> int: - if ( - not isinstance(items_per_thread, Integral) - or isinstance(items_per_thread, bool) - or items_per_thread < 1 - ): - raise ValueError("items_per_thread must be a positive integer") - return int(items_per_thread) - - -def _normalize_logical_warp_threads(threads_in_warp: Any) -> int: - """Require a power-of-two width that fits within one physical warp.""" - - if ( - not isinstance(threads_in_warp, Integral) - or isinstance(threads_in_warp, bool) - or threads_in_warp not in _SUPPORTED_LOGICAL_WARP_THREADS - ): - supported = ", ".join( - str(value) for value in sorted(_SUPPORTED_LOGICAL_WARP_THREADS) - ) - raise ValueError( - "Warp Load/Store requires threads_in_warp in " - f"{{{supported}}}; got {threads_in_warp!r}" - ) - return int(threads_in_warp) - - -def _base_parameters(kind: WarpLoadStoreKind) -> list[Any]: - """Describe scratch, the memory pointer, and the per-thread item array. - - Load and store take the pointer first but reverse the input/output roles. - Both write through arguments instead of returning an array or pointer. - """ - - if kind is WarpLoadStoreKind.LOAD: - return [ - TempStorageParameter(), - Pointer(_T, name="src", is_array_pointer=True, restrict=True), - Array( - _T, - _ITEMS_PER_THREAD, - name="dst", - is_output=True, - is_return=False, - ), - ] - return [ - TempStorageParameter(), - Pointer( - _T, - name="dst", - is_output=True, - is_return=False, - is_array_pointer=True, - restrict=True, - ), - Array(_T, _ITEMS_PER_THREAD, name="src"), - ] - - -def _normalize_optional_binding( - value: bool | ArgumentBinding, - *, - name: str, -) -> ArgumentBinding: - """Use booleans to select signatures and keep explicit bindings.""" - - if isinstance(value, ArgumentBinding): - return value - if not isinstance(value, bool): - raise TypeError(f"{name} must be a bool or ArgumentBinding") - return ArgumentBinding.runtime() if value else ArgumentBinding.omitted() - - -def _with_pointer_offset( - parameters: list[Any], - offset: ArgumentBinding, -) -> tuple[Any, ...]: - """Append an element offset for the memory pointer after scratch storage. - - The descriptor targets argument zero after the leading storage parameter - is removed: ``src`` for loads and ``dst`` for stores. Lowering applies - the offset to that pointer before invoking CUB. - """ - - static_value = offset.value if offset.kind is BindingKind.STATIC else None - return ( - *parameters, - PointerOffset( - INT64, - name="offset", - pointer_arg_index=0, - static_value=static_value, - ), - ) - - -@dataclass(frozen=True) -class WarpLoadStoreSemantics: - """Normalized call variants for one physical or logical warp's tile. - - Construct this description with :func:`make_warp_load_store_semantics` - to validate options and build the parameter descriptors consumed by a - backend. The tile contains ``threads_in_warp * items_per_thread`` items. - - Attributes - ---------- - kind : WarpLoadStoreKind - Load from memory into per-thread items, or store those items to - memory. - dtype : object - Element type used by the CUB primitive and its per-thread items. - algorithm : WarpLoadStoreAlgorithm - Memory-access strategy and per-thread arrangement. - items_per_thread : int - Number of items owned by each participating thread. - threads_in_warp : int - Logical warp width: one of 1, 2, 4, 8, 16, or 32. A width of 32 - represents a physical warp. - valid_items : ArgumentBinding - Omitted, static, or runtime valid-item count for a guarded call. - oob_default : ArgumentBinding - Omitted, static, or runtime value for invalid load items. - has_full_tile : bool - Whether the described variants include an unguarded full-tile call. - This records signature availability, not a runtime bounds check. - pointer_offset : ArgumentBinding - Omitted, static, or runtime element offset applied to the memory - pointer. A runtime binding denotes an effective offset supplied by - the backend; group planning accounts for the warp's tile origin. - parameters : tuple of tuples - Parameter descriptors for each selected call variant. Every variant - starts with temporary storage, the memory pointer, and the - per-thread array, followed by any bounds, default, and offset - descriptors. Static descriptors are embedded during lowering rather - than passed as runtime arguments. - """ - - kind: WarpLoadStoreKind - dtype: Any - algorithm: WarpLoadStoreAlgorithm - items_per_thread: int - threads_in_warp: int - valid_items: ArgumentBinding - oob_default: ArgumentBinding - has_full_tile: bool - pointer_offset: ArgumentBinding - parameters: tuple[tuple[Any, ...], ...] - - @property - def has_valid_items(self) -> bool: - return self.valid_items.kind is not BindingKind.OMITTED - - @property - def has_oob_default(self) -> bool: - return self.oob_default.kind is not BindingKind.OMITTED - - @property - def has_pointer_offset(self) -> bool: - return self.pointer_offset.kind is not BindingKind.OMITTED - - @property - def requires_runtime_effective_offset(self) -> bool: - """Whether lowering must supply the effective element offset.""" - - return self.pointer_offset.kind is BindingKind.RUNTIME - - @property - def method_name(self) -> str: - return self.kind.value.title() - - @property - def algorithm_cpp(self) -> str: - return _algorithm_cpp_map(self.kind)[self.algorithm] - - @property - def semantic_key(self) -> tuple[Any, ...]: - """Identify the bindings and selected call variants.""" - - return ( - f"warp_{self.kind.value}", - semantic_token(self.dtype), - self.algorithm.value, - self.items_per_thread, - self.threads_in_warp, - self.valid_items.semantic_key, - self.oob_default.semantic_key, - self.has_full_tile, - self.pointer_offset.semantic_key, - semantic_token(self.parameters), - ) - - -def make_warp_load_store_semantics( - *, - kind: str | WarpLoadStoreKind, - dtype: Any, - items_per_thread: int, - algorithm: str | WarpLoadStoreAlgorithm, - threads_in_warp: int = _MAX_WARP_THREADS, - valid_items: bool | ArgumentBinding = False, - oob_default: bool | ArgumentBinding = False, - include_full_tile: bool = False, - include_pointer_offset: bool | ArgumentBinding = False, -) -> WarpLoadStoreSemantics: - """Describe CUB Warp Load/Store call variants and their scalar bindings. - - Normalize the algorithm and optional arguments, validate static bounds, - and build descriptors for the selected full-tile and guarded calls. - Each invocation moves one logical warp's tile; the caller must provide - complete participation by that warp. This builds metadata for later - specialization and lowering. - - Parameters - ---------- - kind : str or WarpLoadStoreKind - ``"load"`` or ``"store"``. - dtype : object - Element type carried into the CUB specialization and used for any - out-of-bounds default. Type interpretation belongs to the backend. - items_per_thread : int - Positive number of items per thread. Booleans are rejected. - algorithm : str or WarpLoadStoreAlgorithm - Algorithm enum, its lowercase value, or the matching fully - qualified CUB enumerator, such as ``"::cub::WARP_LOAD_TRANSPOSE"``. - See :class:`WarpLoadStoreAlgorithm` for the per-thread arrangements. - threads_in_warp : int, optional - Logical warp width in ``{1, 2, 4, 8, 16, 32}``, defaulting to the - physical warp width of 32. Booleans are rejected. - valid_items : bool or ArgumentBinding, optional - ``False`` omits the guarded signature; ``True`` selects a runtime - signed 32-bit count. An explicit binding can omit the count, embed - a static count, or request a runtime count. The count covers the - first valid items of this warp's tile, not items per thread. - Static counts must fit signed 32-bit arithmetic and lie between - zero and ``threads_in_warp * items_per_thread``, inclusive. - Runtime counts must satisfy the same bounds at the call site. - oob_default : bool or ArgumentBinding, optional - Value assigned to invalid items by a guarded load. ``False`` omits - the default; ``True`` adds a runtime value of ``dtype``. Explicit - bindings select omitted, static, or runtime values. Static defaults - must be finite numeric scalars representable as C++ literals. - Requires ``valid_items`` and is invalid for stores. With no default, - invalid load items are unspecified; their prior values need not be - preserved. Guarded stores leave memory outside the valid range - untouched. - include_full_tile : bool, optional - Add an unguarded full-tile variant alongside a guarded variant. - Setting this requires ``valid_items``. When ``valid_items`` is - omitted, the full-tile variant is already selected by default. - include_pointer_offset : bool or ArgumentBinding, optional - Select element offsets into the memory pointer. ``False`` produces - only variants without an offset; ``True`` produces both those - variants and variants with a runtime signed 64-bit offset. An - explicit static or runtime binding instead produces only variants - with that offset. An omitted binding produces only variants without - it. Static offsets must be nonnegative signed 64-bit integers. - Runtime offsets use signed 64-bit arithmetic; the caller must - establish pointer validity and bounds. - Applying an offset does not change the meaning of ``valid_items``: - the count starts at the adjusted pointer. - - Returns - ------- - WarpLoadStoreSemantics - Normalized bindings and parameter descriptors for each requested - call variant. The descriptors retain template dependencies until - :func:`make_warp_load_store_specialization` binds them. - - Raises - ------ - TypeError - An optional binding is neither a bool nor an ``ArgumentBinding``; - a static count or offset is not an integer or is a boolean; or a - static default is not a supported numeric scalar. - ValueError - The kind, algorithm, item count, or warp width is unsupported; a - static count or offset is outside its supported range; a static - default is nonfinite or its integer literal exceeds 64-bit range; - or the requested guarded, default, and full-tile variants conflict. - - Notes - ----- - The offset describes an adjustment to the supplied pointer. This - builder does not calculate the logical warp's index or add its tile - origin. Group planning requests a runtime effective offset, and the - backend combines the group-instance tile origin with the user's offset. - That planning path records nonnegative, overflow-safe runtime offset - bounds as caller preconditions. - """ - - pointer_offset_overload_cohort = isinstance(include_pointer_offset, bool) - kind = WarpLoadStoreKind(kind) - items_per_thread = _normalize_items_per_thread(items_per_thread) - threads_in_warp = _normalize_logical_warp_threads(threads_in_warp) - algorithm = _normalize_algorithm(kind, algorithm) - valid_items = _normalize_optional_binding(valid_items, name="valid_items") - valid_items = normalize_i32_binding(valid_items, name="valid_items") - if valid_items.kind is BindingKind.STATIC: - tile_items = items_per_thread * threads_in_warp - if not 0 <= int(valid_items.value) <= tile_items: - raise ValueError( - "static valid_items must be between zero and the warp tile " - f"size ({tile_items})" - ) - oob_default = _normalize_optional_binding(oob_default, name="oob_default") - pointer_offset = _normalize_optional_binding( - include_pointer_offset, - name="include_pointer_offset", - ) - pointer_offset = normalize_i64_binding( - pointer_offset, name="pointer offset" - ) - if ( - pointer_offset.kind is BindingKind.STATIC - and int(pointer_offset.value) < 0 - ): - raise ValueError("static pointer offset must be nonnegative") - if kind is WarpLoadStoreKind.STORE and ( - oob_default.kind is not BindingKind.OMITTED - ): - raise ValueError("oob_default is only valid for WarpLoad") - if ( - oob_default.kind is not BindingKind.OMITTED - and valid_items.kind is BindingKind.OMITTED - ): - raise ValueError("oob_default requires a valid_items signature") - if include_full_tile and valid_items.kind is BindingKind.OMITTED: - raise ValueError("include_full_tile requires a valid_items signature") - - base = _base_parameters(kind) - has_full_tile = valid_items.kind is BindingKind.OMITTED or include_full_tile - methods: list[tuple[Any, ...]] = [] - if has_full_tile and ( - pointer_offset.kind is BindingKind.OMITTED - or pointer_offset_overload_cohort - ): - methods.append(tuple(base)) - if valid_items.kind is not BindingKind.OMITTED: - num_valid_items = i32_parameter(valid_items, name="num_valid_items") - assert num_valid_items is not None - partial = [*base, num_valid_items] - if oob_default.kind is BindingKind.RUNTIME: - partial.append(Value(dtype, name="oob_default")) - elif oob_default.kind is BindingKind.STATIC: - partial.append( - CxxFunction( - cxx_scalar_literal(oob_default.value, name="oob_default"), - dtype, - name="oob_default", - ) - ) - if ( - pointer_offset.kind is BindingKind.OMITTED - or pointer_offset_overload_cohort - ): - methods.append(tuple(partial)) - if pointer_offset.kind is not BindingKind.OMITTED: - methods.append(_with_pointer_offset(partial, pointer_offset)) - if pointer_offset.kind is not BindingKind.OMITTED and has_full_tile: - methods.append(_with_pointer_offset(base, pointer_offset)) - - return WarpLoadStoreSemantics( - kind=kind, - dtype=dtype, - algorithm=algorithm, - items_per_thread=items_per_thread, - threads_in_warp=threads_in_warp, - valid_items=valid_items, - oob_default=oob_default, - has_full_tile=has_full_tile, - pointer_offset=pointer_offset, - parameters=tuple(methods), - ) - - -def make_warp_load_store_specialization( - *, - kind: str | WarpLoadStoreKind, - dtype: Any, - items_per_thread: int, - algorithm: str | WarpLoadStoreAlgorithm, - threads_in_warp: int = _MAX_WARP_THREADS, - valid_items: bool | ArgumentBinding = False, - oob_default: bool | ArgumentBinding = False, - include_full_tile: bool = False, - include_pointer_offset: bool | ArgumentBinding = False, -) -> Algorithm: - """Bind CUB template arguments for physical or logical Warp Load/Store. - - Normalize and validate the options with - :func:`make_warp_load_store_semantics`, then construct an ``Algorithm`` - binding ``T``, ``ITEMS_PER_THREAD``, ``ALGORITHM``, and - ``LOGICAL_WARP_THREADS``. The parameters and validation rules are those - of that semantics builder. - - Parameters - ---------- - kind : str or WarpLoadStoreKind - ``"load"`` or ``"store"``. - dtype : object - Element type bound to CUB's ``T``. - items_per_thread : int - Positive number of elements per thread, excluding booleans. - algorithm : str or WarpLoadStoreAlgorithm - CUB strategy, accepted in the forms documented by - :func:`make_warp_load_store_semantics`. - threads_in_warp : int, optional - Logical warp width in ``{1, 2, 4, 8, 16, 32}``, defaulting to 32. - valid_items : bool or ArgumentBinding, optional - Guarded-call count binding, bounded by this warp's tile size. - oob_default : bool or ArgumentBinding, optional - Optional invalid-item default for guarded loads. - include_full_tile : bool, optional - Include an unguarded overload alongside guarded overloads. - include_pointer_offset : bool or ArgumentBinding, optional - Select memory-pointer offset overloads. Boolean and explicit-binding - selection follow :func:`make_warp_load_store_semantics`. - - Returns - ------- - Algorithm - Bound CUB template arguments, header and method, selected parameter - variants, and effective-offset metadata for group lowering. - Its tile stride is ``threads_in_warp * items_per_thread`` elements. - - Notes - ----- - Specialization constructs a description ready for backend materialization. - It does not compile wrappers, allocate temporary storage, or calculate a - runtime tile origin. When group planning supplies a runtime effective - offset, lowering adds the group-instance tile origin to the user's - offset before invoking the specialized wrapper. - """ - - call = make_warp_load_store_semantics( - kind=kind, - dtype=dtype, - items_per_thread=items_per_thread, - algorithm=algorithm, - threads_in_warp=threads_in_warp, - valid_items=valid_items, - oob_default=oob_default, - include_full_tile=include_full_tile, - include_pointer_offset=include_pointer_offset, - ) - title = call.kind.value.title() - return Algorithm( - struct_name=f"Warp{title}", - method_name=title, - c_name=f"warp_{call.kind.value}", - includes=(f"cub/warp/warp_{call.kind.value}.cuh",), - template_parameters=_TEMPLATE_PARAMETERS, - parameters=call.parameters, - template_arguments={ - "T": dtype, - "ITEMS_PER_THREAD": call.items_per_thread, - "ALGORITHM": call.algorithm_cpp, - "LOGICAL_WARP_THREADS": call.threads_in_warp, - }, - metadata={ - "scope": "warp", - "primitive": call.kind.value, - "algorithm": call.algorithm.value, - "valid_items": call.has_valid_items, - "oob_default": call.has_oob_default, - "full_tile": call.has_full_tile, - "pointer_offset": call.has_pointer_offset, - "requires_runtime_effective_offset": ( - call.requires_runtime_effective_offset - ), - "effective_offset_origin": "group_instance", - "effective_offset_stride": ( - call.threads_in_warp * call.items_per_thread - ), - }, - ) - - -def make_warp_load_specialization(**kwargs: Any) -> Algorithm: - """Build a WarpLoad specialization with the shared builder's options. - - See :func:`make_warp_load_store_specialization`; ``kind`` is fixed to - ``WarpLoadStoreKind.LOAD``. - """ - - return make_warp_load_store_specialization( - kind=WarpLoadStoreKind.LOAD, **kwargs - ) - - -def make_warp_store_specialization(**kwargs: Any) -> Algorithm: - """Build a WarpStore specialization with the shared builder's options. - - See :func:`make_warp_load_store_specialization`; ``kind`` is fixed to - ``WarpLoadStoreKind.STORE`` and ``oob_default`` must be omitted. - """ - - return make_warp_load_store_specialization( - kind=WarpLoadStoreKind.STORE, **kwargs - ) - - -__all__ = [ - "WarpLoadAlgorithm", - "WarpLoadStoreAlgorithm", - "WarpLoadStoreKind", - "WarpLoadStoreSemantics", - "WarpStoreAlgorithm", - "make_warp_load_specialization", - "make_warp_load_store_semantics", - "make_warp_load_store_specialization", - "make_warp_store_specialization", -] diff --git a/python/cuda_coop/cuda/coop/_headers/__init__.py b/python/cuda_coop/cuda/coop/_headers/__init__.py deleted file mode 100644 index c019c0afe972..000000000000 --- a/python/cuda_coop/cuda/coop/_headers/__init__.py +++ /dev/null @@ -1,334 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Select the CCCL and CUDA headers used to compile cooperative primitives. - -CCCL headers come from an explicit root, this package's source checkout, or -the installed wheel. CUDA runtime headers are selected separately. Keeping -these choices separate prevents a compiler from using the toolkit's bundled -CUB in place of the CCCL headers that belong to ``cuda.coop``. -""" - -from __future__ import annotations - -import atexit -import os -import sys -from collections.abc import Iterable -from contextlib import ExitStack -from dataclasses import dataclass -from functools import cache -from importlib.resources import as_file, files -from pathlib import Path - -_CUB_PROBE = Path("cub/version.cuh") -_CUDA_PROBE = Path("cuda_runtime.h") -_INSTALLED_HEADER_CONTEXTS = ExitStack() -atexit.register(_INSTALLED_HEADER_CONTEXTS.close) - - -class HeaderResolutionError(RuntimeError): - """Raised when cuda.coop cannot resolve a coherent CCCL header set.""" - - -@dataclass(frozen=True) -class CoopIncludePaths: - """Keep the selected header paths and a description of their source. - - ``cccl`` lists the include roots in search order. ``cuda`` contains at - most one CUDA runtime include directory. If it is empty, :meth:`as_tuple` - raises an error because compilation needs CUDA headers. ``origin`` names - the selected CCCL source for diagnostics. - """ - - cccl: tuple[Path, ...] - cuda: tuple[Path, ...] - origin: str - - def as_tuple(self) -> tuple[Path, ...]: - """Require CUDA headers and put the selected CCCL roots first. - - This order gives the selected CCCL headers priority over the copies - shipped with the CUDA toolkit. - """ - - if not self.cuda: - raise HeaderResolutionError( - "Unable to locate one CUDA include directory containing " - "cuda_runtime.h. Configure CUDA_PATH, CUDA_HOME, or CUDA_ROOT " - "before compiling cuda.coop providers." - ) - return (*self.cccl, *self.cuda) - - -def _unique_existing_dirs( - paths: Iterable[Path | str | None], -) -> tuple[Path, ...]: - """Resolve existing directories and remove duplicates in input order.""" - - result: list[Path] = [] - seen: set[Path] = set() - for raw_path in paths: - if raw_path is None: - continue - path = Path(raw_path).expanduser().resolve() - if not path.is_dir() or path in seen: - continue - seen.add(path) - result.append(path) - return tuple(result) - - -def _source_checkout_paths(root: Path) -> tuple[Path, ...]: - """Find the separate library roots in a CCCL source checkout. - - Return an empty tuple if the CUB marker is absent. Once that marker is - present, require all library roots so an incomplete checkout is an error. - """ - - if not (root / "cub" / _CUB_PROBE).is_file(): - return () - candidates = ( - root / "thrust", - root / "cub", - root / "libcudacxx" / "include", - ) - missing = tuple( - path.relative_to(root) for path in candidates if not path.is_dir() - ) - if missing: - rendered = ", ".join(path.as_posix() for path in missing) - raise HeaderResolutionError( - f"{root} looks like a CCCL source checkout but is missing source " - f"include roots: {rendered}." - ) - return _unique_existing_dirs(candidates) - - -def _packaged_header_paths(root: Path) -> tuple[Path, ...]: - """Find a bundled include tree at ``root`` or directly below it. - - A CUDA runtime header in the same tree identifies a toolkit include - directory, which cannot serve as the package's CCCL header bundle. - """ - - candidates = ( - root, - root / "include", - ) - for candidate in candidates: - if not (candidate / _CUB_PROBE).is_file(): - continue - if (candidate / "cuda_runtime.h").is_file(): - raise HeaderResolutionError( - f"{candidate} is a CUDA toolkit include directory, not a " - "cuda-coop header bundle." - ) - return _unique_existing_dirs((candidate,)) - return () - - -def _configured_cccl_paths(root: Path) -> tuple[Path, ...]: - """Require an explicit root to contain a checkout or header bundle. - - Invalid configuration raises an error instead of selecting another source. - """ - - resolved_root = root.expanduser().resolve() - paths = _source_checkout_paths(resolved_root) - if paths: - return paths - paths = _packaged_header_paths(resolved_root) - if paths: - return paths - raise HeaderResolutionError( - f"{resolved_root} is neither a CCCL source checkout nor a cuda-coop " - "header bundle." - ) - - -def _find_source_checkout(start: Path) -> tuple[Path, tuple[Path, ...]] | None: - """Find the checkout that contains this Python source package. - - Search above ``start``, normally a module's file path. A wheel installed - inside an unrelated checkout must still use its own bundled headers. - """ - - current = start.expanduser().resolve() - if current.is_file(): - current = current.parent - for candidate in (current, *current.parents): - if not _belongs_to_source_package(current, candidate): - continue - paths = _source_checkout_paths(candidate) - if paths: - return candidate, paths - return None - - -def _belongs_to_source_package(start: Path, root: Path) -> bool: - """Return whether ``start`` belongs to this checkout's source package.""" - - package_root = root / "python" / "cuda_coop" - try: - start.relative_to(package_root) - except ValueError: - return False - return True - - -@cache -def _installed_header_root() -> Path: - """Keep the wheel's header resource available for the process lifetime. - - Resource access can extract files to a temporary directory. The shared - context stack keeps that directory alive for later compiler invocations. - """ - - resource = files(__package__).joinpath("include") - return Path(_INSTALLED_HEADER_CONTEXTS.enter_context(as_file(resource))) - - -def _installed_include_paths() -> CoopIncludePaths: - """Require the wheel's CUB bundle and discover CUDA headers separately.""" - - cccl = _packaged_header_paths(_installed_header_root()) - - if not any((path / _CUB_PROBE).is_file() for path in cccl): - raise HeaderResolutionError( - "The installed cuda-coop wheel does not contain " - "its bundled CUB headers." - ) - return CoopIncludePaths( - cccl=cccl, - cuda=_cuda_include_paths(), - origin="cuda-coop wheel header bundle", - ) - - -def _select_cuda_include_path( - paths: Iterable[Path | str | None], -) -> tuple[Path, ...]: - """Select only the first directory that contains CUDA runtime headers.""" - - candidates = _unique_existing_dirs(paths) - for candidate in candidates: - if (candidate / _CUDA_PROBE).is_file(): - return (candidate,) - return () - - -def _cuda_include_paths() -> tuple[Path, ...]: - """Find one CUDA include directory without mixing toolkit installations. - - Prefer Pathfinder's result, then the configured toolkit roots, then the - conventional Unix installation. Return an empty tuple if none is usable. - """ - - try: - from cuda.pathfinder import find_nvidia_header_directory - - cuda_include = find_nvidia_header_directory("cudart") - except (ImportError, RuntimeError): - cuda_include = None - - # Pathfinder applies its own search order. Use its result alone so headers - # from another toolkit cannot enter the same provider compilation. - discovered = _select_cuda_include_path((cuda_include,)) - if discovered: - return discovered - - configured = tuple( - Path(root).expanduser() / "include" - for env_name in ("CUDA_PATH", "CUDA_HOME", "CUDA_ROOT") - if (root := os.environ.get(env_name)) - ) - fallback = ( - (Path("/usr/local/cuda/include"),) if sys.platform != "win32" else () - ) - return _select_cuda_include_path((*configured, *fallback)) - - -def _validate_required_headers( - required_headers: Iterable[str], - *, - cccl_paths: tuple[Path, ...], - origin: str, -) -> None: - """Check required relative paths only within the selected CCCL roots. - - Reject parent traversal and absolute paths. A missing header is an error. - Searching CUDA roots could substitute headers from another CCCL version. - """ - - missing: list[str] = [] - for header in dict.fromkeys(required_headers): - relative = Path(header) - if relative.is_absolute() or ".." in relative.parts: - raise HeaderResolutionError( - f"Invalid CCCL header path: {header!r}." - ) - if not any( - (include_path / relative).is_file() for include_path in cccl_paths - ): - missing.append(header) - if missing: - raise HeaderResolutionError( - f"The selected {origin} is missing required CCCL headers " - f"({', '.join(missing)}). cuda.coop does not fall back to CUDA " - "Toolkit CCCL headers." - ) - - -def resolve_include_paths( - *, - start: Path, - configured_roots: Iterable[Path | str | None] = (), - required_headers: Iterable[str] = (), -) -> CoopIncludePaths: - """Select one CCCL header source and discover CUDA runtime headers. - - The first nonempty entry in ``configured_roots`` takes precedence. Without - one, use the checkout containing ``start`` if it is inside this package's - source directory. Otherwise, use the installed wheel's header bundle. - An invalid explicit root raises :class:`HeaderResolutionError`. - - ``required_headers`` contains paths relative to the CCCL include roots. - Require each file in the selected source; do not fill gaps with toolkit - headers. The returned object can have no CUDA include directory. Its - :meth:`CoopIncludePaths.as_tuple` method requires one before compilation. - """ - - configured = next((Path(root) for root in configured_roots if root), None) - if configured is not None: - cccl = _configured_cccl_paths(configured) - origin = f"configured CCCL root {configured.expanduser().resolve()}" - paths = CoopIncludePaths( - cccl=cccl, - cuda=_cuda_include_paths(), - origin=origin, - ) - elif source := _find_source_checkout(start): - root, cccl = source - paths = CoopIncludePaths( - cccl=cccl, - cuda=_cuda_include_paths(), - origin=f"CCCL source checkout {root}", - ) - else: - paths = _installed_include_paths() - - _validate_required_headers( - required_headers, - cccl_paths=paths.cccl, - origin=paths.origin, - ) - return paths - - -__all__ = [ - "CoopIncludePaths", - "HeaderResolutionError", - "resolve_include_paths", -] diff --git a/python/cuda_coop/cuda/coop/_headers/_identity.py b/python/cuda_coop/cuda/coop/_headers/_identity.py deleted file mode 100644 index a673311e54c8..000000000000 --- a/python/cuda_coop/cuda/coop/_headers/_identity.py +++ /dev/null @@ -1,393 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Identify the headers that a compiler cache entry depends on. - -The include search order and file contents both affect compilation. Clean Git -trees provide a fast content identifier. Other directories need a recursive -content hash, including the targets of symbolic links. File timestamps alone -cannot detect every header change. -""" - -from __future__ import annotations - -import hashlib -import os -import re -import subprocess -import time -from collections.abc import Iterable -from dataclasses import dataclass -from pathlib import Path - - -class HeaderIdentityError(RuntimeError): - """Raised when an include root cannot be fingerprinted safely.""" - - -@dataclass(frozen=True) -class IncludeRootIdentity: - """Record one include root's path, content digest, and hashing method. - - ``method`` distinguishes a clean Git tree from a recursive file walk. - ``duration_ns`` measures discovery and hashing work for performance - reports; it is not part of the cache identity. - """ - - path: str - method: str - digest: str - duration_ns: int - - -@dataclass(frozen=True) -class IncludeDirsIdentity: - """Combine include roots in the order the compiler searches them. - - ``digest`` includes each root's path, method, and content digest. - Reordering roots can select a different header with the same name, so - order matters. The walk count and elapsed time help diagnose the cost of - cache lookups. - """ - - roots: tuple[IncludeRootIdentity, ...] - digest: str - recursive_walks: int - duration_ns: int - - -def _find_git_root(root: Path) -> Path | None: - """Find the nearest enclosing Git checkout or worktree.""" - - for candidate in (root, *root.parents): - if (candidate / ".git").exists(): - return candidate - return None - - -def _run_git(repo: Path, *args: str) -> subprocess.CompletedProcess[bytes]: - """Run a Git query with literal paths and return its status.""" - - return subprocess.run( - ["git", "--literal-pathspecs", "-C", str(repo), *args], - check=False, - stdout=subprocess.PIPE, - stderr=subprocess.DEVNULL, - ) - - -def _git_include_root_identities( - repo: Path, - roots: tuple[Path, ...], -) -> dict[Path, str]: - """Use Git tree IDs only where the working files match tracked content. - - Check the requested roots together for edits, extra files, and index flags - that can hide changes. If those checks fail, return no identities so the - caller hashes files instead. Exclude roots with tracked symbolic links: - Git records the link text but does not identify its target's contents. - """ - - relatives = tuple(root.relative_to(repo) for root in roots) - pathspecs = tuple( - "." if not relative.parts else relative.as_posix() - for relative in relatives - ) - status = _run_git( - repo, - "status", - "--porcelain=v1", - "-z", - "--untracked-files=all", - "--ignored=matching", - "--ignore-submodules=none", - "--", - *pathspecs, - ) - if status.returncode != 0 or status.stdout: - return {} - - index_entries = _run_git(repo, "ls-files", "-v", "-z", "--", *pathspecs) - if index_entries.returncode != 0 or not index_entries.stdout: - return {} - for entry in index_entries.stdout.split(b"\0"): - if entry and entry[:1] != b"H": - # `S` marks skip-worktree; lowercase marks assume-unchanged. - return {} - - staged_entries = _run_git( - repo, "ls-files", "--stage", "-z", "--", *pathspecs - ) - if staged_entries.returncode != 0 or not staged_entries.stdout: - return {} - symlinked_roots: set[Path] = set() - for entry in staged_entries.stdout.split(b"\0"): - if not entry: - continue - metadata, separator, path = entry.partition(b"\t") - fields = metadata.split() - if separator != b"\t" or len(fields) != 3 or fields[2] != b"0": - return {} - if fields[0] != b"120000": - continue - tracked_path = path.decode("utf-8", errors="surrogateescape") - for root, relative in zip(roots, relatives): - relative_path = relative.as_posix() - if ( - not relative.parts - or tracked_path == relative_path - or tracked_path.startswith(f"{relative_path}/") - ): - symlinked_roots.add(root) - - result: dict[Path, str] = {} - non_root_relatives = tuple( - relative for relative in relatives if relative.parts - ) - tree_oids: dict[str, bytes] = {} - if non_root_relatives: - tree = _run_git( - repo, - "ls-tree", - "-d", - "-z", - "HEAD", - "--", - *(relative.as_posix() for relative in non_root_relatives), - ) - if tree.returncode != 0: - return {} - for record in tree.stdout.split(b"\0"): - if not record: - continue - metadata, separator, path = record.partition(b"\t") - fields = metadata.split() - if separator != b"\t" or len(fields) != 3 or fields[1] != b"tree": - return {} - tree_oids[path.decode("utf-8", errors="surrogateescape")] = fields[ - 2 - ] - - repository_tree: bytes | None = None - if any(not relative.parts for relative in relatives): - tree = _run_git(repo, "rev-parse", "--verify", "HEAD^{tree}") - if tree.returncode != 0: - return {} - repository_tree = tree.stdout.strip() - - for root, relative in zip(roots, relatives): - tree_oid = ( - repository_tree - if not relative.parts - else tree_oids.get(relative.as_posix()) - ) - if ( - tree_oid is None - or len(tree_oid) not in (40, 64) - or re.fullmatch(rb"[0-9a-f]+", tree_oid) is None - ): - return {} - if root in symlinked_roots: - continue - digest = hashlib.sha256() - digest.update(b"git-tree-v1\0") - digest.update(tree_oid) - result[root] = digest.hexdigest() - return result - - -def _recursive_include_root_identity(root: Path) -> str: - """Hash names and contents in a stable order, following symbolic links. - - Include link text and reachable target contents because either can change - what the compiler reads. Mark broken links and directory cycles, and - reject special files whose contents cannot be hashed as ordinary headers. - """ - - digest = hashlib.sha256() - digest.update(b"recursive-content-v3\0") - - def update_path(relative: Path) -> None: - """Prefix each path with its length to keep names distinct.""" - - encoded = relative.as_posix().encode("utf-8", errors="surrogateescape") - digest.update(str(len(encoded)).encode("ascii")) - digest.update(b":") - digest.update(encoded) - digest.update(b"\0") - - def hash_file(path: Path) -> None: - """Hash the file in chunks and record its size and digest.""" - - content_digest = hashlib.sha256() - content_length = 0 - with path.open("rb") as header: - while chunk := header.read(1024 * 1024): - content_length += len(chunk) - content_digest.update(chunk) - digest.update(str(content_length).encode("ascii")) - digest.update(b":") - digest.update(content_digest.digest()) - - def visit( - directory: Path, - relative: Path, - ancestors: frozenset[Path], - ) -> None: - """Visit entries by name and stop cycles through directory links.""" - - resolved_directory = directory.resolve() - if resolved_directory in ancestors: - digest.update(b"directory-cycle\0") - return - ancestors = ancestors | {resolved_directory} - with os.scandir(directory) as scanned: - entries = sorted(scanned, key=lambda entry: os.fsencode(entry.name)) - for entry in entries: - if entry.name in {".git", "__pycache__"}: - continue - entry_path = Path(entry.path) - entry_relative = relative / entry.name - update_path(entry_relative) - if entry.is_symlink(): - digest.update(b"symlink\0") - target = os.readlink(entry.path) - digest.update( - target.encode("utf-8", errors="surrogateescape") - if isinstance(target, str) - else target - ) - digest.update(b"\0") - try: - target_path = entry_path.resolve(strict=True) - except OSError: - digest.update(b"broken\0") - continue - if target_path.is_dir(): - digest.update(b"target-directory\0") - visit(target_path, entry_relative, ancestors) - elif target_path.is_file(): - digest.update(b"target-file\0") - hash_file(target_path) - else: - raise HeaderIdentityError( - f"Unsupported special include path: {target_path}" - ) - elif entry.is_dir(follow_symlinks=False): - digest.update(b"directory\0") - visit(entry_path, entry_relative, ancestors) - elif entry.is_file(follow_symlinks=False): - digest.update(b"file\0") - hash_file(entry_path) - else: - raise HeaderIdentityError( - f"Unsupported special include path: {entry_path}" - ) - - visit(root, Path(), frozenset()) - return digest.hexdigest() - - -def include_dirs_identity(include_dirs: Iterable[str]) -> IncludeDirsIdentity: - """Build a cache identity for the compiler's ordered include directories. - - Resolve every directory before hashing it. Use Git tree IDs where possible - and walk the remaining directories, then combine the results in input - order. Raise :class:`HeaderIdentityError` for missing directories or - unreadable files so the cache key cannot omit those headers. - - Timing fields report this call's work and do not affect the digest. - """ - - started_ns = time.perf_counter_ns() - resolved_roots = tuple( - Path(include_dir).expanduser().resolve() for include_dir in include_dirs - ) - for root in resolved_roots: - if not root.is_dir(): - raise HeaderIdentityError( - f"Provider include path is not a directory: {root}" - ) - if not resolved_roots: - digest = hashlib.sha256(b"ordered-include-roots-v3\0none").hexdigest() - return IncludeDirsIdentity( - roots=(), - digest=digest, - recursive_walks=0, - duration_ns=max(0, time.perf_counter_ns() - started_ns), - ) - - root_durations_ns = dict.fromkeys(resolved_roots, 0) - roots_by_repo: dict[Path, list[Path]] = {} - for root in resolved_roots: - if (repo := _find_git_root(root)) is not None: - roots_by_repo.setdefault(repo, []).append(root) - git_identities: dict[Path, str] = {} - for repo, roots in roots_by_repo.items(): - git_started_ns = time.perf_counter_ns() - try: - repo_identities = _git_include_root_identities(repo, tuple(roots)) - except OSError: - repo_identities = {} - git_identities.update(repo_identities) - git_duration_ns = max(0, time.perf_counter_ns() - git_started_ns) - per_root_duration_ns = git_duration_ns // len(roots) - for root in roots: - root_durations_ns[root] += per_root_duration_ns - - root_identities: list[IncludeRootIdentity] = [] - recursive_walks = 0 - for root in resolved_roots: - method = "git-tree" - root_digest = git_identities.get(root) - if root_digest is None: - method = "recursive-content" - recursive_walks += 1 - recursive_started_ns = time.perf_counter_ns() - try: - root_digest = _recursive_include_root_identity(root) - except OSError as exc: - raise HeaderIdentityError( - f"Failed fingerprinting provider include path: {root}" - ) from exc - root_durations_ns[root] += max( - 0, - time.perf_counter_ns() - recursive_started_ns, - ) - root_identities.append( - IncludeRootIdentity( - path=str(root), - method=method, - digest=root_digest, - duration_ns=root_durations_ns[root], - ) - ) - - aggregate_digest = hashlib.sha256() - aggregate_digest.update(b"ordered-include-roots-v3\0") - for index, identity in enumerate(root_identities): - aggregate_digest.update(str(index).encode("ascii")) - aggregate_digest.update(b"\0") - aggregate_digest.update( - identity.path.encode("utf-8", errors="surrogateescape") - ) - aggregate_digest.update(b"\0") - aggregate_digest.update(identity.method.encode("ascii")) - aggregate_digest.update(b"\0") - aggregate_digest.update(identity.digest.encode("ascii")) - aggregate_digest.update(b"\0") - return IncludeDirsIdentity( - roots=tuple(root_identities), - digest=aggregate_digest.hexdigest(), - recursive_walks=recursive_walks, - duration_ns=max(0, time.perf_counter_ns() - started_ns), - ) - - -__all__ = [ - "HeaderIdentityError", - "IncludeDirsIdentity", - "IncludeRootIdentity", - "include_dirs_identity", -] diff --git a/python/cuda_coop/cuda/coop/_headers/_toolkit.py b/python/cuda_coop/cuda/coop/_headers/_toolkit.py deleted file mode 100644 index 55c0f30cb887..000000000000 --- a/python/cuda_coop/cuda/coop/_headers/_toolkit.py +++ /dev/null @@ -1,632 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Load CUDA compiler libraries from the toolkit selected by the headers. - -NVRTC compiles the generated CUDA source, and nvJitLink links the result. -Their libraries must come from the toolkit selected by the CUDA headers. NVRTC -must match the headers' major and minor versions. nvJitLink may use a newer -minor version of the same major release. - -Loaded libraries remain in the process. A lock serializes selection and -loading, and the first successful load fixes the toolkit root for later -requests. This also prevents a retry after a partial failure from mixing -toolkit installations. -""" - -from __future__ import annotations - -import ctypes -import os -import re -import threading -from collections.abc import Iterable -from dataclasses import dataclass -from pathlib import Path - -_CUDART_VERSION = re.compile( - r"^\s*#\s*define\s+CUDART_VERSION\s+(\d+)", re.MULTILINE -) -_PRELOAD_LOCK = threading.RLock() -_EXACT_LIBRARY_HANDLES: dict[str, object] = {} -_PROCESS_TOOLKIT_SELECTION: tuple[str, tuple[int, int]] | None = None - - -@dataclass(frozen=True) -class ToolkitCompilerLibraries: - """Report loaded library paths and their checked versions. - - ``toolkit_version`` comes from the CUDA runtime header. The NVRTC and - nvJitLink versions come from the loaded library handles. Keep both so - callers can check later compiler use against the same selection. - """ - - toolkit_root: str - nvrtc_path: str - nvrtc_builtins_path: str - nvjitlink_path: str - toolkit_version: tuple[int, int] - nvrtc_version: tuple[int, int] - nvjitlink_version: tuple[int, int] - - -@dataclass(frozen=True) -class _ToolkitRootCandidates: - """List complete library candidates within one toolkit installation. - - Each NVRTC candidate is paired with its adjacent builtins library. These - paths have been found on disk but have not yet been loaded or validated. - """ - - toolkit_root: Path - nvrtc_pairs: tuple[tuple[Path, Path], ...] - nvjitlink: tuple[Path, ...] - - -@dataclass(frozen=True) -class _ToolkitLibraryLayout: - """Map a header directory to the matching compiler-library directories. - - A toolkit can be one directory tree or several sibling Python packages. - ``toolkit_root`` identifies the selected installation in either layout. - """ - - toolkit_root: Path - nvrtc_dirs: tuple[Path, ...] - nvjitlink_dirs: tuple[Path, ...] - - -def _cuda_include_dirs( - include_dirs: Iterable[str | os.PathLike[str]], -) -> tuple[Path, ...]: - """Keep distinct CUDA include directories in compiler search order.""" - - result: list[Path] = [] - for raw_path in include_dirs: - path = Path(raw_path).expanduser().resolve() - if (path / "cuda_runtime_api.h").is_file() and path not in result: - result.append(path) - return tuple(result) - - -def _toolkit_version( - cuda_include_dirs: tuple[Path, ...], -) -> tuple[int, int] | None: - """Read the header versions and require all selected roots to agree. - - Return ``None`` when there are no CUDA include roots. Raise an error for - a malformed version macro or conflicting versions. - """ - - versions: set[tuple[int, int]] = set() - for include_dir in cuda_include_dirs: - header = include_dir / "cuda_runtime_api.h" - try: - contents = header.read_text(encoding="utf-8") - except UnicodeError as exc: - raise RuntimeError( - f"failed decoding CUDA Toolkit version header: {header}" - ) from exc - match = _CUDART_VERSION.search(contents) - if match is None: - raise RuntimeError( - "failed parsing CUDART_VERSION from CUDA Toolkit " - f"version header: {header}" - ) - encoded = int(match.group(1)) - versions.add((encoded // 1000, (encoded % 1000) // 10)) - if len(versions) > 1: - rendered = ", ".join( - f"{major}.{minor}" for major, minor in sorted(versions) - ) - raise RuntimeError( - "resolved CUDA include directories disagree on Toolkit version: " - f"{rendered}" - ) - return next(iter(versions), None) - - -def _library_dirs(root: Path) -> tuple[Path, ...]: - """List existing library directories for supported toolkit layouts.""" - - result: list[Path] = [] - names = ("lib", "lib64", "bin") - if os.name == "nt": - names += ("bin/x64", "bin/x86_64") - for name in names: - candidate = root / name - if candidate.is_dir(): - result.append(candidate) - return tuple(result) - - -def _toolkit_library_layout(include_dir: Path) -> _ToolkitLibraryLayout: - """Locate compiler libraries beside the selected CUDA headers. - - CUDA 12 wheels split runtime, compiler, and linker files into sibling - packages. Treat their shared ``nvidia`` directory as the installation - root; do not search other Python environments for missing components. - """ - - toolkit_root = include_dir.parent.resolve() - if ( - include_dir.name == "include" - and toolkit_root.name == "cuda_runtime" - and toolkit_root.parent.name == "nvidia" - ): - # CUDA 12 wheels install Toolkit components into sibling namespace - # packages below one ``site-packages/nvidia`` anchor. Treat that anchor - # as the Toolkit root without searching any other site-packages tree. - toolkit_root = toolkit_root.parent - return _ToolkitLibraryLayout( - toolkit_root=toolkit_root, - nvrtc_dirs=_library_dirs(toolkit_root / "cuda_nvrtc"), - nvjitlink_dirs=_library_dirs(toolkit_root / "nvjitlink"), - ) - - library_dirs = _library_dirs(toolkit_root) - return _ToolkitLibraryLayout( - toolkit_root=toolkit_root, - nvrtc_dirs=library_dirs, - nvjitlink_dirs=library_dirs, - ) - - -def _library_names(kind: str, major: int) -> tuple[str, ...]: - """Name platform libraries for the requested major release.""" - - if os.name == "nt": - version = major * 10 - if kind == "nvrtc": - return (f"nvrtc64_{version}_0.dll",) - return (f"nvJitLink_{version}_0.dll",) - if kind == "nvrtc": - return (f"libnvrtc.so.{major}",) - return (f"libnvJitLink.so.{major}",) - - -def _nvrtc_builtins_names(major: int, minor: int) -> tuple[str, ...]: - """Name the NVRTC builtins library for the exact header version.""" - - if os.name == "nt": - return (f"nvrtc-builtins64_{major}{minor}.dll",) - return (f"libnvrtc-builtins.so.{major}.{minor}",) - - -def _toolkit_root_candidates( - include_dir: Path, - *, - major: int, - minor: int, -) -> tuple[_ToolkitRootCandidates | None, str]: - """Find a complete compiler-library set below one CUDA toolkit root. - - Pair NVRTC with builtins in the same directory, then look for nvJitLink in - that installation. Return either the candidates or a missing-file - diagnostic. This stage checks paths only; loading and version checks - happen later. - """ - - layout = _toolkit_library_layout(include_dir) - nvrtc_pairs: list[tuple[Path, Path]] = [] - unpaired_nvrtc: list[Path] = [] - for library_dir in layout.nvrtc_dirs: - builtins = next( - ( - library_dir / name - for name in _nvrtc_builtins_names(major, minor) - if (library_dir / name).is_file() - ), - None, - ) - for name in _library_names("nvrtc", major): - nvrtc = library_dir / name - if not nvrtc.is_file(): - continue - if builtins is None: - unpaired_nvrtc.append(nvrtc) - else: - nvrtc_pairs.append((nvrtc, builtins)) - - nvjitlink = tuple( - candidate - for library_dir in layout.nvjitlink_dirs - for name in _library_names("nvJitLink", major) - if (candidate := library_dir / name).is_file() - ) - - missing: list[str] = [] - if not nvrtc_pairs: - expected_builtins = ", ".join(_nvrtc_builtins_names(major, minor)) - if unpaired_nvrtc: - rendered = ", ".join(str(path) for path in unpaired_nvrtc) - missing.append( - "NVRTC candidates without adjacent builtins " - f"({expected_builtins}): {rendered}" - ) - else: - expected_nvrtc = ", ".join(_library_names("nvrtc", major)) - missing.append( - "NVRTC and adjacent builtins " - f"({expected_nvrtc}; {expected_builtins})" - ) - if not nvjitlink: - expected_nvjitlink = ", ".join(_library_names("nvJitLink", major)) - missing.append(f"nvJitLink ({expected_nvjitlink})") - if missing: - return None, f"{layout.toolkit_root}: missing {'; '.join(missing)}" - return ( - _ToolkitRootCandidates( - toolkit_root=layout.toolkit_root, - nvrtc_pairs=tuple(nvrtc_pairs), - nvjitlink=nvjitlink, - ), - "", - ) - - -def _claim_process_toolkit(root: Path, version: tuple[int, int]) -> None: - """Fix the toolkit selection after the first successful library load. - - Call while holding ``_PRELOAD_LOCK``. Later claims must match this root - and version because loaded libraries remain in the process. - """ - - global _PROCESS_TOOLKIT_SELECTION - - selection = (str(root.resolve()), version) - if _PROCESS_TOOLKIT_SELECTION is None: - _PROCESS_TOOLKIT_SELECTION = selection - return - if _PROCESS_TOOLKIT_SELECTION == selection: - return - selected_root, selected_version = _PROCESS_TOOLKIT_SELECTION - raise RuntimeError( - "CUDA compiler libraries are already process-global from Toolkit " - f"{selected_root} ({selected_version[0]}.{selected_version[1]}); " - f"refusing to mix Toolkit {selection[0]} ({version[0]}.{version[1]})" - ) - - -def _exact_candidate_path(candidate: Path, toolkit_root: Path) -> str: - """Require the resolved library path to stay inside the toolkit root.""" - - exact_path = Path(os.path.realpath(candidate)) - try: - exact_path.relative_to(toolkit_root.resolve()) - except ValueError as exc: - raise RuntimeError( - f"resolved CUDA Toolkit library escapes {toolkit_root}: " - f"{candidate} -> {exact_path}" - ) from exc - return str(exact_path) - - -def _load_exact_candidate(candidate: Path, *, toolkit_root: Path) -> str: - """Load a library from the selected root and retain its exact handle. - - Call while holding ``_PRELOAD_LOCK``. Retained handles let version checks - query these exact libraries, not another copy found by name. - """ - - exact_path = _exact_candidate_path(candidate, toolkit_root) - if exact_path not in _EXACT_LIBRARY_HANDLES: - handle = ctypes.CDLL( - str(candidate), - mode=getattr(ctypes, "RTLD_GLOBAL", 0), - ) - _EXACT_LIBRARY_HANDLES[exact_path] = handle - return exact_path - - -def _preload_toolkit_root( - candidates: _ToolkitRootCandidates, - *, - version: tuple[int, int], - builtins_failures: list[tuple[Path, OSError]], -) -> tuple[str, str, str] | None: - """Load builtins, NVRTC, and nvJitLink from the selected installation. - - Call while holding ``_PRELOAD_LOCK``. A builtins load failure permits the - next candidate. Once builtins load, claim the toolkit before loading NVRTC - so a later failure cannot lead to a retry with another installation. - Return the loaded paths, or ``None`` if every builtins candidate failed. - """ - - selected_nvrtc: tuple[str, str] | None = None - for nvrtc, builtins in candidates.nvrtc_pairs: - try: - builtins_path = _load_exact_candidate( - builtins, - toolkit_root=candidates.toolkit_root, - ) - except OSError as exc: - builtins_failures.append((builtins, exc)) - continue - - # A successfully loaded builtins library is now process-global. Claim - # its Toolkit before loading anything else so no later attempt can mix - # roots after a partial failure. - _claim_process_toolkit(candidates.toolkit_root, version) - try: - nvrtc_path = _load_exact_candidate( - nvrtc, - toolkit_root=candidates.toolkit_root, - ) - except OSError as exc: - raise RuntimeError( - "failed loading same-root NVRTC after its builtins library was " - f"made process-global: {nvrtc}: {exc}" - ) from exc - selected_nvrtc = nvrtc_path, builtins_path - break - - if selected_nvrtc is None: - return None - - nvjitlink_failures: list[tuple[Path, OSError]] = [] - for nvjitlink in candidates.nvjitlink: - try: - nvjitlink_path = _load_exact_candidate( - nvjitlink, - toolkit_root=candidates.toolkit_root, - ) - except OSError as exc: - nvjitlink_failures.append((nvjitlink, exc)) - continue - return (*selected_nvrtc, nvjitlink_path) - - rendered = "; ".join( - f"{candidate}: {error}" for candidate, error in nvjitlink_failures - ) - raise RuntimeError( - "failed loading all resolved CUDA Toolkit nvJitLink candidates below " - f"{candidates.toolkit_root}: {rendered}" - ) from nvjitlink_failures[-1][1] - - -def _exact_library_handle(path: str, *, kind: str) -> object: - """Get the retained library handle without another loader search.""" - - exact_path = os.path.realpath(path) - try: - return _EXACT_LIBRARY_HANDLES[exact_path] - except KeyError as exc: - raise RuntimeError( - f"the exact {kind} handle was not retained for validation: " - f"{exact_path}" - ) from exc - - -def _nvrtc_version(path: str) -> tuple[int, int]: - """Query the exact loaded NVRTC while ``_PRELOAD_LOCK`` is held.""" - - handle = _exact_library_handle(path, kind="NVRTC") - try: - version = handle.nvrtcVersion - except AttributeError as exc: - raise RuntimeError( - f"loaded NVRTC does not export nvrtcVersion: {path}" - ) from exc - version.argtypes = ( - ctypes.POINTER(ctypes.c_int), - ctypes.POINTER(ctypes.c_int), - ) - version.restype = ctypes.c_int - major = ctypes.c_int() - minor = ctypes.c_int() - result = version(ctypes.byref(major), ctypes.byref(minor)) - if result != 0: - raise RuntimeError(f"nvrtcVersion failed with result {result}: {path}") - return major.value, minor.value - - -def _nvjitlink_version(path: str) -> tuple[int, int]: - """Query the exact loaded nvJitLink while ``_PRELOAD_LOCK`` is held.""" - - handle = _exact_library_handle(path, kind="nvJitLink") - try: - version = handle.nvJitLinkVersion - except AttributeError as exc: - raise RuntimeError( - f"loaded nvJitLink does not export nvJitLinkVersion: {path}" - ) from exc - version.argtypes = ( - ctypes.POINTER(ctypes.c_uint), - ctypes.POINTER(ctypes.c_uint), - ) - version.restype = ctypes.c_int - major = ctypes.c_uint() - minor = ctypes.c_uint() - result = version(ctypes.byref(major), ctypes.byref(minor)) - if result != 0: - raise RuntimeError( - f"nvJitLinkVersion failed with result {result}: {path}" - ) - return major.value, minor.value - - -def preload_toolkit_compiler_libraries( - include_dirs: Iterable[str | os.PathLike[str]], -) -> ToolkitCompilerLibraries: - """Load compiler libraries from the toolkit selected by the CUDA headers. - - ``include_dirs`` is the compiler's ordered include path. Its first CUDA - runtime header directory selects the installation. Require NVRTC, its - builtins library, and nvJitLink there; a later include root cannot replace - an incomplete first installation. - - Loading is serialized and can fix the process's toolkit selection even if - a later library fails. Successful calls return exact library paths and - versions after checking the headers and the libraries selected by - Pathfinder. Incompatible paths, versions, or later toolkit selections - raise an error. - """ - - # Keep CUDA bindings out of root import. Pathfinder is loaded only once a - # provider actually needs compiler libraries. - from cuda.pathfinder import load_nvidia_dynamic_lib - - cuda_include_dirs = _cuda_include_dirs(include_dirs) - if not cuda_include_dirs: - raise RuntimeError( - "no resolved include directory contains cuda_runtime_api.h; " - "cannot select same-root CUDA compiler libraries" - ) - toolkit_version = _toolkit_version(cuda_include_dirs) - assert toolkit_version is not None - major, minor = toolkit_version - - # The first CUDA header root is authoritative because it is the one the - # compiler's ordered include search will use. A complete library set under - # a later include root must not silently replace it, even at the same - # reported Toolkit version. - primary_include = cuda_include_dirs[0] - primary_root = _toolkit_library_layout(primary_include).toolkit_root - candidate, diagnostic = _toolkit_root_candidates( - primary_include, - major=major, - minor=minor, - ) - if candidate is None: - other_roots = tuple( - _toolkit_library_layout(include_dir).toolkit_root - for include_dir in cuda_include_dirs[1:] - if _toolkit_library_layout(include_dir).toolkit_root != primary_root - ) - ignored = ( - "; later ordered CUDA header roots are not eligible: " - + ", ".join(str(root) for root in other_roots) - if other_roots - else "" - ) - raise RuntimeError( - "resolved CUDA headers require NVRTC, nvrtc-builtins, and " - f"nvJitLink from one CUDA Toolkit root: {diagnostic}{ignored}" - ) - candidates = [candidate] - - with _PRELOAD_LOCK: - if _PROCESS_TOOLKIT_SELECTION is not None: - selected_root, selected_version = _PROCESS_TOOLKIT_SELECTION - if selected_version != toolkit_version: - raise RuntimeError( - "CUDA compiler libraries are already process-global from " - f"Toolkit {selected_root} ({selected_version[0]}." - f"{selected_version[1]}); resolved headers report " - f"{major}.{minor}" - ) - candidates = [ - candidate - for candidate in candidates - if str(candidate.toolkit_root) == selected_root - ] - if not candidates: - raise RuntimeError( - "CUDA compiler libraries are already process-global from " - f"Toolkit {selected_root}; " - "resolved headers select different Toolkit roots" - ) - - builtins_failures: list[tuple[Path, OSError]] = [] - selected: tuple[_ToolkitRootCandidates, tuple[str, str, str]] | None = ( - None - ) - for candidate in candidates: - paths = _preload_toolkit_root( - candidate, - version=toolkit_version, - builtins_failures=builtins_failures, - ) - if paths is not None: - selected = candidate, paths - break - if selected is None: - rendered = "; ".join( - f"{candidate}: {error}" - for candidate, error in builtins_failures - ) - raise RuntimeError( - "failed loading all same-root CUDA Toolkit NVRTC builtins " - f"candidates: {rendered}" - ) from builtins_failures[-1][1] - - root_candidates, (nvrtc_path, builtins_path, nvjitlink_path) = selected - exact_paths = { - "nvrtc": nvrtc_path, - "nvJitLink": nvjitlink_path, - } - loaded = { - kind: load_nvidia_dynamic_lib(kind) - for kind in ("nvrtc", "nvJitLink") - } - actual_paths = { - kind: os.path.realpath(result.abs_path) - for kind, result in loaded.items() - } - for kind, exact_path in exact_paths.items(): - if actual_paths[kind] != exact_path: - raise RuntimeError( - f"resolved CUDA headers expect {exact_path}, " - f"but the process uses {actual_paths[kind]} for {kind}" - ) - - nvrtc_version = _nvrtc_version(nvrtc_path) - nvjitlink_version = _nvjitlink_version(nvjitlink_path) - libraries = ToolkitCompilerLibraries( - toolkit_root=str(root_candidates.toolkit_root), - nvrtc_path=nvrtc_path, - nvrtc_builtins_path=builtins_path, - nvjitlink_path=nvjitlink_path, - toolkit_version=toolkit_version, - nvrtc_version=nvrtc_version, - nvjitlink_version=nvjitlink_version, - ) - validate_nvrtc_version(libraries, nvrtc_version) - validate_nvjitlink_version(libraries, nvjitlink_version) - return libraries - - -def validate_nvrtc_version( - libraries: ToolkitCompilerLibraries, - actual_version: tuple[int, int], -) -> None: - """Require NVRTC to match the selected CUDA headers exactly.""" - - expected = libraries.toolkit_version - if actual_version == expected: - return - raise RuntimeError( - "resolved CUDA headers report Toolkit " - f"{expected[0]}.{expected[1]}, but loaded NVRTC " - f"{libraries.nvrtc_path} reports " - f"{actual_version[0]}.{actual_version[1]}" - ) - - -def validate_nvjitlink_version( - libraries: ToolkitCompilerLibraries, - actual_version: tuple[int, int], -) -> None: - """Require nvJitLink to be the same major as the headers and no older.""" - - expected = libraries.toolkit_version - if actual_version[0] == expected[0] and actual_version[1] >= expected[1]: - return - raise RuntimeError( - "resolved CUDA headers report Toolkit " - f"{expected[0]}.{expected[1]}, but loaded nvJitLink " - f"{libraries.nvjitlink_path} reports " - f"{actual_version[0]}.{actual_version[1]}; " - "nvJitLink must use the same major release " - "and be no older than the headers" - ) - - -__all__ = [ - "ToolkitCompilerLibraries", - "preload_toolkit_compiler_libraries", - "validate_nvjitlink_version", - "validate_nvrtc_version", -] diff --git a/python/cuda_coop/cuda/coop/_registration.py b/python/cuda_coop/cuda/coop/_registration.py deleted file mode 100644 index aa7e16036750..000000000000 --- a/python/cuda_coop/cuda/coop/_registration.py +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Load a compiler adapter when the host requests registration. - -The adapter installs its compiler hooks during import. Python's import cache -makes repeated registration calls safe and keeps compiler imports out of the -common API's normal import path. -""" - -from __future__ import annotations - -import importlib -from typing import Literal - - -def 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 - :mod:`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. - """ - - if backend not in ("numba-cuda-mlir", "numba_cuda_mlir"): - raise ValueError( - f"Unsupported cuda.coop backend {backend!r}; " - "expected 'numba-cuda-mlir' or 'numba_cuda_mlir'." - ) - - module_name = "cuda.coop.numba_mlir" - try: - importlib.import_module(module_name) - except ModuleNotFoundError as error: - if error.name != module_name: - raise - raise ImportError( - "This cuda-coop installation does not include the Numba-CUDA-MLIR " - "adapter. Install a cuda-coop version with Numba-CUDA-MLIR support." - ) from error diff --git a/python/cuda_coop/cuda/coop/_typing.py b/python/cuda_coop/cuda/coop/_typing.py deleted file mode 100644 index 3588508de7fd..000000000000 --- a/python/cuda_coop/cuda/coop/_typing.py +++ /dev/null @@ -1,104 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Provide compiler-neutral types for public runtime annotations. - -The adjacent stubs retain overload and dtype relationships. These aliases -also let documentation tools and runtime introspection resolve signatures -without importing an optional GPU compiler. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias, TypeVar - -if TYPE_CHECKING: - import numpy -else: - try: - import numpy - except ModuleNotFoundError as exc: - if exc.name != "numpy": - raise - numpy = None - -from ._core.api._payload import ( - TempStorageLike, - ThreadDataLike, - _ReadableThreadDataLike, -) - -BlockLoadStoreAlgorithm: TypeAlias = Literal[ - "direct", - "striped", - "vectorize", - "transpose", - "warp_transpose", - "warp_transpose_timesliced", -] - -WarpLoadStoreAlgorithm: TypeAlias = Literal[ - "direct", - "striped", - "vectorize", - "transpose", -] - - -class CompilerScalarLike(Protocol): - """Describe a compiler scalar without importing its concrete type.""" - - width: int - - @property - def dtype(self) -> object: - """Return this value's compiler dtype.""" - - def ir_value(self) -> object: - """Return this scalar's compiler IR value.""" - - -class CompilerIntegerLike(CompilerScalarLike, Protocol): - """Compiler scalar carrying the signedness metadata of an integer.""" - - signed: bool - - -if not TYPE_CHECKING and numpy is None: - CommonNumericScalar: TypeAlias = int | float | CompilerScalarLike - IntegerValue: TypeAlias = int | CompilerIntegerLike -else: - CommonNumericScalar: TypeAlias = ( - int - | float - | numpy.int8 - | numpy.uint8 - | numpy.int16 - | numpy.uint16 - | numpy.int32 - | numpy.uint32 - | numpy.int64 - | numpy.uint64 - | numpy.float32 - | numpy.float64 - | CompilerScalarLike - ) - - if TYPE_CHECKING: - IntegerValue: TypeAlias = int | numpy.integer[Any] | CompilerIntegerLike - else: - IntegerValue: TypeAlias = int | numpy.integer | CompilerIntegerLike - -ValidItems: TypeAlias = IntegerValue - -_CommonNumericT = TypeVar("_CommonNumericT", bound=CommonNumericScalar) - -CommonThreadDataLike = _ReadableThreadDataLike - -__all__ = [ - "CommonThreadDataLike", - "TempStorageLike", - "ThreadDataLike", - "_CommonNumericT", -] diff --git a/python/cuda_coop/cuda/coop/_typing.pyi b/python/cuda_coop/cuda/coop/_typing.pyi deleted file mode 100644 index 41cde9d4d4c1..000000000000 --- a/python/cuda_coop/cuda/coop/_typing.pyi +++ /dev/null @@ -1,167 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe common API types without requiring a compiler package. - -Protocols let type checkers recognize the attributes supplied by each backend. -They describe compatible values; they do not make a new Python class usable in -a GPU kernel. The compiler must still know how to lower each concrete value. -""" - -from typing import Any, Literal, Protocol, TypeAlias, TypeVar - -import numpy - -_ItemT = TypeVar("_ItemT") - -ThreadLevel: TypeAlias = Literal[ - "thread", - "gpu_thread", - "warp", - "block", - "cluster", - "grid", -] -ThreadGroupKind: TypeAlias = Literal[ - "thread", - "warp", - "block", - "cluster", - "grid", - "threads_within_warp", - "warps_within_block", -] -SynchronizableGroupKind: TypeAlias = Literal[ - "thread", - "warp", - "block", - "cluster", - "threads_within_warp", - "warps_within_block", -] -BlockLoadStoreAlgorithm: TypeAlias = Literal[ - "direct", - "striped", - "vectorize", - "transpose", - "warp_transpose", - "warp_transpose_timesliced", -] -WarpLoadStoreAlgorithm: TypeAlias = Literal[ - "direct", - "striped", - "vectorize", - "transpose", -] -LoadStoreAlgorithm: TypeAlias = BlockLoadStoreAlgorithm | WarpLoadStoreAlgorithm -TempStorageSharing: TypeAlias = Literal["shared", "exclusive"] - -class CompilerScalarLike(Protocol): - """Describe a compiler scalar without importing its concrete type.""" - - width: int - - @property - def dtype(self) -> object: - """Return this value's compiler dtype.""" - def ir_value(self) -> object: - """Return this scalar's compiler IR value.""" - -class CompilerIntegerLike(CompilerScalarLike, Protocol): - """Compiler scalar carrying the signedness metadata of an integer.""" - - signed: bool - -CommonNumericScalar: TypeAlias = ( - int - | float - | numpy.int8 - | numpy.uint8 - | numpy.int16 - | numpy.uint16 - | numpy.int32 - | numpy.uint32 - | numpy.int64 - | numpy.uint64 - | numpy.float32 - | numpy.float64 - | CompilerScalarLike -) - -_CommonNumericT = TypeVar("_CommonNumericT", bound=CommonNumericScalar) -_ReadableItemT_co = TypeVar( - "_ReadableItemT_co", bound=CommonNumericScalar, covariant=True -) -ScalarValue: TypeAlias = ( - bool | int | float | complex | numpy.number | CompilerScalarLike -) -IntegerValue: TypeAlias = int | numpy.integer[Any] | CompilerIntegerLike -TraceInteger: TypeAlias = int | numpy.integer[Any] -ValidItems: TypeAlias = IntegerValue - -class ThreadDataLike(Protocol[_ItemT]): - """Common mutable, indexable per-thread payload contract. - - Concrete compiler backends may attach additional helpers and metadata, but - common operations rely only on this payload shape and item access - contract. Structural type compatibility does not register arbitrary user - classes with a compiler; kernels must use payloads that their active - backend recognizes. - """ - - items_per_thread: int - dtype: object | None - - def __len__(self) -> int: - """Return the number of logical items owned by this thread.""" - - def __getitem__(self, index: int, /) -> _ItemT: - """Return one thread-local item.""" - - def __setitem__(self, index: int, value: _ItemT, /) -> None: - """Replace one thread-local item.""" - -class CommonThreadDataLike(Protocol[_ReadableItemT_co]): - """Describe readable per-thread items with common numeric types. - - Operations that only read a payload use this protocol. Mutable operations - use ``ThreadDataLike``, which also requires item assignment. - """ - - items_per_thread: int - dtype: object | None - - def __len__(self) -> int: - """Return the number of items owned by this thread.""" - - def __getitem__(self, index: int, /) -> _ReadableItemT_co: - """Return one numeric register value supported by the common API.""" - -class TempStorageLike(Protocol): - """Describe the attributes of an explicit scratch descriptor. - - ``size_in_bytes`` and ``alignment`` may be ``None`` so the compiler - derives them from the calls that use the descriptor. ``auto_sync`` - selects automatic reuse barriers. ``sharing`` selects one shared region - or separate slices per call site. The descriptor holds no buffer. A - supported compiler allocates the memory and rejects a size that is too - small for its uses. - """ - - size_in_bytes: int | None - alignment: int | None - auto_sync: bool - sharing: TempStorageSharing - -__all__ = [ - "BlockLoadStoreAlgorithm", - "LoadStoreAlgorithm", - "TempStorageLike", - "TempStorageSharing", - "ThreadDataLike", - "ThreadGroupKind", - "ThreadLevel", - "WarpLoadStoreAlgorithm", - "_CommonNumericT", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/__init__.py b/python/cuda_coop/cuda/coop/numba_mlir/__init__.py deleted file mode 100644 index 273051be9f0a..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/__init__.py +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Expose cooperative operations and activate their Numba-CUDA-MLIR planner. - -Import this qualified API on the host before compiling kernels. Importing it -loads the supported compiler runtime and registers the whole-function planner. -Load, Store, and ThreadData are kernel constructs. Group descriptors can also -be created on the host and used as kernel globals. Construct TempStorage -inside each kernel. During compilation, the compiler rebuilds the descriptor -from its compile-time constant arguments and validates them. The compiler -rejects a descriptor that comes from a module global. The ``local`` and -``shared`` array namespaces are resolved and cached on first access. -""" - -import importlib -from typing import TYPE_CHECKING - -from .._core.api import TempStorageLike, ThreadDataLike -from ._compiler._activation import _initialize_runtime_hooks -from ._group._load_store import load, store -from ._temp_storage import TempStorage -from ._thread_data import ThreadData -from ._thread_group import ( - Hierarchy, - ThreadGroup, - ThreadHierarchy, - this_block, - this_cluster, - this_grid, - this_thread, - this_warp, -) - -if TYPE_CHECKING: - from ._thread_data import local, shared - -__all__ = [ - "Hierarchy", - "TempStorage", - "TempStorageLike", - "ThreadData", - "ThreadDataLike", - "ThreadGroup", - "ThreadHierarchy", - "load", - "local", - "shared", - "store", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -] - - -def __getattr__(name): - """Resolve and cache the runtime array namespaces on first access.""" - if name in {"local", "shared"}: - value = getattr( - importlib.import_module(f"{__name__}._thread_data"), name - ) - globals()[name] = value - return value - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -def __dir__(): - return sorted(__all__) - - -_initialize_runtime_hooks() diff --git a/python/cuda_coop/cuda/coop/numba_mlir/__init__.pyi b/python/cuda_coop/cuda/coop/numba_mlir/__init__.pyi deleted file mode 100644 index ccd195ed6a1b..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/__init__.pyi +++ /dev/null @@ -1,37 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from .._core.api import TempStorageLike, ThreadDataLike -from ._group._load_store import load, store -from ._temp_storage import TempStorage -from ._thread_data import ThreadData, local, shared -from ._thread_group import ( - Hierarchy, - ThreadGroup, - ThreadHierarchy, - this_block, - this_cluster, - this_grid, - this_thread, - this_warp, -) - -__all__ = [ - "Hierarchy", - "TempStorage", - "TempStorageLike", - "ThreadData", - "ThreadDataLike", - "ThreadGroup", - "ThreadHierarchy", - "load", - "local", - "shared", - "store", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/__init__.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/__init__.py deleted file mode 100644 index e21683ac824b..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -__all__: tuple[str, ...] = () diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_activation.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_activation.py deleted file mode 100644 index fc309edcb7b4..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_activation.py +++ /dev/null @@ -1,221 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Activate the Numba backend after checking its compiler runtime. - -Load the compiler package, check its supported version, and import the CUDA -compiler interface before registering ``CoopWholeFunctionPlanner``. The -registration entry point acquires ``_activation_lock`` itself and imports all -planner dependencies before registering the planner as its final step. Failed -imports can therefore be retried without leaving a partially registered pass. - -Runtime-load and version diagnostics include installation details, and wrapped -import failures retain their original causes. Planner-dependency import errors -propagate unchanged. Explicit activation raises these failures; automatic -backend discovery reports them as warnings. -""" - -from __future__ import annotations - -import importlib -from threading import RLock -from types import ModuleType -from typing import cast - -from ._numba_mlir_compat import ( - NumbaMlirBackendImportError, - _require_numba_mlir_version, - _runtime_requirement, -) - -assert __package__ is not None -_BACKEND_PACKAGE = __package__.removesuffix("._compiler") -_activation_lock = RLock() - - -def _load_runtime() -> ( - tuple[ModuleType, None] | tuple[None, NumbaMlirBackendImportError] -): - """Load the Numba CUDA compiler and classify activation failures. - - Import ``numba_cuda_mlir``, validate its version, then import - ``numba_cuda_mlir.cuda``. Classify import failures to distinguish - a missing runtime, an incomplete or shadowed package, and a dependency - that fails while either module is imported. A supported Numba-CUDA-MLIR - installation includes ``.cuda``; its absence is an import-path or broken - installation diagnostic, not a supported compiler configuration. For - example, a local ``numba_cuda_mlir`` directory can shadow the installed - package while distribution metadata still reports a supported version. - This checks the Python compiler installation; it does not check GPU - availability, driver support, or the installed CUDA Toolkit. - - Return import failures with their original exception as the cause for - ``_require_runtime`` to raise. Version-validation failures raise directly, - before importing the ``.cuda`` submodule. In either case, explicit - registration and qualified imports propagate an ``ImportError``. - Automatic registration catches it and warns so the common import works. - - Returns - ------- - runtime : module or None - The ``numba_cuda_mlir.cuda`` module on success, otherwise ``None``. - error : NumbaMlirBackendImportError or None - Failure with a reason code and import details, otherwise ``None``. - Exactly one of the two return values is non-``None``. - - Raises - ------ - NumbaMlirBackendImportError - The runtime version cannot be established, falls outside the supported - range, or the version checker's ``packaging`` dependency is missing. - """ - - try: - runtime = importlib.import_module("numba_cuda_mlir") - except ImportError as exc: - missing = getattr(exc, "name", None) - if missing == "numba_cuda_mlir": - return ( - None, - NumbaMlirBackendImportError( - "backend-runtime-missing", - "cuda.coop.numba_mlir requires a compatible " - f"Numba-CUDA-MLIR runtime. {_runtime_requirement()}", - cause=exc, - missing=missing, - ), - ) - return ( - None, - NumbaMlirBackendImportError( - "transitive-runtime-import-failed", - "cuda.coop.numba_mlir found the Numba-CUDA-MLIR runtime, but " - f"importing it failed at dependency {missing!r}. " - f"{_runtime_requirement()}", - cause=exc, - missing=missing, - ), - ) - except Exception as exc: # noqa: BLE001 - preserve backend import context - return ( - None, - NumbaMlirBackendImportError( - "transitive-runtime-import-failed", - "cuda.coop.numba_mlir found the Numba-CUDA-MLIR runtime, but " - "importing it failed with " - f"{type(exc).__name__}. {_runtime_requirement()}", - cause=exc, - exception_type=type(exc).__name__, - ), - ) - - _require_numba_mlir_version(runtime) - - try: - cuda_module = importlib.import_module("numba_cuda_mlir.cuda") - except ImportError as exc: - missing = getattr(exc, "name", None) - if missing == "numba_cuda_mlir.cuda": - return ( - None, - NumbaMlirBackendImportError( - "conflicting-backend-runtime", - "cuda.coop.numba_mlir found a package named " - "'numba_cuda_mlir', but it does not provide the CUDA " - "compiler runtime at 'numba_cuda_mlir.cuda'. Remove the " - f"conflicting package. {_runtime_requirement(runtime)}", - cause=exc, - missing=missing, - ), - ) - return ( - None, - NumbaMlirBackendImportError( - "transitive-runtime-import-failed", - "cuda.coop.numba_mlir found the Numba-CUDA-MLIR runtime, but " - f"its CUDA compiler failed to import dependency {missing!r}. " - f"{_runtime_requirement(runtime)}", - cause=exc, - missing=missing, - ), - ) - except Exception as exc: # noqa: BLE001 - preserve backend import context - return ( - None, - NumbaMlirBackendImportError( - "transitive-runtime-import-failed", - "cuda.coop.numba_mlir found the Numba-CUDA-MLIR runtime, but " - "its CUDA compiler failed to import with " - f"{type(exc).__name__}. {_runtime_requirement(runtime)}", - cause=exc, - exception_type=type(exc).__name__, - ), - ) - return cuda_module, None - - -_cuda_module: ModuleType | None = None - - -def _require_runtime() -> ModuleType: - """Return the validated CUDA compiler module, loading it on first use. - - Cache successful loads only. Explicit activation propagates the classified - import or version error. Automatic registration catches it and warns. - - Returns - ------- - ModuleType - The ``numba_cuda_mlir.cuda`` module. - - Raises - ------ - NumbaMlirBackendImportError - The runtime is absent, unsupported, or fails to import. The diagnostic - retains the reason and original exception established by the loader. - """ - - global _cuda_module - if _cuda_module is None: - runtime, error = _load_runtime() - if error is not None: - raise error - _cuda_module = cast(ModuleType, runtime) - return _cuda_module - - -def _initialize_runtime_hooks() -> None: - """Register the planner after all compiler support has loaded. - - Serialize this backend's activation attempts and require a supported - runtime before importing its implementation. Those imports do not register - planner or rewrite hooks. Registering the single planner as the final step - therefore keeps failed imports from leaving partially activated - cooperative hooks; a later attempt can reuse successfully loaded modules. - The runtime deduplicates repeated registration of the same planner class. - - Raises - ------ - NumbaMlirBackendImportError - The runtime cannot be loaded or its version is unsupported. - ImportError - A compiler integration dependency cannot be imported. The original - import error propagates; automatic registration converts it to a - warning, while explicit activation leaves it visible to the caller. - """ - - with _activation_lock: - _require_runtime() - planner_module = importlib.import_module( - f"{_BACKEND_PACKAGE}._compiler._planner" - ) - from numba_cuda_mlir.extending import register_planner - - register_planner(planner_module.CoopWholeFunctionPlanner) - - -__all__ = [ - "_initialize_runtime_hooks", - "_require_runtime", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_artifacts.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_artifacts.py deleted file mode 100644 index 06d5bfbd2e6a..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_artifacts.py +++ /dev/null @@ -1,66 +0,0 @@ -# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Prepare compiled provider files for linking. - -The linker consumes binary inputs by filename after provider construction. -Temporary-file helpers therefore keep those paths alive for the owner to clean -up later. -""" - -import os -import tempfile -from collections import namedtuple -from typing import IO - -version = namedtuple("version", ("major", "minor")) - - -def make_binary_tempfile(content: bytes, suffix: str) -> IO[bytes]: - """Write a persistent link-input file and return its closed handle. - - Compiler link inputs are consumed by filename after provider construction, - so flush and close the file before returning while leaving its path - intact. The caller must arrange successful-file cleanup, normally through - an ``Invocable`` or shared-bundle owner. If writing raises, remove the - partially written file before propagating the exception. - - Parameters - ---------- - content : bytes - Binary image to write without buffering. - suffix : str - Filename suffix used by the linker to recognize the image format. - - Returns - ------- - IO[bytes] - Closed ``NamedTemporaryFile`` wrapper whose ``name`` identifies the - persistent file. It is not an open stream for subsequent writes. - """ - - with tempfile.NamedTemporaryFile( - mode="w+b", suffix=suffix, buffering=0, delete=False - ) as tmp: - try: - tmp.write(content) - except Exception: - name = tmp.name - tmp.close() - try: - os.unlink(name) - except FileNotFoundError: - pass - raise - return tmp - - -def check_in(name, arg, values): - """Validate a small closed-set compiler option.""" - - if arg not in values: - raise ValueError(f"{name} must be in {values} ; got {name} = {arg}") - - -__all__ = ["check_in", "make_binary_tempfile", "version"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_caching.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_caching.py deleted file mode 100644 index 7f24ea8e0c0d..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_caching.py +++ /dev/null @@ -1,397 +0,0 @@ -# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Reuse compilation results across processes when their inputs are unchanged. - -NVRTC compilation can be expensive even when the resulting device code is -small. The decorator in this module stores results under a key derived from -the function's identity and arguments, allowing later processes to reuse the -same artifact. The compilation caller supplies source, target, options, and -compiler/header identity as ordinary values; the cache does not inspect -compiler objects or discover dependencies itself. - -``CUDA_COOP_ENABLE_CACHE`` enables persistent caching when set to ``"1"`` or -any value other than empty, ``"0"``, ``"false"``, ``"no"``, or ``"off"``. -Comparison ignores surrounding whitespace and letter case. Unset disables the -disk cache. This flag does not control the separate in-memory -LRU cache on NVRTC compilation. - -``XDG_CACHE_HOME`` selects the cache parent directory on non-Windows systems; -the default is ``~/.cache``. ``LOCALAPPDATA`` selects it on Windows; the -default is ``~/AppData/Local``. Only the variable for the current platform is -read. Its value must be an absolute path: unset, empty, and relative values -use the platform default. Cache files live under a ``cccl`` subdirectory, -partitioned further by callable identity. - -The enable flag and selected cache location are captured when this module is -imported. Set them before backend initialization imports this module; later -environment changes do not reconfigure its wrappers. Directory or write -``OSError`` failures disable disk caching for the rest of the process, while -unreadable entries simply miss. Compilation proceeds without persistent -caching in those cases. -""" - -import hashlib -import json -import os -import tempfile -from base64 import b64decode, b64encode -from collections.abc import Callable -from functools import wraps -from typing import ParamSpec, TypeVar, cast - -_P = ParamSpec("_P") -_R = TypeVar("_R") - -_FALSE_CACHE_VALUES = frozenset(("", "0", "false", "no", "off")) -_CACHE_ENV_VALUE = os.environ.get("CUDA_COOP_ENABLE_CACHE") -_ENABLE_CACHE = ( - _CACHE_ENV_VALUE is not None - and _CACHE_ENV_VALUE.strip().lower() not in _FALSE_CACHE_VALUES -) -_CACHE_USABLE = _ENABLE_CACHE -_CACHE_SCHEMA_VERSION = 6 -_CACHE_MISS = object() - - -def _cache_location() -> str: - """Select the platform cache parent, accepting only absolute overrides.""" - - if os.name == "nt": - cache_home = os.environ.get("LOCALAPPDATA", "") - fallback = ("AppData", "Local") - else: - cache_home = os.environ.get("XDG_CACHE_HOME", "") - fallback = (".cache",) - if not os.path.isabs(cache_home): - cache_home = os.path.join(os.path.expanduser("~"), *fallback) - return os.path.join(cache_home, "cccl") - - -_CACHE_LOCATION = _cache_location() - - -def _json_cache_key(value: object) -> object: - """Convert supported key values into a tagged JSON-serializable tree. - - Plain JSON loses distinctions such as tuples versus lists and cannot - encode bytes. Tag these containers, preserve a tuple subclass's qualified - type name, and encode bytes as base64. Represent dictionaries as sorted - key/value pairs so supported non-string keys can participate without JSON - coercing them to strings. Scalar values retain JSON's native encoding. - - Parameters - ---------- - value : object - A scalar (``None``, bool, int, float, or str), bytes, or a recursively - supported tuple, list, or dictionary. Dictionary entries are ordered - by ``repr`` of their keys. Nested containers - are supported; cycles are not. - - Returns - ------- - object - JSON-compatible representation used to hash arguments. This is a key - encoding, not a general-purpose serialization format for cache values. - - Raises - ------ - TypeError - A value has no supported key representation. The disk-cache wrapper - treats this as a request to compute without caching. Examples include - a ``set``, a ``pathlib.Path``, or a custom object, including one - nested inside an otherwise supported container. - RecursionError - A container contains a cycle. The disk-cache wrapper - does not catch this error. - """ - - if value is None or isinstance(value, (bool, int, float, str)): - return value - if isinstance(value, bytes): - return { - "__cuda_coop_numba_mlir_cache_type__": "bytes", - "data": b64encode(value).decode("ascii"), - } - if isinstance(value, tuple): - value_type = f"{type(value).__module__}.{type(value).__qualname__}" - return { - "__cuda_coop_numba_mlir_cache_type__": value_type, - "items": [_json_cache_key(item) for item in value], - } - if isinstance(value, list): - return { - "__cuda_coop_numba_mlir_cache_type__": "builtins.list", - "items": [_json_cache_key(item) for item in value], - } - if isinstance(value, dict): - return { - "__cuda_coop_numba_mlir_cache_type__": "builtins.dict", - "items": [ - (_json_cache_key(key), _json_cache_key(item)) - for key, item in sorted( - value.items(), key=lambda entry: repr(entry[0]) - ) - ], - } - raise TypeError(f"Unsupported disk cache key value: {value!r}") - - -def json_hash(*args: object, **kwargs: object) -> str: - """Hash positional and keyword arguments using the current cache schema. - - Prefix the serialized argument tree with the schema version so changes to - the cache format invalidate older entries. Keyword names are part of the - key, and their insertion order does not affect it. This function does not - add callable identity; ``disk_cache`` passes that identity as an argument. - - Parameters - ---------- - *args : object - Positional key components accepted by ``_json_cache_key``. - **kwargs : object - Named key components with the same serialization restrictions. - - Returns - ------- - str - Hexadecimal SHA-256 digest of the versioned argument representation. - - Raises - ------ - TypeError - An argument cannot be represented by the cache key encoder. - """ - - hasher = hashlib.sha256() - hasher.update(f"v{_CACHE_SCHEMA_VERSION}:".encode()) - payload = json.dumps( - _json_cache_key((args, kwargs)), - separators=(",", ":"), - sort_keys=True, - ) - hasher.update(payload.encode("utf-8")) - return hasher.hexdigest() - - -def _cache_identity_path(cache_identity: str) -> str: - """Create a separate cache directory for this callable's identity.""" - - identity_hash = hashlib.sha256(cache_identity.encode("utf-8")).hexdigest() - path = os.path.join(_CACHE_LOCATION, identity_hash) - os.makedirs(path, exist_ok=True) - return path - - -def _cache_value_type(value: object) -> str: - if isinstance(value, bytes): - return "bytes" - return f"{type(value).__module__}.{type(value).__qualname__}" - - -def _encode_cache_value(value: object) -> object: - """Keep LTO bytes and nested layout tuples together in one JSON entry. - - Tag bytes and recurse through tuples so a result such as - ``(image, ((size, alignment), ...))`` survives a process restart with its - types intact. Other values pass through to JSON unchanged; this encoder - does not recursively tag the contents of lists or dictionaries. - """ - - if isinstance(value, bytes): - return { - "__cuda_coop_numba_mlir_cache_type__": "bytes", - "data": b64encode(value).decode("ascii"), - } - if isinstance(value, tuple): - return { - "__cuda_coop_numba_mlir_cache_type__": "tuple", - "items": [_encode_cache_value(item) for item in value], - } - return value - - -def _decode_cache_value(value: object) -> object: - """Restore bytes and tuple structure from the cache's tagged values. - - Strict base64 decoding rejects malformed byte payloads. The outer reader - handles unusable entries and checks their schema and top-level result - type; this helper does not validate a particular compiler result shape. - """ - - if ( - isinstance(value, dict) - and value.get("__cuda_coop_numba_mlir_cache_type__") == "bytes" - ): - return b64decode(value["data"].encode("ascii"), validate=True) - if ( - isinstance(value, dict) - and value.get("__cuda_coop_numba_mlir_cache_type__") == "tuple" - ): - return tuple(_decode_cache_value(item) for item in value["items"]) - return value - - -def _read_cache(path: str | os.PathLike[str]) -> object: - """Read one cache entry, treating unusable files as cache misses. - - Validate the schema version and the decoded value's recorded type before - returning it. This catches stale formats and values whose JSON round trip - changes their top-level type. It does not authenticate cached data or - validate the value against a particular compiler invocation. - - Parameters - ---------- - path : str or os.PathLike[str] - JSON cache entry to open. The function does not modify or remove it. - - Returns - ------- - object - Decoded value on success, otherwise the unique ``_CACHE_MISS`` - sentinel for an unreadable, malformed, stale, or type-inconsistent - entry. ``None`` is a valid cached result and is distinct from a miss. - """ - - try: - with open(path, encoding="utf-8") as f: - cached = json.load(f) - if not isinstance(cached, dict): - return _CACHE_MISS - if cached.get("version") != _CACHE_SCHEMA_VERSION: - return _CACHE_MISS - value = _decode_cache_value(cached["value"]) - if cached.get("value_type") != _cache_value_type(value): - return _CACHE_MISS - return value - except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError): - return _CACHE_MISS - - -def _write_cache(path: str | os.PathLike[str], value: object) -> None: - """Serialize a result and atomically replace its cache entry. - - Write a schema and top-level type tag alongside the value, encoding bytes - as base64 and recursively tagging tuples. Use a temporary file in the - destination directory, flush and fsync it, then replace the destination - so readers do not see a partially written JSON document. A tuple result, - such as an LTO image with its layouts, is one entry, so a reader gets both - parts or neither. Concurrent writers may replace the same entry. On a - write or replacement failure, attempt to remove the temporary file and - propagate the error for ``disk_cache`` to handle. - - Parameters - ---------- - path : str or os.PathLike[str] - Destination entry. Its parent directory must already exist. - value : object - Result to persist: bytes, recursively supported tuples, or a - JSON-serializable value. A later read also requires the decoded - top-level type to match the saved type tag. - - Raises - ------ - OSError - Creating, writing, syncing, or replacing the entry fails. - TypeError or ValueError - The result cannot be serialized as a cache entry. - """ - - cached = { - "version": _CACHE_SCHEMA_VERSION, - "value_type": _cache_value_type(value), - "value": _encode_cache_value(value), - } - fd, tmp_path = tempfile.mkstemp( - prefix=f".{os.path.basename(path)}.", - dir=os.path.dirname(path), - text=True, - ) - try: - with os.fdopen(fd, "w", encoding="utf-8") as f: - json.dump(cached, f) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp_path, path) - except Exception: - try: - os.unlink(tmp_path) - except OSError: - pass - raise - - -def disk_cache(func: Callable[_P, _R]) -> Callable[_P, _R]: - """Decorate a computation with best-effort persistent result caching. - - Cache files are namespaced by the callable's module and qualified name; - the entry key includes that identity, arguments, and the cache schema. - Callers must therefore include every compilation input that affects the - result in the arguments. Callable source changes do not invalidate entries - by themselves. The enable flag is read from ``CUDA_COOP_ENABLE_CACHE`` - when this module is imported. - - Key values may be JSON scalars, bytes, tuples, lists, or dictionaries of - supported values. A ``set``, ``pathlib.Path``, or custom object has no key - encoding, even when nested inside a list or dictionary. Those values cause - the wrapper to call ``func`` without reading or writing the cache. A path - represented as a string, for example, can be cached; a ``Path`` object - cannot. See ``_json_cache_key`` for the full encoding rules. - - Results that cannot be encoded skip the write. Directory or write - ``OSError`` failures disable caching for every wrapper in this module for - the rest of the process; unreadable entries simply miss. Exceptions from - the wrapped computation propagate without being cached. There is no lock - around computation, so concurrent misses can compute the same result - before atomic publication. - - Parameters - ---------- - func : callable - Computation whose result can be reused for the same serialized inputs. - Arguments must be accepted by ``_json_cache_key`` to use the cache; - results must round-trip through the cache value encoding. - - Returns - ------- - callable - Wrapper preserving ``func`` metadata and forwarding its arguments. - Returns a decoded cache hit or the newly computed result. - """ - - cache_identity = f"{func.__module__}.{func.__qualname__}" - - @wraps(func) - def cacher(*args: _P.args, **kwargs: _P.kwargs) -> _R: - """Read a usable entry or compute and try to save the result.""" - - global _CACHE_USABLE - - if not _CACHE_USABLE: - return func(*args, **kwargs) - - try: - key = json_hash(cache_identity, *args, **kwargs) - path = os.path.join(_cache_identity_path(cache_identity), key) - except (TypeError, ValueError): - return func(*args, **kwargs) - except OSError: - _CACHE_USABLE = False - return func(*args, **kwargs) - - if os.path.isfile(path): - cached = _read_cache(path) - if cached is not _CACHE_MISS: - return cast(_R, cached) - - result = func(*args, **kwargs) - try: - _write_cache(path, result) - except (TypeError, ValueError): - pass - except OSError: - _CACHE_USABLE = False - return result - - return cacher diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_descriptor_provenance.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_descriptor_provenance.py deleted file mode 100644 index 5fc018054672..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_descriptor_provenance.py +++ /dev/null @@ -1,253 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Find where ``ThreadData`` and ``TempStorage`` values came from in the IR. - -These objects describe per-thread payloads and shared scratch for the -compiler; they do not become ordinary Python objects in a GPU kernel. To -replace them with arrays and storage pointers, the compiler must recover their -constructors and options even when the kernel assigns aliases, casts values, -or joins branches. A "reaching definition" is an assignment that may supply a -variable's value; "provenance" here means following those assignments back to -their sources. - -The helpers trace descriptor sources, bind ``TempStorage`` constructor -options, and collect the element types written through payload aliases. Group -planning uses these facts to select an operation implementation; call -rewriting uses them to materialize payloads and scratch storage. Each caller -supplies its own definition and constant lookups because the available IR -facts differ between those phases. The shared traversal keeps their treatment -of aliases and control-flow joins consistent without evaluating arbitrary -kernel code. -""" - -from __future__ import annotations - -from collections.abc import Callable, Iterable, Iterator -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from numba_cuda_mlir.numba_cuda.core import ir - - from .._temp_storage import TempStorage -else: - from numba_cuda_mlir.numbair_transforms import ir - - -def descriptor_definitions( - value: object, - definitions: Callable[[ir.Var], Iterable[object]], - *, - seen: set[str] | None = None, -) -> Iterator[tuple[str | None, object]]: - """Trace a possible descriptor to the assignments that supply its value. - - Follow aliases, casts, iterator unpacking, and phi inputs using the - caller's definition lookup, so the same traversal works before and after - single static assignment (SSA) construction. A cycle contributes no leaf. - Concrete non-descriptors, including ``None``, remain leaves: callers need - to reject paths that mix such values with descriptors rather than - accepting one valid constructor. Sibling paths are independent, so a - shared constructor may be yielded more than once. - - Parameters - ---------- - value : object - IR value to trace. A non-variable is yielded unchanged with no owner. - definitions : callable - Lookup accepting an IR variable and returning all its reaching - definitions, including multiple definitions before SSA construction. - seen : set of str, optional - Variable names on the current recursion path. This traversal copies - the set before extending it and does not mutate the supplied set. - - Yields - ------ - owner : str or None - Name of the variable whose definition is the leaf, not necessarily the - original alias. ``None`` denotes a non-variable input. - leaf : object - Definition reached after following the supported forwarding forms. No - descriptor recognition or constructor validation is performed. - """ - - if not isinstance(value, ir.Var): - yield None, value - return - if seen is None: - seen = set() - if value.name in seen: - return - seen = {*seen, value.name} - for definition in definitions(value): - if isinstance(definition, ir.Var): - yield from descriptor_definitions( - definition, definitions, seen=seen - ) - elif isinstance(definition, ir.Expr) and definition.op in { - "cast", - "exhaust_iter", - }: - yield from descriptor_definitions( - definition.value, definitions, seen=seen - ) - elif isinstance(definition, ir.Expr) and definition.op == "phi": - for incoming in getattr(definition, "incoming_values", ()): - yield from descriptor_definitions( - incoming, definitions, seen=seen - ) - else: - yield value.name, definition - - -def temp_storage_constructor( - call: ir.Expr, - constant: Callable[..., Any], - *, - syntax_error: type[Exception] = TypeError, -) -> TempStorage: - """Build a storage descriptor from an IR constructor call. - - Share argument binding and descriptor validation between group planning - and provider rewriting while letting each phase supply its own constant - resolver. This creates the host-side descriptor only; it does not allocate - scratch or modify the call. Constructor defaults and value validation come - from ``TempStorage`` itself. - - Parameters - ---------- - call : ir.Expr - Call already recognized as a ``TempStorage`` constructor. - constant : callable - Resolve an argument as ``constant(value, name=parameter_name)``. The - callback controls constant specialization and its diagnostics. - syntax_error : type of Exception, optional - Exception class for unsupported call syntax, duplicate arguments, or - unknown keywords. Defaults to ``TypeError``; exceptions from the - resolver and descriptor constructor propagate unchanged. - - Returns - ------- - TempStorage - Descriptor with resolved, validated constructor arguments. - - Raises - ------ - TypeError - Invalid call syntax when ``syntax_error`` has its default value, or an - invalid descriptor option type. - ValueError - A descriptor option has an invalid value. - """ - - from .._temp_storage import TempStorage - - if ( - getattr(call, "vararg", None) is not None - or getattr(call, "varkwarg", None) is not None - ): - raise syntax_error( - "TempStorage does not accept *args or **kwargs; pass " - "size_in_bytes, alignment, auto_sync, and sharing explicitly." - ) - if len(call.args) > 1: - raise syntax_error( - "TempStorage accepts only size_in_bytes positionally; " - "alignment, auto_sync, and sharing are keyword-only." - ) - refs = dict(zip(("size_in_bytes",), call.args)) - for name, value in call.kws: - if name in refs: - raise syntax_error( - f"TempStorage got multiple values for argument {name!r}" - ) - refs[name] = value - unexpected = sorted( - refs.keys() - {"size_in_bytes", "alignment", "auto_sync", "sharing"} - ) - if unexpected: - raise syntax_error( - f"TempStorage got unexpected keyword(s): {', '.join(unexpected)}" - ) - return TempStorage( - **{name: constant(value, name=name) for name, value in refs.items()} - ) - - -def payload_write_dtypes( - func_ir: ir.FunctionIR, - payload: object, - dtype: Callable[[ir.Var], object], -) -> Iterator[object]: - """Yield known element dtypes written through a payload's aliases. - - Build an alias set to a fixed point across the entire function, following - assignments, casts, iterator unpacking, and phi inputs in both directions. - This lets inference start at either a constructor or a later alias. The - scan is flow-insensitive: it collects writes throughout the function, - including all connected phi inputs, without checking path feasibility or - write order. The caller decides whether the collected dtypes agree. - - Parameters - ---------- - func_ir : ir.FunctionIR - Function with alias assignments and element writes; not modified. - payload : object - Variable whose aliases are searched; non-variables yield nothing. - dtype : callable - Lookup accepting the variable assigned to an element and returning its - dtype, or ``None`` when unknown. - - Yields - ------ - object - Known dtype for each ``SetItem`` or ``StaticSetItem`` write through - the alias set, in the representation supplied by ``dtype``. This - helper does not normalize the callback's results. Unknown dtypes are - skipped; duplicates are retained. An empty result does not establish - that the payload has no writes. - """ - - if not isinstance(payload, ir.Var): - return - aliases = {payload.name} - changed = True - while changed: - changed = False - for block in func_ir.blocks.values(): - for statement in block.body: - if not isinstance(statement, ir.Assign): - continue - definition = statement.value - sources = () - if isinstance(definition, ir.Var): - sources = (definition,) - elif isinstance(definition, ir.Expr): - if definition.op in {"cast", "exhaust_iter"}: - sources = (definition.value,) - elif definition.op == "phi": - sources = getattr(definition, "incoming_values", ()) - source_names = { - source.name - for source in sources - if isinstance(source, ir.Var) - } - if statement.target.name in aliases or source_names & aliases: - additions = {statement.target.name, *source_names} - aliases - if additions: - aliases.update(additions) - changed = True - - for block in func_ir.blocks.values(): - for statement in block.body: - if not isinstance(statement, (ir.SetItem, ir.StaticSetItem)): - continue - if ( - isinstance(statement.target, ir.Var) - and statement.target.name in aliases - and isinstance(statement.value, ir.Var) - ): - value_dtype = dtype(statement.value) - if value_dtype is not None: - yield value_dtype diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_errors.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_errors.py deleted file mode 100644 index 5cecd5385f06..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_errors.py +++ /dev/null @@ -1,321 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Report unsupported group, payload, and storage choices during planning. - -Each exception names the kernel argument or missing fact that prevents a -cooperative operation from being lowered. Keeping their message construction -here lets the planner's branches focus on the condition being checked and -keeps related diagnostics consistent. These errors leave planning with a -diagnostic; they do not request another pass or literal specialization. -""" - -from textwrap import fill - -from ._group_planner_support import GroupRewriteError - - -def _wrap_diagnostic(message: str) -> str: - """Wrap a diagnostic without splitting names or hyphenated terms.""" - - return fill( - message, width=80, break_long_words=False, break_on_hyphens=False - ) - - -class InvalidGroupSelectorError(ValueError): - """The common API received a selector outside its supported choices. - - Point callers to the qualified API when they need backend-only controls. - """ - - def __init__(self, operation, parameter, choices): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.{operation} {parameter} must be one of: " - f"{choices}; use a backend-qualified import for " - "backend-only controls", - ) - ) - - -class CyclicArrayProvenanceError(GroupRewriteError): - """A payload cycle has no concrete source that establishes its form. - - The planner cannot choose scalar or array rules from a backedge alone. - """ - - def __init__(self, operation): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir.{operation} could not resolve " - "cyclic array provenance to a concrete scalar or array value", - ) - ) - - -class InconsistentArrayExtentError(GroupRewriteError): - """Payload aliases imply different item counts for one provider call.""" - - def __init__(self): - super().__init__( - _wrap_diagnostic( - "cuda.coop.numba_mlir array aliases have inconsistent " - "items_per_thread extents", - ) - ) - - -class InconsistentTupleExtentError(GroupRewriteError): - """The same tuple element reaches payloads with different item counts.""" - - def __init__(self): - super().__init__( - _wrap_diagnostic( - "cuda.coop.numba_mlir tuple projections have inconsistent " - "items_per_thread extents", - ) - ) - - -class NonConstantThreadGroupError(GroupRewriteError): - """The group operand cannot be reconstructed as a planning descriptor. - - Provider selection needs the group before Numba type inference. - """ - - def __init__(self, operation): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir.{operation} requires a compile-time " - f"ThreadGroup from this_*()", - ) - ) - - -class NonConstantGroupArgumentError(GroupRewriteError): - """A specialization control cannot be recovered as a constant. - - This is an unsupported value after constant resolution. A non-literal - kernel argument instead raises ``ForceLiteralArg`` to request a retry. - """ - - def __init__(self, value_name): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir group arguments that shape provider " - f"specialization must be compile-time constants; got " - f"{value_name!r}", - ) - ) - - -class InconsistentLoopTupleExtentError(GroupRewriteError): - """A merged tuple element has conflicting known payload item counts.""" - - def __init__(self): - super().__init__( - _wrap_diagnostic( - "cuda.coop.numba_mlir loop-carried tuple payloads have " - "inconsistent items_per_thread extents", - ) - ) - - -class InconsistentLoopPayloadExtentError(GroupRewriteError): - """A merged payload has conflicting known per-thread item counts.""" - - def __init__(self): - super().__init__( - _wrap_diagnostic( - "cuda.coop.numba_mlir loop-carried payloads have " - "inconsistent items_per_thread extents", - ) - ) - - -class EscapingGroupDescriptorError(GroupRewriteError): - """A ThreadGroup or ThreadHierarchy value would be needed at runtime. - - Group planning erases these host descriptions after consuming their - supported uses. Report the remaining variable names before changing IR. - """ - - def __init__(self, names): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir ThreadGroup/ThreadHierarchy values " - f"are compile-time descriptors and may only feed this_*(), " - f"group_by(), or group-first primitives; descriptor use " - f"involving {names!r} would escape to runtime", - ) - ) - - -class InvalidLoadStoreAlgorithmError(ValueError): - """A Load/Store algorithm is outside the choices for this group kind.""" - - def __init__(self, operation, choices, group_kind=None): - group = "" if group_kind is None else f" for {group_kind} groups" - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir.{operation} " - f"algorithm{group} must be one of: {choices}", - ) - ) - - -class UnknownBlockDimensionError(GroupRewriteError): - """Provider selection lacks the exact block size needed for its shape.""" - - def __init__(self, operation): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir.{operation} requires an exact block " - f"dimension before provider selection", - ) - ) - - -class UnknownLoadStoreExtentError(GroupRewriteError): - """Provider selection cannot determine the per-thread payload size.""" - - def __init__(self, operation): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir.{operation} requires a static " - f"items_per_thread extent before provider selection", - ) - ) - - -class UnknownLoadStoreDtypeError(GroupRewriteError): - """Provider selection cannot determine the payload's element type.""" - - def __init__(self, operation): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir.{operation} could not infer " - "a dtype before provider selection", - ) - ) - - -class MemoryDtypeMismatchError(TypeError): - """The memory operand and payload require different element types. - - Load/Store transfers elements with matching types; it does not insert - an element conversion between these operands. - """ - - def __init__(self, operation, memory_dtype, payload_dtype): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir.{operation} memory dtype " - f"{memory_dtype} does not match " - f"payload dtype {payload_dtype}", - ) - ) - - -class NonConstantTempStorageError(GroupRewriteError): - """A storage argument cannot be recovered as a compile-time descriptor. - - The provider needs its storage options before scratch is materialized. - """ - - def __init__(self, operation): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir.{operation} temp_storage must " - "resolve to a compile-time TempStorage descriptor", - ) - ) - - -class UnsupportedLoadStoreTargetError(GroupRewriteError): - """The shared plan selected a target this backend cannot lower.""" - - def __init__(self, target): - super().__init__( - _wrap_diagnostic( - "cuda.coop.numba_mlir Load/Store received an unsupported " - f"lowering target {target!r}", - ) - ) - - -class DefaultDtypeMismatchError(TypeError): - """A runtime load fill value has a different type from the payload. - - Its existing type must match; contextual coercion applies to untyped - static literals, not arbitrary runtime operands. - """ - - def __init__(self, value_dtype, payload_dtype): - super().__init__( - "cuda.coop.numba_mlir.load runtime oob_default dtype " - f"{value_dtype}\n" - f"does not match payload dtype {payload_dtype}" - ) - - -class UnknownLoadStoreProviderError(GroupRewriteError): - """The plan's CUB implementation has no matching provider factory.""" - - def __init__(self, provenance): - super().__init__( - _wrap_diagnostic( - f"cuda.coop.numba_mlir Load/Store received an unknown CUB " - f"implementation provenance {provenance!r}", - ) - ) - - -class CommonLoadPayloadError(TypeError): - """The common load API received an output without a ThreadData origin. - - Qualified Numba calls also support local-array destinations. - """ - - def __init__(self): - super().__init__( - _wrap_diagnostic( - "cuda.coop.load requires output to be a fixed-size " - "ThreadData payload in the common API; " - "use cuda.coop.numba_mlir for backend-qualified " - "local-array payload support", - ) - ) - - -class CommonStorePayloadError(TypeError): - """The common store API received an unsupported payload form. - - Its payload must be a scalar or ThreadData. Qualified Numba calls also - support local arrays. - """ - - def __init__(self): - super().__init__( - _wrap_diagnostic( - "cuda.coop.store accepts only a scalar or fixed-size " - "ThreadData value payload in the common API; use " - "cuda.coop.numba_mlir for backend-qualified local-array " - "payload support", - ) - ) - - -class UnsupportedLoadStoreGroupError(NotImplementedError): - """Load/Store has no supported implementation for the requested group. - - Include the shared plan's reason when it supplies one. - """ - - def __init__(self, operation, reason=None): - message = f"cuda.coop.numba_mlir.{operation} cannot lower this group" - if reason is not None: - message += f": {reason}" - super().__init__(_wrap_diagnostic(message)) diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_load_store.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_load_store.py deleted file mode 100644 index 0c94015856ea..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_load_store.py +++ /dev/null @@ -1,891 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Choose and prepare the implementation of a group ``load()`` or ``store()``. - -The group planner delegates these two operations here after resolving which -threads participate. Recover the per-thread item count and element type, check -memory and payload compatibility, and classify optional controls as omitted, -static, or runtime values. Pass those facts to the compiler-neutral planner, -then translate its decision into a call to the matching private CUB factory. - -The returned IR carries both the runtime operands and the complete lowering -plan. Later in the same whole-function planner, call rewriting specializes the -factory and supplies payload arrays and any shared scratch it needs. Warp -operations may also need an offset for their particular warp within the block; -this module emits that arithmetic because it depends on runtime thread -indices. It builds replacement statements without installing them, so the -caller can check every group descriptor use before it changes any block. -""" - -import inspect -import operator -from typing import Any, cast - -import numba_cuda_mlir.numba_cuda.types as numba_types -from numba_cuda_mlir import cuda as _cuda_module - -from cuda.coop._core import ( - ArgumentBinding, - BindingKind, - GroupLoadStoreAlgorithm, - GroupLoadStoreKind, - GroupLoadStoreSemantics, - GroupLoweringPlan, - GroupLoweringTarget, - StorageOwnership, - ThreadGroup, - make_group_primitive_call, - plan_group_primitive, -) - -from ._group_errors import ( - CommonLoadPayloadError, - CommonStorePayloadError, - DefaultDtypeMismatchError, - InvalidLoadStoreAlgorithmError, - MemoryDtypeMismatchError, - NonConstantTempStorageError, - UnknownBlockDimensionError, - UnknownLoadStoreDtypeError, - UnknownLoadStoreExtentError, - UnknownLoadStoreProviderError, - UnsupportedLoadStoreGroupError, - UnsupportedLoadStoreTargetError, -) -from ._group_planner_support import GroupRewriteError, ir -from ._group_planning import GroupPlanningContext -from ._load_store_algorithms import ( - _BLOCK_LOAD_STORE_ALGORITHMS, - _WARP_LOAD_STORE_ALGORITHMS, - _resolve_algorithm, -) -from ._operations import ( - RewriteOperationSpecification, - register_group_primitive, - register_rewrite_operation, -) -from ._parameters import _validate_common_numeric_dtype, coerce_static_scalar -from ._rewrite_load_store import ( - analyze_load_store_match, - infer_load_store_payload, - prepare_load_store_runtime_args, - validate_load_store_runtime_controls, -) - -_SUPPORTED_WARP_WIDTHS = frozenset({1, 2, 4, 8, 16, 32}) - - -def _load_store_algorithm( - value: object, - *, - operation: str, - group_kind: str, -) -> str: - """Resolve an algorithm from the choices for the selected group kind. - - Physical and logical warps use the warp choices; block groups use the - block choices. Keep the group kind in an invalid-choice diagnostic so - callers can see why an algorithm is unavailable for this request. - """ - - algorithm_scope = ( - "warp" if group_kind in {"warp", "threads_within_warp"} else group_kind - ) - allowed = ( - _WARP_LOAD_STORE_ALGORITHMS - if algorithm_scope == "warp" - else _BLOCK_LOAD_STORE_ALGORITHMS - ) - try: - return _resolve_algorithm( - value, allowed, f"cuda.coop.numba_mlir.{operation}" - ) - except ValueError: - choices = ", ".join(sorted(allowed)) - raise InvalidLoadStoreAlgorithmError( - operation, choices, group_kind - ) from None - - -# Match the selected implementation to a provider operation. Each key contains -# the library, include path, C++ type, and method recorded by shared planning. -_CUB_PLAN_ROUTES = { - ( - "CUB", - "cub/block/block_load.cuh", - "cub::BlockLoad", - "Load", - ): "load", - ( - "CUB", - "cub/block/block_store.cuh", - "cub::BlockStore", - "Store", - ): "store", - ( - "CUB", - "cub/warp/warp_load.cuh", - "cub::WarpLoad", - "Load", - ): "load", - ( - "CUB", - "cub/warp/warp_store.cuh", - "cub::WarpStore", - "Store", - ): "store", -} - - -class _LoadStorePlanning: - """Plan registered Load and Store operations with shared rules. - - A compiler family is a module supplying planning and rewrite hooks for one - or more operation names. Load and store share this implementation because - they share payload, dtype, and tile rules; each has its own registration. - """ - - def __init__(self, context: GroupPlanningContext) -> None: - self._context = context - - def _validate_common_arguments( - self, operation: str, bound: inspect.BoundArguments - ) -> None: - """Normalize selectors in the bound arguments of a common API call. - - Update ``algorithm`` in place. Rewriting bypasses the Python wrapper, - so compiled calls must apply its selector restrictions here. - """ - - bound.arguments["algorithm"] = self._context.validate_common_selector( - operation, - "algorithm", - bound.arguments["algorithm"], - _BLOCK_LOAD_STORE_ALGORITHMS, - ) - - def _scope_factory( - self, - plan: GroupLoweringPlan, - operation: str, - ) -> tuple[Any, dict[str, Any]]: - """Choose the registered factory and its group-dimension arguments. - - Use the plan's block/warp target and storage ownership to select the - matching private factory. Warp factories also need the logical group - width. Both forms need exact enclosing block dimensions, including - when the selected algorithm uses no scratch. - - Return ``(factory, kwargs)`` without invoking the factory. Reject - missing block dimensions, unsupported warp topology, or an unknown - target before building provider-call IR. - """ - - group = plan.resolved_group - assert group.hierarchy is not None - block_dim = group.hierarchy.block_dim - if block_dim is None: - raise UnknownBlockDimensionError(operation) - - from .._lowering import _load_store - - assert plan.temp_storage is not None - storage_free = plan.temp_storage.ownership is StorageOwnership.NONE - if plan.target is GroupLoweringTarget.CUB_BLOCK: - factory_name = ( - operation if storage_free else f"_{operation}_with_storage" - ) - factory_kwargs = {"threads_per_block": block_dim} - elif plan.target is GroupLoweringTarget.CUB_WARP: - if ( - plan.topology is None - or plan.topology.group_kind - not in {"warp", "threads_within_warp"} - or plan.topology.logical_width not in _SUPPORTED_WARP_WIDTHS - ): - raise GroupRewriteError( - "cuda.coop.numba_mlir Warp Load/Store requires a " - "power-of-two lowering width in [1, 32]" - ) - factory_name = ( - f"warp_{operation}" - if storage_free - else f"_warp_{operation}_with_storage" - ) - factory_kwargs = { - "threads_per_block": block_dim, - "threads_in_warp": plan.topology.logical_width, - } - else: - raise UnsupportedLoadStoreTargetError(plan.target.value) - return ( - getattr(_load_store, factory_name), - factory_kwargs, - ) - - def _planning_oob_default( - self, - value: Any, - *, - payload_dtype: numba_types.Type, - ) -> ArgumentBinding: - """Validate the ``load()`` operation's out-of-bounds fill value. - - Classify ``oob_default`` as omitted, static, or runtime. A runtime - value with a known dtype must match the payload exactly; an unknown - dtype is left for later validation. Static values retain their - original dtype information during coercion, preserving the distinction - between Python literals whose type is determined by the operation and - values whose numeric width is already established. - - Parameters - ---------- - value : ir.Var or object - Bound ``oob_default`` argument, including ``None`` for omission. - payload_dtype : numba_types.Type - Normalized element dtype selected for the load. - - Returns - ------- - ArgumentBinding - Omitted or runtime binding, or a static binding containing the - validated, coerced default. - - Raises - ------ - DefaultDtypeMismatchError - A runtime default has a known dtype different from the payload. - TypeError - A known scalar dtype is unsupported or incompatible with the - payload. - ValueError - A static default is nonfinite or outside the payload dtype range. - """ - - binding = self._context.planning_binding(value) - if binding.kind is BindingKind.OMITTED: - return binding - if binding.kind is BindingKind.RUNTIME: - value_dtype = self._context.dtype(value) - if value_dtype is None: - return binding - value_dtype = _validate_common_numeric_dtype( - value_dtype, - operation="load", - parameter="oob_default", - ) - if value_dtype != payload_dtype: - raise DefaultDtypeMismatchError(value_dtype, payload_dtype) - return binding - - scalar = binding.value - resolved, provenance = self._context.try_static_scalar_provenance(value) - assert resolved and provenance is not None - scalar = coerce_static_scalar( - scalar, - payload_dtype, - operation="load", - parameter="oob_default", - source_dtype=provenance.dtype, - ) - return ArgumentBinding.static(scalar) - - def _planning_items_per_thread(self, operation: str, payload: Any) -> int: - """Recover a fixed payload extent for provider specialization. - - A scalar store uses one item per thread. Array payloads must have a - known extent, and a load always needs an array destination. This - determines the item count; dtype and common API payload rules are - checked separately. - """ - - is_array = self._context.is_array(operation, payload) - if not is_array: - if operation == "load": - raise TypeError( - "cuda.coop.numba_mlir.load output " - "must be a fixed-size local array" - ) - return 1 - extent = self._context.array_extent(payload) - if extent is None: - raise UnknownLoadStoreExtentError(operation) - return extent - - def _plan_load_store( - self, - *, - operation: str, - group: ThreadGroup, - bound: inspect.BoundArguments, - ) -> GroupLoweringPlan: - """Build a supported core plan from a bound group memory call. - - Recover payload shape and element type before selecting a provider. - Memory dtype takes precedence when known, but a known payload dtype - must match it. For a ``store()`` operation with an untyped array - payload, inspect values written to its elements first; for a static - scalar payload, check whether its value and original dtype can be used - with the destination dtype. A ``load()`` operation records the - inferred output dtype in the planning context for later group calls. - - Classify optional scalar controls by provenance, parse caller storage, - and pass compiler-neutral load/store semantics plus exact launch facts - to the core planner. Warp groups reject explicit storage here. The - result selects topology, storage, synchronization, and implementation - metadata; it does not yet construct provider-call IR. The dtype cache - may already be updated if a later validation fails. - - Parameters - ---------- - operation : {"load", "store"} - Public operation being planned. - group : ThreadGroup - Group already resolved against this kernel's launch facts. - bound : inspect.BoundArguments - Public call arguments with defaults applied; values may be IR - variables or host constants. The mapping is read without - modification. - - Returns - ------- - GroupLoweringPlan - Supported core plan containing the inferred semantic operation and - the selected provider implementation. - - Raises - ------ - GroupRewriteError - Payload extent or dtype is unknown, or storage provenance is - invalid. - TypeError - A payload or scalar control has an unsupported type, or known - payload and memory dtypes disagree. - ValueError - A scalar or algorithm is invalid, or explicit storage is supplied - for a warp group. - ForceLiteralArg - A shape, selector, or storage option needs literal specialization. - NotImplementedError - The core planner cannot lower the requested group operation. - """ - - from .._lowering._core import NumbaMlirCoreAdapter - - payload_name = "output" if operation == "load" else "value" - memory_name = "source" if operation == "load" else "destination" - payload = bound.arguments[payload_name] - payload_is_array = self._context.is_array(operation, payload) - items_per_thread = self._planning_items_per_thread(operation, payload) - payload_dtype = self._context.dtype(payload) - if operation == "store" and payload_dtype is None: - payload_dtype = self._context.payload_write_dtype(payload) - memory_dtype = self._context.dtype(bound.arguments[memory_name]) - dtype = memory_dtype if memory_dtype is not None else payload_dtype - if dtype is None: - raise UnknownLoadStoreDtypeError(operation) - dtype = _validate_common_numeric_dtype(dtype, operation=operation) - if operation == "store" and not payload_is_array: - resolved, provenance = self._context.try_static_scalar_provenance( - payload - ) - if resolved: - assert provenance is not None - coerce_static_scalar( - provenance.value, - dtype, - operation="store", - parameter="value", - source_dtype=provenance.dtype, - ) - payload_dtype = dtype - if payload_dtype is not None: - _validate_common_numeric_dtype(payload_dtype, operation=operation) - if ( - payload_dtype is not None - and memory_dtype is not None - and payload_dtype != memory_dtype - ): - raise MemoryDtypeMismatchError( - operation, memory_dtype, payload_dtype - ) - if operation == "load": - self._context.record_thread_data_dtype(payload, dtype) - - oob_default = ( - self._planning_oob_default( - bound.arguments["oob_default"], - payload_dtype=dtype, - ) - if operation == "load" - else ArgumentBinding.omitted() - ) - - algorithm = _load_store_algorithm( - self._context.constant(bound.arguments["algorithm"]), - operation=operation, - group_kind=group.kind, - ) - temp_storage_value = bound.arguments["temp_storage"] - if group.kind in {"warp", "threads_within_warp"} and not ( - self._context.is_none(temp_storage_value) - ): - raise ValueError( - f"cuda.coop.numba_mlir.{operation} temp_storage is not " - "supported for Warp groups; omit it so the implementation " - "can provide per-group storage" - ) - storage_options: dict[str, Any] = {} - if not self._context.is_none(temp_storage_value): - descriptor = self._context.temp_storage(temp_storage_value) - if descriptor is None: - raise NonConstantTempStorageError(operation) - size_in_bytes, alignment, auto_sync, sharing = descriptor - storage_options = { - "storage_ownership": StorageOwnership.CALLER, - "storage_sharing": sharing, - "storage_size_in_bytes": size_in_bytes, - "storage_alignment": alignment, - "storage_auto_sync": auto_sync, - } - - semantics = GroupLoadStoreSemantics( - kind=GroupLoadStoreKind(operation), - dtype=NumbaMlirCoreAdapter().core_dtype(dtype), - items_per_thread=items_per_thread, - algorithm=GroupLoadStoreAlgorithm(algorithm), - valid_items=self._context.planning_binding( - bound.arguments["valid_items"] - ), - oob_default=oob_default, - offset=self._context.planning_binding(bound.arguments["offset"]), - **storage_options, - ) - plan = plan_group_primitive( - make_group_primitive_call(group, semantics), - self._context.launch, - ) - try: - return plan.require_supported() - except NotImplementedError as exc: - raise UnsupportedLoadStoreGroupError(operation, exc) from exc - - @staticmethod - def _plan_provider_operation(plan: GroupLoweringPlan) -> str: - """Route the selected CUB implementation to its provider operation. - - Match the plan's library, header, C++ type, and method against the - known routes. This checks the implementation the core actually - selected. The caller then verifies that it matches the requested load - or store. - """ - - if plan.target not in { - GroupLoweringTarget.CUB_BLOCK, - GroupLoweringTarget.CUB_WARP, - }: - raise UnsupportedLoadStoreTargetError(plan.target.value) - assert plan.provenance is not None - try: - return _CUB_PLAN_ROUTES[plan.provenance.semantic_key] - except KeyError as exc: - raise UnknownLoadStoreProviderError( - plan.provenance.semantic_key - ) from exc - - @staticmethod - def _planned_argument( - binding: ArgumentBinding, - runtime_value: Any, - ) -> Any: - """Forward runtime operands and keep static binding descriptors. - - Factories read any plain non-None value as a runtime operand. Keep a - static control wrapped in its ArgumentBinding so the factory embeds - it as a constant. - """ - - return runtime_value if binding.kind is BindingKind.RUNTIME else binding - - def _warp_group_effective_offset( - self, - statements: list[ir.Assign], - *, - inst: ir.Assign, - plan: GroupLoweringPlan, - binding: ArgumentBinding, - runtime_value: Any, - items_per_thread: int, - ) -> ir.Var: - """Build IR that computes the starting element for this warp's tile. - - Each physical or logical warp must access its own tile within the - block's data. Compute the current thread's x-fastest linear rank as - ``x + block_x * (y + block_y * z)``. Divide by the logical group - width, then multiply by ``width * items_per_thread``. Add the user's - base offset to that tile origin. The arithmetic remains runtime IR - because the warp's position in the block depends on ``threadIdx``, - even when the user's base offset is static. For a one-dimensional - block, the linear rank is simply ``threadIdx.x``. - - Parameters - ---------- - statements : list of ir.Assign - Output list extended in place with constants, CUDA attribute - reads, and arithmetic assignments in dependency order. - inst : ir.Assign - Original public call; supplies scope and source locations. - plan : GroupLoweringPlan - Plan with a physical or logical warp topology, supported - power-of-two width, and exact enclosing block dimensions. - binding : ArgumentBinding - Planned user offset. Omission supplies zero; a static binding - supplies its value; a runtime binding selects ``runtime_value``. - runtime_value : object - Original offset operand, used only for a runtime binding. - items_per_thread : int - Planned payload extent used to determine the group's tile size. - - Returns - ------- - ir.Var - Variable holding the provider call's effective element offset. - - Raises - ------ - GroupRewriteError - The plan lacks exact warp topology or block dimensions. - """ - - topology = plan.topology - participation = plan.participation - if ( - topology is None - or topology.group_kind not in {"warp", "threads_within_warp"} - or topology.logical_width not in _SUPPORTED_WARP_WIDTHS - or participation is None - or participation.exact_block_dim is None - ): - raise GroupRewriteError( - "Warp Load/Store requires exact power-of-two topology" - ) - scope = inst.target.scope - loc = inst.loc - - def new_var(stem: str) -> ir.Var: - return self._context.new_var(scope, loc, f"warp_tile_{stem}") - - def value_var(value: Any, stem: str) -> ir.Var: - return self._context.value_var( - statements, - scope=scope, - loc=loc, - stem=f"warp_tile_{stem}", - value=value, - ) - - def binary( - function: Any, lhs: ir.Var, rhs: ir.Var, stem: str - ) -> ir.Var: - result = new_var(stem) - statements.append( - ir.Assign(ir.Expr.binop(function, lhs, rhs, loc), result, loc) - ) - return result - - module = value_var(_cuda_module, "cuda") - thread_idx = new_var("thread_idx") - statements.append( - ir.Assign( - ir.Expr.getattr(module, "threadIdx", loc), thread_idx, loc - ) - ) - - def component(axis: str) -> ir.Var: - result = new_var(f"thread_idx_{axis}") - statements.append( - ir.Assign(ir.Expr.getattr(thread_idx, axis, loc), result, loc) - ) - return result - - block_dim = participation.exact_block_dim - linear_rank = component("x") - if block_dim[1] > 1 or block_dim[2] > 1: - y = component("y") - z = component("z") - yz = binary( - operator.mul, - value_var(block_dim[1], "block_y"), - z, - "linear_yz", - ) - yz = binary(operator.add, y, yz, "linear_y") - yz = binary( - operator.mul, - value_var(block_dim[0], "block_x"), - yz, - "linear_x_stride", - ) - linear_rank = binary(operator.add, linear_rank, yz, "linear_rank") - group_index = binary( - operator.floordiv, - linear_rank, - value_var(topology.logical_width, "group_width"), - "group_index", - ) - tile_origin = binary( - operator.mul, - group_index, - value_var( - topology.logical_width * items_per_thread, - "group_tile_items", - ), - "origin", - ) - if binding.kind is BindingKind.OMITTED: - base_offset = 0 - elif binding.kind is BindingKind.STATIC: - base_offset = binding.value - else: - base_offset = runtime_value - return binary( - operator.add, - tile_origin, - value_var(base_offset, "base_offset"), - "effective_offset", - ) - - def _lower_load_store( - self, - inst: ir.Assign, - *, - operation: str, - group: ThreadGroup, - bound: inspect.BoundArguments, - is_common_root: bool, - ) -> list[Any]: - """Build replacement IR for a public ``load()`` or ``store()`` call. - - Validate common-API payload restrictions, obtain a semantic lowering - plan, and select the matching private factory by provider provenance - and scope. Translate planned dtype, shape, algorithm, and scalar - bindings into the provider's arguments. Warp plans that require a tile - origin first emit the effective-offset arithmetic. Other static - controls retain their ``ArgumentBinding`` objects; runtime controls - retain their original IR operands. - - The final provider call carries the complete plan for the subsequent - provider rewrite. Planning may record inferred payload dtypes, but - this method only returns replacement statements; the whole-function - planner installs them after validating all descriptor uses. - - Parameters - ---------- - inst : ir.Assign - Public call assignment whose result target and location are - preserved. - operation : {"load", "store"} - Registered public operation to lower. - group : ThreadGroup - Group resolved against the current kernel launch. - bound : inspect.BoundArguments - Public call arguments with defaults applied. Runtime operands are - forwarded; the argument mapping is not modified. - is_common_root : bool - Whether the call uses the common API. Enforce its payload - contract and retain its operation marker on the provider call. - - Returns - ------- - list of object - Ordered IR statements computing any warp offset, materializing - factory arguments, and calling the selected provider with the - original target. - - Raises - ------ - GroupRewriteError - The payload is unresolved, provenance selects the wrong - operation, or the provider contract conflicts with the plan. - TypeError - A common-API payload is unsupported or known dtypes disagree. - ValueError - A scalar, algorithm, or explicit storage argument is invalid. - ForceLiteralArg - A planning argument needs literal specialization. - NotImplementedError - The requested group operation has no supported lowering. - """ - - from .._lowering._core import NumbaMlirCoreAdapter - - if is_common_root: - if operation == "load": - if not self._context.is_thread_data( - operation, "output", bound.arguments["output"] - ): - raise CommonLoadPayloadError() - else: - value = bound.arguments["value"] - if self._context.is_array(operation, value) and ( - not self._context.is_thread_data(operation, "value", value) - ): - raise CommonStorePayloadError() - plan = self._plan_load_store( - operation=operation, - group=group, - bound=bound, - ) - planned_operation = self._plan_provider_operation(plan) - if planned_operation != operation: - raise GroupRewriteError( - f"cuda.coop.numba_mlir.{operation} selected the " - f"{planned_operation!r} provider" - ) - factory, factory_kwargs = self._scope_factory(plan, planned_operation) - assert plan.implementation is not None - factory_kwargs.update( - { - "algorithm": plan.implementation.metadata["algorithm"], - "dtype": NumbaMlirCoreAdapter().normalize_dtype( - plan.implementation.template_arguments["T"] - ), - "items_per_thread": plan.implementation.template_arguments[ - "ITEMS_PER_THREAD" - ], - } - ) - if is_common_root: - factory_kwargs["_common_root_operation"] = operation - semantics = cast(GroupLoadStoreSemantics, plan.call.operation) - requires_runtime_effective_offset = bool( - plan.implementation.metadata.get( - "requires_runtime_effective_offset", False - ) - ) - statements: list[Any] = [] - for public_name, factory_name in ( - ("valid_items", "num_valid_items"), - ("oob_default", "oob_default"), - ("offset", "offset"), - ): - binding = getattr(semantics, public_name) - if public_name == "offset" and requires_runtime_effective_offset: - continue - if binding.kind is BindingKind.OMITTED: - continue - factory_kwargs[factory_name] = self._planned_argument( - binding, - bound.arguments[public_name], - ) - if requires_runtime_effective_offset: - factory_kwargs["offset"] = self._warp_group_effective_offset( - statements, - inst=inst, - plan=plan, - binding=semantics.offset, - runtime_value=bound.arguments["offset"], - items_per_thread=semantics.items_per_thread, - ) - if operation == "store": - factory_kwargs["_group_root_store"] = True - if not self._context.is_none(bound.arguments["temp_storage"]): - factory_kwargs["temp_storage"] = bound.arguments["temp_storage"] - if operation == "load": - runtime_args = [ - bound.arguments["source"], - bound.arguments["output"], - ] - else: - value = bound.arguments["value"] - runtime_args = [bound.arguments["destination"], value] - statements.extend( - self._context.rewrite_call( - inst, - lowering_plan=plan, - factory=factory, - args=runtime_args, - kwargs=factory_kwargs, - ) - ) - return statements - - -def _lower_registered_load_store( - context: GroupPlanningContext, *args: Any, **kwargs: Any -) -> list[Any]: - return _LoadStorePlanning(context)._lower_load_store(*args, **kwargs) - - -def _validate_registered_common_arguments( - context: GroupPlanningContext, - operation: str, - bound: inspect.BoundArguments, -) -> None: - _LoadStorePlanning(context)._validate_common_arguments(operation, bound) - - -for _operation in ("load", "store"): - register_group_primitive( - _operation, - lower=_lower_registered_load_store, - validate_common_arguments=_validate_registered_common_arguments, - ) -del _operation - -_COMMON_REWRITE_KWARGS = frozenset( - { - "algorithm", - "dim", - "dtype", - "items_per_thread", - "num_valid_items", - "offset", - "threads_per_block", - "threads_in_warp", - "_common_root_operation", - } -) -register_rewrite_operation( - "load", - RewriteOperationSpecification( - factory_namespaces=frozenset({"block", "warp"}), - dtype_factory_kwargs=frozenset({"dtype"}), - runtime_arg_counts=frozenset({2, 3, 4}), - runtime_factory_kwargs=("num_valid_items", "oob_default"), - runtime_factory_kw_prerequisites=(("oob_default", "num_valid_items"),), - allowed_factory_kwargs=_COMMON_REWRITE_KWARGS | {"oob_default"}, - required_factory_kwargs=frozenset({"threads_per_block", "dtype"}), - accepts_temp_storage=True, - scalar_binding_kwargs=frozenset({"num_valid_items", "oob_default"}), - runtime_offset_kwarg="offset", - infer_payload=infer_load_store_payload, - analyze_match=analyze_load_store_match, - prepare_runtime_args=prepare_load_store_runtime_args, - validate_runtime_controls=validate_load_store_runtime_controls, - ), -) -register_rewrite_operation( - "store", - RewriteOperationSpecification( - factory_namespaces=frozenset({"block", "warp"}), - dtype_factory_kwargs=frozenset({"dtype"}), - runtime_arg_counts=frozenset({2, 3}), - runtime_factory_kwargs=("num_valid_items",), - runtime_factory_kw_prerequisites=(), - allowed_factory_kwargs=_COMMON_REWRITE_KWARGS | {"_group_root_store"}, - required_factory_kwargs=frozenset({"threads_per_block", "dtype"}), - accepts_temp_storage=True, - scalar_binding_kwargs=frozenset({"num_valid_items"}), - runtime_offset_kwarg="offset", - infer_payload=infer_load_store_payload, - analyze_match=analyze_load_store_match, - prepare_runtime_args=prepare_load_store_runtime_args, - validate_runtime_controls=validate_load_store_runtime_controls, - ), -) - - -__all__: tuple[str, ...] = () diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planner.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planner.py deleted file mode 100644 index f8d6e9ebda08..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planner.py +++ /dev/null @@ -1,2192 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Replace kernel group descriptions with calls to chosen implementations. - -A kernel can describe its participating threads with ``this_block()`` and -``group_by()``, then pass that group to operations such as ``load()``. Those -Python descriptions cannot execute on the device. This module recovers them -from the kernel's IR, resolves their sizes against the configured launch, and -asks the operation's registered planning code to choose an implementation. -These compile-time ``ThreadGroup`` and ``ThreadHierarchy`` values are called -group descriptors here. ``ThreadData`` supplies the per-thread payload and -is lowered separately to a local array. - -``CoopWholeFunctionPlanner`` invokes this work after device-helper inlining -and before type inference. ``has_group_markers`` first determines whether -group resolution is needed; ``_GroupPlanning`` then requests launch facts and -runs ``_GroupCallPlanner``. The latter builds replacements, checks that group -descriptors have no remaining runtime uses, and substitutes calls to private -provider factories. A provider factory is a host callable that specializes the -selected implementation for types, item counts, and launch dimensions. - -The same whole-function planner repairs the changed IR before its next phase -materializes those provider calls and allocates payloads and scratch storage. -This module does not register a separate Numba pass. Its analysis is limited -to supported constructors and expressions; ordinary Numba type inference -handles the rest of the kernel afterward. -""" - -import inspect -from enum import Enum -from numbers import Integral -from typing import TYPE_CHECKING, Any, cast - -import numba_cuda_mlir.numba_cuda.types as _numba_types -from numba_cuda_mlir.cuda.local import array as _cuda_local_array -from numba_cuda_mlir.extending import require_launch_config -from numba_cuda_mlir.numba_cuda.core.errors import ForceLiteralArg - -import cuda.coop._core.api as _common_api -import cuda.coop._core.api._dispatch as _common_dispatch -from cuda.coop._core import ( - LaunchFactOrigin, - LaunchFacts, - ThreadGroup, - ThreadHierarchy, - resolve_thread_group, -) - -from .._temp_storage import TempStorage -from .._thread_data import ThreadData -from ._group_errors import ( - CyclicArrayProvenanceError, - EscapingGroupDescriptorError, - InconsistentArrayExtentError, - InconsistentLoopPayloadExtentError, - InconsistentLoopTupleExtentError, - InconsistentTupleExtentError, - InvalidGroupSelectorError, - NonConstantGroupArgumentError, - NonConstantThreadGroupError, -) -from ._group_planner_support import ( - _COMMON_GROUP_CONSTRUCTORS, - _GROUP_CONSTRUCTORS, - _NAME_COUNTER, - GroupRewriteError, - _group_operation_name, - _is_common_root_operation, - ir, -) -from ._group_planning import GroupPlanningContext -from ._operations import group_primitive -from ._scalar_provenance import ( - try_resolve_static_scalar, - try_resolve_static_scalar_provenance, -) - -if TYPE_CHECKING: - from ._planner import CoopWholeFunctionPlanner - - -class _GroupCallPlanner: - """Resolve group descriptions before ordinary type inference. - - A group such as ``this_block()`` describes which threads cooperate; it - is not an object that the device should construct. For one compiler - attempt, this planner recovers those descriptions and payload facts - from IR, asks each operation's registered planner to choose a - provider, then replaces the public calls with calls that the next - rewriting phase understands. - - Analysis and mutation are separate. Helpers collect descriptor - assignments and replacement statements while the original assignment - definitions remain available. ``run()`` installs the replacements only - after every operation and remaining descriptor use has passed - validation. Create a fresh planner for each attempt; its caches and - pending replacements belong to that IR. - - Parameters - ---------- - state : compiler state - Numba state containing ``func_ir`` and argument types in ``args``. - Device helpers must already be inlined into this function. - launch_config : dict of str to object - Normalized configuration returned by ``require_launch_config``. - Its exact dimensions let group resolution choose a supported - topology. - - Attributes - ---------- - func_ir : ir.FunctionIR - Function being analyzed and, after successful planning, rewritten. - launch : LaunchFacts - Shared-core representation of the configured launch, including - where its dimensions came from. - descriptor_assigns : set of ir.Assign - Group and hierarchy construction, alias, and cast assignments - whose compile-time meaning has been consumed. Their targets become - ``None``. - dead_func_names : set of str - Callable-variable names consumed by descriptor or operation - rewriting. Their old assignments also become ``None`` so Numba - need not type the removed Python callables. - replacements : dict of ir.Assign to list - Original operation assignments mapped to pending provider-call IR. - Building this map does not yet replace the function's block - bodies. - context : GroupPlanningContext - Restricted analysis interface passed to operation planners. It - exposes group, payload, scalar, and storage facts and retains - inferred payload dtypes for later calls in this planning attempt. - """ - - def __init__(self, state, launch_config: dict[str, Any]) -> None: - self.state = state - self.func_ir = state.func_ir - self.launch_config = launch_config - self.launch = self._make_launch_facts(launch_config) - # Pending edits refer to the original IR until all calls validate. - self.dead_func_names: set[str] = set() - self.descriptor_assigns: set[ir.Assign] = set() - self.replacements: dict[ir.Assign, list[Any]] = {} - # Host-side thread groups and hierarchies, keyed by IR variable name. - self._group_cache: dict[str, ThreadGroup] = {} - self._hierarchy_cache: dict[str, ThreadHierarchy] = {} - self.context = GroupPlanningContext(self) - - @staticmethod - def _make_launch_facts(config: dict[str, Any]) -> LaunchFacts: - """Convert the configured launch to shared group-planning facts. - - The shared resolver consumes ``LaunchFacts`` rather than Numba - metadata. Record the configured block, grid, and optional cluster - dimensions with origins identifying the backend launch configuration. - Supplying a cluster is also the evidence for ``cluster_launch``. This - launch interface does not establish cooperative-grid launch support, - so that fact stays false. - - Parameters - ---------- - config : dict of str to object - Normalized launch configuration supplied by the compiler - dispatcher. Other entries, such as dynamic shared-memory bytes, - are not topology facts and do not enter this record. - - Returns - ------- - LaunchFacts - Dimensions and their origins for resolving a requested thread - group. - """ - - block = config.get("block") - grid = config.get("grid") - cluster = config.get("cluster") - cluster_launch = cluster is not None - origins = [ - LaunchFactOrigin( - fact="exact_block_dim", - source="numba_cuda_mlir_launch_config", - verified=True, - ), - LaunchFactOrigin( - fact="exact_grid_dim", - source="numba_cuda_mlir_launch_config", - verified=True, - ), - LaunchFactOrigin( - fact="cluster_launch", - source="numba_cuda_mlir_launch_config", - verified=True, - ), - ] - if cluster is not None: - origins.append( - LaunchFactOrigin( - fact="exact_cluster_dim", - source="numba_cuda_mlir_launch_config", - verified=True, - ) - ) - return LaunchFacts( - exact_block_dim=block, - exact_grid_dim=grid, - exact_cluster_dim=cluster, - cluster_launch=cluster_launch, - cooperative_launch=False, - provenance=tuple(origins), - ) - - def _definition(self, value: Any) -> Any: - """Look up the expression that supplies a variable's value. - - A definition is the right-hand side of an assignment, such as the - ``ThreadData(...)`` call assigned to a payload variable. Call and - descriptor recognition use this helper when they need a single source. - For an IR variable ``value``, return Numba's recorded definition, or - ``None`` if it is missing or ambiguous. Other inputs pass through - unchanged. Use ``_all_definitions`` when several assignments may supply - the value after a branch or loop. - """ - - if not isinstance(value, ir.Var): - return value - try: - return self.func_ir.get_definition(value) - except KeyError: - return None - - def _all_definitions(self, value: ir.Var) -> tuple[Any, ...]: - """Find every recorded assignment that may supply a variable's value. - - An assignment is a reaching definition for a use if execution can get - from that assignment to the use without replacing the value. For - example, two branches can assign different ``ThreadData`` objects to - the same payload variable. Array analysis must check both sources - before choosing a provider. This helper conservatively collects all - recorded sources for the name; it does not prove that each assignment - can reach the particular use being analyzed. - - ``value`` is an IR variable; its ``name`` indexes Numba's definition - table. Return the recorded right-hand-side expressions as a tuple, - falling back to ``_definition`` if the table has no entry. An empty - tuple means no definition was found, so callers have no source to - classify. - """ - - definitions = getattr(self.func_ir, "_definitions", {}).get( - value.name, () - ) - if definitions: - return tuple(definitions) - definition = self._definition(value) - return () if definition is None else (definition,) - - def _callable(self, value: Any) -> Any: - """Recover the Python object behind a call target or attribute chain. - - Resolve globals, captured values, constants, and their attributes - without calling the resulting object. Operation recognition compares - the returned object's identity with registered constructors and - operations. Return ``None`` for an unresolved base or failed attribute - lookup. Despite this helper's name, the recovered object is not - required to be callable. - """ - - current = self._definition(value) - attrs: list[str] = [] - while isinstance(current, ir.Expr) and current.op == "getattr": - attrs.append(current.attr) - current = self._definition(current.value) - if isinstance(current, (ir.Global, ir.FreeVar, ir.Const)): - obj = current.value - elif callable(current): - obj = current - else: - return None - try: - for attr in reversed(attrs): - obj = getattr(obj, attr) - except (AttributeError, ImportError): - return None - return obj - - def _reject_literal_unroll_value(self, value: Any, parameter: str) -> None: - """Reject constant controls that depend on literal unrolling. - - Group planning needs shapes and selectors before the literal-unroll - pass has expanded its iterations. Trace all reaching definitions and - expression operands for a recognized ``literal_unroll`` call instead - of trying to resolve an iteration value prematurely. Unrelated - literal-unroll loops are allowed. Cycles terminate the search without - establishing a dependency. - - Parameters - ---------- - value : ir.Var or object - Argument whose IR dependencies are inspected. Non-variables have - no dependencies to inspect. - parameter : str - Description of the shape, selector, or other compile-time control - used in the diagnostic. - - Returns - ------- - None - The value has no detected dependency on a recognized unroll call. - - Raises - ------ - GroupRewriteError - A dependency reaches the standard or vendored Numba literal-unroll - marker; use constant controls or an ordinary loop with a fixed - cooperative shape. - """ - - def depends_on_unroll(current, seen): - if not isinstance(current, ir.Var) or current.name in seen: - return False - seen = {*seen, current.name} - for definition in self._all_definitions(current): - if isinstance(definition, ir.Var): - if depends_on_unroll(definition, seen): - return True - elif isinstance(definition, ir.Expr): - if definition.op == "call": - function = self._callable(definition.func) - if getattr( - function, "__name__", None - ) == "literal_unroll" and getattr( - function, "__module__", None - ) in { - "numba.misc.special", - "numba_cuda_mlir.numba_cuda.misc.special", - }: - return True - if any( - depends_on_unroll(source, seen) - for source in definition.list_vars() - ): - return True - return False - - if depends_on_unroll(value, set()): - raise GroupRewriteError( - "cuda.coop.numba_mlir does not support literal_unroll values " - f"that determine {parameter}. Write separate " - "cooperative calls with explicit constant shapes/selectors, " - "or use an ordinary loop with a fixed cooperative shape." - ) - - def _reject_literal_unroll_constructors(self) -> None: - """Check constructor controls before building any group replacements. - - Inspect recognized group, hierarchy, payload, and storage - constructors. Their arguments must be usable before Numba expands - ``literal_unroll``; ``_reject_literal_unroll_value`` supplies the - diagnostic for a detected dependency. Ordinary loops and unrelated - unrolls remain available. - """ - - constructors = { - *_GROUP_CONSTRUCTORS, - ThreadHierarchy, - ThreadData, - _common_api.ThreadData, - TempStorage, - _common_api.TempStorage, - } - for block in self.func_ir.blocks.values(): - for inst in block.body: - call = getattr(inst, "value", None) - if not isinstance(call, ir.Expr) or call.op != "call": - continue - function = self._callable(call.func) - if function not in constructors: - continue - for argument in (*call.args, *(value for _, value in call.kws)): - self._reject_literal_unroll_value( - argument, f"{function.__name__} arguments" - ) - - def _constant(self, value: Any) -> Any: - """Resolve a required compile-time argument or request specialization. - - Try direct argument and constant definitions first, then reconstruct - hierarchy or group descriptors, and finally use Numba constant - inference. A non-literal kernel argument requests dispatcher - specialization rather than being treated as an unsupported value. - Descriptor reconstruction may populate the planner's caches. This is - the required-constant path; callers classifying optional runtime - controls use ``_try_static_scalar`` instead. - - Parameters - ---------- - value : ir.Var or object - IR variable to resolve, or an already-resolved value returned - unchanged. - - Returns - ------- - object - Compile-time value, including ``None`` or a reconstructed - descriptor. - - Raises - ------ - ForceLiteralArg - A function argument must be recompiled with a literal type. - NonConstantGroupArgumentError - The value cannot be resolved by descriptor or constant inference. - GroupRewriteError - A recognized descriptor has unsupported call syntax or arguments. - """ - - if not isinstance(value, ir.Var): - return value - definition = self._definition(value) - if isinstance(definition, ir.Arg): - position = definition.index - argtype = self.state.args[position] - if not isinstance(argtype, _numba_types.Literal): - raise ForceLiteralArg({position}) - return argtype.literal_value - if isinstance(definition, (ir.Global, ir.FreeVar, ir.Const)): - return definition.value - hierarchy = self._hierarchy(value) - if hierarchy is not None: - return hierarchy - group = self._group(value) - if group is not None: - return group - try: - return self.func_ir.infer_constant(value) - except Exception as exc: - raise NonConstantGroupArgumentError(value.name) from exc - - def _try_constant(self, value: Any) -> tuple[bool, Any]: - """Probe for a constant without requesting dispatcher specialization. - - Literal arguments and arguments known to be ``None`` resolve directly. - Other inputs use ``_constant``; its specialization requests and group - resolution failures become an unresolved result. This broad constant - probe may use Numba constant inference. Use ``_try_static_scalar`` - when an operation needs evidence that a scalar was explicitly supplied - as a compile-time value. - - Parameters - ---------- - value : ir.Var or object - Argument to inspect, usually an IR variable or an already-resolved - Python default. Unknown runtime values leave this probe unresolved. - - Returns - ------- - resolved : bool - Whether a constant was recovered without another compiler attempt. - value : object - Recovered value, or ``None`` on failure. Check ``resolved`` to - distinguish failure from a known ``None``. - """ - if isinstance(value, ir.Var): - definition = self._definition(value) - if isinstance(definition, ir.Arg): - argtype = self.state.args[definition.index] - if isinstance(argtype, _numba_types.Literal): - return (True, argtype.literal_value) - if isinstance(argtype, _numba_types.NoneType) or ( - isinstance(argtype, _numba_types.Omitted) - and argtype.value is None - ): - return (True, None) - return (False, None) - try: - return (True, self._constant(value)) - except (ForceLiteralArg, GroupRewriteError): - return (False, None) - - def _try_static_scalar( - self, - value: Any, - ) -> tuple[bool, Any]: - """Resolve an explicitly static scalar with its known numeric width. - - Delegate to the shared provenance walker with this function's recorded - definitions and argument types. Constants and literal arguments must - agree across all paths; runtime expressions are not evaluated to turn - them into constants. The result is ``(resolved, value)``; an - unresolved value is ``(False, None)``. No specialization is requested. - """ - - return try_resolve_static_scalar( - value, - definitions=self._all_definitions, - argument_type=lambda index: ( - self.state.args[index] - if 0 <= index < len(self.state.args) - else None - ), - ) - - def _try_static_scalar_provenance(self, value: Any) -> tuple[bool, Any]: - """Resolve a static scalar and retain its dtype provenance. - - Use the same conservative traversal as ``_try_static_scalar``, - retaining the provenance record so callers can distinguish an untyped - Python literal from an explicitly typed value. Return ``(True, - provenance)`` when every path agrees, otherwise ``(False, None)``. A - known ``None`` is contained in a provenance record and is distinct - from failure. - """ - - return try_resolve_static_scalar_provenance( - value, - definitions=self._all_definitions, - argument_type=lambda index: ( - self.state.args[index] - if 0 <= index < len(self.state.args) - else None - ), - ) - - def _bind(self, function: Any, call: ir.Expr) -> inspect.BoundArguments: - """Bind IR arguments to the API signature and apply its defaults. - - Explicit arguments remain IR values; defaults are ordinary Python - values. This lets operation planners use parameter names independently - of the caller's positional or keyword spelling. Argument unpacking is - rejected because its contents are not available to this binding step. - - Parameters - ---------- - function : callable - Recognized constructor or operation whose signature defines the - call. - call : ir.Expr - Call expression to validate and bind. - - Returns - ------- - inspect.BoundArguments - Named arguments with Python defaults applied, ready for planning. - - Raises - ------ - GroupRewriteError - The call uses ``*args`` or ``**kwargs``. - TypeError - Arguments do not match the signature. The diagnostic includes the - function name and the IR call's source location. - """ - if call.vararg is not None or call.varkwarg is not None: - raise GroupRewriteError( - "cuda.coop.numba_mlir group calls do not support *args/**kwargs" - ) - try: - bound = inspect.signature(function).bind( - *call.args, **dict(call.kws) - ) - except TypeError as exc: - raise TypeError( - f"{function.__name__}() {exc}\n{call.loc.strformat()}" - ) from None - bound.apply_defaults() - return bound - - def _validate_common_selector( - self, - operation: str, - parameter: str, - value: Any, - allowed: frozenset[str], - *, - allow_none: bool = False, - ) -> Any: - """Validate selector tokens for a rewritten common-API call. - - Rewriting recognizes a public operation by identity and does not - execute its host wrapper. Repeat the selector validation here so - compiled calls retain that API's accepted strings and diagnostics. The - value must be constant and independent of pending literal unrolling. - - Parameters - ---------- - operation : str - Public operation name used in diagnostics. - parameter : str - Selector parameter name used in diagnostics. - value : object - IR argument or an already-resolved default. - allowed : frozenset of str - Accepted normalized selector tokens. - allow_none : bool, optional - Whether ``None`` is an accepted omission value. - - Returns - ------- - str or None - Token with whitespace stripped, lowercased, and hyphens replaced - by underscores, or an explicitly allowed ``None``. - - Raises - ------ - TypeError - The value is not a string, or is an enum instead of a string - token. - InvalidGroupSelectorError - The normalized token is not supported. - ForceLiteralArg - Resolving the selector requires a literal function argument. - GroupRewriteError - The selector cannot be resolved or depends on literal unrolling. - """ - self._reject_literal_unroll_value(value, f"{operation} {parameter}") - token = self._constant(value) - if token is None and allow_none: - return None - if not isinstance(token, str) or isinstance(token, Enum): - raise TypeError( - f"cuda.coop.{operation} {parameter} must be a string" - ) - token = token.strip().lower().replace("-", "_") - if token not in allowed: - choices = ", ".join(sorted(allowed)) - raise InvalidGroupSelectorError(operation, parameter, choices) - return token - - def _hierarchy(self, value: Any) -> ThreadHierarchy | None: - """Reconstruct a hierarchy descriptor for host-side planning. - - Accept an existing ``ThreadHierarchy``, or follow an IR variable - through aliases, casts, constants, and the recognized constructor. - Validate a constructor's signature and cache the resulting host - descriptor by variable name. Return ``None`` for an unrecognized - definition; descriptor arguments requiring a group or hierarchy are - diagnosed by their caller. - """ - - if isinstance(value, ThreadHierarchy): - return value - if not isinstance(value, ir.Var): - return None - cached = self._hierarchy_cache.get(value.name) - if cached is not None: - return cached - definition = self._definition(value) - if isinstance(definition, ir.Var): - return self._hierarchy(definition) - if isinstance(definition, ir.Expr) and definition.op == "cast": - return self._hierarchy(definition.value) - if isinstance(definition, (ir.Global, ir.FreeVar, ir.Const)): - if isinstance(definition.value, ThreadHierarchy): - return definition.value - return None - if not isinstance(definition, ir.Expr) or definition.op != "call": - return None - function = self._callable(definition.func) - if function is not ThreadHierarchy: - return None - self._bind(ThreadHierarchy, definition) - hierarchy = ThreadHierarchy() - self._hierarchy_cache[value.name] = hierarchy - return hierarchy - - def _group(self, value: object) -> ThreadGroup | None: - """Reconstruct a group descriptor from a supported IR definition. - - Follow aliases and casts, accept existing host descriptors, and - interpret registered constructors and ``group_by`` calls on recognized - parents. Constructor identity matters; matching a callable's name is - insufficient. Arguments are resolved through ``_constant``, which may - request literal specialization. Common-API constructors retain - ``common_root`` provenance so later validation applies the common API - contract. - - Cache newly constructed descriptors by variable name for this planner. - This describes the requested group; launch-dependent resolution - belongs to ``_resolve_group``. Unlike marker detection, this routine - requires a single resolvable definition and does not merge phi inputs. - - Parameters - ---------- - value : object - Existing ``ThreadGroup`` or IR variable expected to name one. - Other inputs are allowed as probes and return ``None``; the - caller decides whether an unrecognized group is an error. - - Returns - ------- - ThreadGroup or None - Reconstructed or cached descriptor, or ``None`` if the definition - is not a recognized group expression. - - Raises - ------ - ForceLiteralArg - A constructor or subgroup argument needs literal specialization. - TypeError - A recognized constructor call does not match its Python signature. - GroupRewriteError - A group constructor or subgroup call uses unsupported syntax, - or its arguments cannot be resolved as compile-time values. - """ - - if isinstance(value, ThreadGroup): - return value - if not isinstance(value, ir.Var): - return None - cached = self._group_cache.get(value.name) - if cached is not None: - return cached - definition = self._definition(value) - if isinstance(definition, ir.Var): - return self._group(definition) - if isinstance(definition, (ir.Global, ir.FreeVar, ir.Const)): - if isinstance(definition.value, ThreadGroup): - return definition.value - return None - if isinstance(definition, ir.Expr) and definition.op == "cast": - return self._group(definition.value) - if not isinstance(definition, ir.Expr) or definition.op != "call": - return None - function = self._callable(definition.func) - if function in _GROUP_CONSTRUCTORS: - bound = self._bind(function, definition) - args = [] - kwargs = {} - parameters = inspect.signature(function).parameters - for name, parameter in parameters.items(): - argument = self._constant(bound.arguments[name]) - if parameter.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD: - args.append(argument) - elif parameter.kind is inspect.Parameter.KEYWORD_ONLY: - kwargs[name] = argument - group = _GROUP_CONSTRUCTORS[function](*args, **kwargs) - if function in _COMMON_GROUP_CONSTRUCTORS: - assert group.hierarchy is not None - group = group.with_hierarchy( - group.hierarchy, source="common_root" - ) - self._group_cache[value.name] = group - return group - function_definition = self._definition(definition.func) - if ( - isinstance(function_definition, ir.Expr) - and function_definition.op == "getattr" - and (function_definition.attr == "group_by") - ): - parent = self._group(function_definition.value) - if parent is None: - return None - if definition.vararg is not None or definition.varkwarg is not None: - raise GroupRewriteError( - "ThreadGroup.group_by does not support *args/**kwargs" - ) - raw_args = tuple(definition.args) - raw_kwargs = dict(definition.kws) - unknown_kwargs = set(raw_kwargs) - {"count", "exhaustive"} - if unknown_kwargs: - names = ", ".join(sorted(unknown_kwargs)) - raise GroupRewriteError( - f"ThreadGroup.group_by got unexpected keyword(s): {names}" - ) - if len(raw_args) > 1: - raise GroupRewriteError( - "ThreadGroup.group_by accepts one positional count argument" - ) - if raw_args and "count" in raw_kwargs: - raise GroupRewriteError( - "ThreadGroup.group_by received count more than once" - ) - if raw_args: - count_arg = raw_args[0] - elif "count" in raw_kwargs: - count_arg = raw_kwargs["count"] - else: - raise GroupRewriteError("ThreadGroup.group_by requires count") - self._reject_literal_unroll_value(count_arg, "group_by count") - exhaustive_arg = raw_kwargs.get("exhaustive", True) - self._reject_literal_unroll_value( - exhaustive_arg, "group_by exhaustive" - ) - count_value = self._constant(count_arg) - exhaustive = self._constant(exhaustive_arg) - group = parent.group_by(count_value, exhaustive=exhaustive) - self._group_cache[value.name] = group - return group - return None - - def _resolve_group( - self, - group: ThreadGroup, - *, - feature: str, - through_level: str | None = None, - ) -> ThreadGroup: - """Resolve a group request against the configured launch. - - Descriptor reconstruction records what the kernel requested. - Resolution fills the launch-dependent hierarchy and verifies that the - requested composition can be represented. Preserve common-API origin - information so operation validation can still enforce that contract - afterward. - - Parameters - ---------- - group : ThreadGroup - Reconstructed or supplied group description. - feature : str - Operation or group-method name prepended to unsupported - diagnostics. - through_level : str, optional - Hierarchy level through which resolution is required, for example - when a query needs enclosing group counts. - - Returns - ------- - ThreadGroup - Supported group with launch-dependent topology resolved. - - Raises - ------ - NotImplementedError - Shared group resolution reports an unsupported request. The - message identifies the Numba operation that needed the group. - """ - - resolution = resolve_thread_group( - group, self.launch, through_level=through_level - ) - try: - resolved = resolution.require_supported() - except NotImplementedError as exc: - raise NotImplementedError( - f"cuda.coop.numba_mlir.{feature} {exc}" - ) from exc - if group.source == "common_root": - assert resolved.hierarchy is not None - resolved = resolved.with_hierarchy( - resolved.hierarchy, source="common_root" - ) - return resolved - - def _is_none(self, value: Any) -> bool: - """Check whether an optional operation argument is explicitly omitted. - - Operation planners call this through ``GroupPlanningContext.is_none`` - for arguments such as ``temp_storage``. ``value`` is an IR operand or - Python default. Return true only when constant resolution establishes - ``None``; an unknown runtime value is not treated as omission. This - probe does not request another compiler specialization. - """ - - resolved, constant = self._try_constant(value) - return resolved and constant is None - - @staticmethod - def _merge_array_states(states: tuple[bool | None, ...]) -> bool | None: - """Combine the array checks for all possible sources of an operand. - - The array walkers call this when a branch or loop gives a variable more - than one possible value. An operand can be treated as an array only if - no concrete source contradicts that choice. A loop that keeps a payload - created before the loop is valid: revisiting the loop variable adds no - new evidence, but must not discard the known constructor either. - - Parameters - ---------- - states : tuple of bool or None - Results for the alternative sources. ``True`` means a supported - array constructor was found; ``False`` means that source is - unsupported; ``None`` means traversal revisited a variable on the - current recursion path, before finding a constructor there. - - Returns - ------- - bool or None - ``False`` if any source is unsupported or there are no definitions. - Otherwise ``True`` if a constructor was found, or ``None`` if every - source only led back into the same cycle. Thus ``(True, None)`` is - accepted, while ``(True, False)`` is rejected. - """ - - if not states or any(state is False for state in states): - return False - if any(state is True for state in states): - return True - return None - - def _is_array_tuple_item( - self, - value: Any, - index: int, - *, - seen: set[str], - thread_data_only: bool = False, - ) -> bool | None: - """Check whether a selected tuple element holds an array payload. - - An operation can receive a payload taken from a tuple, for example - ``payloads[0]``. ``_is_array_definition`` calls this helper for such - indexing expressions. Follow every possible source of the tuple so a - branch cannot hide an unsupported value in the selected position. - Track the tuple and index together: revisiting ``payloads[0]`` is a - cycle, while inspecting another element is a separate question. - - Parameters - ---------- - value : ir.Var or object - IR variable containing the tuple. Other values return ``False``. - index : int - Tuple position to inspect, with Python's zero-based and negative - indexing rules. This is an element index, not a byte offset. - seen : set of str - Variable names and ``"name[index]"`` keys already visited on this - recursion path. The selected key is added in place; each possible - definition receives a copy so one branch cannot hide another. - thread_data_only : bool, optional - Accept only ``ThreadData`` constructors when true. Otherwise, - recognized CUDA local-array constructors are also accepted. - - Returns - ------- - bool or None - ``True`` when at least one supported constructor is found and no - source is rejected; ``False`` for a missing or unsupported source; - ``None`` if traversal only revisits variables in a cycle. The - caller combines this result with the operand's other sources. - """ - - if not isinstance(value, ir.Var): - return False - seen_key = f"{value.name}[{index}]" - if seen_key in seen: - return None - seen.add(seen_key) - return self._merge_array_states( - tuple( - self._is_array_tuple_item_definition( - definition, - index, - seen=set(seen), - thread_data_only=thread_data_only, - ) - for definition in self._all_definitions(value) - ) - ) - - def _is_array_tuple_item_definition( - self, - definition: Any, - index: int, - *, - seen: set[str], - thread_data_only: bool, - ) -> bool | None: - """Check one tuple source for the selected array payload. - - ``_is_array_tuple_item`` calls this once per recorded tuple definition. - Follow tuple copies, casts, and the IR used to unpack an iterator; each - preserves which position the caller selected. A phi expression chooses - between values arriving from different branches or loop iterations, so - its inputs must all be checked. When a ``build_tuple`` expression is - reached, classify the selected element with ``_is_array_value``. - - Parameters - ---------- - definition : object - One right-hand-side IR value or expression that supplies the tuple. - index : int - Zero-based tuple position; negative indices count from the end. - seen : set of str - Variable names and tuple-position keys visited on this recursion - path. Independent branches receive copies to avoid skipping work. - thread_data_only : bool - Whether the selected payload must originate from ``ThreadData``. - False also permits a recognized CUDA local-array constructor. - - Returns - ------- - bool or None - The selected element's array classification. ``False`` means an - unsupported source or invalid index. ``None`` means a recursive - path supplied no constructor evidence; ``True`` means a supported - constructor was found without a conflicting source. - """ - - if isinstance(definition, ir.Var): - return self._is_array_tuple_item( - definition, index, seen=seen, thread_data_only=thread_data_only - ) - if not isinstance(definition, ir.Expr): - return False - if definition.op in {"cast", "exhaust_iter"}: - return self._is_array_tuple_item( - definition.value, - index, - seen=seen, - thread_data_only=thread_data_only, - ) - if definition.op == "phi": - incoming_values = getattr(definition, "incoming_values", ()) - return self._merge_array_states( - tuple( - self._is_array_tuple_item( - incoming, - index, - seen=set(seen), - thread_data_only=thread_data_only, - ) - for incoming in incoming_values - ) - ) - if definition.op == "build_tuple": - items = tuple(getattr(definition, "items", ())) - if not -len(items) <= index < len(items): - return False - return self._is_array_value( - items[index], seen=set(seen), thread_data_only=thread_data_only - ) - if definition.op != "call": - return False - return False - - def _is_array_value( - self, - value: object, - *, - seen: set[str] | None = None, - thread_data_only: bool = False, - ) -> bool | None: - """Determine whether a value comes from a supported per-thread array. - - A ``store()`` operation can consume a scalar or a per-thread array, - whereas ``load()`` needs an array destination. Their validation uses - this result to select the appropriate shape and argument rules. A - variable may instead name a scalar, a kernel argument, an unrelated - object, or a loop-carried value whose origin is not yet known. A false - result only rules out a recognized array origin; further validation - must establish whether the value is a supported scalar. - - Trace aliases, casts, phi inputs, and constant tuple projections to - ``ThreadData`` or qualified local-array constructors. The three-state - result lets a loop backedge coexist with a known constructor: a cycle - alone is unresolved, a constructor plus cycles is accepted, and any - concrete unrecognized path rejects the value. This classifies the - payload form only; dtype and extent are inferred separately. - - Parameters - ---------- - value : object - IR variable to classify, or any other value to probe. - Non-variables return ``False``. - seen : set of str, optional - Recursion-path variable names and tuple-projection keys. The - current name is added in place; branches receive separate copies. - thread_data_only : bool, optional - If true, accept only common or qualified ``ThreadData`` - constructors. Otherwise, also accept the recognized CUDA - local-array constructor. - - Returns - ------- - bool or None - ``True`` if at least one constructor is found and no path is - rejected; ``False`` for an unsupported definition, missing - definitions, or a non-variable; ``None`` when only cyclic paths - remain. The operand validation wrappers turn that last state into - a provenance diagnostic. - """ - - if not isinstance(value, ir.Var): - return False - if seen is None: - seen = set() - if value.name in seen: - return None - seen.add(value.name) - return self._merge_array_states( - tuple( - self._is_array_definition( - definition, - seen=set(seen), - thread_data_only=thread_data_only, - ) - for definition in self._all_definitions(value) - ) - ) - - def _is_array_definition( - self, definition: Any, *, seen: set[str], thread_data_only: bool - ) -> bool | None: - """Check whether one assignment supplies a supported array payload. - - ``_is_array_value`` calls this for each possible source of an operation - operand. Tracing the assignment back to its constructor tells Load/Store - planning whether it has a per-thread array or must apply scalar rules. - For a value selected from a tuple, resolve the index before examining - that tuple element. Branch and loop merges inspect all alternatives; - an unsupported source cannot be hidden by another valid source. - - Parameters - ---------- - definition : object - A recorded right-hand-side IR value or expression, such as a - constructor call, variable alias, cast, phi, or tuple lookup. - seen : set of str - Variable names and tuple-position keys already visited on the - current recursion path. Each branch receives its own copy. - thread_data_only : bool - Accept common and qualified ``ThreadData`` constructors only when - true; false also accepts the recognized CUDA local-array call. - - Returns - ------- - bool or None - ``True`` for a supported constructor or alternatives that agree; - ``False`` for an unsupported source or unresolved tuple index; - ``None`` when following aliases or loop inputs only leads back to - a variable already being examined. Operand validation reports that - unresolved cycle if no other path supplies a constructor. - """ - - if isinstance(definition, ir.Var): - return self._is_array_value( - definition, seen=seen, thread_data_only=thread_data_only - ) - if not isinstance(definition, ir.Expr): - return False - if definition.op == "cast": - return self._is_array_value( - definition.value, seen=seen, thread_data_only=thread_data_only - ) - if definition.op == "phi": - incoming_values = getattr(definition, "incoming_values", ()) - return self._merge_array_states( - tuple( - self._is_array_value( - incoming, - seen=set(seen), - thread_data_only=thread_data_only, - ) - for incoming in incoming_values - ) - ) - if definition.op in {"getitem", "static_getitem"}: - index = getattr(definition, "index", None) - if isinstance(index, ir.Var): - resolved, index = self._try_constant(index) - if not resolved: - return False - if isinstance(index, Integral) and (not isinstance(index, bool)): - return self._is_array_tuple_item( - definition.value, - int(index), - seen=set(seen), - thread_data_only=thread_data_only, - ) - return False - if definition.op != "call": - return False - function = self._callable(definition.func) - if function in {ThreadData, _common_api.ThreadData}: - return True - if function is _cuda_local_array: - return not thread_data_only - return False - - @staticmethod - def _new_var(scope: Any, loc: ir.Loc, stem: str) -> ir.Var: - """Name a temporary needed by a generated provider call. - - ``_rewritten_call`` and operation planners use these variables for the - selected callable and its arguments. The shared counter avoids name - collisions when several operations are rewritten in the same function. - - Parameters - ---------- - scope : ir.Scope - Lexical scope of the original operation's result variable. - loc : ir.Loc - Source location to retain for compiler diagnostics. - stem : str - Readable name fragment, such as ``"factory"`` or ``"arg0"``, - identifying the temporary's purpose in an IR dump. - - Returns - ------- - ir.Var - A fresh variable. Its defining assignment is created by the caller. - """ - - return ir.Var( - scope, f"__cuda_coop_group_{stem}_{next(_NAME_COUNTER)}__", loc - ) - - def _value_var( - self, - statements: list[Any], - *, - scope: Any, - loc: ir.Loc, - stem: str, - value: Any, - ) -> ir.Var: - """Supply an IR variable for one argument of a generated provider call. - - Numba call expressions refer to variables, but planning also produces - Python values such as an item count or a provider's metadata object. - ``_rewritten_call`` and operation planners use this helper to bind those - values before emitting a call. Existing variables need no assignment. - For a Python value, append a constant or global assignment to the - caller's pending statement list. ``run()`` installs that list only after - all cooperative operations pass validation. - - Parameters - ---------- - statements : list of IR statements - Pending replacement statements, in execution order. A new - assignment is appended here when ``value`` needs a variable. - scope : ir.Scope - Scope for a new variable, taken from the original call's result. - loc : ir.Loc - Original source location for the variable and assignment. - stem : str - Name fragment describing the argument, such as ``"arg0"`` or a - keyword name. A counter makes the complete name distinct. - value : object - Existing IR variable or Python value to pass. ``None``, booleans, - integers, floats, strings, and tuples use ``ir.Const``; other - objects use ``ir.Global`` so later compiler phases can recover - the object. - - Returns - ------- - ir.Var - The unchanged input variable or the newly assigned temporary. - No function block is modified by this helper. - """ - - if isinstance(value, ir.Var): - return value - result = self._new_var(scope, loc, stem) - if value is None or isinstance(value, (bool, int, float, str, tuple)): - rhs = ir.Const(value, loc) - else: - rhs = ir.Global(result.name, value, loc) - statements.append(ir.Assign(rhs, result, loc)) - return result - - def _rewritten_call( - self, - inst: ir.Assign, - *, - factory: Any, - args: list[Any], - kwargs: dict[str, Any], - common_root_operation: str | None = None, - ) -> list[Any]: - """Build replacement IR calling the selected provider. - - Bind the host factory and any non-variable arguments to fresh - temporaries, then emit a call at the original source location. - Existing argument variables are reused, and the replacement preserves - the original result target. Nothing is installed in a block until - ``run()`` accepts the whole rewrite. - - Parameters - ---------- - inst : ir.Assign - Original operation assignment whose target, scope, and location - are retained by the replacement. - factory : callable - Selected private provider callable for the next rewriting phase. - args : list of object - Positional IR variables or Python values supplied to the provider. - kwargs : dict of str to object - Named IR variables or Python values supplied to the provider. - common_root_operation : str, optional - Public operation identity forwarded to provider validation. It - keeps common-API rules available after the public wrapper has been - removed. - - Returns - ------- - list - Ordered assignments defining the callable and arguments and making - the replacement call. - """ - - statements: list[Any] = [] - scope = inst.target.scope - loc = inst.loc - if common_root_operation is not None: - kwargs = dict(kwargs) - kwargs.setdefault("_common_root_operation", common_root_operation) - function_var = self._new_var(scope, loc, "factory") - statements.append( - ir.Assign( - ir.Global(function_var.name, factory, loc), function_var, loc - ) - ) - rewritten_args = [ - self._value_var( - statements, scope=scope, loc=loc, stem=f"arg{idx}", value=value - ) - for idx, value in enumerate(args) - ] - rewritten_kwargs = tuple( - ( - ( - name, - self._value_var( - statements, scope=scope, loc=loc, stem=name, value=value - ), - ) - for name, value in kwargs.items() - ) - ) - statements.append( - ir.Assign( - ir.Expr.call( - function_var, rewritten_args, rewritten_kwargs, loc - ), - inst.target, - loc, - ) - ) - return statements - - def _array_operand_state(self, operation: str, value: Any) -> bool: - """Tell an operation planner whether its operand is a per-thread array. - - ``GroupPlanningContext.is_array`` calls this to validate an operand. - For example, Store needs to choose between a scalar value and an array - of values per thread, while Load requires an array destination. Trace - the operand to its constructor before making that choice. If traversal - only finds a cycle, report that the payload's source could not be - established rather than treating it as a scalar. - - Parameters - ---------- - operation : str - Public operation name, such as ``"load"`` or ``"store"``, used - to identify the call if source tracing fails. - value : object - Operand's IR variable, or another value to probe. Recognized - sources are ``ThreadData`` and CUDA local-array constructors. - - Returns - ------- - bool - Whether the operand has a supported array source on all concrete - paths. ``False`` leaves scalar or invalid-value checks to the - operation planner; it does not establish that the value is scalar. - - Raises - ------ - CyclicArrayProvenanceError - All traced sources only revisit variables already being examined, - so none establishes how the payload was created. - """ - - state = self._is_array_value(value) - if state is None: - raise CyclicArrayProvenanceError(operation) - return state - - def _thread_data_operand_state( - self, operation: str, parameter: str, value: Any - ) -> bool: - """Check the payload constructor required by a common-API parameter. - - ``GroupPlanningContext.is_thread_data`` calls this when an operation - needs a ``ThreadData`` payload. For example, common ``load`` requires - ``ThreadData`` for its output even though the qualified backend also - accepts CUDA local arrays. Follow copies and control-flow alternatives - so that the same rule applies when the constructor is elsewhere in the - kernel. - - Parameters - ---------- - operation : str - Public operation name used in a source-tracing diagnostic. - parameter : str - Public parameter name, such as ``"output"`` or ``"value"``, - identifying which operand could not be traced. - value : object - IR operand to trace to a common or qualified ``ThreadData`` call. - - Returns - ------- - bool - ``True`` for a recognized ``ThreadData`` source on all concrete - paths. ``False`` lets the caller report its parameter-specific - error for an unsupported source. - - Raises - ------ - GroupRewriteError - The search only finds cyclic references and cannot establish the - payload's constructor. - """ - - state = self._is_array_value(value, thread_data_only=True) - if state is None: - raise GroupRewriteError( - f"cuda.coop.{operation} could not " - f"resolve {parameter} payload provenance" - ) - return state - - def _array_extent( - self, value: Any, *, seen: set[str] | None = None - ) -> int | None: - """Recover how many elements one thread holds in an array payload. - - This element count is the payload's *extent*. For example, - ``ThreadData(items_per_thread=4)`` has extent four, independent of its - element dtype or the number of threads in the block. Operation planners - call ``GroupPlanningContext.array_extent`` to obtain this count before - specializing a CUB provider whose item count is a template argument. - - Follow the assignments that can supply ``value`` through copies, casts, - branch or loop merges, and constant tuple indices. Read the count from - ``ThreadData`` or a CUDA local-array constructor with an integer shape. - All known counts must agree: one provider cannot use different array - sizes on different paths through the kernel. Paths with no known count - are ignored here; the separate array-kind check must establish that the - operand comes from supported constructors. - - Parameters - ---------- - value : ir.Var or object - Payload variable whose element count is needed. Non-variables - provide no count and return ``None``. - seen : set of str, optional - Variable names and tuple-position keys already visited on this - recursion path. Add the current name in place and copy the set - when following alternative assignments to avoid infinite loops. - - Returns - ------- - int or None - The single known number of elements per thread, or ``None`` if no - count is found. Counts are integral and exclude booleans; the - operation's later shape validation checks positivity. - - Raises - ------ - GroupRewriteError - Known counts disagree, or a constructor's dimension depends on - literal unrolling that has not run yet. - ForceLiteralArg - A constructor dimension needs a kernel argument specialized to a - compile-time value. - """ - - if not isinstance(value, ir.Var): - return None - if seen is None: - seen = set() - if value.name in seen: - return None - seen.add(value.name) - extents: set[int] = set() - for definition in self._all_definitions(value): - extent = self._array_extent_definition(definition, seen=set(seen)) - if extent is not None: - extents.add(extent) - if len(extents) > 1: - raise InconsistentArrayExtentError() - return next(iter(extents), None) - - def _array_extent_tuple_item( - self, value: Any, index: int, *, seen: set[str] - ) -> int | None: - """Find the per-thread element count of a payload selected from a tuple. - - ``_array_extent_definition`` calls this for indexing such as - ``payloads[0]``. The tuple may have been copied or assigned on several - branches, so inspect each possible assignment before selecting one - count for provider specialization. The tuple length is unrelated to - the payload's extent: a tuple of two four-element payloads has length - two, while either selected payload has extent four. - - Parameters - ---------- - value : ir.Var or object - IR variable containing the tuple. Other values have no known - payload count and return ``None``. - index : int - Payload position within the tuple, using Python's zero-based and - negative indexing rules. - seen : set of str - Variable names and tuple-position keys visited on this recursion - path. Add ``"name[index]"`` in place and copy the set for each - alternative assignment so loops terminate without hiding branches. - - Returns - ------- - int or None - The unique known element count per thread for the selected payload. - Unknown sources and repeated visits contribute no count; ``None`` - means none was found. Array-kind validation is a separate check. - - Raises - ------ - InconsistentTupleExtentError - Possible definitions of this tuple give the selected payload - different known element counts. - """ - - if not isinstance(value, ir.Var): - return None - seen_key = f"{value.name}[{index}]" - if seen_key in seen: - return None - seen.add(seen_key) - extents = { - extent - for definition in self._all_definitions(value) - if ( - extent := self._array_extent_tuple_item_definition( - definition, index, seen=set(seen) - ) - ) - is not None - } - if len(extents) > 1: - raise InconsistentTupleExtentError() - return next(iter(extents), None) - - def _array_extent_tuple_item_definition( - self, definition: Any, index: int, *, seen: set[str] - ) -> int | None: - """Read one tuple assignment for the selected payload's element count. - - ``_array_extent_tuple_item`` calls this for each possible tuple source. - A *reaching definition* is an assignment whose value can arrive at the - operation along a branch or loop path. The caller conservatively checks - all recorded sources, without proving each reaches this use. Trace back - to the payload constructor because CUB provider specialization needs a - fixed number of elements per thread. That number is the payload's - *extent*, not the tuple length or its size in bytes. - - Copies, casts, and iterator-unpacking IR preserve the selected position. - A phi expression combines values from different control-flow paths, - so known counts on its inputs must agree. At a tuple construction, - ``_array_extent`` examines the selected payload itself. - - Parameters - ---------- - definition : object - One right-hand-side IR value or expression that supplies the tuple. - index : int - Payload position in the tuple; negative indices count from the end. - seen : set of str - Variable names and tuple-position keys already visited while - tracing this path. Separate phi inputs receive copies of the set. - - Returns - ------- - int or None - Known number of elements per thread in the selected payload. - ``None`` means the source or index is unsupported, its count is - unknown, or following it only revisits an existing recursion path. - The caller checks payload kind separately. - - Raises - ------ - InconsistentLoopTupleExtentError - A branch or loop merge supplies different known counts for the - selected tuple position. - """ - - if isinstance(definition, ir.Var): - return self._array_extent_tuple_item(definition, index, seen=seen) - if not isinstance(definition, ir.Expr): - return None - if definition.op in {"cast", "exhaust_iter"}: - return self._array_extent_tuple_item( - definition.value, index, seen=seen - ) - if definition.op == "phi": - extents = { - extent - for incoming in getattr(definition, "incoming_values", ()) - if ( - extent := self._array_extent_tuple_item( - incoming, index, seen=set(seen) - ) - ) - is not None - } - if len(extents) > 1: - raise InconsistentLoopTupleExtentError() - return next(iter(extents), None) - if definition.op == "build_tuple": - items = tuple(getattr(definition, "items", ())) - if not -len(items) <= index < len(items): - return None - return self._array_extent(items[index], seen=set(seen)) - if definition.op != "call": - return None - return None - - def _array_extent_definition( - self, definition: Any, *, seen: set[str] - ) -> int | None: - """Read one payload assignment to recover its per-thread element count. - - ``_array_extent`` calls this for every assignment that might supply an - operand. Provider selection needs the constructor's fixed element - count, even when the operation receives a copied value or one selected - from a tuple. Follow these intermediate expressions until a recognized - ``ThreadData`` or CUDA local-array constructor is reached, then resolve - its ``items_per_thread`` or integer ``shape`` argument. - - Parameters - ---------- - definition : object - One right-hand-side IR value or expression that supplies a payload: - a constructor call, alias, cast, phi, or tuple lookup. - seen : set of str - Variable names and tuple-position keys already visited on this - recursion path. Copies let each phi input be examined independently. - - Returns - ------- - int or None - The constructor's number of elements per thread, as an integer - excluding booleans. ``None`` means no usable count was found. In - particular, an unsupported source, nonconstant tuple index, or - noninteger local-array shape supplies no count. This does not - establish that every source is an array or that the count is - positive. - - Raises - ------ - InconsistentLoopPayloadExtentError - A phi's inputs have different known element counts, so one provider - specialization cannot describe every path. - GroupRewriteError - A constructor dimension depends on an unexpanded literal-unroll - value. Errors from nested tuple analysis also propagate. - ForceLiteralArg - Resolving a constructor dimension requires the dispatcher to retry - compilation with a literal kernel argument. - """ - - if isinstance(definition, ir.Var): - return self._array_extent(definition, seen=seen) - if not isinstance(definition, ir.Expr): - return None - if definition.op == "cast": - return self._array_extent(definition.value, seen=seen) - if definition.op == "phi": - extents = { - extent - for incoming in getattr(definition, "incoming_values", ()) - if (extent := self._array_extent(incoming, seen=set(seen))) - is not None - } - if len(extents) > 1: - raise InconsistentLoopPayloadExtentError() - return next(iter(extents), None) - if definition.op in {"getitem", "static_getitem"}: - index = getattr(definition, "index", None) - if isinstance(index, ir.Var): - resolved, index = self._try_constant(index) - if not resolved: - return None - if isinstance(index, Integral) and (not isinstance(index, bool)): - return self._array_extent_tuple_item( - definition.value, int(index), seen=set(seen) - ) - return None - if definition.op != "call": - return None - function = self._callable(definition.func) - if function in {ThreadData, _common_api.ThreadData}: - bound = self._bind(function, definition) - extent_argument = bound.arguments["items_per_thread"] - self._reject_literal_unroll_value(extent_argument, "payload extent") - try: - extent = self._constant(extent_argument) - except GroupRewriteError: - return None - if isinstance(extent, Integral) and (not isinstance(extent, bool)): - return int(extent) - return None - if function is _cuda_local_array: - shape_ref = ( - definition.args[0] - if definition.args - else dict(definition.kws).get("shape") - ) - if shape_ref is None: - return None - self._reject_literal_unroll_value(shape_ref, "payload extent") - try: - extent = self._constant(shape_ref) - except GroupRewriteError: - return None - if isinstance(extent, Integral) and (not isinstance(extent, bool)): - return int(extent) - return None - return None - - def _lower_root_operation( - self, inst: ir.Assign, call: ir.Expr, function: Any, operation: str - ) -> None: - """Validate a public call and queue its registered provider rewrite. - - Bind the public signature, reconstruct and resolve its group, and - apply common-API group and argument rules when the call came from that - API. Then delegate operation-specific choices to its registration - using ``GroupPlanningContext``. Record the returned statements and - mark the old callable variable for removal; the original IR remains - available while other calls are planned. - - Parameters - ---------- - inst : ir.Assign - Assignment to replace after the complete function passes - validation. - call : ir.Expr - Public operation call stored in ``inst``. - function : callable - Recognized public callable whose signature and API origin apply. - operation : str - Canonical name used to look up the registered operation planner. - - Raises - ------ - GroupRewriteError - Call arguments or group reconstruction are invalid, no planner is - registered, or operation-specific validation rejects the request. - TypeError - Arguments do not match the public signature or accepted parameter - types. - ForceLiteralArg - Planning needs a function argument specialized as a literal. - NotImplementedError - Group resolution or the operation planner finds no supported - lowering. - """ - - bound = self._bind(function, call) - if bound.arguments.get("kwargs"): - names = ", ".join(sorted(bound.arguments["kwargs"])) - raise GroupRewriteError( - f"cuda.coop.numba_mlir.{operation} " - f"got unexpected keyword(s): {names}" - ) - group = self._group(bound.arguments["group"]) - if group is None: - raise NonConstantThreadGroupError(operation) - is_common_root = _is_common_root_operation(function, operation) - if is_common_root: - _common_dispatch._validate_common_operation_group(operation, group) - group = self._resolve_group(group, feature=operation) - registration = group_primitive(operation) - if registration is None: - raise GroupRewriteError( - f"cuda.coop.numba_mlir operation {operation!r} has no planner" - ) - if ( - is_common_root - and registration.validate_common_arguments is not None - ): - registration.validate_common_arguments( - self.context, operation, bound - ) - replacement = registration.lower( - self.context, - inst, - operation=operation, - group=group, - bound=bound, - is_common_root=is_common_root, - ) - self.dead_func_names.add(call.func.name) - self.replacements[inst] = replacement - - def _mark_descriptor_calls(self) -> None: - """Mark thread-group and hierarchy assignments for removal. - - Recognize hierarchy and group constructors by callable identity, and - ``group_by`` only when its receiver resolves to a group. Then - repeatedly follow aliases and casts until every reachable descriptor - assignment is marked. Record consumed callable names too. This is - bookkeeping only: ``_validate_group_descriptor_references`` must - approve removal before blocks change. - """ - - for block in self.func_ir.blocks.values(): - for inst in block.body: - if not isinstance(inst, ir.Assign): - continue - call = inst.value - if not isinstance(call, ir.Expr) or call.op != "call": - continue - function = self._callable(call.func) - if ( - function is ThreadHierarchy - or function in _GROUP_CONSTRUCTORS - ): - self.descriptor_assigns.add(inst) - self.dead_func_names.add(call.func.name) - continue - definition = self._definition(call.func) - if ( - isinstance(definition, ir.Expr) - and definition.op == "getattr" - and (definition.attr == "group_by") - and (self._group(definition.value) is not None) - ): - self.descriptor_assigns.add(inst) - self.dead_func_names.add(call.func.name) - descriptor_names = { - inst.target.name for inst in self.descriptor_assigns - } - changed = True - while changed: - changed = False - for block in self.func_ir.blocks.values(): - for inst in block.body: - if not isinstance(inst, ir.Assign): - continue - source = inst.value - if isinstance(source, ir.Expr) and source.op == "cast": - source = source.value - if ( - isinstance(source, ir.Var) - and source.name in descriptor_names - and (inst.target.name not in descriptor_names) - ): - self.descriptor_assigns.add(inst) - descriptor_names.add(inst.target.name) - changed = True - - def _validate_group_descriptor_references(self) -> None: - """Reject runtime references to thread groups or hierarchies. - - Here, a descriptor is a ``ThreadGroup`` from ``this_block()``, - ``this_warp()``, or another group constructor or ``group_by()`` call, - or a ``ThreadHierarchy`` value. These describe which threads cooperate. - ``ThreadData`` provides the per-thread payload array and is handled - separately. - - A group or hierarchy value may feed another group/hierarchy assignment, - a primitive call with a queued replacement, or a method lookup marked - for removal. Returning that value or passing it to an unrelated call - would require a device object after its construction has been erased. - Such uses raise ``EscapingGroupDescriptorError`` naming the variables. - Run this after collecting replacements, before rewriting block bodies. - """ - - descriptor_names = { - inst.target.name for inst in self.descriptor_assigns - } - if not descriptor_names: - return - for block in self.func_ir.blocks.values(): - for inst in block.body: - used_names = { - value.name for value in inst.list_vars() - } & descriptor_names - if isinstance(inst, ir.Assign): - used_names.discard(inst.target.name) - if not used_names: - continue - if inst in self.descriptor_assigns or inst in self.replacements: - continue - if isinstance(inst, ir.Assign): - value = inst.value - if ( - isinstance(value, ir.Expr) - and value.op == "getattr" - and isinstance(value.value, ir.Var) - and (value.value.name in used_names) - and (inst.target.name in self.dead_func_names) - ): - continue - names = ", ".join(sorted(used_names)) - raise EscapingGroupDescriptorError(names) - - def run(self) -> bool: - """Plan group operations and replace their descriptors in the IR. - - ``_GroupPlanning._resolve_groups`` creates this planner and calls - ``run()`` during ``CoopWholeFunctionPlanner.run``. Numba has already - inlined device helpers, making their cooperative calls visible in the - kernel, and ordinary type inference has not started. The caller has - obtained the configured launch dimensions so this phase can resolve - thread groups and choose provider specializations. - - First reject unsupported literal-unroll dependencies and identify - assignments that construct ``ThreadHierarchy`` or groups, including - ``group_by()`` calls and descriptor aliases or casts. These describe - compile-time choices and are marked for removal once their uses have - been consumed. Ask each registered operation family to build - replacement statements, then check the remaining descriptor uses - before replacing any block bodies. Returning a descriptor or passing - it to an unrelated runtime call, for example, prevents its removal and - raises an error. - - Accepted descriptor assignments and obsolete callable assignments - become ``None`` assignments, preserving their targets; public - operations become private provider calls carrying their lowering - plans. This removes the group markers so the same whole-function - planner can proceed to provider rewriting. Planning updates this - instance's caches, dtype facts, and replacement bookkeeping; block - bodies change only after all calls and descriptor uses pass. - - Returns - ------- - bool - ``True`` when block bodies were rewritten, including - descriptor-only cleanup. ``False`` when no descriptor, call, or - callable needs replacing. ``_resolve_groups`` forwards this result - to ``CoopWholeFunctionPlanner.run``: a true result triggers IR - repair before provider rewriting, which runs in either case. The - outer planner reports changes from either phase to Numba, which - repairs the final IR and proceeds to the next registered planner. - This boolean does not request another run of group planning. - - Raises - ------ - GroupRewriteError - A group call is invalid or a descriptor escapes its compile-time - uses. - ForceLiteralArg - Planning requires specialization of a function argument. - NotImplementedError - The requested group or operation has no supported lowering. - """ - - self._reject_literal_unroll_constructors() - self._mark_descriptor_calls() - for block in self.func_ir.blocks.values(): - for inst in block.body: - if not isinstance(inst, ir.Assign): - continue - call = inst.value - if not isinstance(call, ir.Expr) or call.op != "call": - continue - function = self._callable(call.func) - operation = _group_operation_name(function) - if operation is not None: - self._lower_root_operation(inst, call, function, operation) - continue - self._validate_group_descriptor_references() - if not ( - self.descriptor_assigns or self.replacements or self.dead_func_names - ): - return False - for block in self.func_ir.blocks.values(): - rewritten: list[Any] = [] - for inst in block.body: - replacement = self.replacements.get(inst) - if replacement is not None: - rewritten.extend(replacement) - continue - if isinstance(inst, ir.Assign) and ( - inst in self.descriptor_assigns - or inst.target.name in self.dead_func_names - ): - rewritten.append( - ir.Assign( - ir.Const(None, inst.loc), inst.target, inst.loc - ) - ) - continue - rewritten.append(inst) - block.body = rewritten - return True - - -def has_group_markers(func_ir: ir.FunctionIR) -> bool: - """Return whether the current function IR still needs group planning. - - One recognized call anywhere in the function is enough: a group - constructor such as ``this_block()``, ``ThreadHierarchy()``, a registered - public group operation such as ``load()`` or ``store()``, or - ``group_by()`` on a recognized group descriptor. For ``group_by()``, trace - the receiver through aliases, casts, control-flow merges, and earlier - subgroup calls to distinguish group descriptors from unrelated objects - with that method. - - Inspect only the supplied IR. Calls inside device helpers become visible - here after inlining; this scan does not visit their bodies. ``ThreadData`` - and ``TempStorage`` constructors alone do not count as group markers. - - The single whole-function planner uses a positive result in its group - resolution phase to request launch metadata and resolve the calls. - Successful resolution removes group descriptors and replaces public - operations with private provider calls, which do not count as group - markers. The following ``CoopSinglePhaseRewrite`` helper also checks for - unresolved markers before materializing those provider calls. - - Detection does not validate the calls or resolve their launch dimensions; - those checks belong to the group planner. - - Parameters - ---------- - func_ir : ir.FunctionIR - Current function IR with its recorded-definition lookup available. - Scanned without modifying its blocks or resolving launch metadata. - - Returns - ------- - bool - Whether at least one recognized group-planning marker remains. - """ - # Detection needs definition lookup only, before launch facts exist. - analyzer = object.__new__(_GroupCallPlanner) - analyzer.func_ir = func_ir - analyzer._group_cache = {} - analyzer._hierarchy_cache = {} - - def is_group_descriptor(value: Any, seen: set[str]) -> bool: - if isinstance(value, ThreadGroup): - return True - if not isinstance(value, ir.Var) or value.name in seen: - return False - seen = {*seen, value.name} - for definition in analyzer._all_definitions(value): - if isinstance(definition, ir.Var) and is_group_descriptor( - definition, seen - ): - return True - if isinstance(definition, (ir.Global, ir.FreeVar, ir.Const)): - if isinstance(definition.value, ThreadGroup): - return True - continue - if not isinstance(definition, ir.Expr): - continue - if definition.op == "cast" and is_group_descriptor( - definition.value, seen - ): - return True - if definition.op == "phi" and any( - is_group_descriptor(incoming, seen) - for incoming in getattr(definition, "incoming_values", ()) - ): - return True - if definition.op != "call": - continue - function = analyzer._callable(definition.func) - if function in _GROUP_CONSTRUCTORS: - return True - function_definition = analyzer._definition(definition.func) - if ( - isinstance(function_definition, ir.Expr) - and function_definition.op == "getattr" - and function_definition.attr == "group_by" - and is_group_descriptor(function_definition.value, seen) - ): - return True - return False - - for block in func_ir.blocks.values(): - for inst in block.body: - if not isinstance(inst, ir.Assign): - continue - value = inst.value - if not isinstance(value, ir.Expr) or value.op != "call": - continue - function = analyzer._callable(value.func) - if ( - function is ThreadHierarchy - or function in _GROUP_CONSTRUCTORS - or _group_operation_name(function) is not None - ): - return True - function_definition = analyzer._definition(value.func) - if ( - isinstance(function_definition, ir.Expr) - and function_definition.op == "getattr" - and (function_definition.attr == "group_by") - and is_group_descriptor(function_definition.value, set()) - ): - return True - return False - - -class _GroupPlanning: - """Add the group-resolution phase to ``CoopWholeFunctionPlanner``. - - This is a mixin: the whole-function planner inherits its methods rather - than creating a separate ``_GroupPlanning`` object. The owning planner - supplies ``state`` and ``is_device_function`` through Numba's - ``WholeFunctionPlanner`` base class. Its ``run()`` method calls - ``_resolve_groups()`` after device helpers have been inlined and before - provider-call rewriting starts. - - The method first checks whether any group calls remain, then requests - launch dimensions and delegates to a fresh ``_GroupCallPlanner``. Calls - in a standalone device function must instead be inlined into a kernel, - whose configured launch supplies those dimensions. The returned boolean - tells the owner whether to rebuild IR analysis before rewriting provider - calls. - """ - - def _resolve_groups(self) -> bool: - """Resolve group calls using launch metadata and report IR changes. - - Request launch facts only when group markers remain in the function. - Device helpers must first be inlined into a kernel so their groups - can be resolved against that kernel's configured launch. - - ``require_launch_config()`` explicitly declares that generated code - depends on launch values. Numba-CUDA-MLIR already has those values in - its dispatcher, but does not automatically expose them as compiler - specialization inputs. This opt-in lets kernels that do not need - launch facts retain generic dispatch, avoiding additional launch - normalization and specialization lookups. The helper runs during - compilation; the later launch cost comes from selecting a matching - specialization. - - The runtime helper reads the options shared by compiler passes in - ``state.metadata["targetoptions"]``. If ``"__launch_config__"`` - already holds a dictionary containing ``grid`` and ``block``, it - returns that dictionary. Otherwise, it gets the per-attempt tracker - from ``state.metadata`` using the runtime's - ``_LAUNCH_CONFIG_TRACKER_METADATA_KEY`` (``"launch_config_tracker"``). - This slot holds a temporary coordination object. The tracker lazily - normalizes the available ``grid``, ``block``, ``sharedmem``, and - ``cluster`` values and marks them as required. The helper caches the - result in ``targetoptions["__launch_config__"]`` for later passes. - - Compilation continues in the same attempt. The dispatcher observes - the tracker's required flag and keys the compiled result by launch - values so later launches reuse a matching specialization. The runtime - removes the temporary tracker from the final compilation metadata. - Without an active configuration or tracker, the helper raises an - error requiring compilation through a configured kernel launch. - - Pass the returned configuration to ``_GroupCallPlanner`` to resolve - the group topology and replace public operations with provider calls. - - Returns - ------- - bool - Whether group descriptors or public operations were replaced in - ``state.func_ir``. A function without group markers returns - ``False`` without requesting launch metadata. - - Raises - ------ - GroupRewriteError - Cooperative calls remain in a standalone device function, or group - planning finds invalid calls or escaping descriptors. - ForceLiteralArg - A group or operation needs a compile-time argument value. The - dispatcher consumes this compiler signal and retries with the - requested literal specialization. - RuntimeError - Numba has no configured launch metadata for this compilation. - Public compiler entry points translate the compiler's internal - signal to a diagnostic requiring a configured kernel launch. - NotImplementedError - The requested group or operation has no supported lowering. - """ - planner = cast("CoopWholeFunctionPlanner", self) - if not has_group_markers(planner.state.func_ir): - return False - if planner.is_device_function: - function_name = planner.state.func_ir.func_id.func_qualname - raise GroupRewriteError( - "cuda.coop.numba_mlir cooperative calls in device function " - f"{function_name!r} must be inlined into a kernel. " - "Standalone primitive helpers and primitives inside " - "standalone callbacks are unsupported; use inline='always' " - "for a kernel helper or move the cooperative calls " - "into the kernel." - ) - launch_config = require_launch_config(planner.state) - return _GroupCallPlanner(planner.state, launch_config).run() - - -__all__ = [ - "GroupRewriteError", - "_GroupCallPlanner", - "_GroupPlanning", - "has_group_markers", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planner_support.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planner_support.py deleted file mode 100644 index a62046678077..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planner_support.py +++ /dev/null @@ -1,121 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Identify the public group calls understood by the Numba planner. - -Common and backend-qualified constructors use different Python callables. For -example, ``cuda.coop.this_block`` and ``cuda.coop.numba_mlir.this_block`` both -describe a block, and the tables here map either callable to the backend's -descriptor constructor. The planner also records whether the kernel used the -common API, whose accepted group kinds and arguments must remain consistent -across backends. - -That distinction still matters after the public call is replaced by a private -provider call. For example, common ``load`` requires a ``ThreadData`` output; -the Numba-CUDA-MLIR qualified API also accepts a CUDA local array. The planner -must retain the call's API origin to check the appropriate payload rule. -Operation lookup compares the actual registered Python objects, so a user -function named ``load`` is not mistaken for ``cuda.coop.load``. - -The planner helpers also share an IR import, temporary-name counter, and base -exception here. Keeping these definitions separate lets group planning and its -operation modules use them without importing each other during initialization. -""" - -from __future__ import annotations - -from itertools import count -from typing import TYPE_CHECKING - -import cuda.coop._core.api as _common_api -import cuda.coop._core.api._dispatch as _common_dispatch - -from .. import _thread_group as _thread_groups -from ._operations import group_operation_name - -if TYPE_CHECKING: - from numba_cuda_mlir.numba_cuda.core import ( - ir as ir, # noqa: PLC0414 - Re-export for typing. - ) -else: - from numba_cuda_mlir.numbair_transforms import ( - ir as ir, # noqa: PLC0414 - Re-export for typing. - ) - -_NAME_COUNTER = count() -_GROUP_CONSTRUCTORS = { - _thread_groups.this_thread: _thread_groups.this_thread, - _thread_groups.this_warp: _thread_groups.this_warp, - _thread_groups.this_block: _thread_groups.this_block, - _thread_groups.this_cluster: _thread_groups.this_cluster, - _thread_groups.this_grid: _thread_groups.this_grid, - _common_api.this_thread: _thread_groups.this_thread, - _common_api.this_warp: _thread_groups.this_warp, - _common_api.this_block: _thread_groups.this_block, - _common_api.this_cluster: _thread_groups.this_cluster, - _common_api.this_grid: _thread_groups.this_grid, -} -_COMMON_GROUP_CONSTRUCTORS = frozenset( - { - _common_api.this_thread, - _common_api.this_warp, - _common_api.this_block, - _common_api.this_cluster, - _common_api.this_grid, - } -) - - -class GroupRewriteError(Exception): - """A recognized group call cannot be lowered with the available facts. - - Report a planning failure to the compiler. ``ForceLiteralArg`` is the - separate compiler signal for retrying with a specialized argument. - """ - - -def _group_operation_name(function: object) -> str | None: - """Find the operation planner associated with a public Python callable. - - Marker detection and call rewriting pass the object recovered from the - kernel's call target as ``function``. Return its registered name, such as - ``"load"`` or ``"store"``, or ``None`` if it is not a recognized common - or qualified operation. Comparing callable identity preserves aliases - such as ``my_load = coop.load`` without accepting unrelated functions - that happen to have the same name. - """ - - operation = group_operation_name(function) - if operation is None: - operation = _common_dispatch._common_group_operation_name(function) - return operation - - -def _is_common_root_operation(function: object, operation: str) -> bool: - """Tell call rewriting whether to apply the common operation's rules. - - ``_lower_root_operation`` uses this before replacing a public call with a - backend provider. Since rewriting bypasses the common Python wrapper, - the planner must check that API's supported group kinds, selector values, - and payload types itself. - - ``function`` is the Python object recovered from the call target; - ``operation`` is its expected registered name, such as ``"load"``. - Return ``True`` only when that object is the common API marker for the - named operation. Qualified operations and unrelated callables return - ``False``. - """ - - return _common_dispatch._common_group_operation_name(function) == operation - - -__all__ = [ - "_COMMON_GROUP_CONSTRUCTORS", - "_GROUP_CONSTRUCTORS", - "_NAME_COUNTER", - "GroupRewriteError", - "_group_operation_name", - "_is_common_root_operation", - "ir", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planning.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planning.py deleted file mode 100644 index 9f61879e39df..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_planning.py +++ /dev/null @@ -1,1143 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Provide operation planners with group, payload, and storage facts. - -Each supported operation needs some of the same information: the participating -threads, a payload's element type and item count, compile-time scalar -controls, and the options of a ``TempStorage`` descriptor. -``GroupPlanningContext`` exposes those queries to operation-specific code -without direct access to all of ``_GroupCallPlanner``'s bookkeeping. - -This context belongs to the group-resolution phase of the single -whole-function planner, before provider calls are materialized. It follows IR -assignments to recover facts that ordinary type inference has not established -yet and records element types inferred by earlier ``load()`` operations for -later calls. Once an operation has chosen a provider, it also checks that the -provider's storage and synchronization declarations agree with the group plan -and builds the call IR that carries that plan to the rewriting phase. -""" - -from __future__ import annotations - -from collections.abc import Callable -from numbers import Integral -from typing import TYPE_CHECKING, Any - -import numba_cuda_mlir.numba_cuda.types as _numba_types -from numba_cuda_mlir import cuda as _cuda_module -from numba_cuda_mlir.cuda.local import array as _cuda_local_array - -import cuda.coop._core.api as _common_api -from cuda.coop._core import ( - ArgumentBinding, - GroupLoweringPlan, - StorageOwnership, - SynchronizationScope, -) - -from .._temp_storage import TempStorage -from .._thread_data import ThreadData -from ._descriptor_provenance import ( - descriptor_definitions, - payload_write_dtypes, - temp_storage_constructor, -) -from ._group_planner_support import GroupRewriteError, ir -from ._operations import ( - _GROUP_LOWERING_PLAN_KWARG, - StorageABI, - expected_storage_reuse_barrier, - factory_operation, - provider_synchronization_matches, -) -from ._parameters import ( - _python_scalar_dtype, - normalize_dtype_param, -) -from ._scalar_provenance import ( - cuda_index_dtype, - scalar_call_dtype, - scalar_expression_dtype, -) - -if TYPE_CHECKING: - from ._group_planner import _GroupCallPlanner - - -class GroupPlanningContext: - """Share one planner's analysis with operation implementations. - - An operation planner needs launch dimensions, payload types and extents, - and static scalar controls before ordinary Numba typing runs. This - context exposes those queries and builds provider-call replacements. - The owning planner retains the IR and decides when to install them. - - The context retains inferred ThreadData dtypes at constructor sites so - a later operation can use the type established by an earlier load. - Aliases of the same constructor share this fact. Loop dtype candidates - have a shorter lifetime: they exist for one query, seed a type cycle, - and must pass a second query in which all definitions resolve and agree. - See ``_loop_dtype`` for that two-stage analysis. - - Parameters - ---------- - planner : _GroupCallPlanner - Active planner supplying IR definitions, argument types, and launch - facts. Create one context for each planner attempt. - """ - - __slots__ = ( - "__loop_dtypes", - "__planner", - "__seed_loop_dtypes", - "__thread_data_dtypes", - ) - - def __init__(self, planner: _GroupCallPlanner) -> None: - self.__planner = planner - self.__thread_data_dtypes: dict[int, Any] = {} - self.__loop_dtypes: dict[str, Any] = {} - self.__seed_loop_dtypes = False - - @property - def launch(self) -> Any: - """Return launch dimensions and their sources for this attempt.""" - - return self.__planner.launch - - def _definition(self, value: Any) -> Any: - return self.__planner._definition(value) - - def _all_definitions(self, value: ir.Var) -> tuple[Any, ...]: - return self.__planner._all_definitions(value) - - def _callable(self, value: Any) -> Any: - return self.__planner._callable(value) - - def constant(self, value: Any) -> Any: - """Require a constant control without literal-unroll dependencies. - - The planner may request literal specialization for a kernel argument. - Use ``try_static_scalar`` when a runtime scalar is also permitted. - """ - - self.__planner._reject_literal_unroll_value( - value, "a compile-time argument" - ) - return self.__planner._constant(value) - - def try_constant(self, value: Any) -> tuple[bool, Any]: - """Probe for a constant without requesting another compiler attempt. - - Return ``(resolved, value)``. Numba constant inference may resolve it; - use ``try_static_scalar`` when the source must be explicitly static. - """ - - return self.__planner._try_constant(value) - - def try_static_scalar(self, value: Any) -> tuple[bool, Any]: - """Recover an explicitly static scalar and keep its known width. - - Return ``(resolved, value)`` without evaluating runtime expressions or - requesting specialization. A false result means a runtime operand. - """ - - return self.__planner._try_static_scalar(value) - - def try_static_scalar_provenance(self, value: Any) -> tuple[bool, Any]: - """Recover a static scalar together with its recorded dtype. - - Return ``(resolved, provenance)``. The record distinguishes an untyped - Python literal from a value that already has a compiler dtype. - """ - - return self.__planner._try_static_scalar_provenance(value) - - def bind(self, function: Any, call: ir.Expr) -> Any: - """Bind IR arguments and defaults to the public signature.""" - - return self.__planner._bind(function, call) - - def validate_common_selector( - self, - operation: str, - parameter: str, - value: Any, - allowed: frozenset[str], - *, - allow_none: bool = False, - ) -> Any: - """Apply the common API's selector rules after its wrapper is removed. - - Resolve the selector as a constant, normalize its spelling, and check - ``allowed``. ``allow_none`` controls whether omission is valid. - """ - - return self.__planner._validate_common_selector( - operation, - parameter, - value, - allowed, - allow_none=allow_none, - ) - - def is_none(self, value: Any) -> bool: - return self.__planner._is_none(value) - - def is_array(self, operation: str, value: Any) -> bool: - """Check for a supported per-thread payload constructor. - - Accept ThreadData and CUDA local-array constructors. The common API's - ThreadData-only rule is checked separately by ``is_thread_data``. An - unresolved cycle raises a diagnostic; a false result still needs - scalar validation by the caller. - """ - - return self.__planner._array_operand_state(operation, value) - - def is_thread_data( - self, operation: str, parameter: str, value: Any - ) -> bool: - """Check the payload-origin restriction required by the common API. - - Follow aliases to ThreadData constructors. An unresolved cycle raises - a diagnostic that names the operation and parameter. - """ - - return self.__planner._thread_data_operand_state( - operation, - parameter, - value, - ) - - def array_extent(self, value: Any) -> int | None: - """Recover the unique known per-thread item count, or ``None``. - - Known counts must agree. Unknown paths contribute no count, so the - caller must validate the payload's form separately. - """ - - return self.__planner._array_extent(value) - - def new_var(self, scope: Any, loc: ir.Loc, stem: str) -> ir.Var: - return self.__planner._new_var(scope, loc, stem) - - def value_var( - self, - statements: list[Any], - *, - scope: Any, - loc: ir.Loc, - stem: str, - value: Any, - ) -> ir.Var: - return self.__planner._value_var( - statements, - scope=scope, - loc=loc, - stem=stem, - value=value, - ) - - @staticmethod - def _validate_provider_contract( - lowering_plan: GroupLoweringPlan, - factory: Callable[..., Any], - *, - runtime_temp_storage_supplied: bool | None = None, - ) -> None: - """Check the selected provider against planned storage and execution. - - This is the boundary between compiler-neutral planning and the private - provider call. Require complete topology, participation, - synchronization, and storage facts, then compare them with the - registered factory's ABI and scopes. - - Some algorithms need temporary shared memory to exchange values - between threads. A block ``load()`` operation with - ``algorithm="transpose"`` uses one scratch region for the block; a - warp ``load()`` operation with that algorithm needs a separate region - for each participating physical or logical warp. For such plans, - require exact block dimensions, groups that cover the block, and one - shared-memory slice per group instance. Direct algorithms for - ``load()`` and ``store()`` need no scratch and skip those checks. An - explicit ``temp_storage`` argument is supported only for a single - block-scoped instance. - - Normally the provider's declared reuse barrier must match the plan. - Caller-owned storage with ``auto_sync=False`` also permits the - provider's execution-scope barrier declaration: the pointer rewrite - bypasses its allocating wrapper and controls synchronization itself. - This exception does not apply to implementation-owned storage. No IR - is mutated here. - - Parameters - ---------- - lowering_plan : GroupLoweringPlan - Supported plan with storage and execution requirements to check. - factory : callable - Selected host-side provider factory registered with operation - metadata. It is not invoked here. - runtime_temp_storage_supplied : bool or None, optional - Whether the proposed provider call supplies ``temp_storage``. When - the plan needs scratch, this flag must agree with caller - ownership. ``None`` skips this argument-presence check only. - - Returns - ------- - None - The provider metadata and supported storage contracts agree. - - Raises - ------ - TypeError - ``lowering_plan`` is not a ``GroupLoweringPlan``. - GroupRewriteError - The plan is unsupported or incomplete, the provider is - unregistered, or its ABI, scopes, ownership, or layout disagree. - """ - - if not isinstance(lowering_plan, GroupLoweringPlan): - raise TypeError("lowering_plan must be a GroupLoweringPlan") - if lowering_plan.unsupported is not None: - raise GroupRewriteError( - "cuda.coop.numba_mlir cannot select a provider for an " - "unsupported lowering plan" - ) - metadata = factory_operation(factory) - if metadata is None: - raise GroupRewriteError( - "cuda.coop.numba_mlir lowering plan selected an unregistered " - "provider factory" - ) - topology = lowering_plan.topology - participation = lowering_plan.participation - synchronization = lowering_plan.synchronization - storage = lowering_plan.temp_storage - if ( - topology is None - or participation is None - or synchronization is None - or storage is None - ): - raise GroupRewriteError( - "cuda.coop.numba_mlir provider selection requires complete " - "group topology, participation, synchronization, and storage " - "contracts" - ) - if topology.execution_scope is SynchronizationScope.GROUP: - raise GroupRewriteError( - "cuda.coop.numba_mlir provider execution scope 'group' " - "has no storage or synchronization emitter" - ) - storage_bearing = storage.ownership is not StorageOwnership.NONE - if storage_bearing: - if storage.address_space != "shared": - raise GroupRewriteError( - "cuda.coop.numba_mlir storage-bearing providers require " - "shared-address-space TempStorage" - ) - if ( - storage.instances != topology.instances - or storage.instance_index != topology.instance_index - ): - raise GroupRewriteError( - "cuda.coop.numba_mlir TempStorage layout disagrees with " - "its group topology" - ) - exact_block_dim = participation.exact_block_dim - if exact_block_dim is None: - raise GroupRewriteError( - "cuda.coop.numba_mlir storage-bearing providers require " - "exact block dimensions" - ) - block_threads = ( - exact_block_dim[0] * exact_block_dim[1] * exact_block_dim[2] - ) - if topology.logical_width * topology.instances != block_threads: - raise GroupRewriteError( - "cuda.coop.numba_mlir group topology does not cover the " - "exact block dimensions" - ) - caller_owned = storage.ownership is StorageOwnership.CALLER - if ( - storage_bearing - and runtime_temp_storage_supplied is not None - and (caller_owned != runtime_temp_storage_supplied) - ): - raise GroupRewriteError( - "cuda.coop.numba_mlir TempStorage ownership disagrees with " - "the provider call arguments" - ) - if caller_owned and ( - topology.execution_scope is not SynchronizationScope.BLOCK - or topology.instances != 1 - ): - if topology.execution_scope is SynchronizationScope.WARP: - raise GroupRewriteError( - "cuda.coop.numba_mlir caller-owned TempStorage is not " - "supported for warp-scoped cooperative primitives; omit " - "temp_storage so the implementation can provide one " - "aligned slice per group instance" - ) - raise GroupRewriteError( - "cuda.coop.numba_mlir caller-owned TempStorage is supported " - "only for single-instance block-scoped cooperative primitives" - ) - expected_reuse_barrier = expected_storage_reuse_barrier( - topology, storage - ) - if synchronization.storage_reuse_barrier is not expected_reuse_barrier: - raise GroupRewriteError( - "cuda.coop.numba_mlir TempStorage automatic synchronization " - "disagrees with the planned storage-reuse barrier" - ) - expected = { - "storage_abi": ( - StorageABI.LEADING_POINTER - if storage_bearing - else StorageABI.NONE - ), - "execution_scope": topology.execution_scope, - "synchronization_scope": synchronization.storage_reuse_barrier, - } - mismatches = [ - f"{name}={getattr(metadata, name).value!r} (plan {planned.value!r})" - for name, planned in expected.items() - if name != "synchronization_scope" - if getattr(metadata, name) is not planned - ] - planned_synchronization = expected["synchronization_scope"] - if not provider_synchronization_matches( - metadata, topology, synchronization, storage - ): - mismatches.append( - "synchronization_scope=" - f"{metadata.synchronization_scope.value!r} " - f"(plan {planned_synchronization.value!r})" - ) - if mismatches: - details = ", ".join(mismatches) - raise GroupRewriteError( - f"cuda.coop.numba_mlir provider {metadata.operation!r} " - f"metadata disagrees with its lowering plan: {details}" - ) - - def rewrite_call( - self, - inst: ir.Assign, - *, - lowering_plan: GroupLoweringPlan, - factory: Callable[..., Any], - args: list[Any], - kwargs: dict[str, Any], - common_root_operation: str | None = None, - ) -> list[Any]: - """Build a provider call carrying the validated group-lowering plan. - - Check the provider ABI and storage contract before embedding the plan - in its reserved keyword argument. The later provider rewrite consumes - this metadata, so it does not have to reconstruct the public group - semantics. The returned assignments materialize non-IR arguments and - invoke the factory with the original result target. The caller - installs them into the function; this method does not replace the - original instruction. - - Parameters - ---------- - inst : ir.Assign - Original public call assignment; supplies the result target, - scope, and source location for generated statements. - lowering_plan : GroupLoweringPlan - Supported plan to validate and attach to the provider call. - factory : callable - Registered host-side provider factory selected by the operation - family. Embedded as the generated call target, not invoked here. - args : list of object - Positional provider arguments, as existing IR variables or host - values. - kwargs : dict of str to object - Provider keyword arguments. Copied before plan metadata is added; - presence of ``temp_storage`` is checked against planned ownership. - common_root_operation : str or None, optional - Common API operation name to retain for downstream validation. - When present, supplies the private marker unless ``kwargs`` - already has it. - - Returns - ------- - list of object - Ordered argument-materialization and call assignments replacing - ``inst``. - - Raises - ------ - GroupRewriteError - Provider-contract validation fails or ``kwargs`` uses the reserved - lowering-plan keyword. - """ - - self._validate_provider_contract( - lowering_plan, - factory, - runtime_temp_storage_supplied="temp_storage" in kwargs, - ) - if _GROUP_LOWERING_PLAN_KWARG in kwargs: - raise GroupRewriteError( - "cuda.coop.numba_mlir family lowering " - "used a reserved provider keyword" - ) - kwargs = { - **kwargs, - _GROUP_LOWERING_PLAN_KWARG: lowering_plan, - } - return self.__planner._rewritten_call( - inst, - factory=factory, - args=args, - kwargs=kwargs, - common_root_operation=common_root_operation, - ) - - def planning_binding(self, value: Any) -> ArgumentBinding: - """Classify a scalar control from its explicit static provenance. - - Use explicit static provenance rather than general constant inference. - A runtime expression remains a runtime binding even if another - compiler analysis could fold it. The original runtime operand is - retained by the operation family, not inside the returned binding. - - Parameters - ---------- - value : ir.Var or object - Optional scalar control such as ``valid_items``, ``offset``, or a - ``load()`` fill value, represented by an IR variable or a value - already known during compilation. - - Returns - ------- - ArgumentBinding - ``OMITTED`` for statically known ``None``, ``STATIC`` with the - resolved value otherwise, or ``RUNTIME`` when static provenance is - not established. Numeric validity and operation-specific - constraints are checked later. - """ - - resolved, constant = self.try_static_scalar(value) - if not resolved: - return ArgumentBinding.runtime() - if constant is None: - return ArgumentBinding.omitted() - return ArgumentBinding.static(constant) - - @staticmethod - def _dtype_from_numba_type(value: object) -> _numba_types.Type | None: - """Get an element dtype from a compiler type before IR typing runs. - - ``dtype`` uses this for an already known Numba type; - ``_dtype_definition`` uses it for kernel argument types. An array - contributes its element dtype, another Numba type is normalized - directly, and an object that is not a Numba type returns ``None``. - """ - - if isinstance(value, _numba_types.Array): - value = value.dtype - elif not isinstance(value, _numba_types.Type): - return None - return normalize_dtype_param(value) - - @staticmethod - def _one_dtype(candidates: set[Any], *, message: str) -> Any | None: - """Return the sole known dtype collected by ``_complete_dtype``. - - ``candidates`` contains distinct, resolved dtypes. An empty set - returns ``None``; more than one dtype raises ``GroupRewriteError`` - with the caller's ``message`` so it identifies the failing query. - """ - - if len(candidates) > 1: - raise GroupRewriteError(message) - return next(iter(candidates), None) - - def _complete_dtype( - self, - candidates: Any, - *, - message: str, - ) -> Any | None: - """Require agreement among the dtype candidates for one query. - - Variable, tuple, and phi queries call this with an iterable of - normalized dtypes or ``None`` entries for unknown sources. - During loop discovery, omit unknown candidates so a known entry value - can seed a cycle. The later strict query requires every candidate to - be known. Return ``None`` if information is incomplete. After that - check, raise ``GroupRewriteError`` with ``message`` when the remaining - dtypes disagree. - """ - - resolved = list(candidates) - if self.__seed_loop_dtypes: - # Discovery can use a known loop-entry type before the backedge - # resolves. The final pass requires every candidate to resolve. - resolved = [dtype for dtype in resolved if dtype is not None] - if not resolved or any(dtype is None for dtype in resolved): - return None - return self._one_dtype(set(resolved), message=message) - - def _loop_dtype(self, value: ir.Var) -> Any | None: - """Resolve a loop's type cycle without accepting an unknown producer. - - Group planning runs before ordinary type inference. A value - initialized from a typed array may then feed a computation whose - result becomes the next iteration's input. Following that input - recursively reaches the same variable before its dtype is known. - Returning ``None`` at every such backedge would discard the useful - type supplied by the array. - - First propagate candidate dtypes from known definitions until they - stop changing. During this discovery pass only, a join may omit - unresolved inputs; a recursive backedge can use its variable's - candidate. Then repeat the query with strict joins: every recorded - definition must resolve and agree, including the computation on the - backedge. An opaque helper remains unknown, and a conflicting dtype is - still an error. - - Candidates live only for this query and are cleared on failure too. - They are neither IR annotations nor permanent facts for later calls. - This solves supported type-preserving cycles; it does not implement - general type promotion or infer arbitrary device-helper return types. - - Parameters - ---------- - value : ir.Var - Variable passed to ``dtype`` without an existing recursion path. - Its recorded assignments may include a loop's initial value and - values produced on later iterations. - - Returns - ------- - numba_types.Type or None - Dtype verified on all paths, or ``None`` if any path is unknown. - - Raises - ------ - GroupRewriteError - Known recorded definitions require different dtypes. - """ - self.__seed_loop_dtypes = True - try: - while True: - previous = self.__loop_dtypes.copy() - self.dtype(value, seen=set()) - if self.__loop_dtypes == previous: - break - self.__seed_loop_dtypes = False - return self.dtype(value, seen=set()) - finally: - self.__seed_loop_dtypes = False - self.__loop_dtypes.clear() - - def record_thread_data_dtype( - self, value: Any, dtype: _numba_types.Type - ) -> None: - """Record a producer's dtype at the payload's constructor sites. - - Group planning precedes the provider rewrite that materializes - payloads. A load into untyped ``ThreadData`` therefore records its - inferred dtype here so subsequent group calls can recover it. Follow - descriptor aliases, casts, phi inputs, and constant tuple projections - to constructor calls, keying the cache by call-expression identity so - aliases share the fact. Explicit constructor dtypes and earlier - inferred dtypes must agree. - - Only recognized constructors reached by this traversal are updated; - unresolved tuple projections and other leaves contribute no cache - entry. This updates the planning context, not constructor arguments in - the IR. - - Parameters - ---------- - value : ir.Var - Producer's output payload, possibly reached through supported - aliases or tuple projections. - dtype : numba_types.Type - Normalized element dtype inferred by the producer. - - Returns - ------- - None - Any recognized constructor sites now carry the inferred dtype. - - Raises - ------ - GroupRewriteError - A reached constructor already has a different explicit or inferred - dtype. Cache entries recorded before the conflict are not rolled - back. - """ - - def payload_definitions(current, seen): - for _, definition in descriptor_definitions( - current, self._all_definitions, seen=seen - ): - if not isinstance(definition, ir.Expr) or definition.op not in { - "getitem", - "static_getitem", - }: - yield definition - continue - index = definition.index - if isinstance(index, ir.Var): - resolved, index = self.try_constant(index) - if not resolved: - continue - if not isinstance(index, Integral) or isinstance(index, bool): - continue - index = int(index) - next_seen = {*seen, current.name} - for packed in payload_definitions(definition.value, next_seen): - if ( - isinstance(packed, ir.Expr) - and packed.op == "build_tuple" - and -len(packed.items) <= index < len(packed.items) - ): - yield from payload_definitions( - packed.items[index], next_seen - ) - - for definition in payload_definitions(value, set()): - if ( - isinstance(definition, ir.Expr) - and definition.op == "call" - and self._callable(definition.func) - in {ThreadData, _common_api.ThreadData} - ): - previous = self._dtype_definition(definition, seen=set()) - if previous is not None and previous != dtype: - raise GroupRewriteError( - "cuda.coop.numba_mlir ThreadData aliases have " - "inconsistent dtypes" - ) - self.__thread_data_dtypes[id(definition)] = dtype - - def _tuple_dtype( - self, - value: Any, - index: int, - *, - seen: set[str], - ) -> Any | None: - """Infer one tuple element across recorded container definitions. - - ``_dtype_definition`` calls this for a constant tuple index so a - payload packed into a tuple retains its element dtype when unpacked. - Track ``variable[index]`` separately from the container variable. This - allows one projection to recurse without hiding a different element. - Return ``None`` for an unresolved projection or cycle. Candidate - agreement uses ``_complete_dtype``, including its temporary - loop-discovery rule. - - Parameters - ---------- - value : ir.Var or object - Variable holding the tuple; other objects return ``None``. - index : int - Tuple position, with Python's negative-index convention. This is - not an index into the payload array stored at that position. - seen : set of str - Recursion-path keys. Add the current projection in place, then - give each recorded definition its own copy. - """ - - if not isinstance(value, ir.Var): - return None - seen_key = f"{value.name}[{index}]" - if seen_key in seen: - return None - seen.add(seen_key) - return self._complete_dtype( - ( - self._tuple_dtype_definition( - definition, - index, - seen=set(seen), - ) - for definition in self._all_definitions(value) - ), - message=( - "cuda.coop.numba_mlir tuple " - "projections have inconsistent dtypes" - ), - ) - - def _tuple_dtype_definition( - self, - definition: Any, - index: int, - *, - seen: set[str], - ) -> Any | None: - """Follow one tuple definition to the selected element's dtype. - - ``_tuple_dtype`` calls this for each recorded source of a tuple. - Aliases, casts, iterator unpacking, and phi inputs retain the index. A - built tuple delegates its selected element to ``dtype``. Unknown forms - and out-of-range indices contribute no dtype; known conflicts at a phi - join raise ``GroupRewriteError``. - - Parameters - ---------- - definition : object - Assignment source, usually an IR variable or expression. - index : int - Position within the tuple, accepting Python negative indices. - seen : set of str - Variable and tuple-projection keys already on this recursion path. - Each phi input receives a separate copy. - - Returns - ------- - numba_types.Type or None - Selected element's normalized dtype, or ``None`` if unknown. - """ - - if isinstance(definition, ir.Var): - return self._tuple_dtype(definition, index, seen=seen) - if not isinstance(definition, ir.Expr): - return None - if definition.op in {"cast", "exhaust_iter"}: - return self._tuple_dtype(definition.value, index, seen=seen) - if definition.op == "phi": - return self._complete_dtype( - ( - self._tuple_dtype(incoming, index, seen=set(seen)) - for incoming in getattr(definition, "incoming_values", ()) - ), - message=( - "cuda.coop.numba_mlir loop-carried tuple payloads have " - "inconsistent dtypes" - ), - ) - if definition.op == "build_tuple": - items = tuple(getattr(definition, "items", ())) - if not -len(items) <= index < len(items): - return None - return self.dtype(items[index], seen=seen) - return None - - def _dtype_definition( - self, definition: Any, *, seen: set[str] - ) -> Any | None: - """Extract dtype evidence from one supported IR definition. - - Use argument types, constants, payload constructors, selected scalar - operators, casts, and CUDA index attributes. Follow aliases and joins - through the context's dtype queries. A tuple projection can have its - own element type. Array indexing uses the source's element type. - - ThreadData constructors use an explicit dtype or a dtype recorded by a - producer. An unrelated call stays unknown unless the scalar-cast - helper recognizes it. This limited analysis supplies provider - selection before ordinary Numba typing; it does not execute kernel - expressions. - - ``dtype`` calls this for each recorded source of its variable. - ``definition`` is that assignment source; ``seen`` contains variable - and tuple-projection keys already on the recursion path. Return a - normalized dtype or ``None`` when the source supplies no known dtype. - """ - - if isinstance(definition, ir.Var): - return self.dtype(definition, seen=seen) - if isinstance(definition, ir.Arg): - if not 0 <= definition.index < len(self.__planner.state.args): - return None - return self._dtype_from_numba_type( - self.__planner.state.args[definition.index] - ) - if isinstance(definition, (ir.Global, ir.FreeVar, ir.Const)): - return _python_scalar_dtype(definition.value) - if not isinstance(definition, ir.Expr): - return None - if definition.op in {"cast", "exhaust_iter"}: - return self.dtype(definition.value, seen=seen) - if definition.op == "phi": - return self._complete_dtype( - ( - self.dtype(incoming, seen=set(seen)) - for incoming in getattr(definition, "incoming_values", ()) - ), - message=( - "cuda.coop.numba_mlir payload " - "aliases have inconsistent dtypes" - ), - ) - if definition.op in {"getitem", "static_getitem"}: - index = getattr(definition, "index", None) - if isinstance(index, ir.Var): - resolved, index = self.try_constant(index) - if not resolved: - return self.dtype(definition.value, seen=seen) - if isinstance(index, Integral) and not isinstance(index, bool): - tuple_dtype = self._tuple_dtype( - definition.value, - int(index), - seen=set(seen), - ) - if tuple_dtype is not None: - return tuple_dtype - return self.dtype(definition.value, seen=seen) - if definition.op in {"binop", "inplace_binop", "unary"}: - return scalar_expression_dtype( - definition, lambda value: self.dtype(value, seen=set(seen)) - ) - if definition.op == "getattr": - return cuda_index_dtype( - definition, self._attribute_chain, _cuda_module - ) - if definition.op != "call": - return None - function = self._callable(definition.func) - if function in {ThreadData, _common_api.ThreadData}: - bound = self.bind(function, definition) - resolved, dtype = self.try_constant(bound.arguments["dtype"]) - if resolved and dtype is not None: - return normalize_dtype_param(dtype) - return self.__thread_data_dtypes.get(id(definition)) - if function is _cuda_local_array: - if len(definition.args) >= 2: - resolved, dtype = self.try_constant(definition.args[1]) - if resolved: - return normalize_dtype_param(dtype) - dtype_ref = dict(definition.kws).get("dtype") - if dtype_ref is not None: - resolved, dtype = self.try_constant(dtype_ref) - if resolved: - return normalize_dtype_param(dtype) - return None - return scalar_call_dtype( - function, - definition.args, - lambda value: self.dtype(value, seen=set(seen)), - ) - - def _attribute_chain( - self, value: Any - ) -> tuple[Any, tuple[str, ...]] | None: - """Recover an attribute path and its constant root without calling it. - - CUDA index recognition compares the root object's identity and the - attribute names. Return ``None`` when the root cannot be recovered. - """ - - attributes: list[str] = [] - current = self._definition(value) - while isinstance(current, ir.Expr) and current.op == "getattr": - attributes.append(current.attr) - current = self._definition(current.value) - if not isinstance(current, (ir.Global, ir.FreeVar, ir.Const)): - return None - attributes.reverse() - return current.value, tuple(attributes) - - def dtype( - self, value: object, *, seen: set[str] | None = None - ) -> _numba_types.Type | None: - """Infer a normalized dtype from facts known during group planning. - - Use argument types, scalar constants and operators, supported CUDA - index attributes, local-array constructors, and ``ThreadData`` - declarations or recorded producer dtypes. Follow aliases, casts, phi - inputs, and tuple projections; array indexing contributes the source - element dtype. This is a limited pre-typing analysis, not full Numba - type inference. - - Loop backedges can use a candidate established by a known incoming - definition, provided every recorded definition then resolves to the - same type. Unknown definitions and unseeded cycles return ``None``; - fully known but inconsistent paths are rejected. - - Parameters - ---------- - value : object - IR variable to inspect, or a Numba compiler type to normalize - directly. Array types contribute their element dtype. Other inputs - are accepted as probes and return ``None``; a raw Python scalar is - not treated as an IR constant by this entry point. - seen : set of str, optional - Recursion-path variable names and tuple-projection keys. The - current name is added in place; definitions are visited with - separate copies. - - Returns - ------- - numba_types.Type or None - Normalized compiler dtype when every relevant path is known and - agrees, or ``None`` when inference is incomplete. - - Raises - ------ - GroupRewriteError - Fully resolved aliases or tuple projections disagree on the dtype. - """ - - if not isinstance(value, ir.Var): - return self._dtype_from_numba_type(value) - if seen is None: - inferred = self.dtype(value, seen=set()) - return inferred if inferred is not None else self._loop_dtype(value) - if value.name in seen: - # Only the loop-resolution query supplies a candidate here. - return self.__loop_dtypes.get(value.name) - seen.add(value.name) - inferred = self._complete_dtype( - ( - self._dtype_definition( - definition, - seen=set(seen), - ) - for definition in self._all_definitions(value) - ), - message=( - "cuda.coop.numba_mlir payload aliases have inconsistent dtypes" - ), - ) - if self.__seed_loop_dtypes and inferred is not None: - self.__loop_dtypes[value.name] = inferred - return inferred - - def payload_write_dtype(self, payload: Any) -> Any | None: - """Infer a payload dtype from values assigned through its aliases. - - The Store planner calls this when an untyped ``ThreadData`` payload - has no declared or producer-inferred dtype. ``payload`` is its IR - variable, possibly reached through aliases. - - Inspect the known types of element writes across the function. All - known types must agree or ``TypeError`` is raised. Unknown writes - contribute no evidence; ``None`` means no known write dtype was found. - The scan does not prove that every element is initialized or that all - paths write. - """ - - inferred = None - for value_dtype in payload_write_dtypes( - self.__planner.func_ir, payload, self.dtype - ): - if inferred is None: - inferred = value_dtype - elif inferred != value_dtype: - raise TypeError( - "cuda.coop.numba_mlir could not infer one consistent " - "dtype from payload writes" - ) - return inferred - - def temp_storage( - self, - value: Any, - *, - seen: set[str] | None = None, - ) -> tuple[int | None, int | None, bool, str] | None: - """Recover one storage contract from recorded descriptor definitions. - - The Load/Store planner uses this to validate an explicit storage - operand before choosing and recording its provider's storage plan. - Parse each recognized constructor with the planning constant resolver, - then require its normalized contract to agree with the others. A - concrete non-descriptor path, including a ``None`` initializer, - invalidates a value that also reaches a descriptor. Backedges - contribute no new leaf. - - Equivalent constructors may merge when automatic synchronization is - used. With ``auto_sync=False``, all aliases must reach exactly one - call expression: merging separately constructed regions would lose the - origin needed to reason about caller-managed synchronization. This - checks provenance and options, not backing storage or capacity. - - Parameters - ---------- - value : ir.Var or object - Value expected to name a storage descriptor. Non-variables return - ``None`` without parsing. - seen : set of str, optional - Recursion-path names passed to ``descriptor_definitions``. The - supplied set is not mutated. - - Returns - ------- - tuple or None - ``(size_in_bytes, alignment, auto_sync, sharing)`` for the unique - contract, or ``None`` if no recognized constructor is reached. The - first two fields may be ``None`` to defer size or alignment - selection. - - Raises - ------ - GroupRewriteError - Contracts conflict, descriptor and non-descriptor paths mix, - multiple constructor sites use manual synchronization, or call - syntax is invalid. - ForceLiteralArg - A constructor option requires literal argument specialization. - """ - - if not isinstance(value, ir.Var): - return None - candidates = set() - sites: set[int] = set() - non_descriptor = False - for _, definition in descriptor_definitions( - value, self._all_definitions, seen=seen - ): - if not ( - isinstance(definition, ir.Expr) - and definition.op == "call" - and self._callable(definition.func) - in {TempStorage, _common_api.TempStorage} - ): - non_descriptor = True - continue - descriptor = temp_storage_constructor( - definition, - lambda argument, *, name: self.constant(argument), - syntax_error=GroupRewriteError, - ) - candidates.add( - ( - descriptor.size_in_bytes, - descriptor.alignment, - descriptor.auto_sync, - descriptor.sharing, - ) - ) - sites.add(id(definition)) - if len(candidates) > 1: - raise GroupRewriteError( - "cuda.coop.numba_mlir TempStorage aliases have " - "inconsistent contracts" - ) - descriptor = next(iter(candidates), None) - if descriptor is not None and non_descriptor: - raise GroupRewriteError( - "cuda.coop.numba_mlir TempStorage variables must be bound to " - f"a TempStorage descriptor on every path; {value.name!r} is " - "also bound to a non-descriptor value such as None. " - "Remove the None initializer or construct the descriptor " - "unconditionally." - ) - if descriptor is not None and descriptor[2] is False and len(sites) > 1: - raise GroupRewriteError( - "cuda.coop.numba_mlir TempStorage with auto_sync=False must " - f"be constructed at exactly one site; {value.name!r} reaches " - f"{len(sites)} constructor sites. The compiler cannot verify " - "caller synchronization when it merges these regions. " - "Construct the descriptor once or set auto_sync=True." - ) - return descriptor - - -__all__ = ["GroupPlanningContext"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_rewriting.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_rewriting.py deleted file mode 100644 index 8666b670ae8f..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_group_rewriting.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Expose payload and scalar facts to operation-specific call rewriting. - -After group resolution has selected implementations, -``CoopSinglePhaseRewrite`` uses primitive-specific hooks to infer factory -arguments and prepare runtime operands. ``GroupRewriteContext`` gives those -hooks access to the active rewrite's array descriptions, compiler types, and -recorded ``ThreadData`` element types. It forwards queries and updates to the -rewrite so all operations use the same analysis and caches. - -This context is used during the second phase of ``CoopWholeFunctionPlanner``. -It does not register a separate Numba rewrite or own the IR; its returned -facts include only what the active helper knows before type inference. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -from ._rewrite_support import _ThreadDataSpecification, ir - -if TYPE_CHECKING: - from ._rewrite import CoopSinglePhaseRewrite - - -class GroupRewriteContext: - """Share the active call rewrite's analysis with primitive-specific hooks. - - Parameters - ---------- - rewrite : CoopSinglePhaseRewrite - Active rewrite supplying IR facts to primitive-family hooks. - """ - - __slots__ = ("__rewrite",) - - def __init__(self, rewrite: CoopSinglePhaseRewrite) -> None: - self.__rewrite = rewrite - - def thread_data(self, value: ir.Var) -> _ThreadDataSpecification | None: - """Return known payload facts, or ``None`` when none can be recovered. - - A description can have an unknown dtype or extent. It can also - describe a native local array; it does not prove that the value came - from public ``ThreadData``. - """ - - return self.__rewrite._resolve_thread_data_specification(value) - - def array(self, value: ir.Var) -> _ThreadDataSpecification | None: - """Return known dtype and extent facts for local or shared arrays.""" - - return self.__rewrite._resolve_array_specification_from_var( - value, seen=set() - ) - - def dtype(self, value: ir.Var) -> Any | None: - """Return the dtype inferred for an IR value, when known.""" - - return self.__rewrite._resolve_var_dtype(value) - - def numba_type(self, value: ir.Var) -> Any | None: - """Return the compiler type assigned to an IR value, when known.""" - - return self.__rewrite._resolve_var_numba_type(value) - - def infer_thread_data_write_dtype(self, value: ir.Var) -> Any | None: - """Infer an untyped ``ThreadData`` dtype from its element writes.""" - - return self.__rewrite._infer_thread_data_dtype_from_writes(value) - - def record_thread_data_dtype(self, value: ir.Var, dtype: Any) -> None: - """Record a dtype inferred for an otherwise untyped payload.""" - - self.__rewrite._record_inferred_thread_data_dtype(value, dtype) - - def static_scalar_provenance(self, value: Any) -> Any: - """Resolve a static scalar with its known source dtype. - - Return ``_UNRESOLVED`` when the analysis cannot establish static - origin. A provenance record with value ``None`` remains valid. - """ - - return self.__rewrite._resolve_static_scalar_provenance(value) - - -__all__ = ["GroupRewriteContext"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_layout.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_layout.py deleted file mode 100644 index 5a21d9851a7a..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_layout.py +++ /dev/null @@ -1,128 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Read C++ size and alignment from names produced by NVRTC. - -The compiler knows a provider's scratch layout after template instantiation. -A small variable template puts ``sizeof`` and ``alignof`` values into its -symbol name. Its address is registered as an NVRTC name expression. After -compilation, ``nvrtcGetLoweredName`` returns the mangled name containing both -values. The same compilation builds the provider image, so no separate -metadata program is compiled or linked. The layout stays tied to the -compiled source and target. -""" - -import hashlib -import re - - -def prepare_layout_queries( - cpp: str, layout_types: tuple[str, ...] -) -> tuple[str, str, tuple[str, ...]]: - """Append a probe that encodes each requested layout in a template name. - - The probe's two ``unsigned long long`` template arguments are the byte - size and alignment. Its name includes a digest of the ordered type names - so the same request produces the same source and cache identity. This - function only prepares source; NVRTC later evaluates the expressions. - - Parameters - ---------- - cpp : str - Provider translation unit containing the type definitions or aliases. - layout_types : tuple of str - Non-empty C++ type expressions visible at the end of ``cpp``. The - returned queries keep their order and duplicates, so each layout - matches its requested type. - - Returns - ------- - tuple - ``(source, symbol, queries)`` with the probe appended to the source, - its symbol name, and one address expression per requested type. An - empty tuple returns ``(cpp, "", ())`` and preserves the source. - - Raises - ------ - ValueError - A type expression is empty or contains only whitespace. Other C++ - syntax and type errors are left to NVRTC compilation. - """ - - if not layout_types: - return cpp, "", () - if any(not storage_type.strip() for storage_type in layout_types): - raise ValueError("layout types must be non-empty") - digest = hashlib.sha256(repr(layout_types).encode("utf-8")).hexdigest() - symbol = f"cuda_coop_layout_probe_{digest}" - source = ( - f"{cpp.rstrip()}\n\n" - "template \n" - f"__device__ unsigned char {symbol} = 0;\n" - ) - queries = tuple( - f"&{symbol}<(sizeof({storage_type})), (alignof({storage_type}))>" - for storage_type in layout_types - ) - return source, symbol, queries - - -def decode_layout_name( - lowered_name: bytes | str, *, symbol: str, expression: str -) -> tuple[int, int]: - """Decode one probe's compiler-generated name into a byte layout. - - Accept only the expected symbol followed by two unsigned-long-long - template values in NVRTC's mangled-name encoding. This is a decoder for - ``prepare_layout_queries`` probes, not a general C++ demangler. The - function validates size and alignment before returning them for use in - scratch allocation. - - Parameters - ---------- - lowered_name : bytes or str - Name returned by ``nvrtcGetLoweredName`` after compilation. Bytes - must be ASCII. Trailing NUL characters are ignored for both bytes - and str inputs. - symbol : str - Exact probe symbol returned by ``prepare_layout_queries``. - expression : str - Original NVRTC name expression, included in failure diagnostics. - - Returns - ------- - tuple of int - ``(size, alignment)`` in bytes. Both values are positive. Alignment - is a power of two, and size is a multiple of alignment. - - Raises - ------ - ValueError - The name has an unexpected symbol or encoding, or the decoded layout - fails these checks. Non-ASCII bytes also fail decoding. - """ - - if isinstance(lowered_name, bytes): - lowered_name = lowered_name.decode("ascii") - match = re.fullmatch( - rf"_Z{len(symbol)}{re.escape(symbol)}ILy([0-9]+)ELy([0-9]+)EE", - lowered_name.rstrip("\0"), - ) - if match is None: - raise ValueError( - "NVRTC returned an unexpected lowered layout-probe name for " - f"{expression!r}: {lowered_name!r}" - ) - size, alignment = (int(value) for value in match.groups()) - if ( - size <= 0 - or alignment <= 0 - or alignment & (alignment - 1) - or size % alignment - ): - raise ValueError( - f"Invalid storage layout for {expression!r}: " - f"size={size}, alignment={alignment}" - ) - return size, alignment diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_load_store_algorithms.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_load_store_algorithms.py deleted file mode 100644 index 637f0ba1c436..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_load_store_algorithms.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Keep Numba Load/Store selectors consistent with the shared core. - -Allowed names come from the core algorithm enums. Factories use the -storage-free subset to check that an algorithm without shared scratch -uses the factory without a scratch pointer. Callers pass strings, which -are normalized; enum objects are rejected. -""" - -from enum import Enum - -from cuda.coop._core.block.load_store import BlockLoadStoreAlgorithm -from cuda.coop._core.group.load_store import ( - _STORAGE_FREE_ALGORITHMS as _CORE_STORAGE_FREE_ALGORITHMS, -) -from cuda.coop._core.warp.load_store import WarpLoadStoreAlgorithm - -_BLOCK_LOAD_STORE_ALGORITHMS = frozenset( - item.value for item in BlockLoadStoreAlgorithm -) -_WARP_LOAD_STORE_ALGORITHMS = frozenset( - item.value for item in WarpLoadStoreAlgorithm -) -_STORAGE_FREE_ALGORITHMS = frozenset( - item.value for item in _CORE_STORAGE_FREE_ALGORITHMS -) - - -def _resolve_algorithm( - algorithm, allowed_algorithms, primitive_name: str -) -> str: - """Normalize a selector and check the primitive's supported choices. - - Ignore surrounding whitespace and letter case; treat hyphens as - underscores. Reject enum objects even when they inherit from ``str`` so - the API has one consistent selector form. - """ - - if not isinstance(algorithm, str) or isinstance(algorithm, Enum): - raise TypeError(f"{primitive_name} algorithm must be a string") - token = algorithm.strip().lower().replace("-", "_") - if token in allowed_algorithms: - return token - choices = ", ".join(sorted(allowed_algorithms)) - raise ValueError( - f"Unsupported {primitive_name} algorithm {algorithm!r}; expected one " - f"of: {choices}" - ) - - -__all__ = [ - "_BLOCK_LOAD_STORE_ALGORITHMS", - "_STORAGE_FREE_ALGORITHMS", - "_WARP_LOAD_STORE_ALGORITHMS", - "_resolve_algorithm", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_numba_mlir_compat.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_numba_mlir_compat.py deleted file mode 100644 index e5ab30b097a4..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_numba_mlir_compat.py +++ /dev/null @@ -1,225 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Validate the optional Numba backend's version before compiler integration. - -.. rubric:: Why package requirements and an activation check both exist - -The Numba extras in ``pyproject.toml`` declare the supported -``numba-cuda-mlir`` version range. When an application installs -``cuda-coop[numba-cuda-mlir-cu12]`` or ``cuda-coop[numba-cuda-mlir-cu13]``, -its package installer can select a compatible compiler or report a dependency -conflict before the application runs. The lower bound must name a released -compiler containing every API and behavior that this backend requires. Raise -that bound when adopting a newer compiler API, rather than adding a fallback -for older releases. - -The base ``cuda-coop`` distribution deliberately has no compiler dependencies. -Applications may install it with ``pip install cuda-coop`` and manage their -compiler separately. Installing an extra applies additional requirements; -merely having Numba installed does not activate those requirements during a -base-package installation. An independently installed compiler can therefore -be outside our supported range, even though installing ``cuda-coop`` succeeds. -The activation check gives that application a useful error containing the -required range, detected version, and installation instructions. - -.. rubric:: Why validation belongs to backend activation - -An application using another backend, such as CUTLASS, should not need to -upgrade an unrelated Numba installation. We check only when activating -``cuda.coop.numba_mlir``, through its qualified import or -``coop.register("numba-cuda-mlir")``. Activation checks the version before -importing the CUDA compiler integration and registering the planner. A failed -check therefore reports an environment problem before kernel compilation, -without adding checks to individual operations or the kernel launch path. - -The common ``cuda.coop`` import does not inspect every installed compiler. Its -automatic registration considers runtimes the application already imported. -If that optional activation fails, the common import emits a warning and -continues; explicit activation raises the error. Merely installing an old -Numba version has no effect on an application using another backend. - -.. rubric:: Why there is no feature-by-feature compatibility layer - -The supported version range is the compiler contract. Checking that an -attribute exists cannot establish the semantics of launch specialization, -inlining, or generated code. Accepting older releases based on selected -attributes would create a second, less precise support policy and defer -failures to whichever operation first needs an unavailable feature. We use -the required compiler APIs directly after checking the version. - -Likewise, catching AttributeError or TypeError around compiler operations and -turning every failure into an upgrade instruction can hide bugs in our own -integration or in a supported compiler. Import failures retain their original -cause, and compilation errors retain their specific meaning. In particular, -a missing launch-metadata API is different from a supported API reporting -that a particular compilation has no configured launch metadata. The latter -is a compilation/usage error, not evidence of an old compiler. - -.. rubric:: Maintenance - -Keep the range below aligned with both Numba extras in ``pyproject.toml``. Use -packaging's version semantics, including local and development versions, -instead of a string prefix check. Already-installed prereleases within the -range are accepted; prereleases below the minimum or at the excluded next -series remain outside it. A missing or unparsable version cannot establish -support and is reported explicitly. Compiler API loading, planner -registration, and any registration failure handling belong to the activation -code, rather than to a registry of capabilities in this module. -""" - -from __future__ import annotations - -import importlib.metadata -from types import ModuleType - -_REQUIRED_RUNTIME_VERSION = ">=0.5.0,<0.6" -_RUNTIME_INSTALL_HINT = ( - "python -m pip install --upgrade 'cuda-coop[numba-cuda-mlir-cu12]' " - "for CUDA 12 or 'cuda-coop[numba-cuda-mlir-cu13]' for CUDA 13" -) - - -class NumbaMlirBackendImportError(ImportError): - """Report a Numba backend activation failure to the application. - - Explicit registration and qualified backend imports propagate this error; - applications can catch it as ``ImportError``. Automatic registration - catches activation failures and emits ``CudaCoopAutoRegistrationWarning`` - instead, allowing the common import to complete. This is a diagnostic - exception, rather than an internal signal for compiler control flow. - - ``backend`` identifies the compiler, ``reason_code`` classifies the - failure, and ``details`` carries diagnostic context. When supplied, the - original import failure is preserved as ``__cause__``. - - Parameters - ---------- - reason_code : str - Category used to distinguish activation failures. - message : str - Human-readable failure description and recovery guidance. - cause : BaseException or None, optional - Original failure to retain as the explicit exception cause. - **details : object - Named diagnostic values, such as detected and required versions, - stored in the exception's ``details`` dictionary. - """ - - def __init__( - self, - reason_code: str, - message: str, - *, - cause: BaseException | None = None, - **details: object, - ) -> None: - super().__init__(message) - self.backend = "numba-cuda-mlir" - self.reason_code = reason_code - self.details = details - if cause is not None: - self.__cause__ = cause - - -def _detected_version(runtime: ModuleType | None) -> str | None: - """Prefer the imported runtime's version, then distribution metadata.""" - - version = getattr(runtime, "__version__", None) - if isinstance(version, str) and version: - return version - try: - return importlib.metadata.version("numba-cuda-mlir") - except importlib.metadata.PackageNotFoundError: - return None - - -def _runtime_requirement(runtime: ModuleType | None = None) -> str: - """Describe the supported range and detected compiler installation. - - Parameters - ---------- - runtime : module, optional - Imported compiler package whose ``__version__`` is preferred. If it - is ``None`` or lacks a usable version, use distribution metadata. - - Returns - ------- - str - Required version range, detected version or an unavailable-version - notice, and an installation command for the supported compiler. - """ - - version = _detected_version(runtime) - detected = "the installed numba-cuda-mlir version could not be determined" - if version is not None: - detected = f"detected numba-cuda-mlir=={version}" - return ( - f"cuda.coop.numba_mlir requires numba-cuda-mlir" - f"{_REQUIRED_RUNTIME_VERSION}; {detected}. Install with " - f"{_RUNTIME_INSTALL_HINT}." - ) - - -def _require_numba_mlir_version(runtime: ModuleType) -> None: - """Reject unsupported installations before accessing compiler APIs. - - Activation calls this after importing the top-level compiler package and - before loading its CUDA integration. Import ``packaging`` here so a bare - ``cuda-coop`` installation can use the common API without that dependency. - The declared range accepts installed prereleases within its bounds; it - does not admit versions below the minimum or in the excluded next series. - - Parameters - ---------- - runtime : module - Imported ``numba_cuda_mlir`` package. Prefer its nonempty - ``__version__`` string, falling back to distribution metadata when - that attribute is unavailable. - - Raises - ------ - NumbaMlirBackendImportError - ``packaging`` is missing, the compiler version cannot be determined - or parsed, or it lies outside the supported range. The diagnostic - includes installation instructions and relevant version details. - """ - - # Packaging is a backend dependency. Import it here so a missing runtime - # still gets its own diagnostic, including in a base-only installation. - try: - from packaging.specifiers import SpecifierSet - from packaging.version import InvalidVersion, Version - except ModuleNotFoundError as exc: - if exc.name != "packaging": - raise - raise NumbaMlirBackendImportError( - "backend-dependency-missing", - "cuda.coop.numba_mlir requires packaging to validate its compiler " - f"version. Install with {_RUNTIME_INSTALL_HINT}.", - cause=exc, - missing=exc.name, - ) from exc - - version = _detected_version(runtime) - try: - supported = version is not None and SpecifierSet( - _REQUIRED_RUNTIME_VERSION - ).contains(Version(version), prereleases=True) - except InvalidVersion: - supported = False - if not supported: - raise NumbaMlirBackendImportError( - "unsupported-runtime-version", - _runtime_requirement(runtime), - detected_version=version, - required_version=_REQUIRED_RUNTIME_VERSION, - ) - - -__all__ = [ - "NumbaMlirBackendImportError", - "_require_numba_mlir_version", - "_runtime_requirement", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_nvrtc.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_nvrtc.py deleted file mode 100644 index 9ab3b13b8f6e..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_nvrtc.py +++ /dev/null @@ -1,739 +0,0 @@ -# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Compile generated CUDA wrappers with an identified toolkit and header set. - -The backend generates C++ wrappers around CUB operations, but those wrappers -must become device code before a Numba kernel can call them. This module -resolves the compiler libraries and headers, invokes NVRTC, and returns an -LTO IR image or PTX text for later linking. It does not launch GPU work. - -The same source can produce different code with different headers, compiler -versions, targets, or options. ``CompileContext`` records the header and -library dependencies; the target and options also participate in cache keys -and private provider symbol identities. Context resolution preloads the chosen -toolkit's libraries before importing the CUDA bindings. ``compile`` dumps -source before cache lookup, so developers can inspect generated code even -when compilation is reused. - -``compile_with_layouts`` reads C++ scratch sizes and alignments from the same -NVRTC program that emits the provider image. An appended probe encodes each -value in a compiler-generated name. The image and its ordered layouts are -cached as one result, so a cache hit supplies both without another compile -or a separate PTX link. -""" - -from __future__ import annotations - -import functools -import hashlib -import os -import re -from collections.abc import Iterator -from contextlib import contextmanager -from contextvars import ContextVar -from dataclasses import dataclass -from pathlib import Path -from typing import Any, TypeAlias, cast - -from cuda.coop._core._source_dump import dump_source -from cuda.coop._headers import CoopIncludePaths, resolve_include_paths -from cuda.coop._headers._identity import include_dirs_identity -from cuda.coop._headers._toolkit import ( - preload_toolkit_compiler_libraries, - validate_nvrtc_version, -) - -from ._artifacts import check_in, version -from ._caching import disk_cache -from ._layout import decode_layout_name, prepare_layout_queries - -_LayoutResult: TypeAlias = tuple[bytes, tuple[tuple[int, int], ...]] - -_REQUIRED_HEADERS = ( - "cub/block/block_load.cuh", - "cub/block/block_store.cuh", - "cuda/barrier", - "cuda/devices", - "cuda/functional", - "cuda/hierarchy", - "cuda/std/cstdint", - "cuda/std/functional", - "cuda/std/type_traits", -) - - -@dataclass(frozen=True) -class CompileContext: - """Record header and toolkit inputs shared by provider compilations. - - Exact library paths and versions identify the selected compiler and - linker. ``include_dirs`` preserves header search order, while - ``header_identity`` identifies their contents. Source, target, and - compiler options enter each compilation's cache key separately. - """ - - toolkit_root: str - toolkit_version: tuple[int, int] - nvrtc_path: str - nvrtc_builtins_path: str - nvjitlink_path: str - nvrtc_version: version - nvjitlink_version: tuple[int, int] - include_dirs: tuple[str, ...] - header_identity: str - - @property - def symbol_suffix(self) -> str: - """Return a short digest of every toolkit and header field. - - Raw C-ABI helpers add it to their symbols to distinguish code built - with different headers or toolkits. - """ - - digest = hashlib.sha256() - values: tuple[object, ...] = ( - self.toolkit_root, - self.toolkit_version, - self.nvrtc_path, - self.nvrtc_builtins_path, - self.nvjitlink_path, - self.nvrtc_version, - self.nvjitlink_version, - self.include_dirs, - self.header_identity, - ) - for value in values: - digest.update(repr(value).encode("utf-8", errors="surrogateescape")) - digest.update(b"\0") - return digest.hexdigest()[:16] - - -_CHIP_PATTERN = re.compile(r"sm_([1-9][0-9]*[0-9])([af]?)") - - -@dataclass(frozen=True) -class CompileTarget: - """Describe the device target and line-information option of a kernel. - - ``cc`` is the compute capability encoded as major times ten plus minor. - ``arch_suffix`` is ``"a"`` or ``"f"`` for architecture- or - family-specific targets such as ``sm_90a`` or ``sm_100f``, otherwise - empty. ``lineinfo`` requests source-line information for generated code. - Generated device code must use this target so that it matches the kernel - it is linked into, not whichever device happens to be current. - """ - - cc: int - arch_suffix: str = "" - lineinfo: bool = False - - @classmethod - def from_targetoptions( - cls, targetoptions: dict[str, Any] - ) -> CompileTarget | None: - """Read the target that Numba-CUDA-MLIR resolved for a kernel. - - The compiler stores the resolved ``chip`` (for example ``sm_90`` or - ``sm_100f``) in the target options before any planner runs; explicit - ``cc=`` arguments are already mapped to it. Return ``None`` when no - chip has been resolved. - - ``lineinfo=True`` requests line information. ``debug=True`` does not: - the kernel is then linked with ``-g``, and nvJitLink rejects such an - LTO link when more than one input carries debug or line information. - Generated code is usually several inputs, so it stays without both. - """ - - chip = targetoptions.get("chip") - if chip is None: - return None - match = _CHIP_PATTERN.fullmatch(chip) if isinstance(chip, str) else None - if match is None: - raise RuntimeError( - "cuda.coop.numba_mlir requires a resolved CUDA target of the " - f"form 'sm_XY[a|f]'; got chip={chip!r}" - ) - return cls( - cc=int(match.group(1)), - arch_suffix=match.group(2), - lineinfo=bool(targetoptions.get("lineinfo", False)) - and not targetoptions.get("debug", False), - ) - - @property - def arch(self) -> str: - """Return the architecture spelling used in generated symbols.""" - - return f"{self.cc}{self.arch_suffix}" - - def nvrtc_kwargs(self) -> dict[str, Any]: - """Return keywords for ``compile`` and ``compiler_identity``.""" - - return { - "cc": self.cc, - "arch_suffix": self.arch_suffix, - "lineinfo": self.lineinfo, - } - - -@dataclass(frozen=True) -class CompilerIdentity: - """Identify the target and options shared by compatible providers. - - The options include ordered header paths. Callers retain the corresponding - ``CompileContext`` separately for its library and header-content identity. - """ - - cc: int - rdc: bool - code: str - compiler_options: tuple[bytes, ...] - arch_suffix: str = "" - lineinfo: bool = False - - @property - def target(self) -> CompileTarget: - return CompileTarget( - cc=self.cc, arch_suffix=self.arch_suffix, lineinfo=self.lineinfo - ) - - -def _load_nvrtc(): - """Import CUDA bindings after the selected toolkit is preloaded.""" - - import cuda.bindings.nvrtc as _nvrtc_bindings - - return _nvrtc_bindings - - -def _nvrtc_version(nvrtc: Any) -> version: - """Read the bound NVRTC library's version or report its query failure.""" - - err, major, minor = nvrtc.nvrtcVersion() - if err != nvrtc.nvrtcResult.NVRTC_SUCCESS: - raise RuntimeError(f"nvrtcVersion error: {err}") - return version(int(major), int(minor)) - - -def _check_nvrtc(err, prog, *, nvrtc=None): - """Raise an NVRTC failure with the program's compilation log. - - ``compile_impl`` checks compilation, name-query, and image-retrieval - results here so failures include NVRTC's diagnostic text. Keep the - original result code even if fetching the log also fails; a secondary - logging error must not hide the operation that failed. - - Parameters - ---------- - err : nvrtcResult - Result code returned by the operation being checked. Success returns - immediately without fetching a log. - prog : nvrtcProgram - Live program whose log can explain the failure. The caller owns its - lifetime and destroys it after diagnostics have been collected. - nvrtc : module or None - Bindings for the already selected NVRTC library. If omitted, import - the bindings; toolkit selection and preloading must have happened first. - - Raises - ------ - RuntimeError - NVRTC reported failure. The message includes its result code and either - the program log or the reason that the log could not be read. - """ - - nvrtc = _load_nvrtc() if nvrtc is None else nvrtc - if err == nvrtc.nvrtcResult.NVRTC_SUCCESS: - return - original_err = err - log_err, logsize = nvrtc.nvrtcGetProgramLogSize(prog) - if log_err != nvrtc.nvrtcResult.NVRTC_SUCCESS: - raise RuntimeError( - f"NVRTC error: {original_err}; failed to get log size: {log_err}" - ) - log = bytearray(logsize) - log_result = nvrtc.nvrtcGetProgramLog(prog, log) - log_err = log_result[0] if isinstance(log_result, tuple) else log_result - if log_err != nvrtc.nvrtcResult.NVRTC_SUCCESS: - raise RuntimeError( - f"NVRTC error: {original_err}; failed to get log: {log_err}" - ) - rendered = bytes(log).rstrip(b"\0").decode("ascii", errors="replace") - raise RuntimeError(f"NVRTC error: {original_err}: {rendered}") - - -def _dump_source(cpp, cc, code, compiler_options): - """Include the target and output options in the saved source identity.""" - - return dump_source( - cpp, - backend="numba_mlir", - identity=(cc, code, compiler_options), - ) - - -def _include_options(include_dirs: tuple[str, ...]) -> list[bytes]: - return [os.fsencode(f"--include-path={path}") for path in include_dirs] - - -def _compiler_options( - *, - cc: int, - rdc: bool, - code: str, - include_dirs: tuple[str, ...], - arch_suffix: str = "", - lineinfo: bool = False, -) -> tuple[bytes, ...]: - """Return the exact ordered NVRTC option set used for cache identity.""" - - check_in("rdc", rdc, [True, False]) - check_in("code", code, ["lto", "ptx"]) - check_in("arch_suffix", arch_suffix, ["", "a", "f"]) - options = [ - b"--std=c++17", - *_include_options(include_dirs), - f"--gpu-architecture=compute_{cc}{arch_suffix}".encode("ascii"), - ] - if rdc: - options.append(b"--relocatable-device-code=true") - if code == "lto": - options.append(b"-dlto") - if lineinfo: - options.append(b"-lineinfo") - options.append(b"-DCCCL_DISABLE_BF16_SUPPORT") - return tuple(options) - - -def compiler_identity( - *, - context: CompileContext, - cc: int, - rdc: bool, - code: str, - arch_suffix: str = "", - lineinfo: bool = False, -) -> CompilerIdentity: - """Identify target and options that qualify provider symbols. - - The options include the context's ordered header paths. Providers also use - this identity to group compatible specializations into one compilation. The - complete identity also needs the context's library and header-content - identities, which callers retain separately. - """ - - return CompilerIdentity( - cc=int(cc), - rdc=bool(rdc), - code=str(code), - compiler_options=_compiler_options( - cc=int(cc), - rdc=bool(rdc), - code=str(code), - include_dirs=context.include_dirs, - arch_suffix=str(arch_suffix), - lineinfo=bool(lineinfo), - ), - arch_suffix=str(arch_suffix), - lineinfo=bool(lineinfo), - ) - - -@functools.lru_cache(maxsize=8) -@disk_cache -def compile_impl( - cpp: str, - cc: int, - rdc: bool, - code: str, - toolkit_root: str, - toolkit_version: tuple[int, int], - nvrtc_path: str, - nvrtc_builtins_path: str, - nvjitlink_path: str, - nvrtc_version: version, - nvjitlink_version: tuple[int, int], - include_dirs: tuple[str, ...], - header_identity: str, - compiler_options: tuple[bytes, ...], - layout_queries: tuple[str, ...] = (), - layout_symbol: str = "", - arch_suffix: str = "", - lineinfo: bool = False, -) -> bytes | str | _LayoutResult: - """Compile one source unit using a complete, explicit cache identity. - - The memory and disk cache decorators key all arguments, including toolkit - paths and header identity that are not otherwise read by the function - body. Those fields prevent artifacts from different compiler installations - or header sets from sharing an entry. Callers must resolve and preload the - matching compiler context before entering this function. - - On a cache miss, verify that the supplied option tuple matches the request - and that the loaded NVRTC version still matches the context. Compile the - source, retrieve the requested image, and destroy the NVRTC program on - both success and failure. A cleanup error does not replace an earlier - compilation error. Cache hits bypass these body-level checks. - - Register distinct layout expressions before compilation and decode their - names while the program is still alive. Restore duplicate entries in the - caller's order when returning layouts with an LTO image. The image and - layouts form one cache value, tying metadata to the same compilation. - - Parameters - ---------- - cpp : str - Complete CUDA C++ translation unit. - cc : int - Compute capability encoded as major times ten plus minor. - rdc : bool - Whether to enable relocatable device code. - code : {"lto", "ptx"} - Requested output format. - toolkit_root : str - Selected toolkit root, retained in cache identity. - toolkit_version : tuple of int - Selected toolkit version, retained in cache identity. - nvrtc_path, nvrtc_builtins_path, nvjitlink_path : str - Exact compiler-library paths from the preloaded context. - nvrtc_version : version - Expected loaded NVRTC version. - nvjitlink_version : tuple of int - Selected linker version, retained in cache identity. - include_dirs : tuple of str - Ordered header search roots. - header_identity : str - Header-content identity supplied by context resolution. - compiler_options : tuple of bytes - Exact ordered options from ``_compiler_options`` for this request. - layout_queries : tuple of str, optional - Ordered NVRTC name expressions whose lowered names encode storage - layouts. Duplicate expressions are evaluated once and retain their - positions in the returned layouts. - layout_symbol : str, optional - Expected symbol used to validate and decode lowered layout names. - arch_suffix : {"", "a", "f"}, optional - Architecture- or family-specific target suffix appended to ``cc``. - lineinfo : bool, optional - Emit source-line information (``-lineinfo``) for the generated code. - - Returns - ------- - bytes, str, or tuple - LTO image bytes for ``"lto"`` or ASCII-decoded source for ``"ptx"``. - With LTO layout queries, return ``(image, layouts)`` in query order. - - Raises - ------ - RuntimeError - Options or loaded compiler version disagree with the request, or NVRTC - compilation, image retrieval, or program cleanup fails. - ValueError - The output format or relocatable-code option is invalid, or a lowered - layout name does not match the expected encoding. - """ - - # These arguments are only used by the cache decorators to distinguish - # compiler and header contexts. Delete their local bindings to mark them - # as intentionally unused by the compilation body. - del ( - toolkit_root, - toolkit_version, - nvrtc_path, - nvrtc_builtins_path, - nvjitlink_path, - nvjitlink_version, - header_identity, - ) - expected_options = _compiler_options( - cc=cc, - rdc=rdc, - code=code, - include_dirs=include_dirs, - arch_suffix=arch_suffix, - lineinfo=lineinfo, - ) - if compiler_options != expected_options: - raise RuntimeError( - "NVRTC compiler-option identity does not match the " - "requested compile." - ) - nvrtc = _load_nvrtc() - loaded_version = _nvrtc_version(nvrtc) - if loaded_version != nvrtc_version: - raise RuntimeError( - "loaded NVRTC version changed after compile-context resolution: " - f"expected {nvrtc_version}, got {loaded_version}" - ) - - err, prog = nvrtc.nvrtcCreateProgram(cpp.encode(), b"code.cu", 0, [], []) - if err != nvrtc.nvrtcResult.NVRTC_SUCCESS: - raise RuntimeError(f"nvrtcCreateProgram error: {err}") - had_error = False - try: - for expression in dict.fromkeys(layout_queries): - (err,) = nvrtc.nvrtcAddNameExpression( - prog, expression.encode("utf-8") - ) - _check_nvrtc(err, prog, nvrtc=nvrtc) - (err,) = nvrtc.nvrtcCompileProgram( - prog, len(compiler_options), list(compiler_options) - ) - _check_nvrtc(err, prog, nvrtc=nvrtc) - layouts = {} - for expression in dict.fromkeys(layout_queries): - err, lowered_name = nvrtc.nvrtcGetLoweredName( - prog, expression.encode("utf-8") - ) - _check_nvrtc(err, prog, nvrtc=nvrtc) - layouts[expression] = decode_layout_name( - lowered_name, symbol=layout_symbol, expression=expression - ) - if code == "lto": - err, size = nvrtc.nvrtcGetLTOIRSize(prog) - _check_nvrtc(err, prog, nvrtc=nvrtc) - image = bytearray(size) - (err,) = nvrtc.nvrtcGetLTOIR(prog, image) - _check_nvrtc(err, prog, nvrtc=nvrtc) - result = bytes(image) - if layout_queries: - return result, tuple( - layouts[expression] for expression in layout_queries - ) - return result - err, size = nvrtc.nvrtcGetPTXSize(prog) - _check_nvrtc(err, prog, nvrtc=nvrtc) - image = bytearray(size) - (err,) = nvrtc.nvrtcGetPTX(prog, image) - _check_nvrtc(err, prog, nvrtc=nvrtc) - return bytes(image).decode("ascii") - except Exception: - had_error = True - raise - finally: - (destroy_err,) = nvrtc.nvrtcDestroyProgram(prog) - if destroy_err != nvrtc.nvrtcResult.NVRTC_SUCCESS and not had_error: - raise RuntimeError(f"nvrtcDestroyProgram error: {destroy_err}") - - -_COMPILE_CONTEXTS: ContextVar[dict[tuple[str, ...], CompileContext] | None] = ( - ContextVar("cuda_coop_numba_mlir_compile_contexts", default=None) -) - - -@contextmanager -def compile_context_scope() -> Iterator[None]: - """Share lazily resolved contexts during one planner attempt. - - Headers and toolkit files must remain stable until the attempt finishes. - Nested attempts get independent maps and restore their caller on exit. - """ - - token = _COMPILE_CONTEXTS.set({}) - try: - yield - finally: - _COMPILE_CONTEXTS.reset(token) - - -def resolve_shared_compile_context() -> CompileContext: - """Reuse an identity only inside the active attempt's selected roots.""" - - contexts = _COMPILE_CONTEXTS.get() - if contexts is None: - return resolve_compile_context() - include_paths = _resolve_compile_include_paths() - include_dirs = tuple(str(path) for path in include_paths.as_tuple()) - if include_dirs not in contexts: - contexts[include_dirs] = _build_compile_context(include_paths) - return contexts[include_dirs] - - -def resolve_compile_context() -> CompileContext: - """Resolve one header/toolkit identity and preload its compiler libraries. - - ``CUDA_COOP_CCCL_ROOT`` optionally selects a CCCL source checkout or - packaged header bundle instead of automatic source-tree/wheel discovery. - Read it on each call; unset or empty uses discovery. The header resolver - expands ``~`` and resolves relative paths from the current working - directory. A configured root must supply all required CCCL headers; an - invalid or incomplete selection raises rather than falling back. This - variable selects CCCL headers, not the CUDA Toolkit installation. - - Resolve CUDA headers separately, then preload NVRTC, its builtins, and - nvJitLink from that toolkit before importing the CUDA NVRTC bindings. - Check that the loaded NVRTC version matches the selected toolkit. This - ordering keeps wrapper compilation and later linking on the same toolkit. - - Hash the resolved header roots and their contents into the returned - context. The context is used both for artifact cache keys and provider - symbol qualification. Resolution is lazy at its callers; this function - itself is not memoized and may load process-wide compiler libraries. - Algorithms retain their resolved context, so changing - ``CUDA_COOP_CCCL_ROOT`` affects subsequent resolutions, not contexts - already held by providers or supplied explicitly to ``compile``. - - Returns - ------- - CompileContext - Frozen snapshot of exact library paths/versions, ordered include - roots, and header identity for a provider compilation. - """ - - return _build_compile_context(_resolve_compile_include_paths()) - - -def _resolve_compile_include_paths() -> CoopIncludePaths: - """Select and validate ordered header roots for the current request.""" - - return resolve_include_paths( - start=Path(__file__), - configured_roots=(os.environ.get("CUDA_COOP_CCCL_ROOT"),), - required_headers=_REQUIRED_HEADERS, - ) - - -def _build_compile_context(include_paths: CoopIncludePaths) -> CompileContext: - """Preload the selected toolkit and fingerprint the current headers.""" - - include_dirs = tuple(str(path) for path in include_paths.as_tuple()) - libraries = preload_toolkit_compiler_libraries(include_paths.cuda) - nvrtc = _load_nvrtc() - loaded_nvrtc_version = _nvrtc_version(nvrtc) - validate_nvrtc_version(libraries, tuple(loaded_nvrtc_version)) - return CompileContext( - toolkit_root=libraries.toolkit_root, - toolkit_version=libraries.toolkit_version, - nvrtc_path=libraries.nvrtc_path, - nvrtc_builtins_path=libraries.nvrtc_builtins_path, - nvjitlink_path=libraries.nvjitlink_path, - nvrtc_version=loaded_nvrtc_version, - nvjitlink_version=libraries.nvjitlink_version, - include_dirs=include_dirs, - header_identity=include_dirs_identity(include_dirs).digest, - ) - - -def compile( - *, context: CompileContext | None = None, **kwargs: Any -) -> tuple[version, bytes | str | _LayoutResult]: - """Compile provider source with a resolved header and toolkit context. - - Build the ordered options from the selected context, optionally dump the - source through the shared source-dump hook, and pass every context field - to ``compile_impl``. Dumping occurs before cache lookup so source - inspection also works when compilation is reused from memory or disk. - - Parameters - ---------- - context : CompileContext, optional - Previously resolved and preloaded compiler context. ``None`` resolves - one now. Supplying a context reuses its identity; it does not reload - the library paths stored in it. - **kwargs : dict - Required ``cpp`` source string, integer ``cc`` target, boolean - ``rdc``, and ``code`` equal to ``"lto"`` or ``"ptx"``. Optional - ``arch_suffix`` and ``lineinfo`` complete the target; see - ``CompileTarget.nvrtc_kwargs``. These are forwarded to - ``compile_impl`` with the context and generated options. - Internal callers can also supply ``layout_queries`` and - ``layout_symbol``; use ``compile_with_layouts`` to prepare these from - C++ type names. - - Returns - ------- - tuple - ``(nvrtc_version, result)`` with LTO bytes or PTX text according to - ``code``. With LTO layout queries, ``result`` is ``(image, layouts)`` - in query order. The version is taken from the selected context. - - Raises - ------ - RuntimeError - Compiler identity checks or NVRTC compilation fail. - ValueError - An output format or relocatable-code option is invalid. - """ - - context = resolve_compile_context() if context is None else context - compiler_options = _compiler_options( - cc=kwargs["cc"], - rdc=kwargs["rdc"], - code=kwargs["code"], - include_dirs=context.include_dirs, - arch_suffix=kwargs.get("arch_suffix", ""), - lineinfo=kwargs.get("lineinfo", False), - ) - _dump_source(kwargs["cpp"], kwargs["cc"], kwargs["code"], compiler_options) - return context.nvrtc_version, compile_impl( - **kwargs, - toolkit_root=context.toolkit_root, - toolkit_version=context.toolkit_version, - nvrtc_path=context.nvrtc_path, - nvrtc_builtins_path=context.nvrtc_builtins_path, - nvjitlink_path=context.nvjitlink_path, - nvrtc_version=context.nvrtc_version, - nvjitlink_version=context.nvjitlink_version, - include_dirs=context.include_dirs, - header_identity=context.header_identity, - compiler_options=compiler_options, - ) - - -def compile_with_layouts( - *, layout_types: tuple[str, ...], code: str = "lto", **kwargs: Any -) -> tuple[version, _LayoutResult]: - """Cache provider LTO IR and its ordered storage layouts together. - - The same NVRTC program evaluates ``sizeof`` and ``alignof`` and emits - the provider image. Query expressions are part of the cache identity, and - their ordered layouts are stored with the image. A later cache hit can - therefore recover both without compiling or linking a metadata program. - - Parameters - ---------- - layout_types : tuple of str - C++ type names or type expressions visible at the end of the supplied - source. The result has one layout per entry in the same order, with - duplicates preserved. An empty tuple uses ordinary compilation - without adding a probe to the source. - code : {"lto"}, optional - Output format. Only LTO supports the combined image/layout result. - **kwargs : dict - Arguments for ``compile``: ``cpp`` source, integer ``cc`` target, - boolean ``rdc``, and optional resolved ``context``. The source must - make every requested type visible to an appended layout probe. - - Returns - ------- - tuple - ``(nvrtc_version, (image, layouts))``. The image is LTO bytes; each - layout is a ``(size, alignment)`` pair in bytes. For an empty type - tuple, ``layouts`` is also empty. - - Raises - ------ - ValueError - The output format is not LTO, a type expression is empty, or NVRTC - returns an unexpected layout name or invalid size/alignment pair. - RuntimeError - Compiler identity checks or NVRTC compilation fail. - """ - - if code != "lto": - raise ValueError("storage layout queries require LTO compilation") - source, symbol, queries = prepare_layout_queries( - kwargs["cpp"], layout_types - ) - if not queries: - compiler_version, image = compile(code=code, **kwargs) - return compiler_version, (cast(bytes, image), ()) - kwargs["cpp"] = source - return cast( - tuple[version, _LayoutResult], - compile( - code=code, layout_queries=queries, layout_symbol=symbol, **kwargs - ), - ) diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_operations.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_operations.py deleted file mode 100644 index e722114ff5f4..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_operations.py +++ /dev/null @@ -1,784 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Connect public cooperative calls to their compiler implementations. - -A group marker is a public callable such as ``load`` that the group planner -recognizes in kernel IR. A lowering factory is a host-side callable, such as -``_lowering._load_store._load_with_storage``, that accepts specialization -inputs (dtype, block dimensions, algorithm, and scalar bindings) and builds -the compiled provider for that operation. Its usual result is an -``Invocable``: a callable carrying link artifacts and temporary-storage -requirements. During batch collection it returns an uncompiled ``Algorithm`` -specialization instead. - -These registries identify markers and factories by callable identity, so -aliases work without relying on function names. Operation names connect those -identities to family-specific group lowering, argument validation, payload -inference, and runtime-argument preparation. Families load lazily when their -hooks are needed; registration itself does not compile providers or run -cooperative operations. -""" - -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass -from enum import Enum -from importlib import import_module -from threading import RLock -from typing import TYPE_CHECKING, Any, Protocol, TypeVar - -if TYPE_CHECKING: - from cuda.coop._core import ( - GroupTopologyRequirements, - SynchronizationRequirements, - TempStorageRequirements, - ) - - from ._group_rewriting import GroupRewriteContext - from ._rewrite_payload import PayloadInference - -from cuda.coop._core import StorageOwnership, SynchronizationScope - -_CallableT = TypeVar("_CallableT", bound=Callable[..., Any]) - -# The whole-function group planner attaches its exact lowering record to the -# provider marker with this reserved keyword. The provider rewrite -# consumes it before invoking the registered provider factory. -_GROUP_LOWERING_PLAN_KWARG = "__cuda_coop_group_lowering_plan__" - - -class _InferPayloadHook(Protocol): - """Complete factory inputs from planned or inferred payload facts. - - Update the supplied inference state with compatible dtype and shape facts. - That state carries the call's optional lowering plan. The context supplies - the active rewrite's operand and constructor analysis. - """ - - def __call__( - self, - context: GroupRewriteContext, - inference: PayloadInference, - ) -> None: ... - - -class _AnalyzeMatchHook(Protocol): - """Record family facts needed when replacing a provider call. - - Inspect the split arguments, consume private factory markers, and return - metadata for the later runtime-argument hook. - """ - - def __call__( - self, - context: GroupRewriteContext, - *, - op_name: str, - runtime_args: tuple[Any, ...], - factory_kwargs: dict[str, object], - ) -> Any: ... - - -class _PrepareRuntimeArgsHook(Protocol): - """Prepare operands while the rewrite emits a provider call. - - Use the match's family metadata, append preparation IR to ``block``, and - return the provider operands. ``scope`` and ``loc`` identify the generated - variables and statements. - """ - - def __call__( - self, - context: GroupRewriteContext, - block: Any, - *, - match: Any, - runtime_args: list[Any], - scope: Any, - loc: Any, - ) -> list[Any]: ... - - -class _ValidateRuntimeControlsHook(Protocol): - """Check operation-specific scalar controls before provider creation. - - Use available operand types and factory bindings. The hook may leave - unknown types for the later compiler typing pass. - """ - - def __call__( - self, - context: GroupRewriteContext, - *, - op_name: str, - runtime_args: list[Any], - factory_kwargs: dict[str, object], - ) -> None: ... - - -class StorageABI(str, Enum): - """Describe the provider's temporary-storage calling convention. - - ``NONE`` needs no storage operand. ``LEADING_POINTER`` reserves the first - provider argument for scratch memory, supplied separately from the - operation's ordinary operands. - """ - - NONE = "none" - LEADING_POINTER = "leading_pointer" - - -@dataclass(frozen=True) -class FactoryOperation: - """Declare the call contract of a registered provider factory. - - The rewrite uses this record to select argument rules and storage - handling, then checks the resulting invocable against it. Registration - declares the contract; it does not inspect generated device code. - - Attributes - ---------- - operation : str - Operation identifier used to find its argument and rewrite rules. - namespace : str - Provider namespace, such as ``"block"`` or ``"warp"``. - storage_abi : StorageABI - Whether the generated call receives a leading scratch pointer. - execution_scope : SynchronizationScope - Threads that execute the primitive together. - synchronization_scope : SynchronizationScope - Declared synchronization scope: ``NONE`` or the execution scope. - Storage planning decides which reuse barriers to emit. - """ - - operation: str - namespace: str - storage_abi: StorageABI - execution_scope: SynchronizationScope - synchronization_scope: SynchronizationScope - - def __post_init__(self) -> None: - for name in ("operation", "namespace"): - value = getattr(self, name) - if not isinstance(value, str) or not value: - raise ValueError(f"{name} must be a non-empty string") - object.__setattr__(self, "storage_abi", StorageABI(self.storage_abi)) - object.__setattr__( - self, - "execution_scope", - SynchronizationScope(self.execution_scope), - ) - object.__setattr__( - self, - "synchronization_scope", - SynchronizationScope(self.synchronization_scope), - ) - if self.synchronization_scope not in { - SynchronizationScope.NONE, - self.execution_scope, - }: - raise ValueError( - "synchronization_scope must be NONE or match execution_scope" - ) - - -def expected_storage_reuse_barrier( - topology: GroupTopologyRequirements, - storage: TempStorageRequirements, -) -> SynchronizationScope: - """Require a reuse barrier only for automatically synchronized storage.""" - - return ( - topology.execution_scope - if storage.ownership is not StorageOwnership.NONE and storage.auto_sync - else SynchronizationScope.NONE - ) - - -def provider_synchronization_matches( - metadata: FactoryOperation, - topology: GroupTopologyRequirements, - synchronization: SynchronizationRequirements, - storage: TempStorageRequirements, -) -> bool: - """Check the provider's barrier declaration against planned scratch reuse. - - Caller-owned storage may retain its allocating wrapper's declaration when - automatic synchronization is disabled: pointer rewrites bypass that wrapper - and let the caller synchronize. Implementation-owned storage has no such - exception. Callers validate the planned reuse barrier separately. - """ - - planned_barrier = synchronization.storage_reuse_barrier - return metadata.synchronization_scope is planned_barrier or ( - planned_barrier is SynchronizationScope.NONE - and storage.ownership is StorageOwnership.CALLER - and not storage.auto_sync - and metadata.synchronization_scope is topology.execution_scope - ) - - -@dataclass(frozen=True) -class GroupPrimitiveRegistration: - """Connect a public operation to whole-function planning. - - ``lower`` returns replacement IR after the planner has bound public - arguments and resolved the group. ``validate_common_arguments`` checks - calls made through the common ``cuda.coop`` API before that replacement; - backend-qualified calls skip it. ``None`` omits the check. Both hooks - receive the active planning context. - """ - - lower: Callable[..., list[Any]] - validate_common_arguments: Callable[..., None] | None = None - - def __post_init__(self) -> None: - if not callable(self.lower): - raise TypeError("lower must be callable") - if self.validate_common_arguments is not None and not callable( - self.validate_common_arguments - ): - raise TypeError( - "validate_common_arguments must be callable or None" - ) - - -@dataclass(frozen=True) -class RewriteOperationSpecification: - """Define how to interpret and replace one provider call. - - Public group planning first selects a provider factory. This record then - separates its compile-time inputs from device operands and selects the - family hooks used before compiler type inference. A registration can serve - several provider namespaces. - - Attributes - ---------- - factory_namespaces : frozenset of str - Registered provider namespaces accepted for this operation. - dtype_factory_kwargs : frozenset of str - Factory keywords that use dtype resolution and normalization. - runtime_arg_counts : frozenset of int - Accepted positional counts before static controls are removed. The - smallest count gives the number of ordinary operands. - runtime_factory_kwargs : tuple of str - Optional controls after the ordinary operands, in provider order. - Controls supplied by keyword use this same order. - runtime_factory_kw_prerequisites : tuple of tuple of str - ``(control, required_control)`` pairs checked when a keyword - control remains a runtime operand, such as padding needing a count. - Factories remain responsible for validating static bindings. - allowed_factory_kwargs : frozenset of str - Keywords accepted for specialization, including scalar controls. - required_factory_kwargs : frozenset of str - Keywords that must be known after payload and launch inference. - accepts_temp_storage : bool - Whether argument splitting accepts a separate storage operand. - scalar_binding_kwargs : frozenset of str - Optional controls classified as static or runtime bindings. Other - runtime controls are recorded as ``True`` when present. - runtime_offset_kwarg : str or None - Optional offset keyword handled after other runtime controls. A - static offset becomes a binding; a runtime offset is last in the - operand list. ``None`` disables this special case. - infer_payload : callable - Merge operand dtype and shape facts into factory inputs. - analyze_match : callable or None - Inspect a validated match and return family metadata for emission. - prepare_runtime_args : callable or None - Emit operand preparation, such as boxing a scalar Store value. - validate_runtime_controls : callable or None - Check control values and known dtypes before provider creation. - - Notes - ----- - Construction copies collection fields and checks that keyword sets, - positional counts, prerequisite names, and hook types agree. It does not - validate an actual call or compile its provider. - """ - - factory_namespaces: frozenset[str] - dtype_factory_kwargs: frozenset[str] - runtime_arg_counts: frozenset[int] - runtime_factory_kwargs: tuple[str, ...] - runtime_factory_kw_prerequisites: tuple[tuple[str, str], ...] - allowed_factory_kwargs: frozenset[str] - required_factory_kwargs: frozenset[str] - accepts_temp_storage: bool - scalar_binding_kwargs: frozenset[str] - runtime_offset_kwarg: str | None - infer_payload: _InferPayloadHook - analyze_match: _AnalyzeMatchHook | None = None - prepare_runtime_args: _PrepareRuntimeArgsHook | None = None - validate_runtime_controls: _ValidateRuntimeControlsHook | None = None - - def __post_init__(self) -> None: - for name in ( - "factory_namespaces", - "dtype_factory_kwargs", - "runtime_arg_counts", - "allowed_factory_kwargs", - "required_factory_kwargs", - "scalar_binding_kwargs", - ): - object.__setattr__(self, name, frozenset(getattr(self, name))) - object.__setattr__( - self, - "runtime_factory_kwargs", - tuple(self.runtime_factory_kwargs), - ) - raw_prerequisites = tuple(self.runtime_factory_kw_prerequisites) - prerequisites: list[tuple[str, str]] = [] - for prerequisite in raw_prerequisites: - if ( - not isinstance(prerequisite, (tuple, list)) - or len(prerequisite) != 2 - ): - raise TypeError( - "runtime_factory_kw_prerequisites must contain name pairs" - ) - name, required_name = prerequisite - if not isinstance(name, str) or not name: - raise ValueError( - "runtime_factory_kw_prerequisite names " - "must be non-empty strings" - ) - if not isinstance(required_name, str) or not required_name: - raise ValueError( - "runtime_factory_kw_prerequisite names " - "must be non-empty strings" - ) - prerequisites.append((name, required_name)) - object.__setattr__( - self, - "runtime_factory_kw_prerequisites", - tuple(prerequisites), - ) - if not self.factory_namespaces or any( - not isinstance(namespace, str) or not namespace - for namespace in self.factory_namespaces - ): - raise ValueError( - "factory_namespaces must contain non-empty strings" - ) - for field_name in ( - "dtype_factory_kwargs", - "allowed_factory_kwargs", - "required_factory_kwargs", - "scalar_binding_kwargs", - ): - if any( - not isinstance(name, str) or not name - for name in getattr(self, field_name) - ): - raise ValueError(f"{field_name} must contain non-empty strings") - if not self.runtime_arg_counts: - raise ValueError("runtime_arg_counts must not be empty") - if any( - not isinstance(count, int) or isinstance(count, bool) or count < 0 - for count in self.runtime_arg_counts - ): - raise ValueError( - "runtime_arg_counts must contain non-negative integers" - ) - if any( - not isinstance(name, str) or not name - for name in self.runtime_factory_kwargs - ): - raise ValueError( - "runtime_factory_kwargs must contain non-empty strings" - ) - if len(set(self.runtime_factory_kwargs)) != len( - self.runtime_factory_kwargs - ): - raise ValueError("runtime_factory_kwargs must be unique") - runtime_factory_kwargs = frozenset(self.runtime_factory_kwargs) - unknown_runtime_kwargs = ( - runtime_factory_kwargs - self.allowed_factory_kwargs - ) - if unknown_runtime_kwargs: - names = ", ".join(sorted(unknown_runtime_kwargs)) - raise ValueError( - f"runtime_factory_kwargs must be allowed factory kwargs: " - f"{names}" - ) - base_runtime_arg_count = min(self.runtime_arg_counts) - if max(self.runtime_arg_counts) - base_runtime_arg_count > len( - self.runtime_factory_kwargs - ): - raise ValueError( - "runtime_arg_counts require more trailing runtime arguments " - "than runtime_factory_kwargs declares" - ) - prerequisite_names = [name for name, _ in prerequisites] - if len(set(prerequisite_names)) != len(prerequisite_names): - raise ValueError( - "runtime_factory_kw_prerequisite names must be unique" - ) - known_prerequisites = ( - runtime_factory_kwargs | self.allowed_factory_kwargs - ) - for name, required_name in prerequisites: - if name not in runtime_factory_kwargs: - raise ValueError( - f"runtime_factory_kw_prerequisite targets must be " - f"runtime factory kwargs: {name}" - ) - if required_name not in known_prerequisites: - raise ValueError( - f"runtime_factory_kw_prerequisite requirements must be " - f"known factory kwargs: {required_name}" - ) - if name == required_name: - raise ValueError( - "runtime_factory_kw_prerequisites cannot require themselves" - ) - unknown_dtype_kwargs = ( - self.dtype_factory_kwargs - self.allowed_factory_kwargs - ) - if unknown_dtype_kwargs: - names = ", ".join(sorted(unknown_dtype_kwargs)) - raise ValueError( - f"dtype_factory_kwargs must be allowed factory kwargs: {names}" - ) - unknown_required_kwargs = ( - self.required_factory_kwargs - self.allowed_factory_kwargs - ) - if unknown_required_kwargs: - names = ", ".join(sorted(unknown_required_kwargs)) - raise ValueError( - f"required_factory_kwargs must be allowed factory kwargs: " - f"{names}" - ) - unknown_scalar_kwargs = ( - self.scalar_binding_kwargs - runtime_factory_kwargs - ) - if unknown_scalar_kwargs: - names = ", ".join(sorted(unknown_scalar_kwargs)) - raise ValueError( - f"scalar_binding_kwargs must be runtime factory kwargs: {names}" - ) - if self.runtime_offset_kwarg is not None: - if ( - not isinstance(self.runtime_offset_kwarg, str) - or not self.runtime_offset_kwarg - ): - raise ValueError( - "runtime_offset_kwarg must be a non-empty string or None" - ) - if self.runtime_offset_kwarg not in self.allowed_factory_kwargs: - raise ValueError( - "runtime_offset_kwarg must be an allowed factory kwarg" - ) - if self.runtime_offset_kwarg in runtime_factory_kwargs: - raise ValueError( - "runtime_offset_kwarg must not also be " - "a runtime factory kwarg" - ) - if not isinstance(self.accepts_temp_storage, bool): - raise TypeError("accepts_temp_storage must be a bool") - if not callable(self.infer_payload): - raise TypeError("infer_payload must be callable") - for name in ( - "analyze_match", - "prepare_runtime_args", - "validate_runtime_controls", - ): - hook = getattr(self, name) - if hook is not None and not callable(hook): - raise TypeError(f"{name} must be callable or None") - - -_GROUP_OPERATIONS: dict[Callable[..., Any], str] = {} -_GROUP_FAMILY_MODULES: dict[str, str] = {} -_FACTORY_OPERATIONS: dict[Callable[..., Any], FactoryOperation] = {} -_GROUP_PRIMITIVES: dict[str, GroupPrimitiveRegistration] = {} -_REWRITE_OPERATIONS: dict[str, RewriteOperationSpecification] = {} -_GROUP_FAMILY_IMPORT_LOCK = RLock() - - -def group_operation( - operation: str, - *, - family_module: str, -) -> Callable[[_CallableT], _CallableT]: - """Associate a public group marker with its compiler family. - - This decorator records the exact callable object, so the planner - recognizes aliases of the registered function without treating unrelated - functions with the same name as cooperative operations. Record the family - module for lazy loading of planning/rewrite hooks and attach the - backend-member marker used during common API provenance checks. - Registration does not import that family or wrap the decorated function. - - Parameters - ---------- - operation : str - Operation identifier shared by planning and rewrite lookup. - family_module : str - Compiler-family module that registers the operation's hooks. - - Returns - ------- - callable - Decorator that updates registries and callable metadata, then - returns the original function. Equal registrations may repeat. - - Raises - ------ - RuntimeError - The callable already has another operation, or the operation - already has another family module. - """ - - def decorate(function: _CallableT) -> _CallableT: - existing = _GROUP_OPERATIONS.get(function) - if existing is not None and existing != operation: - raise RuntimeError( - f"group marker {function!r} is already registered as " - f"{existing!r}" - ) - _GROUP_OPERATIONS[function] = operation - existing_module = _GROUP_FAMILY_MODULES.get(operation) - if existing_module is not None and existing_module != family_module: - raise RuntimeError( - f"group operation {operation!r} is already assigned to " - f"compiler family {existing_module!r}" - ) - _GROUP_FAMILY_MODULES[operation] = family_module - function.__dict__["__cuda_coop_backend_member__"] = operation - return function - - return decorate - - -def group_operation_name(function: Any) -> str | None: - """Return the operation for an exactly registered group marker.""" - - return _GROUP_OPERATIONS.get(function) - - -def _ensure_group_family_loaded(operation: str) -> None: - """Import an operation's hooks only when a registry lookup needs them. - - Access the qualified backend member first if its marker has not yet - registered a family. A reentrant lock serializes family imports. - An unknown operation leaves no hooks. - """ - - module_name = _GROUP_FAMILY_MODULES.get(operation) - if module_name is None: - backend = import_module("cuda.coop.numba_mlir") - getattr(backend, operation, None) - module_name = _GROUP_FAMILY_MODULES.get(operation) - if module_name is None: - return - with _GROUP_FAMILY_IMPORT_LOCK: - import_module(module_name) - - -def register_group_primitive( - operation: str, - *, - lower: Callable[..., list[Any]], - validate_common_arguments: Callable[..., None] | None = None, -) -> None: - """Register group-call lowering for one public operation. - - A primitive-family module calls this when it is imported. The group - planner later looks up these hooks after resolving the call's thread - group and binding its arguments. Registration connects an operation to - the existing whole-function planner; it does not schedule another pass - or compile a provider. - - Parameters - ---------- - operation : str - Public operation name used to find these hooks, such as ``"load"``. - lower : callable - Hook that receives the planning context, call assignment, resolved - group, and bound arguments, and returns replacement IR statements. - validate_common_arguments : callable or None - Optional check for calls through the common ``cuda.coop`` API. It runs - before ``lower`` with the planning context and bound arguments; - backend-qualified calls skip it. - - Raises - ------ - RuntimeError - Different hooks are already registered under this operation name. - Repeating the same registration is allowed. - """ - - registration = GroupPrimitiveRegistration( - lower=lower, - validate_common_arguments=validate_common_arguments, - ) - existing = _GROUP_PRIMITIVES.get(operation) - if existing is not None and existing != registration: - raise RuntimeError( - f"group primitive {operation!r} is already registered" - ) - _GROUP_PRIMITIVES[operation] = registration - - -def group_primitive(operation: str) -> GroupPrimitiveRegistration | None: - """Return the planning hooks registered for an operation name.""" - - if operation not in _GROUP_PRIMITIVES: - _ensure_group_family_loaded(operation) - return _GROUP_PRIMITIVES.get(operation) - - -def register_rewrite_operation( - operation: str, - specification: RewriteOperationSpecification, -) -> None: - """Register one provider ABI with the shared before-inference rewrite. - - Equal repeated specifications are allowed. Different specifications for - one operation would make its calls ambiguous. - - Parameters - ---------- - operation : str - Operation name used to look up the call grammar and rewrite hooks. - specification : RewriteOperationSpecification - Provider call grammar and hooks for this operation. - - Raises - ------ - TypeError - ``specification`` is not a ``RewriteOperationSpecification``. - RuntimeError - The operation already has a different specification. - """ - - if not isinstance(specification, RewriteOperationSpecification): - raise TypeError("specification must be a RewriteOperationSpecification") - existing = _REWRITE_OPERATIONS.get(operation) - if existing is not None and existing != specification: - raise RuntimeError( - f"rewrite operation {operation!r} is already registered" - ) - _REWRITE_OPERATIONS[operation] = specification - - -def rewrite_operation(operation: str) -> RewriteOperationSpecification | None: - """Return the before-inference registration for one operation. - - Load the owning primitive family when its registration is missing, without - eagerly importing every family. - - Parameters - ---------- - operation : str - Operation name whose call grammar and rewrite hooks are requested. - - Returns - ------- - RewriteOperationSpecification or None - Registered specification, or ``None`` if no registration exists - after loading the owning family. - """ - - if operation not in _REWRITE_OPERATIONS: - _ensure_group_family_loaded(operation) - return _REWRITE_OPERATIONS.get(operation) - - -def register_factory( - function: _CallableT, - *, - operation: str, - namespace: str, - storage_abi: StorageABI, - execution_scope: SynchronizationScope, - synchronization_scope: SynchronizationScope, -) -> _CallableT: - """Register a provider factory and its storage and thread scopes. - - The provider rewrite identifies factories by object identity and uses this - metadata to validate provider calls and materialized invocables. The - factory's import path or function name is not used to infer its ABI. - Registration only declares metadata; the factory and source emitter must - implement the declared storage and synchronization behavior. - - Parameters - ---------- - function : callable - Host-side callable that builds an ``Invocable`` from - specialization keywords, or returns an ``Algorithm`` during batch - collection. See the module overview for the distinction from - public group markers. - operation : str - Non-empty operation identifier, such as ``"load"`` or ``"store"``. - namespace : str - Non-empty provider namespace, such as ``"block"`` or ``"warp"``. - storage_abi : StorageABI - Whether provider calls take a leading scratch pointer. - execution_scope : SynchronizationScope - Scope of threads executing the cooperative operation. - synchronization_scope : SynchronizationScope - Declared synchronization scope: ``NONE`` or the execution scope. - - Returns - ------- - callable - The original factory after registry insertion. Identical repeated - registration is accepted. - - Raises - ------ - TypeError - ``function`` is not callable. - ValueError - Names, enum values, or the scope relationship are invalid. - RuntimeError - This exact factory is already registered with different metadata. - """ - - if not callable(function): - raise TypeError("lowering factory must be callable") - metadata = FactoryOperation( - operation=operation, - namespace=namespace, - storage_abi=storage_abi, - execution_scope=execution_scope, - synchronization_scope=synchronization_scope, - ) - existing = _FACTORY_OPERATIONS.get(function) - if existing is not None and existing != metadata: - raise RuntimeError( - f"lowering factory {function!r} is already registered as " - f"{existing!r}" - ) - _FACTORY_OPERATIONS[function] = metadata - return function - - -def factory_operation(function: Any) -> FactoryOperation | None: - """Return metadata for an exactly registered lowering factory.""" - - return _FACTORY_OPERATIONS.get(function) - - -__all__ = [ - "FactoryOperation", - "GroupPrimitiveRegistration", - "RewriteOperationSpecification", - "StorageABI", - "expected_storage_reuse_barrier", - "factory_operation", - "group_operation", - "group_operation_name", - "group_primitive", - "provider_synchronization_matches", - "register_factory", - "register_group_primitive", - "register_rewrite_operation", - "rewrite_operation", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_parameters.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_parameters.py deleted file mode 100644 index cfbdd689e445..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_parameters.py +++ /dev/null @@ -1,471 +0,0 @@ -# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Normalize specialization inputs and check scalar conversion rules. - -Provider construction and pre-typing payload inference share these helpers so -dimension shapes, dtype spellings, and scalar controls have one -interpretation. Normalization produces compiler types; common-API validation -further restricts them to the common API's numeric profile. Static scalar -conversion preserves explicit dtype provenance, while operator-result -inference delegates promotion rules to the compiler and leaves unresolved -expressions for its later typing pass. -""" - -import math -import operator -from collections import namedtuple -from numbers import Real -from typing import Any, cast - -import numba_cuda_mlir.numba_cuda.types as numba_types -import numpy as np -from numba_cuda_mlir import types as numba_mlir_types - -from cuda.coop._core.dtype_policy import ( - validate_common_numeric_dtype_name, -) - -dim3 = namedtuple("dim3", ("x", "y", "z")) - - -def normalize_dim_param(dim) -> dim3: - """Convert a positive block shape to three integer dimensions. - - Accept a scalar, a tuple of one to three entries, or ``dim3``. Fill - omitted dimensions with one. Each entry must be an integer or an - integer-like value accepted by ``operator.index``. Reject Python booleans - and nonpositive values. Validate device thread-count limits separately. - """ - - if isinstance(dim, dim3): - values = tuple(dim) - elif isinstance(dim, tuple): - if not 1 <= len(dim) <= 3: - raise ValueError( - f"Tuple dimension must have one, two, or three elements; got " - f"{len(dim)}" - ) - values = dim - else: - values = (dim,) - - normalized = [] - for value in values: - if isinstance(value, bool): - raise TypeError("Dimension values must be integers") - try: - value = operator.index(value) - except TypeError as exc: - raise TypeError("Dimension values must be integers") from exc - if value <= 0: - raise ValueError(f"Dimension values must be positive, got {dim!r}") - normalized.append(value) - - normalized.extend([1] * (3 - len(normalized))) - return dim3(*normalized) - - -_NP_DTYPE_TO_NUMBA_MLIR_TYPE = { - np.dtype(np.bool_): numba_types.boolean, - np.dtype(np.int8): numba_types.int8, - np.dtype(np.int16): numba_types.int16, - np.dtype(np.int32): numba_types.int32, - np.dtype(np.int64): numba_types.int64, - np.dtype(np.uint8): numba_types.uint8, - np.dtype(np.uint16): numba_types.uint16, - np.dtype(np.uint32): numba_types.uint32, - np.dtype(np.uint64): numba_types.uint64, - np.dtype(np.float16): numba_types.float16, - np.dtype(np.float32): numba_types.float32, - np.dtype(np.float64): numba_types.float64, - np.dtype(np.complex64): numba_types.complex64, - np.dtype(np.complex128): numba_types.complex128, -} - -_NUMBA_MLIR_TYPE_NAME_ALIASES = { - "bool_": "boolean", - "bool": "boolean", -} - - -def _normalize_numba_mlir_type_name(type_name: str) -> str: - return _NUMBA_MLIR_TYPE_NAME_ALIASES.get(type_name, type_name) - - -def _dtype_from_numpy(np_dtype: np.dtype) -> numba_types.Type: - """Map a NumPy dtype to the compiler's corresponding type. - - Prefer the explicit scalar table, then try the compiler's named type. - Raise ``ValueError`` when neither path supports the dtype. - """ - - canonical = np.dtype(np_dtype) - if canonical in _NP_DTYPE_TO_NUMBA_MLIR_TYPE: - return _NP_DTYPE_TO_NUMBA_MLIR_TYPE[canonical] - - type_name = _normalize_numba_mlir_type_name(canonical.name) - if hasattr(numba_mlir_types, type_name): - resolved = getattr(numba_mlir_types, type_name) - if isinstance(resolved, numba_types.Type): - return resolved - - raise ValueError(f"Unsupported numpy dtype: {canonical}") - - -def normalize_dtype_param( - dtype: object, -) -> "numba_types.Type": - """Convert a supported dtype spelling to a compiler type object. - - Accept existing compiler types, NumPy dtypes and scalar classes, and - supported type-name strings. The Python classes ``int`` and ``float`` - request 32-bit types here; ``bool`` requests boolean and ``complex`` - requests complex128. This interprets a requested type, not the dtype of a - Python scalar value. - - String names may start with ``np.``, ``types.``, or - ``numba_cuda_mlir.types.``. Unknown formats or names raise ``ValueError``. - Common API callers separately check the common numeric profile; - normalization alone does not impose that profile. - """ - - if dtype is bool: - return numba_types.boolean - if dtype is int: - return numba_types.int32 - if dtype is float: - return numba_types.float32 - if dtype is complex: - return numba_types.complex128 - if isinstance(dtype, numba_types.Type): - return dtype - if isinstance(dtype, np.dtype): - return _dtype_from_numpy(dtype) - if isinstance(dtype, type) and issubclass(dtype, np.generic): - return _dtype_from_numpy(np.dtype(dtype)) - if isinstance(dtype, str): - if dtype.startswith("np."): - np_type_name = dtype[3:] - if not hasattr(np, np_type_name): - raise ValueError(f"Invalid numpy dtype: {np_type_name}") - return _dtype_from_numpy(np.dtype(getattr(np, np_type_name))) - - for prefix in ("numba_cuda_mlir.types.", "types."): - if dtype.startswith(prefix): - dtype = dtype[len(prefix) :] - break - - type_name = _normalize_numba_mlir_type_name(dtype) - if hasattr(numba_mlir_types, type_name): - resolved = getattr(numba_mlir_types, type_name) - if isinstance(resolved, numba_types.Type): - return resolved - raise ValueError(f"Invalid Numba-CUDA-MLIR type name: {dtype}") - - raise ValueError(f"Unrecognized dtype format: {dtype}") - - -_NUMBA_MLIR_DTYPE_NAMES = { - numba_mlir_type: np_dtype.name - for np_dtype, numba_mlir_type in _NP_DTYPE_TO_NUMBA_MLIR_TYPE.items() -} - - -def _normalize_common_dtype(dtype: object) -> tuple[numba_types.Type, str]: - """Return a backend dtype and its common normalized name.""" - - dtype = normalize_dtype_param(dtype) - return dtype, _NUMBA_MLIR_DTYPE_NAMES.get(dtype, str(dtype)) - - -def _validate_common_numeric_dtype( - dtype: object, - *, - operation: str, - parameter: str | None = None, -) -> numba_types.Type: - """Return one normalized dtype from the common API's numeric profile.""" - - dtype, dtype_name = _normalize_common_dtype(dtype) - validate_common_numeric_dtype_name( - dtype_name, - operation=operation, - parameter=parameter, - ) - return dtype - - -def _python_scalar_dtype(value: object) -> numba_types.Type | None: - """Infer an ordinary or NumPy scalar's dtype through NumPy. - - Preserve the scalar's inferred width, unlike the 32-bit defaults for - Python type objects in ``normalize_dtype_param``. Return ``None`` for - unsupported values or dtypes. - """ - - if type(value) not in {bool, int, float, complex} and not isinstance( - value, np.generic - ): - return None - try: - return normalize_dtype_param(np.asarray(value).dtype) - except (TypeError, ValueError): - return None - - -def _scalar_cast_dtype(function: Any) -> numba_types.Type | None: - """Recognize the scalar dtype requested by a cast callable. - - Accept compiler type objects and scalar types recognized by ``np.dtype``, - such as ``np.float32``. Return ``None`` for structured, subarray, or - unsupported dtypes. - """ - - if isinstance(function, numba_types.Type): - try: - return normalize_dtype_param(function) - except (TypeError, ValueError): - return None - try: - np_dtype = np.dtype(function) - except (TypeError, ValueError): - return None - if np_dtype.subdtype is not None or np_dtype.fields is not None: - return None - try: - return normalize_dtype_param(np_dtype) - except (TypeError, ValueError): - return None - - -def _scalar_operator_result_dtype( - function: object, *operand_dtypes: object -) -> numba_types.Type | None: - """Infer a scalar expression's result dtype using compiler rules. - - Payload inference uses this probe before the authoritative typing pass. - Resolve the operator through the active Numba-CUDA-MLIR typing context so - integer promotion and result widths follow compiler rules rather than - Python evaluation. Initialize that context lazily when a usable set of - operand dtypes is available. - - All resolution failures become ``None``. This intentionally leaves - unsupported expressions for later typing diagnostics rather than making a - best-effort provenance query reject the kernel. - - Parameters - ---------- - function : callable or None - Operator or scalar function whose result is being inferred. - *operand_dtypes : object - Operand dtypes accepted by ``normalize_dtype_param``. Inference - requires a known dtype for every operand; ``None`` prevents it. - - Returns - ------- - numba_types.Type or None - Normalized result dtype, or ``None`` for missing inputs, no typing - signature, or an exception during normalization or resolution. - """ - - if ( - function is None - or not operand_dtypes - or any(dtype is None for dtype in operand_dtypes) - ): - return None - try: - normalized = tuple( - normalize_dtype_param(dtype) for dtype in operand_dtypes - ) - from numba_cuda_mlir.descriptor import mlir_target - - mlir_target.ensure_initialized() - signature = mlir_target.typing_context.resolve_function_type( - function, - normalized, - {}, - ) - if signature is None: - return None - return normalize_dtype_param(signature.return_type) - except Exception: # noqa: BLE001 - Later typing reports failures. - # This is best-effort provenance, not the authoritative typing pass. - return None - - -def _validate_runtime_integer_dtype( - dtype: object, *, operation: str, parameter: str -) -> numba_types.Integer: - """Check the integer types accepted by Load/Store controls. - - Unwrap literal types, then accept signed integers up to 64 bits or - unsigned integers up to 32 bits. Reject booleans and other types with - ``TypeError``, using ``operation`` and ``parameter`` in the message. This - checks representation; it does not prove runtime value bounds. - """ - - if isinstance(dtype, numba_types.Literal): - dtype = dtype.literal_type - if isinstance(dtype, numba_types.Boolean) or not isinstance( - dtype, numba_types.Integer - ): - raise TypeError( - f"coop {operation} {parameter} must be an integer, not bool or a " - f"noninteger scalar" - ) - if dtype.bitwidth > 64 or (not dtype.signed and dtype.bitwidth > 32): - raise TypeError( - f"coop {operation} {parameter} must be a signed integer up to 64 " - f"bits or an unsigned integer up to 32 bits" - ) - return dtype - - -def coerce_static_scalar( - value: object, - dtype: object, - *, - operation: str, - parameter: str, - source_dtype: object = None, -) -> np.generic: - """Normalize a compile-time scalar without erasing its dtype provenance. - - Ordinary Python ``int``/``float`` literals take the payload's dtype after - range and finiteness checks; float-to-integer conversion is forbidden. - Floating conversion may round a finite in-range literal. NumPy scalars and - values with an explicit compiler ``source_dtype`` instead require an exact - dtype match, even if their numeric value would fit another payload dtype. - This keeps a typed default or Store value from silently changing width. - - The explicit-source path treats ``source_dtype`` as authoritative - provenance: the caller is responsible for pairing it with a value of that - dtype. This helper handles static values only; runtime arguments use - provider typing and generated ABI checks. - - Parameters - ---------- - value : object - Compile-time scalar value to normalize. - dtype : object - Target payload dtype from the common numeric profile. - operation : str - Public operation name used in diagnostics, such as ``"load"``. - parameter : str - Argument name used in diagnostics, such as ``"oob_default"``. - source_dtype : object, optional - Known compiler dtype of ``value``. When absent, a NumPy scalar's - dtype is used; ordinary Python numeric literals remain - contextually typed. - - Returns - ------- - numpy.generic - Scalar in the normalized target dtype. - - Raises - ------ - TypeError - The target is outside the common numeric profile, a value is - boolean or non-numeric, typed provenance differs from the target, - or an ordinary float would be converted to an integer. - ValueError - The dtype cannot be normalized, an ordinary literal is out of - range, or a checked scalar is non-finite. - """ - - target_dtype = _validate_common_numeric_dtype( - dtype, - operation=operation, - parameter=parameter, - ) - target_numpy_dtype = np.dtype(_NUMBA_MLIR_DTYPE_NAMES[target_dtype]) - - if source_dtype is None and isinstance(value, np.generic): - source_dtype = value.dtype - if source_dtype is not None: - try: - normalized_source = normalize_dtype_param(source_dtype) - except (TypeError, ValueError) as exc: - raise TypeError( - f"cuda.coop.{operation} {parameter} must be a numeric scalar" - ) from exc - if normalized_source != target_dtype: - raise TypeError( - f"cuda.coop.{operation} {parameter} dtype {normalized_source}" - f" does not match payload dtype {target_dtype}" - ) - scalar = value.item() if isinstance(value, np.generic) else value - if isinstance(scalar, Real) and not math.isfinite(float(scalar)): - raise ValueError( - f"cuda.coop.{operation} {parameter} must be finite" - ) - return target_numpy_dtype.type(value) - - if isinstance(value, bool) or type(value) is bool: - raise TypeError(f"cuda.coop.{operation} {parameter} must not be bool") - if type(value) not in {int, float}: - raise TypeError( - f"cuda.coop.{operation} {parameter} must be an ordinary Python " - f"numeric literal or an exactly typed NumPy/compiler scalar" - ) - - value = cast(int | float, value) - if np.issubdtype(target_numpy_dtype, np.integer): - if type(value) is float: - raise TypeError( - f"cuda.coop.{operation} {parameter} does not permit " - f"float-to-integer conversion" - ) - bounds = np.iinfo(target_numpy_dtype) - if not bounds.min <= value <= bounds.max: - raise ValueError( - f"cuda.coop.{operation} {parameter} value {value} is outside " - f"the range of {target_numpy_dtype.name}" - ) - return target_numpy_dtype.type(value) - - if type(value) is float and not math.isfinite(value): - raise ValueError(f"cuda.coop.{operation} {parameter} must be finite") - maximum = float(np.finfo(target_numpy_dtype).max) - if not -maximum <= value <= maximum: - raise ValueError( - f"cuda.coop.{operation} {parameter} value {value} is outside the " - f"finite range of {target_numpy_dtype.name}" - ) - with np.errstate(over="ignore", invalid="ignore"): - result = target_numpy_dtype.type(value) - if not np.isfinite(result): - raise ValueError( - f"cuda.coop.{operation} {parameter} value {value} is outside the " - f"finite range of {target_numpy_dtype.name}" - ) - return result - - -def _validate_static_oob_default(value: object, dtype: object) -> np.generic: - """Normalize a static Load default before provider construction.""" - - return coerce_static_scalar( - value, - dtype, - operation="load", - parameter="oob_default", - ) - - -__all__ = [ - "_python_scalar_dtype", - "_scalar_cast_dtype", - "_scalar_operator_result_dtype", - "_validate_common_numeric_dtype", - "_validate_runtime_integer_dtype", - "_validate_static_oob_default", - "coerce_static_scalar", - "dim3", - "normalize_dim_param", - "normalize_dtype_param", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_planner.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_planner.py deleted file mode 100644 index d3b031086f24..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_planner.py +++ /dev/null @@ -1,99 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Connect cooperative-call lowering to Numba's whole-function planner stage. - -Numba invokes ``CoopWholeFunctionPlanner`` after inlining device helpers and -before type inference, when cooperative operations in those helpers are -visible in the caller's IR. The first phase resolves thread groups and chooses -backend providers. The second specializes those providers, allocates payload -and scratch storage, and replaces their calls with compiler-supported IR. -Its driver prepares the whole function before staging shared backing, -rewriting individual blocks, and cleaning unused bindings. Later consumers -can supply an earlier payload's dtype or a descriptor's scratch requirements, -so preparation needs the original calls from every block. - -Keep the phases in this order: provider rewriting needs the decisions made by -group resolution. When the first phase changes IR, rebuild the definitions and -control-flow analysis before the second phase reads them. The final Boolean -result tells Numba whether to repair IR again before continuing compilation. -""" - -from numba_cuda_mlir._whole_function_planners import _planner_registry -from numba_cuda_mlir.extending import WholeFunctionPlanner - -from .._types import kernel_compile_target -from ._group_planner import _GroupPlanning -from ._nvrtc import CompileTarget, compile_context_scope -from ._rewrite import _CallRewriting - - -class CoopWholeFunctionPlanner( - _GroupPlanning, _CallRewriting, WholeFunctionPlanner -): - """Resolve groups before specializing calls and allocating storage. - - One planner instance owns both phases for a compiler attempt. - The mixins keep group resolution and call rewriting in separate modules - while giving both access to the same compiler state and launch metadata. - They do not register additional passes. Keeping the sequence in ``run`` - also puts the required IR repair between the phases in one place: call - rewriting must see the provider calls that group resolution just created. - """ - - def run(self) -> bool: - """Lower cooperative calls while their inlined consumers are visible. - - Numba calls this once per compiler attempt after device-helper - inlining and before type inference. A ``True`` result makes its - planner registry rebuild IR analysis before proceeding to later - planners and compiler passes. It does not ask the registry to repeat - this planner. A request for a literal argument may separately cause - the dispatcher to begin another compiler attempt. Launch metadata is - obtained within the current attempt or reported as unavailable. - - Group resolution supplies the launch-dependent operation and provider - choices that call rewriting needs. Rebuild IR analysis between these - phases when resolution changes the function: otherwise rewriting could - inspect definitions and control-flow facts from before the new - provider calls existed. Numba repairs the final IR after a successful - change report, ready for later compiler passes and type inference. - During call rewriting, keep those original definitions until every - block replacement and the final binding cleanup are complete. - - Both phases generate device code for the kernel's resolved ``chip`` - (including any ``a``/``f`` suffix) and ``lineinfo`` option, not for - the current device. - - Both phases run even when group resolution makes no changes; payload - constructors and private provider calls can still need rewriting. - Planning and provider errors propagate with their original - diagnostics. Compiler specialization requests also propagate so the - dispatcher can retry with the required literal arguments. - - Returns - ------- - bool - ``True`` when either phase changes ``state.func_ir``; ``False`` - when neither phase changes it. The result reports IR changes, not - whether launch metadata was requested or a provider was compiled. - """ - - target = CompileTarget.from_targetoptions( - self.state.metadata.get("targetoptions") or {} - ) - with kernel_compile_target(target), compile_context_scope(): - groups_changed = self._resolve_groups() - if groups_changed: - # Group resolution changes IR statements. Use Numba's repair - # helper to rebuild definitions and control-flow analysis - # before rewriting the new provider calls. The registry - # repairs a planner's changes after run() returns, too late - # for this call to _rewrite_calls(). - _planner_registry._repair_ir(self.state.func_ir) - calls_changed = self._rewrite_calls() - return groups_changed or calls_changed - - -__all__ = ["CoopWholeFunctionPlanner"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite.py deleted file mode 100644 index cbe278a8cc1c..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite.py +++ /dev/null @@ -1,1006 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Lower planned primitives and per-thread payloads to executable Numba IR. - -``CoopWholeFunctionPlanner`` runs after device helpers have been inlined and -before type inference. Its group-resolution phase replaces public primitive -calls with private provider calls. This module specializes those providers -into callable implementations (invocables), supplies their scratch storage, -and replaces ``ThreadData`` constructors with per-thread local arrays. - -``_CallRewriting._rewrite_calls`` owns the function's rewrite lifecycle: -prepare original calls and storage requirements, stage shared backing, rewrite -blocks, then retire compile-time bindings. Preparation must precede every -block replacement: a consumer in a later block can determine the dtype of an -earlier ``ThreadData`` allocation or the requirements of a scratch descriptor. -Block emission still materializes providers as needed and infers standalone -payload dtypes from their writes. - -``ThreadData`` is a per-thread payload; it can also be used in ordinary indexed -computation without a primitive. ``TempStorage`` is an opaque descriptor for -cooperative scratch, with ownership and synchronization rules. Its concrete -layout comes from provider preparation. Keeping these records separate lets -payload allocation and scratch views share facts without confusing their -different lifetimes or ownership. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, cast - -from numba_cuda_mlir import cuda as _cuda_module -from numba_cuda_mlir.extending import require_launch_config - -from cuda.coop._core import GroupLoweringPlan - -from ._operations import _GROUP_LOWERING_PLAN_KWARG, StorageABI -from ._rewrite_arguments import _ArgumentRewrite -from ._rewrite_group_metadata import _GroupMetadataRewrite -from ._rewrite_invocables import _InvocableRewrite -from ._rewrite_launch import _LaunchRewrite -from ._rewrite_payload import _PayloadRewrite -from ._rewrite_provenance import _ProvenanceRewrite -from ._rewrite_storage import _StorageRewrite -from ._rewrite_support import ( - _GLOBAL_NAME_COUNTER, - CoopSinglePhaseRewriteError, - Rewrite, - _DeferredCoopRewrite, - _next_global_name, - _RewriteMatch, - _TempStorageRequirementSummary, - ir, -) - -if TYPE_CHECKING: - from numba_cuda_mlir.numba_cuda.types import Type - from numba_cuda_mlir.numba_cuda.typing.templates import Signature - - from ._planner import CoopWholeFunctionPlanner - from ._rewrite_support import _ThreadDataSpecification - - -class CoopSinglePhaseRewrite( - _ProvenanceRewrite, - _ArgumentRewrite, - _LaunchRewrite, - _GroupMetadataRewrite, - _PayloadRewrite, - _InvocableRewrite, - _StorageRewrite, - Rewrite, -): - """Prepare and rewrite cooperative calls across one function. - - The mixins share payload type/extent information, scratch-storage plans, - and provider caches. This lets argument inference and array allocation - use the same information across calls and control-flow branches. - - The driver calls ``prepare_calls_and_storage`` before any IR changes and - ``begin_rewrite`` to stage shared backing. It then uses ``match``/``apply`` - for each block and calls ``finish_rewrite`` after installing all results. - Original call identity and definitions remain available during emission; - cleanup waits until no unreplaced consumer can need a compile-time binding. - These are steps of one planner, not independently scheduled passes. - """ - - def match( - self, - func_ir: ir.FunctionIR, - block: ir.Block, - typemap: dict[str, Type] | None, - calltypes: dict[ir.Expr, Signature] | None, - ) -> bool: - """Select one block's work using the prepared function's call records. - - ``prepare_calls_and_storage`` must succeed before any block is matched - or replaced. Provider calls are looked up by original assignment - identity; constructors and extent queries are collected for ``apply``. - The rewrite-interface maps ``typemap`` and ``calltypes`` are not read. - """ - - if self._func_ir_identity != id(func_ir): - raise RuntimeError( - "prepare_calls_and_storage must succeed before match" - ) - self._prepare_block(block) - for inst in block.body: - if isinstance(inst, ir.Assign): - self._match_assignment(inst) - return ( - bool(self._matches) - or bool(self._temp_storage_assigns) - or bool(self._thread_data_func_vars) - or bool(self._thread_data_extents) - ) - - def prepare_calls_and_storage(self, func_ir: ir.FunctionIR) -> bool: - """Prepare all calls and storage across the current function's blocks. - - Device helpers have already been inlined into this one function IR. - A ``ThreadData`` dtype can come from a consumer in another block, and - a ``TempStorage`` descriptor's requirements depend on all its consumers. - Collect their constructors and analyze calls while every original - definition is available. Validate storage ownership before preparing - providers, whose concrete layouts supply scratch sizes and alignments. - - Publish completed call records only on success. Missing launch facts - leave the IR unchanged; the driver either retries with a fresh - rewriter or leaves a device helper for its inlined caller. Unresolved - group markers also leave the function for group planning. - """ - - from ._group_planner import has_group_markers - - if has_group_markers(func_ir): - return False - - func_ir_identity = id(func_ir) - if self._func_ir_identity == func_ir_identity: - return True - self._func_ir_identity = None - self._func_matches = {} - self._rewrite_started = False - self._factory_argument_cleanup_candidates = set() - self._payload_callee_cleanup_names = set() - self._func_temp_storage_requirements = {} - self._implicit_temp_storage_requirements = ( - _TempStorageRequirementSummary() - ) - self._func_ir = func_ir - self._thread_data_specifications = {} - self._thread_data_like_vars = set() - self._temp_storage_plans = {} - self._temp_storage_global_plan = None - self._temp_storage_ctor_order = {} - self._temp_storage_ctor_roots = {} - self._implicit_temp_storage_plan = None - self._temp_storage_backing_var = None - self._temp_storage_backing_emitted = False - self._prebundled_specializations = {} - try: - matches = self._collect_function_calls(func_ir) - # Invalid ownership or escaping descriptors must fail before - # provider compilation. The retained calls serve both stages. - storage_uses = self._collect_temp_storage_uses(func_ir, matches) - self._prepare_ltoir_bundle_for_matches(list(matches.values())) - # Sizes and alignments require concrete provider layouts; the - # semantic lowering plans alone do not supply them. - self._func_temp_storage_requirements = ( - self._compute_func_temp_storage_requirements(storage_uses) - ) - except _DeferredCoopRewrite: - self._func_temp_storage_requirements = {} - return False - - self._func_matches = matches - self._func_ir_identity = func_ir_identity - return True - - def begin_rewrite(self) -> None: - """Stage function-wide backing storage before visiting any block. - - Provider layouts are now available. Check shared-memory conflicts - against the original function, then insert the backing in the entry - block so it dominates every scratch view, regardless of visit order. - Preparation and any launch retry must finish before this mutates IR. - Later block matching ignores the generated backing statements, and - block rewriting copies them through with the surrounding statements. - """ - - if self._func_ir_identity != id(self._func_ir): - raise RuntimeError( - "prepare_calls_and_storage must succeed before rewriting" - ) - if self._rewrite_started: - raise RuntimeError("function rewriting has already started") - if self._has_temp_storage_requirements(): - self._stage_temp_storage_backing() - self._rewrite_started = True - - def finish_rewrite(self) -> None: - """Retire compile-time bindings after every replacement is installed. - - A factory input or constructor alias can still feed a call in a later - block. Wait until all calls have been replaced before checking uses - across the function. Keep the original definition table throughout - rewriting: standalone payload inference and constructor-alias cleanup - both need original definitions. The planner registry repairs IR - analysis after this lifecycle completes. - """ - - if not self._rewrite_started: - raise RuntimeError("begin_rewrite must precede finish_rewrite") - self._remove_unused_factory_arguments() - self._clear_unused_payload_callees() - self._factory_argument_cleanup_candidates.clear() - self._payload_callee_cleanup_names.clear() - self._rewrite_started = False - - def _collect_function_calls( - self, func_ir: ir.FunctionIR - ) -> dict[ir.Assign, _RewriteMatch]: - """Analyze every provider call before any block is replaced. - - Record all payload and scratch constructors first so consumers in - other blocks see the same facts. Retain completed call analysis by - original assignment identity for compilation, storage planning and - block emission. Replaced assignments no longer match on later visits. - """ - - matches: dict[ir.Assign, _RewriteMatch] = {} - saved_block_defs = self._block_defs - saved_block = self._block - self._temp_storage_ctor_specifications = {} - self._temp_storage_ctor_order = {} - self._temp_storage_ctor_roots = {} - self._temp_storage_ctor_sites = {} - try: - ctor_order = 0 - for label in sorted(func_ir.blocks): - scan_block = func_ir.blocks[label] - self._block = scan_block - self._block_defs = { - inst.target.name: inst.value - for inst in scan_block.body - if isinstance(inst, ir.Assign) - } - for inst in scan_block.body: - if not isinstance(inst, ir.Assign): - continue - call = inst.value - if not isinstance(call, ir.Expr) or call.op != "call": - continue - if self._is_thread_data_ctor_call(call): - self._thread_data_like_vars.add(inst.target.name) - self._thread_data_specifications[inst.target.name] = ( - self._merge_thread_data_specifications( - self._thread_data_specifications.get( - inst.target.name - ), - self._extract_thread_data_specification(call), - ) - ) - elif self._is_temp_storage_ctor_call(call): - self._record_temp_storage_ctor(inst, call) - self._temp_storage_ctor_order.setdefault( - inst.target.name, ctor_order - ) - ctor_order += 1 - self._validate_temp_storage_ctor_sites() - for label in sorted(func_ir.blocks): - scan_block = func_ir.blocks[label] - self._block = scan_block - self._block_defs = { - inst.target.name: inst.value - for inst in scan_block.body - if isinstance(inst, ir.Assign) - } - for inst in scan_block.body: - if not isinstance(inst, ir.Assign): - continue - call = inst.value - if not isinstance(call, ir.Expr) or call.op != "call": - continue - match = self._analyze_provider_call(inst, call) - if match is not None: - matches[inst] = match - finally: - self._block_defs = saved_block_defs - self._block = saved_block - return matches - - def _prepare_block(self, block: ir.Block) -> None: - """Reset block matches and keep the function-wide storage plan.""" - - self._block = block - self._block_defs = { - inst.target.name: inst.value - for inst in block.body - if isinstance(inst, ir.Assign) - } - self._matches = {} - self._thread_data_extents = {} - self._temp_storage_assigns = set() - self._temp_storage_func_vars = set() - self._thread_data_func_vars = set() - - def _match_assignment(self, inst: ir.Assign) -> None: - """Record a payload query, constructor, or provider call.""" - - match = self._func_matches.get(inst) - if match is not None: - self._matches[inst] = match - return - call = inst.value - if ( - isinstance(call, ir.Expr) - and call.op == "getattr" - and call.attr == "items_per_thread" - ): - if self._is_thread_data_like_var(call.value): - specification = self._resolve_thread_data_specification( - call.value - ) - if ( - specification is not None - and specification.items_per_thread is not None - ): - self._thread_data_extents[inst] = ( - specification.items_per_thread - ) - return - if not isinstance(call, ir.Expr) or call.op != "call": - return - if self._is_temp_storage_ctor_call(call): - self._temp_storage_assigns.add(inst) - self._temp_storage_func_vars.add(call.func.name) - return - if self._is_thread_data_ctor_call(call): - self._thread_data_func_vars.add(call.func.name) - return - - def _analyze_provider_call( - self, inst: ir.Assign, call: ir.Expr - ) -> _RewriteMatch | None: - """Separate compile-time factory inputs from device operands. - - Keep the group lowering plan beside the match. It controls storage - and synchronization during emission, but is not an argument to the - factory. - """ - - target = self._resolve_call_target(call) - if target is None: - return None - op_name = target.operation - # factory_kwargs drives specialization; factory_kw_value_vars tracks - # original IR bindings for cleanup after all calls are rewritten. - ( - runtime_args, - runtime_temp_storage_var, - factory_kwargs, - factory_kw_value_vars, - ) = self._validate_and_split_args( - op_name, call, target.getitem_temp_storage - ) - lowering_plan = cast( - GroupLoweringPlan | None, - factory_kwargs.pop(_GROUP_LOWERING_PLAN_KWARG, None), - ) - family_metadata = self._analyze_family_match( - op_name=op_name, - runtime_args=runtime_args, - factory_kwargs=factory_kwargs, - ) - return _RewriteMatch( - op_name=op_name, - factory=target.factory, - factory_metadata=target.factory_metadata, - func_var_name=target.func_var_name, - func_var_name_extra=target.func_var_name_extra, - runtime_args=runtime_args, - runtime_temp_storage_var=runtime_temp_storage_var, - factory_kwargs=factory_kwargs, - factory_kw_value_vars=factory_kw_value_vars, - loc=inst.loc, - family_metadata=family_metadata, - lowering_plan=lowering_plan, - ) - - def apply(self) -> ir.Block: - """Emit the most recently matched block before type inference. - - Bind provider implementations, then replace statements in source - order. ``begin_rewrite`` has staged shared backing. Payloads become - local arrays. Scratch descriptors become shared-memory views or - ``None`` for storage-free calls. Provider calls receive device operands - and any required reuse barrier. - - Retain cleanup candidates for ``finish_rewrite`` so other blocks can - still use their bindings. Provider materialization and a standalone - payload's final dtype inference can occur here. The caller installs - every returned block before finishing the function rewrite. - """ - - assert self._block is not None - if not self._rewrite_started: - raise RuntimeError("begin_rewrite must precede apply") - ( - call_invocable_globals, - func_var_names_to_clear, - candidate_dead_factory_kw_vars, - ) = self._prepare_call_invocables() - - new_block = ir.Block(self._block.scope, self._block.loc) - for inst in self._block.body: - if inst in self._thread_data_extents: - new_block.append( - ir.Assign( - ir.Const(self._thread_data_extents[inst], inst.loc), - inst.target, - inst.loc, - ) - ) - continue - if ( - isinstance(inst, ir.Assign) - and inst.target.name in func_var_names_to_clear - ): - new_block.append( - ir.Assign(ir.Const(None, inst.loc), inst.target, inst.loc) - ) - continue - if ( - isinstance(inst, ir.Assign) - and inst.target.name in self._temp_storage_func_vars - ): - new_block.append( - ir.Assign(ir.Const(None, inst.loc), inst.target, inst.loc) - ) - continue - if ( - isinstance(inst, ir.Assign) - and isinstance(inst.value, ir.Expr) - and (inst.value.op == "call") - and (self._is_thread_data_ctor_call(inst.value)) - ): - self._emit_thread_data_array(new_block, inst) - continue - match = self._matches.get(inst) - if match is None and inst not in self._temp_storage_assigns: - new_block.append(inst) - continue - if inst in self._temp_storage_assigns: - self._emit_temp_storage(new_block, inst) - continue - assert match is not None - self._emit_provider_call( - new_block, inst, match, call_invocable_globals.get(inst) - ) - - # Only these assignments were eligible for this block's factory-input - # cleanup. Saving their identity avoids deleting a different binding - # of the same name after pre-SSA rebinding. - self._factory_argument_cleanup_candidates.update( - inst - for inst in new_block.body - if isinstance(inst, ir.Assign) - and inst.target.name in candidate_dead_factory_kw_vars - ) - self._payload_callee_cleanup_names.update(self._thread_data_func_vars) - self._state.typingctx.refresh() - return new_block - - def _prepare_call_invocables( - self, - ) -> tuple[dict[ir.Assign, tuple[str, object]], set[str], set[str]]: - """Prepare each call's implementation and track bindings to retire. - - The function-wide scan has already attempted batch compilation. - Reuse those invocables where available; individual materialization - is the fallback. - """ - - call_invocable_globals: dict[ir.Assign, tuple[str, object]] = {} - func_var_names_to_clear: set[str] = set() - candidate_dead_factory_kw_vars: set[str] = set() - # Build or reuse each provider implementation before emitting its call. - # Bind it per call site: two calls through the same Python alias can - # require different specializations. - for match_inst, match in self._matches.items(): - invocable, _ = self._materialize_invocable(match) - candidate_dead_factory_kw_vars.update( - value_var.name for value_var in match.factory_kw_value_vars - ) - global_name = _next_global_name("single_phase") - call_invocable_globals[match_inst] = (global_name, invocable) - func_var_names_to_clear.add(match.func_var_name) - if match.func_var_name_extra is not None: - func_var_names_to_clear.add(match.func_var_name_extra) - return ( - call_invocable_globals, - func_var_names_to_clear, - candidate_dead_factory_kw_vars, - ) - - def _require_thread_data_specification( - self, inst: ir.Assign - ) -> _ThreadDataSpecification: - """Resolve the payload dtype before emitting a local array. - - Primitive operands can supply the dtype during requirement - collection. If they did not, inspect typed indexed writes before - reporting an error. - """ - - thread_data_specification = self._thread_data_specifications.get( - inst.target.name - ) - if ( - thread_data_specification is not None - and thread_data_specification.dtype is None - ): - self._infer_thread_data_dtype_from_writes(inst.target) - thread_data_specification = self._thread_data_specifications.get( - inst.target.name - ) - if ( - thread_data_specification is None - or thread_data_specification.dtype is None - ): - raise CoopSinglePhaseRewriteError( - "Failed to infer dtype for coop.ThreadData(...). " - "Supply type information through a cooperative " - "operation, typed indexed assignments, or an " - "explicit dtype argument." - ) - if thread_data_specification.common_root: - from ._parameters import _validate_common_numeric_dtype - - try: - _validate_common_numeric_dtype( - thread_data_specification.dtype, - operation="ThreadData", - ) - except (TypeError, ValueError) as exc: - raise CoopSinglePhaseRewriteError(str(exc)) from exc - - return thread_data_specification - - def _thread_data_array_arguments( - self, new_block: ir.Block, inst: ir.Assign - ) -> tuple[list[ir.Var], list[tuple[str, ir.Var]]]: - """Emit constant array parameters and adapt constructor arguments. - - ``cuda.local.array`` takes ``shape`` where ``ThreadData`` takes - ``items_per_thread``. Preserve the resolved dtype and explicit - alignment. - """ - - thread_data_specification = self._require_thread_data_specification( - inst - ) - dtype_var = ir.Var( - inst.target.scope, - f"__coop_thread_data_dtype_{next(_GLOBAL_NAME_COUNTER)}__", - inst.loc, - ) - new_block.append( - ir.Assign( - ir.Global( - _next_global_name("thread_data_dtype"), - thread_data_specification.dtype, - inst.loc, - ), - dtype_var, - inst.loc, - ) - ) - assert isinstance(inst.value, ir.Expr) - rewritten_args = list(inst.value.args) - rewritten_kws = list(inst.value.kws) - rewritten_kws = [ - ("shape" if name == "items_per_thread" else name, value) - for name, value in rewritten_kws - if name != "alignment" - ] - if thread_data_specification.items_per_thread is None: - raise CoopSinglePhaseRewriteError( - "Failed to infer items_per_thread for " - "coop.ThreadData(...). The item count must be " - "known at compile time." - ) - items_var = ir.Var( - inst.target.scope, - f"__coop_thread_data_items_{next(_GLOBAL_NAME_COUNTER)}__", - inst.loc, - ) - new_block.append( - ir.Assign( - ir.Const(thread_data_specification.items_per_thread, inst.loc), - items_var, - inst.loc, - ) - ) - if rewritten_args: - rewritten_args[0] = items_var - elif any((name == "shape" for name, _ in rewritten_kws)): - rewritten_kws = [ - (name, items_var if name == "shape" else value) - for name, value in rewritten_kws - ] - else: - rewritten_args.append(items_var) - if len(rewritten_args) >= 2: - rewritten_args[1] = dtype_var - elif any((name == "dtype" for name, _ in rewritten_kws)): - rewritten_kws = [ - (name, dtype_var if name == "dtype" else value) - for name, value in rewritten_kws - ] - elif rewritten_args: - rewritten_args.append(dtype_var) - else: - rewritten_kws.append(("dtype", dtype_var)) - if thread_data_specification.alignment is not None: - alignment_var = ir.Var( - inst.target.scope, - f"__coop_thread_data_alignment_{next(_GLOBAL_NAME_COUNTER)}__", - inst.loc, - ) - new_block.append( - ir.Assign( - ir.Const(thread_data_specification.alignment, inst.loc), - alignment_var, - inst.loc, - ) - ) - rewritten_kws.append(("alignment", alignment_var)) - - return rewritten_args, rewritten_kws - - def _emit_thread_data_array( - self, new_block: ir.Block, inst: ir.Assign - ) -> None: - """Replace a payload constructor with ``cuda.local.array``.""" - - rewritten_args, rewritten_kws = self._thread_data_array_arguments( - new_block, inst - ) - # Bind cuda.local.array at this call site. Other blocks may - # still need the original ThreadData constructor alias, and - # a later scan must recognize this call as already lowered. - array_fn_var = ir.Var( - inst.target.scope, - f"__coop_thread_data_array_{next(_GLOBAL_NAME_COUNTER)}__", - inst.loc, - ) - module_var = ir.Var( - inst.target.scope, - f"__coop_thread_data_module_{next(_GLOBAL_NAME_COUNTER)}__", - inst.loc, - ) - local_var = ir.Var( - inst.target.scope, - f"__coop_thread_data_local_{next(_GLOBAL_NAME_COUNTER)}__", - inst.loc, - ) - new_block.append( - ir.Assign( - ir.Global( - _next_global_name("thread_data_module"), - _cuda_module, - inst.loc, - ), - module_var, - inst.loc, - ) - ) - new_block.append( - ir.Assign( - ir.Expr.getattr(module_var, "local", inst.loc), - local_var, - inst.loc, - ) - ) - new_block.append( - ir.Assign( - ir.Expr.getattr(local_var, "array", inst.loc), - array_fn_var, - inst.loc, - ) - ) - new_block.append( - ir.Assign( - ir.Expr.call( - array_fn_var, - rewritten_args, - tuple(rewritten_kws), - inst.loc, - ), - inst.target, - inst.loc, - ) - ) - - def _emit_temp_storage(self, new_block: ir.Block, inst: ir.Assign) -> None: - """Replace a scratch descriptor with its shared-memory view. - - A descriptor whose consumers need no scratch becomes ``None``. Other - descriptors select their region from the function's shared backing - array. - """ - - ctor_key = self._resolve_temp_storage_ctor_key(inst.target) - if ctor_key is None: - raise CoopSinglePhaseRewriteError( - f"Missing TempStorage metadata for '{inst.target.name}'." - ) - if ctor_key not in self._func_temp_storage_requirements: - # Validation already required a primitive consumer. Its - # provider needs no scratch, so no array is needed here. - new_block.append( - ir.Assign(ir.Const(None, inst.loc), inst.target, inst.loc) - ) - return - plan = self._finalize_temp_storage_plan_for_var(ctor_key) - backing_var = self._temp_storage_backing_var - if backing_var is None: - raise CoopSinglePhaseRewriteError( - "Missing unified TempStorage backing allocation." - ) - self._emit_array_slice( - new_block, - source_var=backing_var, - target_var=inst.target, - start=plan.base_offset, - stop=plan.base_offset + plan.size_in_bytes, - loc=inst.loc, - ) - - def _emit_provider_call( - self, - new_block: ir.Block, - inst: ir.Assign, - match: _RewriteMatch, - call_invocable: tuple[str, object] | None, - ) -> None: - """Emit a provider call with its operands and reuse barrier. - - The provider ABI decides whether to prepend a scratch pointer. The - group and storage plans must agree on automatic synchronization - before emission of the trailing barrier. - """ - - rewritten_runtime_args = self._prepare_family_runtime_args( - new_block, - match=match, - runtime_args=list(match.runtime_args), - scope=inst.target.scope, - loc=match.loc, - ) - runtime_temp_storage_plan = None - if match.factory_metadata.storage_abi is StorageABI.LEADING_POINTER: - if match.runtime_temp_storage_var is not None: - runtime_temp_storage_arg, runtime_temp_storage_plan = ( - self._runtime_temp_storage_arg_for_call( - new_block, - source_var=match.runtime_temp_storage_var, - call_assign=inst, - ) - ) - else: - ( - runtime_temp_storage_arg, - runtime_temp_storage_plan, - ) = self._implicit_temp_storage_arg_for_call( - new_block, - call_assign=inst, - ) - rewritten_runtime_args.insert(0, runtime_temp_storage_arg) - assert isinstance(inst.value, ir.Expr) - call_func = inst.value.func - if call_invocable is not None: - global_name, invocable = call_invocable - call_func = ir.Var( - inst.target.scope, - f"__coop_single_phase_call_{next(_GLOBAL_NAME_COUNTER)}__", - match.loc, - ) - new_block.append( - ir.Assign( - ir.Global(global_name, invocable, match.loc), - call_func, - match.loc, - ) - ) - new_block.append( - ir.Assign( - ir.Expr.call(call_func, rewritten_runtime_args, (), match.loc), - inst.target, - match.loc, - ) - ) - if ( - runtime_temp_storage_plan is not None - and match.lowering_plan is not None - and match.lowering_plan.temp_storage is not None - and ( - match.lowering_plan.temp_storage.auto_sync - is not runtime_temp_storage_plan.auto_sync - ) - ): - # The group lowering plan and runtime storage plan must agree - # on auto_sync before we choose the storage-reuse barrier. - raise CoopSinglePhaseRewriteError( - "cooperative provider TempStorage automatic " - "synchronization disagrees between the group lowering " - "plan and the descriptor." - ) - if ( - runtime_temp_storage_plan is not None - and runtime_temp_storage_plan.auto_sync - ): - synchronization_scope = match.factory_metadata.synchronization_scope - if match.lowering_plan is not None: - planned_synchronization = match.lowering_plan.synchronization - if planned_synchronization is None: - raise CoopSinglePhaseRewriteError( - "cooperative provider storage requires a " - "synchronization contract." - ) - synchronization_scope = ( - planned_synchronization.storage_reuse_barrier - ) - self._emit_temp_storage_auto_sync( - new_block, - scope=inst.target.scope, - loc=inst.loc, - synchronization_scope=synchronization_scope, - lowering_plan=match.lowering_plan, - ) - - def _remove_unused_factory_arguments(self) -> None: - """Remove eligible factory bindings that no rewritten block uses.""" - - if not self._factory_argument_cleanup_candidates: - return - used_var_names: set[str] = set() - for block in self._func_ir.blocks.values(): - for stmt in block.body: - used_var_names.update( - var.name - for var in stmt.list_vars() - if not isinstance(stmt, ir.Assign) - or var.name != stmt.target.name - ) - for block in self._func_ir.blocks.values(): - block.body = [ - stmt - for stmt in block.body - if not ( - isinstance(stmt, ir.Assign) - and stmt in self._factory_argument_cleanup_candidates - and stmt.target.name not in used_var_names - ) - ] - - def _clear_unused_payload_callees(self) -> None: - """Clear references to the ThreadData constructor after lowering calls. - - For example, ``constructor = coop.ThreadData`` must remain while a - call in another block still uses it. Once all calls through that - binding have been replaced with ``cuda.local.array``, replace the - binding with ``None`` so type inference need not type the marker - function. Follow alias chains until no more bindings can be cleared. - - Only clear a surviving, unambiguous original definition. A generated - assignment or another pre-SSA binding with the same name is not a - constructor alias merely because its name is a cleanup candidate. - Payload arrays and computations using them remain in the IR. - """ - - if not self._payload_callee_cleanup_names: - return - blocks = list(self._func_ir.blocks.values()) - original_definitions = getattr(self._func_ir, "_definitions", {}) - bindings: dict[str, ir.Assign] = {} - for block in blocks: - for stmt in block.body: - if not isinstance(stmt, ir.Assign): - continue - definitions = original_definitions.get(stmt.target.name, ()) - if len(definitions) == 1 and definitions[0] is stmt.value: - bindings[stmt.target.name] = stmt - candidates = set(self._payload_callee_cleanup_names) - while candidates: - used_names = set() - for block in blocks: - for stmt in block.body: - used_names.update( - var.name - for var in stmt.list_vars() - if not isinstance(stmt, ir.Assign) - or var.name != stmt.target.name - ) - retired = False - for name in tuple(candidates): - stmt = bindings.get(name) - if stmt is None or name in used_names: - continue - candidates.remove(name) - if isinstance(stmt.value, ir.Var): - candidates.add(stmt.value.name) - stmt.value = ir.Const(None, stmt.loc) - retired = True - if not retired: - break - - -class _CallRewriting: - """Drive provider-call and payload rewriting across the function.""" - - def _rewrite_calls(self) -> bool: - """Prepare, rewrite, and finish one function, retrying launch facts. - - ``CoopWholeFunctionPlanner.run`` invokes this after its - group-resolution step, even when that step made no changes: a kernel - can use ``ThreadData`` without any group operations. Preparation sees - all original inlined consumers before any block is replaced. Once it - succeeds, stage backing storage, rewrite blocks in label order, and - clean up bindings only after installing every replacement. Block - matching never triggers preparation or handles launch deferral. - - Public group resolution normally obtains launch facts before this - method runs. If a remaining private provider call needs dimensions, - request the configured launch and rescan with a fresh rewrite object, - this time reporting unresolved dimensions as errors. This is a local - rescan within the same compiler attempt. ``require_launch_config`` - makes the facts available synchronously; it does not restart Numba - compilation. Literal-argument requests can separately cause the - dispatcher to start another compiler attempt. - - A standalone device function has no kernel launch of its own. Leave - its deferred private calls for rewriting after inlining into a kernel. - - Returns - ------- - bool - Whether a replacement block was installed in ``state.func_ir``. - - Raises - ------ - CoopSinglePhaseRewriteError - Provider arguments, payloads, or storage violate the rewrite - requirements, or dimensions remain unresolved after the rescan. - RuntimeError - A provider needs the configured launch dimensions, but the - runtime has no launch configuration or tracker for this attempt. - """ - - planner = cast("CoopWholeFunctionPlanner", self) - rewrite = CoopSinglePhaseRewrite(planner.state) - modified = False - - def apply_matches() -> None: - """Rewrite a prepared function; leave a deferred one intact.""" - - nonlocal modified - if not rewrite.prepare_calls_and_storage(planner.state.func_ir): - return - rewrite.begin_rewrite() - for label in sorted(planner.state.func_ir.blocks): - block = planner.state.func_ir.blocks[label] - if rewrite.match( - planner.state.func_ir, - block, - planner.state.typemap, - planner.state.calltypes, - ): - planner.state.func_ir.blocks[label] = rewrite.apply() - modified = True - rewrite.finish_rewrite() - - apply_matches() - if ( - rewrite._deferred_launch_dim_inference - and not planner.is_device_function - ): - require_launch_config(planner.state) - # Rebuild inference and storage plans using the now-visible facts. - rewrite = CoopSinglePhaseRewrite( - planner.state, - allow_launch_dim_deferral=False, - ) - apply_matches() - return modified - - -__all__ = [ - "CoopSinglePhaseRewrite", - "CoopSinglePhaseRewriteError", - "_CallRewriting", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_arguments.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_arguments.py deleted file mode 100644 index 8169e4a1b0b6..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_arguments.py +++ /dev/null @@ -1,396 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Separate provider specialization inputs from device-call operands. - -After group resolution selects a provider, argument validation applies its -registered grammar, infers missing payload and launch facts, and distinguishes -static scalar bindings from runtime controls. The result supplies host-side -factory keywords and device-side IR operands to the remaining rewrite helpers; -this module does not replace the call or allocate its storage. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, cast - -from cuda.coop._core import ArgumentBinding, GroupLoweringPlan - -from ._group_rewriting import GroupRewriteContext -from ._operations import _GROUP_LOWERING_PLAN_KWARG, rewrite_operation -from ._rewrite_support import ( - _UNRESOLVED, - CoopSinglePhaseRewriteError, - _DeferredCoopRewrite, - ir, -) - -if TYPE_CHECKING: - from ._rewrite import CoopSinglePhaseRewrite - - -class _ArgumentRewrite: - """Split provider calls using their registered argument rules. - - This mixin gathers static factory values and ordered runtime operands. - ``CoopSinglePhaseRewrite`` supplies provenance, payload, and launch - inference through the other mixins. - """ - - def _validate_and_split_args( - self, op_name: str, call: ir.Expr, getitem_temp_storage: ir.Var | None - ) -> tuple[ - tuple[ir.Var, ...], - ir.Var | None, - dict[str, object], - tuple[ir.Var, ...], - ]: - """Partition provider arguments into static and runtime inputs. - - The operation registry defines positional arity, factory keywords, - scalar bindings, and optional runtime controls. Static scalar controls - become ``ArgumentBinding.static`` values; unresolved controls stay in - the runtime argument list in registry order. A resolved ``None`` omits - an optional control, whereas ``_UNRESOLVED`` preserves its runtime - operand. Temporary storage is returned separately for later - ABI-specific insertion. - - Pass the call's lowering plan to payload inference so Load/Store can - reuse checked dtype and item-count facts while validating explicit - factory inputs. Infer missing inputs from operands and exact launch - metadata, validate runtime controls, and normalize the ``dim`` alias. - The private plan is also carried in the factory dictionary for the - caller to remove before invoking the provider. Record the variables - that supplied compile-time inputs so ``finish_rewrite`` can remove - their unused assignments after all blocks have been rewritten. The - call expression itself is not rewritten here. - - Parameters - ---------- - op_name : str - Registered operation whose argument contract governs the call. - call : ir.Expr - Untyped IR call with positional and keyword operands. - getitem_temp_storage : ir.Var or None - Storage operand discovered in a subscripted callee. Argument - splitting recognizes it, but descriptor-use validation decides - whether that syntax is permitted. - - Returns - ------- - runtime_args : tuple of ir.Var - Provider operands, excluding the leading storage pointer. - runtime_temp_storage : ir.Var or None - Caller scratch operand, or None for implementation-owned storage. - factory_kwargs : dict - Resolved Python specialization values and ``ArgumentBinding`` - descriptors, plus the private lowering plan until the caller - extracts it. - factory_kw_value_vars : tuple of ir.Var - Original IR variables whose assignments may become unused after - rewriting. These are cleanup candidates, never factory inputs. - They do not correspond one-to-one with ``factory_kwargs``: - inferred keywords have no source variable, while omitted controls - and lowering metadata can leave a variable without a keyword. - Runtime controls keep their variables in ``runtime_args``. - - Raises - ------ - CoopSinglePhaseRewriteError - Arguments are invalid, factory values cannot be resolved, or - required inputs remain unavailable when deferral is disabled. - _DeferredCoopRewrite - Internal signal that exact launch metadata is still needed. - ``CoopSinglePhaseRewrite.prepare_calls_and_storage`` catches it and - preserves the call while ``_CallRewriting._rewrite_calls`` obtains - the kernel launch shape and retries within - ``CoopWholeFunctionPlanner``. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - specification = rewrite_operation(op_name) - if specification is None: - raise CoopSinglePhaseRewriteError( - f"unsupported Numba-CUDA-MLIR operation {op_name!r}" - ) - if call.vararg is not None or call.varkwarg is not None: - raise CoopSinglePhaseRewriteError( - "cooperative group calls do not support *args or **kwargs." - ) - runtime_arg_count = len(call.args) - if runtime_arg_count not in specification.runtime_arg_counts: - expected_csv = ", ".join( - str(v) for v in sorted(specification.runtime_arg_counts) - ) - raise CoopSinglePhaseRewriteError( - f"cooperative group operation {op_name!r} expects a " - f"positional runtime argument count in {{{expected_csv}}}; " - f"got {runtime_arg_count}." - ) - base_runtime_arg_count = min(specification.runtime_arg_counts) - runtime_args = list(call.args[:base_runtime_arg_count]) - factory_kw_value_vars: list[ir.Var] = [] - allowed_factory_kwargs = set(specification.allowed_factory_kwargs) - required_factory_kwargs = specification.required_factory_kwargs - seen_factory_kwargs: set[str] = set() - factory_kwargs: dict[str, object] = {} - runtime_temp_storage = getitem_temp_storage - runtime_factory_kwargs = specification.runtime_factory_kwargs - runtime_factory_kw_prerequisites = dict( - specification.runtime_factory_kw_prerequisites - ) - scalar_binding_kwargs = specification.scalar_binding_kwargs - extra_runtime_arg_count = runtime_arg_count - base_runtime_arg_count - positional_runtime_factory_kwargs = set( - runtime_factory_kwargs[:extra_runtime_arg_count] - ) - seen_runtime_factory_kwargs: set[str] = set() - runtime_factory_kw_vars: dict[str, ir.Var] = {} - runtime_offset_var = None - lowering_plan = None - seen_lowering_plan = False - if runtime_factory_kwargs: - if extra_runtime_arg_count > len(runtime_factory_kwargs): - raise CoopSinglePhaseRewriteError( - f"cooperative group operation {op_name!r} received too " - f"many positional runtime arguments." - ) - for index, name in enumerate( - runtime_factory_kwargs[:extra_runtime_arg_count] - ): - value_var = call.args[base_runtime_arg_count + index] - if name in scalar_binding_kwargs: - value = rewrite._resolve_static_scalar_value(value_var) - if value is not _UNRESOLVED: - if value is not None: - factory_kwargs[name] = ( - value - if isinstance(value, ArgumentBinding) - else ArgumentBinding.static(value) - ) - seen_factory_kwargs.add(name) - if isinstance(value_var, ir.Var): - factory_kw_value_vars.append(value_var) - continue - runtime_args.append(value_var) - factory_kwargs[name] = ( - ArgumentBinding.runtime() - if name in scalar_binding_kwargs - else True - ) - seen_factory_kwargs.add(name) - seen_runtime_factory_kwargs.add(name) - for name, value_var in call.kws: - if name == _GROUP_LOWERING_PLAN_KWARG: - if seen_lowering_plan: - raise CoopSinglePhaseRewriteError( - "cooperative group provider marker received " - "duplicate lowering-plan metadata." - ) - seen_lowering_plan = True - lowering_plan = rewrite._resolve_factory_kwarg_value( - op_name, name, value_var - ) - if not isinstance(lowering_plan, GroupLoweringPlan): - raise CoopSinglePhaseRewriteError( - "cooperative group provider marker carries invalid " - "lowering-plan metadata." - ) - if isinstance(value_var, ir.Var): - factory_kw_value_vars.append(value_var) - continue - if name == "temp_storage" and specification.accepts_temp_storage: - if runtime_temp_storage is not None: - raise CoopSinglePhaseRewriteError( - f"cooperative group operation {op_name!r} received " - f"duplicate temp_storage arguments." - ) - if not isinstance(value_var, ir.Var): - raise CoopSinglePhaseRewriteError( - "cooperative group temp_storage must be a variable." - ) - runtime_temp_storage = value_var - continue - if name == specification.runtime_offset_kwarg: - if ( - runtime_offset_var is not None - or name in seen_factory_kwargs - ): - raise CoopSinglePhaseRewriteError( - f"cooperative group operation {op_name!r} received a " - f"duplicate runtime argument {name!r}." - ) - if not isinstance(value_var, ir.Var): - raise CoopSinglePhaseRewriteError( - f"cooperative group runtime argument {name!r} must " - f"be a variable." - ) - value = rewrite._resolve_static_scalar_value(value_var) - if value is not _UNRESOLVED: - if value is not None: - factory_kwargs[name] = ( - value - if isinstance(value, ArgumentBinding) - else ArgumentBinding.static(value) - ) - seen_factory_kwargs.add(name) - factory_kw_value_vars.append(value_var) - continue - runtime_offset_var = value_var - continue - if name in runtime_factory_kwargs: - if ( - name in seen_runtime_factory_kwargs - or name in runtime_factory_kw_vars - or name in seen_factory_kwargs - or name in positional_runtime_factory_kwargs - ): - raise CoopSinglePhaseRewriteError( - f"cooperative group operation {op_name!r} received a " - f"duplicate runtime argument {name!r}." - ) - if not isinstance(value_var, ir.Var): - raise CoopSinglePhaseRewriteError( - f"cooperative group runtime argument {name!r} must " - f"be a variable." - ) - if name in scalar_binding_kwargs: - value = rewrite._resolve_static_scalar_value(value_var) - if value is not _UNRESOLVED: - if value is not None: - factory_kwargs[name] = ( - value - if isinstance(value, ArgumentBinding) - else ArgumentBinding.static(value) - ) - seen_factory_kwargs.add(name) - factory_kw_value_vars.append(value_var) - continue - runtime_factory_kw_vars[name] = value_var - continue - if name not in allowed_factory_kwargs: - allowed = ", ".join( - sorted( - set(allowed_factory_kwargs) - | set(runtime_factory_kwargs) - ) - ) - raise CoopSinglePhaseRewriteError( - f"cooperative group operation {op_name!r} does not " - f"support factory keyword {name!r}. Allowed keywords " - f"are: {allowed}." - ) - if name in seen_factory_kwargs: - raise CoopSinglePhaseRewriteError( - f"cooperative group operation {op_name!r} received a " - f"duplicate factory keyword {name!r}." - ) - seen_factory_kwargs.add(name) - value = rewrite._resolve_factory_kwarg_value( - op_name, name, value_var - ) - if value is _UNRESOLVED: - raise CoopSinglePhaseRewriteError( - f"Failed to evaluate cooperative group operation " - f"{op_name!r} factory argument {name!r} as a " - f"compile-time constant." - ) - factory_kwargs[name] = value - if isinstance(value_var, ir.Var): - factory_kw_value_vars.append(value_var) - for name in runtime_factory_kwargs: - value_var = runtime_factory_kw_vars.get(name) - if value_var is None: - continue - prerequisite = runtime_factory_kw_prerequisites.get(name) - if ( - prerequisite is not None - and prerequisite not in seen_runtime_factory_kwargs - and prerequisite not in runtime_factory_kw_vars - and prerequisite not in seen_factory_kwargs - ): - raise CoopSinglePhaseRewriteError( - f"cooperative group operation {op_name!r} runtime " - f"argument {name!r} requires {prerequisite!r}." - ) - runtime_args.append(value_var) - factory_kwargs[name] = ( - ArgumentBinding.runtime() - if name in scalar_binding_kwargs - else True - ) - seen_factory_kwargs.add(name) - seen_runtime_factory_kwargs.add(name) - if runtime_offset_var is not None: - runtime_args.append(runtime_offset_var) - rewrite._infer_factory_kwargs_from_payload( - op_name, - runtime_args, - allowed_factory_kwargs, - seen_factory_kwargs, - factory_kwargs, - lowering_plan=lowering_plan, - ) - if specification.validate_runtime_controls is not None: - specification.validate_runtime_controls( - GroupRewriteContext(rewrite), - op_name=op_name, - runtime_args=runtime_args, - factory_kwargs=factory_kwargs, - ) - rewrite._canonicalize_dim_factory_alias( - op_name=op_name, - seen_factory_kwargs=seen_factory_kwargs, - factory_kwargs=factory_kwargs, - ) - rewrite._infer_threads_per_block_from_context( - op_name=op_name, - allowed_factory_kwargs=allowed_factory_kwargs, - seen_factory_kwargs=seen_factory_kwargs, - factory_kwargs=factory_kwargs, - ) - missing = required_factory_kwargs - seen_factory_kwargs - if missing: - if "threads_per_block" in missing: - if rewrite._can_defer_launch_dim_inference(): - raise _DeferredCoopRewrite - if not rewrite._allow_launch_dim_deferral: - other_missing = sorted(missing - {"threads_per_block"}) - other_missing_message = ( - " Also missing required factory keywords: " - f"{', '.join(other_missing)}." - if other_missing - else "" - ) - raise CoopSinglePhaseRewriteError( - f"coop operation '{op_name}' could not infer an " - f"exact positive threads_per_block value because " - f"{rewrite._launch_dim_inference_failure_detail()}. " - f"Use a compile-time constant launch shape or pass " - f"explicit threads_per_block.{other_missing_message}" - ) - missing_csv = ", ".join(sorted(missing)) - raise CoopSinglePhaseRewriteError( - f"coop operation '{op_name}' requires explicit factory " - f"keywords: {missing_csv}." - ) - if ( - runtime_temp_storage is not None - and not specification.accepts_temp_storage - ): - raise CoopSinglePhaseRewriteError( - f"cooperative group operation {op_name!r} does not support " - f"runtime temp_storage." - ) - if lowering_plan is not None: - factory_kwargs[_GROUP_LOWERING_PLAN_KWARG] = lowering_plan - return ( - tuple(runtime_args), - runtime_temp_storage, - factory_kwargs, - tuple(factory_kw_value_vars), - ) - - -__all__ = ["_ArgumentRewrite"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_group_metadata.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_group_metadata.py deleted file mode 100644 index b98c68cb6f42..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_group_metadata.py +++ /dev/null @@ -1,130 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Dispatch operation-specific analysis and operand preparation. - -The shared rewrite recognizes providers and manages storage, while registered -primitive families handle details such as scalar Store boxing. Match analysis -records family metadata before provider compilation. Operand preparation uses -that metadata to emit any needed IR while replacing the call. -""" - -from typing import TYPE_CHECKING, cast - -from ._group_rewriting import GroupRewriteContext -from ._operations import rewrite_operation -from ._rewrite_support import CoopSinglePhaseRewriteError, _RewriteMatch, ir - -if TYPE_CHECKING: - from ._rewrite import CoopSinglePhaseRewrite - - -class _GroupMetadataRewrite: - """Connect call rewriting to each operation's analysis and emission.""" - - def _analyze_family_match( - self, - *, - op_name: str, - runtime_args: tuple[ir.Var, ...], - factory_kwargs: dict[str, object], - ) -> object: - """Collect operation-specific facts before compiling a provider. - - Whole-function call analysis invokes this once after separating - specialization inputs from runtime operands. The family - hook can then check details that the generic rewrite cannot infer, - such as whether Store needs to box a scalar in a one-element array. - Its result is saved on the match for operand preparation during - ``apply``. - - Parameters - ---------- - op_name : str - Registered operation whose analysis hook should run. - runtime_args : tuple of ir.Var - Provider operands in call order after argument splitting. - factory_kwargs : dict of str to object - Resolved specialization inputs shared with the argument splitter. - The family hook may update this dictionary before compilation. - - Returns - ------- - object - Family-specific metadata, or ``None`` when no analysis hook exists. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - specification = rewrite_operation(op_name) - if specification is None: - raise CoopSinglePhaseRewriteError( - f"unsupported Numba-CUDA-MLIR operation {op_name!r}" - ) - if specification.analyze_match is None: - return None - return specification.analyze_match( - GroupRewriteContext(rewrite), - op_name=op_name, - runtime_args=runtime_args, - factory_kwargs=factory_kwargs, - ) - - def _prepare_family_runtime_args( - self, - block: ir.Block, - *, - match: _RewriteMatch, - runtime_args: list[ir.Var], - scope: ir.Scope | None, - loc: ir.Loc, - ) -> list[ir.Var]: - """Emit operation-specific operand conversions while replacing a call. - - ``CoopSinglePhaseRewrite.apply`` invokes this after materializing the - provider. For example, Store may need a scalar copied into a local array - before the generated provider can consume it. The hook uses the metadata - collected during matching and may append statements to ``block``. - - Parameters - ---------- - block : ir.Block - Replacement block receiving any operand-conversion statements. - match : _RewriteMatch - Planned call, including its registered operation and family - metadata. - runtime_args : list of ir.Var - Current operands in provider order, before this family's - conversions. - scope : ir.Scope or None - Scope in which the hook creates temporary variables. - loc : ir.Loc - Original call location attached to generated statements and - diagnostics. - - Returns - ------- - list of ir.Var - Operands to pass to the provider. Without a preparation hook, return - the supplied list unchanged. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - specification = rewrite_operation(match.op_name) - if specification is None: - raise CoopSinglePhaseRewriteError( - f"unsupported Numba-CUDA-MLIR operation {match.op_name!r}" - ) - if specification.prepare_runtime_args is None: - return runtime_args - return specification.prepare_runtime_args( - GroupRewriteContext(rewrite), - block, - match=match, - runtime_args=runtime_args, - scope=scope, - loc=loc, - ) - - -__all__ = ["_GroupMetadataRewrite"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_invocables.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_invocables.py deleted file mode 100644 index 5abad5616c8b..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_invocables.py +++ /dev/null @@ -1,304 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Batch provider compilation and reuse callable specializations. - -Whole-function call analysis collects the kernel's distinct provider -specializations before materializing any invocable. On the successful batching -path, their generated C++ shares one translation unit and, on a cache miss, -one NVRTC compilation to LTO IR, avoiding a separate compiler invocation for -each cooperative primitive. Every resulting invocable retains the shared -artifact for linking. - -This reduces provider-compilation overhead; it is not a promise of exactly one -NVRTC invocation per kernel. Cache hits can avoid compilation, equivalent -specializations can collapse to one provider, and unavailable or failed -bundling falls back to individual materialization. Scratch layouts come from -the provider compilation and share its cache entry. Extra support images and -the kernel's own compilation remain separate work. -""" - -from __future__ import annotations - -import hashlib -from collections.abc import Callable -from typing import TYPE_CHECKING, Any, cast - -from .._types import ( - _hash_symbol_value, - collect_specializations, - make_invocable_from_specialization, - prepare_ltoir_bundle, -) -from ._operations import FactoryOperation -from ._rewrite_support import CoopSinglePhaseRewriteError, _RewriteMatch - -if TYPE_CHECKING: - from ._rewrite import CoopSinglePhaseRewrite - - -class _InvocableRewrite: - """Reuse compiled providers across matches and planner retries. - - The mixin validates factory results, caches them in compiler state, and - optionally compiles several specializations together. - """ - - @staticmethod - def _invocable_cache_key( - factory: Callable[..., Any], - factory_metadata: FactoryOperation, - factory_kwargs: dict[str, object], - ) -> tuple[str, tuple[tuple[str, str, str], ...]]: - """Identify a provider specialization within one compiler state. - - Include factory object identity and the registered storage, execution, - and synchronization contracts so distinct providers cannot share an - invocable merely because their operation names match. Keyword order is - irrelevant; each value contributes its Python type and structural - symbol hash. The object identity makes this unsuitable as a persistent - cache key across processes. - - Parameters - ---------- - factory : callable - Registered host-side provider factory whose identity - partitions the cache. It is not invoked here. - factory_metadata : FactoryOperation - Required registration carrying the operation name, namespace, - and ABI and scope contracts for ``factory``. - factory_kwargs : dict of str to object - Resolved specialization inputs, without lowering-plan metadata. - - Returns - ------- - tuple - Provider identity string and sorted keyword-name, value-type, - and value-digest triples used by the rewrite and - compiler-state caches. - """ - - def cache_component(name: str, value: object) -> tuple[str, str, str]: - hasher = hashlib.sha1() - _hash_symbol_value(hasher, value) - value_type = f"{type(value).__module__}.{type(value).__qualname__}" - return (name, value_type, hasher.hexdigest()) - - # This cache belongs to one compiler state. Object identity keeps - # separately registered providers apart even when their operation - # names and specialization arguments are identical. - provider_contract = ( - factory_metadata.namespace, - factory_metadata.storage_abi.value, - factory_metadata.execution_scope.value, - factory_metadata.synchronization_scope.value, - ) - return ( - ( - f"{type(factory).__module__}.{type(factory).__qualname__}:" - f"{id(factory)}:{factory_metadata.operation}:{provider_contract!r}" - ), - tuple( - sorted( - ( - cache_component(name, value) - for name, value in factory_kwargs.items() - ) - ) - ), - ) - - @staticmethod - def _validate_invocable( - invocable: object, factory_metadata: FactoryOperation - ) -> None: - """Check a factory result against its registered provider contract. - - Construction and compiler-cache lookup both use this check before an - invocable enters the rewrite-local cache. Require a callable exposing - link files, then compare its storage ABI and thread scopes with the - registry. This checks declarations, not generated device code. - - Parameters - ---------- - invocable : object - Factory result or cached provider to check. - factory_metadata : FactoryOperation - The factory's registered storage and synchronization contract. - - Raises - ------ - CoopSinglePhaseRewriteError - The result lacks the invocable interface or declares - incompatible storage, execution, or synchronization metadata. - """ - - op_name = factory_metadata.operation - if not callable(invocable) or not hasattr(invocable, "files"): - raise CoopSinglePhaseRewriteError( - f"coop single-phase factory for '{op_name}' did not produce " - f"a coop invocable; got {type(invocable)!r}." - ) - expected_contract = { - "storage_abi": factory_metadata.storage_abi, - "execution_scope": factory_metadata.execution_scope, - "synchronization_scope": factory_metadata.synchronization_scope, - } - mismatches = [] - for name, expected in expected_contract.items(): - observed = getattr(invocable, name, None) - try: - observed = type(expected)(observed) - except (TypeError, ValueError): - pass - if observed != expected: - mismatches.append( - f"{name}={observed!r} (registered {expected.value!r})" - ) - if mismatches: - details = ", ".join(mismatches) - raise CoopSinglePhaseRewriteError( - f"coop provider '{op_name}' returned incompatible metadata: " - f"{details}." - ) - - def _prepare_ltoir_bundle_for_matches( - self, matches: list[_RewriteMatch] - ) -> None: - """Prepare a shared LTO IR bundle when distinct matches permit it. - - Function preparation calls this after analyzing calls and validating - storage uses, before materializing providers. A shared translation - unit lets NVRTC compile all distinct provider bodies together, - avoiding per-primitive compiler startup and - repeated header processing. The resulting invocables share the LTO - artifact; this does not combine the kernel's own compilation with it. - - Collect factory specializations without immediately building each - invocable and deduplicate identical matches. Each collected algorithm - already stores its thread dimensions. Bundling is attempted only for - at least two unique matches. Launch deferral happens during call - analysis, so a retry reaches bundling before any provider is emitted. - - This is an optional compilation optimization. Clear the previous - bundle lookup first; a collection-count mismatch or an import, OS, or - runtime failure leaves it empty so ``_materialize_invocable`` can call - factories individually. Other exceptions propagate. Successful - preparation records specializations by invocable cache key without - replacing function IR. If those specializations coalesce to one - algorithm, bundle preparation may produce no shared bundle; retain the - collected specializations for individual materialization anyway. - - Parameters - ---------- - matches : list of _RewriteMatch - Validated provider calls in whole-function scan order. - - Returns - ------- - None - ``_prebundled_specializations`` holds the prepared - specializations, which may share a bundle. Early exits and - caught failures leave it empty so materialization can invoke - factories directly. - """ - - self._prebundled_specializations = {} - if not matches: - return - matches_by_specialization: dict[ - tuple[str, tuple[tuple[str, str, str], ...]], _RewriteMatch - ] = {} - for match in matches: - key = self._invocable_cache_key( - match.factory, - match.factory_metadata, - match.factory_kwargs, - ) - if key not in matches_by_specialization: - matches_by_specialization[key] = match - # Repeated calls already reuse one invocable; bundling only helps when - # there are distinct specializations to compile together. - if len(matches_by_specialization) < 2: - return - try: - with collect_specializations() as collected: - for match in matches_by_specialization.values(): - _ = match.factory(**match.factory_kwargs) - if len(collected) != len(matches_by_specialization): - return - prepare_ltoir_bundle(collected) - self._prebundled_specializations = dict( - zip(matches_by_specialization, collected) - ) - except (ImportError, OSError, RuntimeError): - self._prebundled_specializations = {} - - def _materialize_invocable( - self, match: _RewriteMatch - ) -> tuple[object, bool]: - """Obtain a callable specialization for one validated match. - - Consult the rewrite-local cache, then the cache in compiler metadata - so a fresh rewrite during launch retries can reuse prior work. On a - miss, construct an invocable from a prepared specialization or - evaluate the factory directly. Newly constructed and compiler-cache - invocables must agree with the provider's storage and synchronization - contracts before use. Successful construction populates both caches. - - Parameters - ---------- - match : _RewriteMatch - Provider factory, registered contract, and resolved - specialization keywords. Lowering-plan metadata has already - been removed. - - Returns - ------- - invocable : object - Callable provider with link files and ABI metadata. - created : bool - True when this call constructed the invocable, including from - a prepared specialization; False for either cache hit. - - Raises - ------ - CoopSinglePhaseRewriteError - Construction fails, or the result breaks the registered contract. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - key = self._invocable_cache_key( - match.factory, - match.factory_metadata, - match.factory_kwargs, - ) - if key in rewrite._invocable_cache: - return (rewrite._invocable_cache[key], False) - compile_cache = rewrite._state.metadata.setdefault( - "__cuda_coop_numba_mlir_invocable_cache__", {} - ) - if key in compile_cache: - invocable = compile_cache[key] - self._validate_invocable(invocable, match.factory_metadata) - rewrite._invocable_cache[key] = invocable - return (invocable, False) - try: - prebundled = self._prebundled_specializations.get(key) - if prebundled is not None: - invocable = make_invocable_from_specialization(prebundled) - else: - invocable = match.factory(**match.factory_kwargs) - except Exception as e: - raise CoopSinglePhaseRewriteError( - f"Failed to evaluate coop single-phase factory at compile " - f"time for '{match.op_name}'." - ) from e - self._validate_invocable(invocable, match.factory_metadata) - rewrite._invocable_cache[key] = invocable - compile_cache[key] = invocable - return (invocable, True) - - -__all__ = ["_InvocableRewrite"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_launch.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_launch.py deleted file mode 100644 index 755f785620d8..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_launch.py +++ /dev/null @@ -1,256 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Reconcile provider block dimensions with the configured kernel launch. - -Provider specialization needs the exact block shape; a launch bound only gives -an upper limit. These helpers read launch metadata, normalize explicit shapes, -and mark work that must wait for the planner to request a launch -configuration. Device helpers retain unresolved calls until inlining supplies -their caller's launch context. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, cast - -from ._parameters import normalize_dim_param -from ._rewrite_support import CoopSinglePhaseRewriteError, _DeferredCoopRewrite - -if TYPE_CHECKING: - from ._rewrite import CoopSinglePhaseRewrite - - -class _LaunchRewrite: - """Resolve exact block dimensions before specializing a provider.""" - - def _infer_threads_per_block_from_context( - self, - *, - op_name: str, - allowed_factory_kwargs: set[str], - seen_factory_kwargs: set[str], - factory_kwargs: dict[str, object], - ) -> None: - """Infer or check the block shape using exact launch metadata. - - Only operations accepting ``threads_per_block`` participate. Compare - normalized three-dimensional shapes when both explicit and launch - values are available; equal thread counts alone do not imply equal - shapes. Malformed explicit dimensions are left for later factory - validation. Launch bounds are never treated as an exact shape. - - An explicit dimension may still need to wait: a device helper or a - kernel with a launch tracker must reconcile it after inlining or after - the whole-function planner requests launch metadata. - - Parameters - ---------- - op_name : str - Operation name used in mismatch diagnostics. - allowed_factory_kwargs : set of str - Keywords accepted by the operation's factory. - seen_factory_kwargs : set of str - Resolved keyword names; receives an inferred block-shape keyword. - factory_kwargs : dict of str to object - Resolved values; updated with inferred ``threads_per_block``. - - Returns - ------- - None - Fill a missing shape from launch metadata; never overwrite one. - - Raises - ------ - CoopSinglePhaseRewriteError - Explicit dimensions disagree with the exact kernel launch shape. - _DeferredCoopRewrite - Internal signal that explicit dimensions need pending launch - metadata. It propagates through argument validation to - ``CoopSinglePhaseRewrite.prepare_calls_and_storage``, which - preserves the IR for ``_CallRewriting._rewrite_calls`` to request - the kernel launch shape and retry within - ``CoopWholeFunctionPlanner``. - """ - - if "threads_per_block" not in allowed_factory_kwargs: - return - threads_per_block = self._infer_threads_per_block_from_launch_config() - if threads_per_block is None: - if ( - "threads_per_block" in seen_factory_kwargs - and self._can_defer_explicit_launch_dim_reconciliation() - ): - raise _DeferredCoopRewrite - return - if "threads_per_block" in seen_factory_kwargs: - explicit_threads_per_block = factory_kwargs["threads_per_block"] - try: - explicit_dim = normalize_dim_param(explicit_threads_per_block) - launch_dim = normalize_dim_param(threads_per_block) - except (TypeError, ValueError): - return - if explicit_dim != launch_dim: - launch_block = self._launch_block_from_context() - raise CoopSinglePhaseRewriteError( - f"cuda.coop factory '{op_name}' received " - f"threads_per_block={explicit_threads_per_block!r}, but " - f"the exact kernel launch block is {launch_block!r}. " - f"Make threads_per_block match the launch block or omit " - f"it to infer the dimension." - ) - return - factory_kwargs["threads_per_block"] = threads_per_block - seen_factory_kwargs.add("threads_per_block") - - def _can_defer_explicit_launch_dim_reconciliation(self) -> bool: - """Record whether explicit dimensions need pending launch metadata. - - Device functions acquire their launch shape from the caller after - inlining. Configured kernels retain a launch tracker until the planner - requests an exact shape. Either case permits deferral when enabled on - this rewrite object; without either signal, this helper returns False. - - This predicate has a side effect: a positive result latches - ``_deferred_launch_dim_inference`` so matching preserves descriptors - and the whole-function planner knows that it must retry. - - Returns - ------- - bool - Whether to wait for launch metadata before comparing dimensions. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - metadata = getattr(rewrite._state, "metadata", {}) or {} - targetoptions = metadata.get("targetoptions", {}) or {} - # Configured kernel launches carry a tracker until the whole-function - # planner requests the exact block. Device functions defer to that - # same planner after inlining into their kernel caller. - should_defer = rewrite._allow_launch_dim_deferral and ( - bool(targetoptions.get("device", False)) - or metadata.get("launch_config_tracker") is not None - ) - self._deferred_launch_dim_inference |= should_defer - return should_defer - - def _launch_block_from_context(self) -> object: - """Read the raw launch block shape, or ``None`` if unavailable.""" - - rewrite = cast("CoopSinglePhaseRewrite", self) - metadata = getattr(rewrite._state, "metadata", {}) or {} - targetoptions = metadata.get("targetoptions", {}) or {} - launch_config = targetoptions.get("__launch_config__") - if not isinstance(launch_config, dict): - return None - return launch_config.get("block") - - def _launch_dim_inference_failure_detail(self) -> str: - """Explain which launch metadata is missing or invalid. - - The argument validator includes this detail in its final diagnostic. A - launch bound is reported separately because it is only a limit. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - metadata = getattr(rewrite._state, "metadata", {}) or {} - targetoptions = metadata.get("targetoptions", {}) or {} - if "__launch_config__" not in targetoptions: - detail = "no __launch_config__ metadata was provided" - else: - launch_config = targetoptions["__launch_config__"] - if not isinstance(launch_config, dict): - detail = ( - f"__launch_config__ metadata is invalid: {launch_config!r}" - ) - elif "block" not in launch_config: - detail = ( - "__launch_config__ metadata contains no block shape: " - f"{launch_config!r}" - ) - else: - detail = ( - "launch metadata reported invalid " - f"block={launch_config['block']!r}" - ) - if "launch_bounds" in targetoptions: - detail += ( - f"; launch_bounds={targetoptions['launch_bounds']!r} " - "is only an upper bound, not an exact launch shape" - ) - return detail - - def _infer_threads_per_block_from_launch_config( - self, - ) -> int | tuple[int, int] | tuple[int, int, int] | None: - """Normalize a launch shape and omit trailing unit dimensions. - - Invalid or missing metadata yields ``None`` so the caller can decide - whether to request a configured launch or report a diagnostic. - """ - - block = self._launch_block_from_context() - if isinstance(block, list): - block = tuple(block) - try: - x, y, z = normalize_dim_param(block) - except (TypeError, ValueError): - return None - if z != 1: - return x, y, z - if y != 1: - return x, y - return x - - def _can_defer_launch_dim_inference(self) -> bool: - """Flag unresolved launch dimensions for the call rewrite to retry.""" - - rewrite = cast("CoopSinglePhaseRewrite", self) - should_defer = ( - rewrite._allow_launch_dim_deferral - and self._infer_threads_per_block_from_launch_config() is None - ) - self._deferred_launch_dim_inference |= should_defer - return should_defer - - @staticmethod - def _canonicalize_dim_factory_alias( - *, - op_name: str, - seen_factory_kwargs: set[str], - factory_kwargs: dict[str, object], - ) -> None: - """Give later dimension inference one keyword for the block shape. - - Argument splitting accepts ``dim`` as an alias of ``threads_per_block``. - It calls this before inferring the launch shape so required-keyword - checks and the provider factory both see the canonical name. Reject a - call that supplies both spellings; later validation checks the dimension - value. - - Parameters - ---------- - op_name : str - Operation name used in the duplicate-keyword diagnostic. - seen_factory_kwargs : set of str - Resolved keyword names. Replace ``dim`` with ``threads_per_block`` - in place when the alias is present. - factory_kwargs : dict of str to object - Corresponding specialization values. Move the alias's value to the - canonical key without changing it. - """ - - if "dim" not in seen_factory_kwargs: - return - if "threads_per_block" in seen_factory_kwargs: - raise CoopSinglePhaseRewriteError( - f"cuda.coop factory '{op_name}' received both " - f"'threads_per_block' and its 'dim' alias; provide only one." - ) - factory_kwargs["threads_per_block"] = factory_kwargs.pop("dim") - seen_factory_kwargs.remove("dim") - seen_factory_kwargs.add("threads_per_block") - - -__all__ = ["_LaunchRewrite"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_load_store.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_load_store.py deleted file mode 100644 index e2b535ae46c5..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_load_store.py +++ /dev/null @@ -1,819 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Apply Load/Store-specific payload and scalar rules during call rewriting. - -Registered hooks merge planned or inferred per-thread payload shape and dtype, -validate optional counts and defaults, and recognize scalar Store operands. -Scalar Stores need a one-element local array for the provider ABI. Emission -checks the exact dtype before boxing so assignment cannot hide a conversion. -Unknown scalar types remain the responsibility of the later compiler typing -pass, and these helpers do not prove runtime bounds or thread uniformity. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from numbers import Integral -from typing import TYPE_CHECKING, NoReturn - -from numba_cuda_mlir import cuda as _cuda_module - -from cuda.coop._core import ( - ArgumentBinding, - BindingKind, - GroupLoadStoreSemantics, -) - -from ._group_rewriting import GroupRewriteContext -from ._rewrite_payload import PayloadInference -from ._rewrite_support import ( - _GLOBAL_NAME_COUNTER, - _UNRESOLVED, - CoopSinglePhaseRewriteError, - _dtype_values_match, - _next_global_name, - _RewriteMatch, - ir, -) - -if TYPE_CHECKING: - from numba_cuda_mlir.numba_cuda.types import Function - - -@dataclass(frozen=True) -class _LoadStoreMatchMetadata: - """Carry the scalar Store boxing decision from analysis to emission. - - A true ``box_root_store_scalar`` asks the operand hook to replace the - scalar with a local array of the inferred dtype and extent. - """ - - box_root_store_scalar: bool = False - - -class _ExactStoreScalar: - """Check a scalar's exact dtype before boxing can introduce a cast. - - Parameters - ---------- - dtype : object - Payload compiler dtype that the scalar must match during typing. - """ - - def __init__(self, dtype: object) -> None: - self.dtype = dtype - self._numba_type = None - - @property - def _numba_type_(self) -> Function: - """Build the compiler function type for an exact-dtype identity check. - - Scalar Store lowering boxes a value into a local array, whose - assignment could silently cast a runtime scalar to the destination - dtype. This callable's overload checks the argument type before that - assignment and returns an inline identity implementation only when the - types match. Compare the scalar types inside literal wrappers. - - The overload template and resulting function type are created lazily - and cached on this validator instance. The callable itself is a - compiler hook; it does not perform the identity operation in Python. - - Returns - ------- - numba_types.Function - Function type with an overload that raises ``TypingError`` for - a dtype mismatch; otherwise, it returns the input unchanged. - """ - - import numba_cuda_mlir.numba_cuda.types as _numba_types - from numba_cuda_mlir.extending import ( - _NumbaCudaMlirOverloadFunctionTemplate, - ) - from numba_cuda_mlir.numba_cuda.core.errors import TypingError - from numba_cuda_mlir.numba_cuda.typing.templates import ( - make_overload_template, - ) - - if self._numba_type is None: - dtype = self.dtype - - def validate(value): - actual_dtype = getattr(value, "literal_type", value) - if actual_dtype != dtype: - raise TypingError( - f"cuda.coop.numba_mlir.store value dtype " - f"{actual_dtype} does not match destination dtype " - f"{dtype}" - ) - - def impl(value): - return value - - return impl - - template = make_overload_template( - self, - validate, - {"no_cpython_wrapper": True, "nopython": True}, - strict=True, - inline="always", - prefer_literal=False, - base=_NumbaCudaMlirOverloadFunctionTemplate, - ) - self._numba_type = _numba_types.Function(template) - return self._numba_type - - def __call__(self, value: object) -> NoReturn: - raise RuntimeError( - "Store scalar validation requires device compilation" - ) - - -class _LoadStoreRewrite: - """Apply Load/Store rules through the operation registry's hooks.""" - - @staticmethod - def _validate_oob_default( - context: GroupRewriteContext, - *, - runtime_args: list[ir.Var], - factory_kwargs: dict[str, object], - ) -> None: - """Check Load padding against the inferred payload dtype. - - Static padding is coerced with its scalar provenance and replaced by a - normalized static binding. Runtime padding must have the payload's - exact supported numeric dtype; it is not implicitly narrowed or - widened. Locate that operand after the two array operands and any - runtime valid-item count. If its type is still unknown, leave - validation for later typing. - - Parameters - ---------- - context : GroupRewriteContext - Access to scalar provenance and available IR variable types. - runtime_args : list of ir.Var - Provider operands in the order produced by argument splitting. - factory_kwargs : dict of str to object - Inferred dtype and control bindings. A static ``oob_default`` - binding is replaced in place by its normalized value. - - Returns - ------- - None - Omitted bindings need no validation. Other bindings are checked - using the dtype information currently available. - - Raises - ------ - CoopSinglePhaseRewriteError - Static padding cannot be converted under the scalar rules, a - runtime operand is missing, or its known dtype is unsupported - or mismatched. - """ - - binding = factory_kwargs.get("oob_default") - if not isinstance(binding, ArgumentBinding) or ( - binding.kind is BindingKind.OMITTED - ): - return - if binding.kind is BindingKind.STATIC: - from ._parameters import coerce_static_scalar - - payload_dtype = factory_kwargs.get("dtype") - if payload_dtype is None: - raise CoopSinglePhaseRewriteError( - "cuda.coop.numba_mlir.load requires an inferred dtype " - "before validating oob_default" - ) - provenance = context.static_scalar_provenance(binding.value) - source_dtype = ( - None - if provenance is _UNRESOLVED or provenance is None - else provenance.dtype - ) - try: - normalized = coerce_static_scalar( - binding.value, - payload_dtype, - operation="load", - parameter="oob_default", - source_dtype=source_dtype, - ) - except (TypeError, ValueError) as exc: - raise CoopSinglePhaseRewriteError(str(exc)) from exc - factory_kwargs["oob_default"] = ArgumentBinding.static(normalized) - return - - runtime_index = 2 - valid_items = factory_kwargs.get("num_valid_items") - if ( - isinstance(valid_items, ArgumentBinding) - and valid_items.kind is BindingKind.RUNTIME - ): - runtime_index += 1 - if runtime_index >= len(runtime_args) or not isinstance( - runtime_args[runtime_index], ir.Var - ): - raise CoopSinglePhaseRewriteError( - "cuda.coop.numba_mlir.load runtime oob_default is missing " - "its runtime value" - ) - - value_var = runtime_args[runtime_index] - value_dtype = context.numba_type(value_var) - if value_dtype is None: - value_dtype = context.dtype(value_var) - if value_dtype is None: - return - - from ._parameters import _validate_common_numeric_dtype - - try: - value_dtype = _validate_common_numeric_dtype( - value_dtype, - operation="load", - parameter="oob_default", - ) - except (TypeError, ValueError) as exc: - raise CoopSinglePhaseRewriteError(str(exc)) from exc - payload_dtype = factory_kwargs.get("dtype") - if payload_dtype is not None and not _dtype_values_match( - value_dtype, - payload_dtype, - ): - raise CoopSinglePhaseRewriteError( - f"cuda.coop.numba_mlir.load runtime oob_default dtype " - f"{value_dtype} does not match payload dtype {payload_dtype}" - ) - - @staticmethod - def _validate_load_store_runtime_controls( - context: GroupRewriteContext, - *, - op_name: str, - runtime_args: list[ir.Var], - factory_kwargs: dict[str, object], - ) -> None: - """Check Load/Store controls before provider creation. - - Static valid-item counts must be integral rather than boolean. Static - offsets must additionally be nonnegative and fit signed 64-bit - storage. For runtime valid-item counts and offsets, check known dtypes - while leaving unknown types for later inference. This method does not - prove runtime value ranges, memory bounds, or group uniformity, nor - does it enforce the static valid-item count's tile bound. - - The operand cursor follows the provider ABI: memory and payload first, - then runtime valid-item count, runtime padding, and optional offset. - Load padding receives its separate dtype/coercion check. - - Parameters - ---------- - context : GroupRewriteContext - Known compiler types and scalar provenance for each operand. - op_name : str - Registered ``"load"`` or ``"store"`` operation. - runtime_args : list of ir.Var - Operands already ordered by argument splitting. - factory_kwargs : dict of str to object - Specialization values and scalar bindings; static Load padding - may be normalized in place. - - Returns - ------- - None - Known control values and types satisfy the checks. - - Raises - ------ - CoopSinglePhaseRewriteError - The operation is unsupported, or a known control is invalid. - """ - - if op_name not in {"load", "store"}: - raise CoopSinglePhaseRewriteError( - f"unsupported Numba-CUDA-MLIR operation {op_name!r}" - ) - - valid_items = factory_kwargs.get("num_valid_items") - if ( - isinstance(valid_items, ArgumentBinding) - and valid_items.kind is BindingKind.STATIC - and ( - isinstance(valid_items.value, bool) - or not isinstance(valid_items.value, Integral) - ) - ): - raise CoopSinglePhaseRewriteError( - f"coop {op_name} valid_items must be an integer" - ) - - offset = factory_kwargs.get("offset") - if ( - isinstance(offset, ArgumentBinding) - and offset.kind is BindingKind.STATIC - ): - if isinstance(offset.value, bool) or not isinstance( - offset.value, Integral - ): - raise CoopSinglePhaseRewriteError( - f"coop {op_name} offset must be an integer" - ) - normalized_offset = int(offset.value) - if normalized_offset < 0: - raise CoopSinglePhaseRewriteError( - "coop load/store static offset must be nonnegative" - ) - if normalized_offset > (1 << 63) - 1: - raise CoopSinglePhaseRewriteError( - "coop load/store static offset must fit " - "a signed 64-bit integer" - ) - - checks: list[tuple[str, int]] = [] - cursor = 2 - if ( - isinstance(valid_items, ArgumentBinding) - and valid_items.kind is BindingKind.RUNTIME - ): - checks.append(("valid_items", cursor)) - cursor += 1 - oob_default = factory_kwargs.get("oob_default") - if ( - isinstance(oob_default, ArgumentBinding) - and oob_default.kind is BindingKind.RUNTIME - ): - cursor += 1 - if len(runtime_args) > cursor: - checks.append(("offset", cursor)) - - from ._parameters import _validate_runtime_integer_dtype - - for parameter, index in checks: - if index >= len(runtime_args) or not isinstance( - runtime_args[index], ir.Var - ): - raise CoopSinglePhaseRewriteError( - f"coop {op_name} {parameter} must be an integer" - ) - value_type = context.numba_type(runtime_args[index]) - if value_type is None: - value_type = context.dtype(runtime_args[index]) - if value_type is None: - continue - try: - _validate_runtime_integer_dtype( - value_type, - operation=op_name, - parameter=parameter, - ) - except TypeError as exc: - raise CoopSinglePhaseRewriteError(str(exc)) from exc - - if op_name == "load": - _LoadStoreRewrite._validate_oob_default( - context, - runtime_args=runtime_args, - factory_kwargs=factory_kwargs, - ) - - @staticmethod - def _infer_load_store_payload( - context: GroupRewriteContext, - inference: PayloadInference, - ) -> None: - """Reconcile Load/Store payload shape and dtype. - - Planned calls reuse their checked dtype and extent while still merging - explicit factory inputs and populating constructor allocation state. - Calls without a Load/Store plan infer these facts from their operands. - - Without a supported Load/Store plan, array payloads supply their - static extent; scalars imply one item per thread. Prefer the memory - element dtype when available, while checking it against any known - payload dtype. For an untyped Store array, inspect typed writes before - recording a destination-derived dtype so an incompatible producer - cannot be hidden by the destination type. This fallback also checks - static scalar coercion against the destination dtype; known runtime - scalar dtypes must match it exactly. - - Both paths record dtypes through payload aliases for later constructor - lowering. This phase updates inference metadata; boxing a scalar into - the provider's array operand happens during runtime-argument emission. - - Parameters - ---------- - context : GroupRewriteContext - Access to payload, dtype, write, and scalar provenance. - inference : PayloadInference - Mutable Load/Store inference state. Runtime operands are - memory then payload; factory keywords and payload dtype caches - may be updated. - - Returns - ------- - None - Known shape and dtype facts are merged into the inference state. - - Raises - ------ - CoopSinglePhaseRewriteError - Known memory and payload dtypes conflict, static scalar - conversion fails, or inferred values violate the supported - numeric contract. - """ - - payload_var, payload_specification = inference.candidate(1) - plan = inference.lowering_plan - if ( - plan is not None - and plan.unsupported is None - and isinstance(plan.call.operation, GroupLoadStoreSemantics) - ): - from .._lowering._core import NumbaMlirCoreAdapter - - semantics = plan.call.operation - dtype = NumbaMlirCoreAdapter().normalize_dtype(semantics.dtype) - inference.infer_kwarg( - "items_per_thread", semantics.items_per_thread - ) - inference.infer_kwarg( - "items_per_thread", - payload_specification.items_per_thread - if payload_specification is not None - else 1, - ) - inference.infer_kwarg("dtype", dtype) - if payload_specification is not None and payload_var is not None: - # Allocation still needs constructor and alias state. The plan - # has already checked memory, payload writes and scalar widths. - context.record_thread_data_dtype(payload_var, dtype) - return - - memory_var = ( - inference.runtime_args[0] if inference.runtime_args else None - ) - memory_dtype = ( - context.dtype(memory_var) - if isinstance(memory_var, ir.Var) - else None - ) - - payload_is_array = payload_specification is not None - if payload_specification is None: - payload_dtype = ( - context.dtype(payload_var) - if isinstance(payload_var, ir.Var) - else None - ) - inference.infer_kwarg("items_per_thread", 1) - inference.infer_kwarg( - "dtype", - memory_dtype if memory_dtype is not None else payload_dtype, - ) - else: - inference.infer_kwarg( - "items_per_thread", payload_specification.items_per_thread - ) - payload_dtype = payload_specification.dtype - if payload_dtype is None and payload_var is not None: - payload_dtype = context.dtype(payload_var) - if ( - inference.op_name == "store" - and payload_dtype is None - and payload_var is not None - ): - payload_dtype = context.infer_thread_data_write_dtype( - payload_var - ) - inferred_dtype = ( - memory_dtype if memory_dtype is not None else payload_dtype - ) - if inferred_dtype is None: - inferred_dtype = inference.factory_value("dtype") - inference.infer_kwarg("dtype", inferred_dtype) - if inferred_dtype is not None and payload_var is not None: - context.record_thread_data_dtype(payload_var, inferred_dtype) - - from ._parameters import _validate_common_numeric_dtype - - if inference.op_name == "store" and not payload_is_array: - provenance = context.static_scalar_provenance(payload_var) - if provenance is not _UNRESOLVED and memory_dtype is not None: - from ._parameters import coerce_static_scalar - - try: - coerce_static_scalar( - provenance.value, - memory_dtype, - operation="store", - parameter="value", - source_dtype=provenance.dtype, - ) - except (TypeError, ValueError) as exc: - raise CoopSinglePhaseRewriteError(str(exc)) from exc - payload_dtype = memory_dtype - - provider_dtype = inference.factory_value("dtype") - - if provider_dtype is not None: - try: - _validate_common_numeric_dtype( - provider_dtype, operation=inference.op_name - ) - except (TypeError, ValueError) as exc: - raise CoopSinglePhaseRewriteError(str(exc)) from exc - if payload_dtype is not None: - try: - _validate_common_numeric_dtype( - payload_dtype, - operation=inference.op_name, - ) - except (TypeError, ValueError) as exc: - raise CoopSinglePhaseRewriteError(str(exc)) from exc - if ( - memory_dtype is not None - and payload_dtype is not None - and not _dtype_values_match(memory_dtype, payload_dtype) - ): - if inference.op_name != "store" or payload_is_array: - raise CoopSinglePhaseRewriteError( - f"cuda.coop.numba_mlir.{inference.op_name} memory dtype " - f"{memory_dtype} does not match payload dtype " - f"{payload_dtype}" - ) - raise CoopSinglePhaseRewriteError( - f"cuda.coop.numba_mlir.store value dtype {payload_dtype} " - f"does not match destination dtype {memory_dtype}" - ) - - -def infer_load_store_payload( - context: GroupRewriteContext, - inference: PayloadInference, -) -> None: - """Expose Load/Store payload inference to the operation registry.""" - - _LoadStoreRewrite._infer_load_store_payload(context, inference) - - -def validate_load_store_runtime_controls( - context: GroupRewriteContext, - *, - op_name: str, - runtime_args: list[ir.Var], - factory_kwargs: dict[str, object], -) -> None: - """Expose Load/Store scalar-control checks to the operation registry.""" - - _LoadStoreRewrite._validate_load_store_runtime_controls( - context, - op_name=op_name, - runtime_args=runtime_args, - factory_kwargs=factory_kwargs, - ) - - -def analyze_load_store_match( - context: GroupRewriteContext, - *, - op_name: str, - runtime_args: tuple[ir.Var, ...], - factory_kwargs: dict[str, object], -) -> _LoadStoreMatchMetadata: - """Consume Load/Store markers and identify scalar boxing. - - Remove the private ``_common_root_operation`` and ``_group_root_store`` - markers from factory keywords so they are not passed to the provider. - The common-API marker requires supported numeric dtypes for memory and - array operands. Scalar inference and the later typing guard handle scalar - Store values, so they need no array payload specification here. - - A public group Store with no recognized per-thread array records a boxing - request for either the common or backend-qualified API. The helper - ``prepare_load_store_runtime_args`` supplies the required array operand. - This analysis emits no runtime statements. - - Parameters - ---------- - context : GroupRewriteContext - Access to operand dtypes and per-thread payload specifications. - op_name : str - Load/Store operation whose private markers are being consumed. - runtime_args : tuple of ir.Var - Ordered provider operands, starting with memory and payload. - factory_kwargs : dict of str to object - Specialization values; private family markers are popped in place. - - Returns - ------- - _LoadStoreMatchMetadata - Whether Store needs a local array to hold its scalar operand. - - Raises - ------ - CoopSinglePhaseRewriteError - Private markers are inconsistent or common-API operand dtypes cannot - be established or are unsupported. - """ - - group_root_store = factory_kwargs.pop("_group_root_store", False) - common_root_operation = factory_kwargs.pop("_common_root_operation", None) - if not isinstance(group_root_store, bool): - raise CoopSinglePhaseRewriteError( - "_group_root_store must be a compile-time bool" - ) - if common_root_operation is not None: - if common_root_operation != op_name: - raise CoopSinglePhaseRewriteError( - "_common_root_operation does not match " - "the rewritten group operation" - ) - from ._parameters import _validate_common_numeric_dtype - - operand_names = ( - ("source", "output") - if op_name == "load" - else ("destination", "value") - ) - operands = list(zip(operand_names, runtime_args)) - if op_name == "store" and len(runtime_args) >= 2: - value_is_array = context.thread_data(runtime_args[1]) - if value_is_array is None: - operands = operands[:1] - for operand_name, operand in operands: - operand_dtype = context.dtype(operand) - if operand_dtype is None: - raise CoopSinglePhaseRewriteError( - f"Failed to infer cuda.coop.{common_root_operation} " - f"{operand_name} dtype for common API validation." - ) - try: - _validate_common_numeric_dtype( - operand_dtype, - operation=op_name, - ) - except (TypeError, ValueError) as exc: - raise CoopSinglePhaseRewriteError(str(exc)) from exc - box_root_store_scalar = False - if group_root_store: - if op_name != "store" or len(runtime_args) < 2: - raise CoopSinglePhaseRewriteError( - "_group_root_store is valid only for root store calls" - ) - box_root_store_scalar = context.thread_data(runtime_args[1]) is None - return _LoadStoreMatchMetadata(box_root_store_scalar=box_root_store_scalar) - - -def prepare_load_store_runtime_args( - context: GroupRewriteContext, - block: ir.Block, - *, - match: _RewriteMatch, - runtime_args: list[ir.Var], - scope: ir.Scope | None, - loc: ir.Loc, -) -> list[ir.Var]: - """Build a local-array payload for a scalar group Store value. - - Use the inferred item count and dtype to allocate the provider's payload, - then assign the scalar to every item. A scalar without static provenance - first passes through ``_ExactStoreScalar`` so compiler typing rejects a - dtype mismatch before array assignment can silently cast it. Static values - have already been checked using scalar coercion rules. - - Calls without a boxing request return their operand list unchanged. For - boxed calls, append preparation statements before the eventual provider - call and replace only its payload operand. - - Parameters - ---------- - context : GroupRewriteContext - Access to static scalar provenance for the Store value. - block : ir.Block - IR block receiving allocation, validation, and item assignments. - match : _RewriteMatch - Match with Load/Store family metadata and inferred factory inputs. - runtime_args : list of ir.Var - Ordered operands, mutated at index one when boxing is required. - scope : ir.Scope or None - Scope assigned to generated variables. - loc : ir.Loc - Source location assigned to generated statements and variables. - - Returns - ------- - list of ir.Var - The supplied list with the scalar payload replaced by a local - array, or the original operands when boxing is unnecessary. - - Raises - ------ - CoopSinglePhaseRewriteError - Family metadata, the Store value, or inferred dtype/extent is - missing or invalid. Later typing rejects runtime scalar mismatches. - """ - - metadata = match.family_metadata - if not isinstance(metadata, _LoadStoreMatchMetadata): - raise CoopSinglePhaseRewriteError("missing Load/Store family metadata") - if not metadata.box_root_store_scalar: - return runtime_args - if len(runtime_args) < 2: - raise CoopSinglePhaseRewriteError( - "root store is missing its value argument" - ) - dtype = match.factory_kwargs.get("dtype") - if dtype is None: - raise CoopSinglePhaseRewriteError( - "root store requires an inferred dtype" - ) - items_per_thread = match.factory_kwargs.get("items_per_thread", 1) - if ( - isinstance(items_per_thread, bool) - or not isinstance(items_per_thread, int) - or items_per_thread < 1 - ): - raise CoopSinglePhaseRewriteError( - "root store requires an inferred positive items_per_thread" - ) - - def new_var(stem: str) -> ir.Var: - return ir.Var( - scope, - f"__coop_root_store_{stem}_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - - module_var = new_var("cuda") - local_var = new_var("local") - array_fn = new_var("array") - shape_var = new_var("shape") - dtype_var = new_var("dtype") - payload = new_var("payload") - block.append( - ir.Assign( - ir.Global(_next_global_name("root_store_cuda"), _cuda_module, loc), - module_var, - loc, - ) - ) - block.append( - ir.Assign(ir.Expr.getattr(module_var, "local", loc), local_var, loc) - ) - block.append( - ir.Assign(ir.Expr.getattr(local_var, "array", loc), array_fn, loc) - ) - block.append(ir.Assign(ir.Const(items_per_thread, loc), shape_var, loc)) - block.append( - ir.Assign( - ir.Global(_next_global_name("root_store_dtype"), dtype, loc), - dtype_var, - loc, - ) - ) - block.append( - ir.Assign( - ir.Expr.call(array_fn, [shape_var, dtype_var], (), loc), - payload, - loc, - ) - ) - value = runtime_args[1] - if context.static_scalar_provenance(value) is _UNRESOLVED: - validator = new_var("validate") - checked_value = new_var("checked_value") - block.append( - ir.Assign( - ir.Global( - _next_global_name("store_scalar"), - _ExactStoreScalar(dtype), - loc, - ), - validator, - loc, - ) - ) - block.append( - ir.Assign( - ir.Expr.call(validator, [value], (), loc), checked_value, loc - ) - ) - value = checked_value - for item_index in range(items_per_thread): - index_var = new_var(f"index_{item_index}") - block.append(ir.Assign(ir.Const(item_index, loc), index_var, loc)) - block.append(ir.SetItem(payload, index_var, value, loc)) - runtime_args[1] = payload - return runtime_args - - -__all__ = [ - "analyze_load_store_match", - "infer_load_store_payload", - "prepare_load_store_runtime_args", - "validate_load_store_runtime_controls", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_payload.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_payload.py deleted file mode 100644 index c03cc0f003cd..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_payload.py +++ /dev/null @@ -1,223 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Complete provider inputs from planned or inferred payload facts. - -Primitive-family hooks use ``PayloadInference`` to merge dtype and extent -values with explicit factory keywords. It carries the current call's lowering -plan, allowing Load/Store to reuse checked facts while retaining constructor -bookkeeping. Direct calls without plans infer those facts from their operands. -Hooks update the pending specialization and payload metadata before provider -creation; they do not replace later compiler type inference or emit array IR. -""" - -from typing import TYPE_CHECKING, cast - -from cuda.coop._core import GroupLoweringPlan - -from ._group_rewriting import GroupRewriteContext -from ._operations import rewrite_operation -from ._parameters import normalize_dtype_param -from ._rewrite_support import ( - CoopSinglePhaseRewriteError, - _ThreadDataSpecification, - ir, -) - -if TYPE_CHECKING: - from ._rewrite import CoopSinglePhaseRewrite - - -class PayloadInference: - """Merge payload facts into one call's pending factory inputs. - - Family hooks share these mutable collections with argument splitting. - ``infer_kwarg`` fills missing inputs and checks existing ones, so a - payload cannot silently replace an explicit specialization choice. - - Parameters - ---------- - context : GroupRewriteContext - Access to operand types, payload descriptions, and scalar origins. - op_name : str - Operation name used in conflict diagnostics. - runtime_args : list of ir.Var - Operands in provider order, without the storage pointer. - allowed_factory_kwargs : set of str - Keywords the operation accepts; other inferred names are ignored. - seen_factory_kwargs : set of str - Names already resolved, updated when inference supplies a value. - factory_kwargs : dict of str to object - Resolved inputs, updated in place without copying the dictionary. - dtype_factory_kwargs : frozenset of str - Names whose values need dtype normalization before comparison. - lowering_plan : GroupLoweringPlan or None - Facts already validated for this call in the current planning attempt. - Direct provider calls have no plan and require full payload inference. - """ - - def __init__( - self, - context: GroupRewriteContext, - op_name: str, - runtime_args: list[ir.Var], - allowed_factory_kwargs: set[str], - seen_factory_kwargs: set[str], - factory_kwargs: dict[str, object], - dtype_factory_kwargs: frozenset[str], - *, - lowering_plan: GroupLoweringPlan | None = None, - ) -> None: - self.context = context - self.op_name = op_name - self.runtime_args = runtime_args - self.allowed_factory_kwargs = allowed_factory_kwargs - self.seen_factory_kwargs = seen_factory_kwargs - self.factory_kwargs = factory_kwargs - self.dtype_factory_kwargs = dtype_factory_kwargs - self.lowering_plan = lowering_plan - - def factory_value(self, name: str) -> object: - return self.factory_kwargs.get(name) - - def _factory_kwarg_matches( - self, name: str, actual: object, expected: object - ) -> bool: - """Compare an inferred value with an existing factory input. - - Normalize dtype keywords when possible so equivalent spellings agree. - Fall back to ordinary equality if normalization fails. - """ - - if name in self.dtype_factory_kwargs: - try: - actual = normalize_dtype_param(actual) - expected = normalize_dtype_param(expected) - except (TypeError, ValueError): - pass - return actual == expected - - def infer_kwarg(self, name: str, value: object) -> None: - """Merge a planned or inferred payload value into factory inputs. - - Ignore unavailable values and keywords the operation does not accept. - An already resolved keyword must agree with the payload; dtype - keywords are compared after normalization when possible so equivalent - dtype spellings do not conflict. Never overwrite a conflicting - explicit value. - - Parameters - ---------- - name : str - Factory keyword to infer or check. - value : object - Planned or inferred payload value, or None when not known. - - Returns - ------- - None - Update ``factory_kwargs`` and ``seen_factory_kwargs`` in place - when this supplies a previously missing keyword. - - Raises - ------ - CoopSinglePhaseRewriteError - The supplied value conflicts with an already resolved keyword. - """ - - if name not in self.allowed_factory_kwargs or value is None: - return - if name in self.seen_factory_kwargs: - if not self._factory_kwarg_matches( - name, self.factory_kwargs[name], value - ): - raise CoopSinglePhaseRewriteError( - f"cooperative group operation {self.op_name!r} factory " - f"argument {name!r} does not match the value inferred " - f"from the payload." - ) - return - self.factory_kwargs[name] = value - self.seen_factory_kwargs.add(name) - - def candidate( - self, index: int - ) -> tuple[ir.Var | None, _ThreadDataSpecification | None]: - """Look up an operand and its available per-thread payload facts. - - Return ``(None, None)`` for an absent or non-variable operand at - ``index``. A valid variable can still have no payload description. - """ - - if not 0 <= index < len(self.runtime_args): - return (None, None) - value = self.runtime_args[index] - if not isinstance(value, ir.Var): - return (None, None) - specification = self.context.thread_data(value) - return (value, specification) - - -class _PayloadRewrite: - """Dispatch payload inference to the registered primitive-family hook.""" - - def _infer_factory_kwargs_from_payload( - self, - op_name: str, - runtime_args: list[ir.Var], - allowed_factory_kwargs: set[str], - seen_factory_kwargs: set[str], - factory_kwargs: dict[str, object], - *, - lowering_plan: GroupLoweringPlan | None = None, - ) -> None: - """Use payload facts to complete a provider's specialization inputs. - - The argument splitter calls the registered family's inference hook - before checking required factory keywords. This lets an operand's dtype - or per-thread element count supply an omitted keyword while detecting a - conflict with an explicit value. ``PayloadInference`` shares the pending - collections with the splitter so those decisions remain visible. - - Parameters - ---------- - op_name : str - Registered operation whose payload inference hook should run. - runtime_args : list of ir.Var - Operands in provider order, without the storage pointer. - allowed_factory_kwargs : set of str - Keywords accepted by this operation; other inferred names are - ignored. - seen_factory_kwargs : set of str - Names already resolved. The hook adds names that it successfully - infers. - factory_kwargs : dict of str to object - Resolved specialization inputs, updated in place by the hook. - lowering_plan : GroupLoweringPlan or None - Semantic facts belonging to this provider call, when planned. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - specification = rewrite_operation(op_name) - if specification is None: - raise CoopSinglePhaseRewriteError( - f"unsupported Numba-CUDA-MLIR operation {op_name!r}" - ) - inference = PayloadInference( - GroupRewriteContext(rewrite), - op_name, - runtime_args, - allowed_factory_kwargs, - seen_factory_kwargs, - factory_kwargs, - specification.dtype_factory_kwargs, - lowering_plan=lowering_plan, - ) - specification.infer_payload(inference.context, inference) - - -__all__ = [ - "PayloadInference", - "_PayloadRewrite", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_provenance.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_provenance.py deleted file mode 100644 index c64a97b90601..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_provenance.py +++ /dev/null @@ -1,2165 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Recover compile-time facts from cooperative values before type inference. - -The call rewrite needs payload shapes, scalar bindings, and scratch ownership -before ordinary compiler typing is available. These helpers follow variable -definitions, aliases, casts, and control-flow merges to recover those facts. -They distinguish unknown provenance from conflicting concrete origins so a -branch or loop rebinding cannot silently select one descriptor's contract. - -Storage constructor identities also determine which calls may reuse an aligned -scratch region. This module builds per-descriptor layouts; -``_rewrite_storage`` combines them into one backing allocation and emits array -views and barriers. -""" - -from __future__ import annotations - -import operator -from typing import TYPE_CHECKING, cast - -import numba_cuda_mlir.numba_cuda.types as numba_types -from numba_cuda_mlir import cuda as _cuda_module -from numba_cuda_mlir.cuda.local import array as _cuda_local_array -from numba_cuda_mlir.cuda.shared import array as _cuda_shared_array -from numba_cuda_mlir.numba_cuda.core.errors import ForceLiteralArg - -from cuda.coop._core import StorageOwnership, SynchronizationScope -from cuda.coop._core import api as _common_api - -from .._temp_storage import TempStorage -from .._thread_data import ThreadData, _normalize_thread_data_alignment -from ._descriptor_provenance import ( - descriptor_definitions, - payload_write_dtypes, - temp_storage_constructor, -) -from ._operations import factory_operation -from ._parameters import normalize_dtype_param -from ._rewrite_support import ( - _INFERENCE_EXCEPTIONS, - _MIN_TEMP_STORAGE_ALIGNMENT, - _UNRESOLVED, - CoopSinglePhaseRewriteError, - Rewrite, - _align_up, - _default_temp_storage_alignment, - _dtype_values_match, - _normalize_temp_storage_alignment, - _phi_incoming_values, - _ResolvedCallTarget, - _RewriteMatch, - _TempStorageCtorSpecification, - _TempStorageGlobalPlan, - _TempStoragePlan, - _TempStorageRequirementSummary, - _TempStorageSlice, - _TempStorageUseRequirement, - _ThreadDataSpecification, - _validate_temp_storage_alignment, - ir, -) -from ._scalar_provenance import ( - StaticScalarProvenance, - cuda_index_dtype, - scalar_call_dtype, - scalar_expression_dtype, - try_resolve_static_scalar, - try_resolve_static_scalar_provenance, -) - -if TYPE_CHECKING: - from .._types import Algorithm - from ._rewrite import CoopSinglePhaseRewrite - - -class _ProvenanceRewrite(Rewrite): - """Keep the facts used while cooperative calls are being replaced. - - Provenance means where an IR value comes from. These queries trace - definitions to recover static values, payload types, and storage owners - before ordinary typing. A missing fact stays unknown; concrete descriptor - conflicts produce diagnostics. - - ``state`` supplies the function IR, argument types, and compiler metadata. - ``allow_launch_dim_deferral`` permits a retry when launch facts are - missing. Mixins share this state, including block lookups, descriptor - tables, storage plans, and provider caches. - """ - - def __init__( - self, - state, - *, - allow_launch_dim_deferral: bool = True, - ) -> None: - super().__init__(state) - self._state = state - self._allow_launch_dim_deferral = allow_launch_dim_deferral - self._func_ir = state.func_ir - self._block: ir.Block | None = None - self._block_defs: dict[str, object] = {} - self._matches: dict[ir.Assign, _RewriteMatch] = {} - self._func_matches: dict[ir.Assign, _RewriteMatch] = {} - self._rewrite_started = False - self._factory_argument_cleanup_candidates: set[ir.Assign] = set() - self._payload_callee_cleanup_names: set[str] = set() - self._temp_storage_assigns: set[ir.Assign] = set() - self._temp_storage_func_vars: set[str] = set() - self._temp_storage_ctor_specifications: dict[ - str, _TempStorageCtorSpecification - ] = {} - self._temp_storage_ctor_order: dict[str, int] = {} - self._temp_storage_ctor_roots: dict[str, str] = {} - self._temp_storage_ctor_sites: dict[str, set[int]] = {} - self._thread_data_func_vars: set[str] = set() - self._thread_data_specifications: dict[ - str, _ThreadDataSpecification - ] = {} - self._thread_data_like_vars: set[str] = set() - self._thread_data_extents: dict[ir.Assign, int] = {} - self._func_ir_identity: int | None = None - self._func_temp_storage_requirements: dict[ - str, _TempStorageRequirementSummary - ] = {} - self._temp_storage_plans: dict[str, _TempStoragePlan] = {} - self._temp_storage_global_plan: _TempStorageGlobalPlan | None = None - self._implicit_temp_storage_requirements = ( - _TempStorageRequirementSummary() - ) - self._implicit_temp_storage_plan: _TempStoragePlan | None = None - self._temp_storage_backing_var: ir.Var | None = None - self._temp_storage_backing_emitted = False - self._arg_type_map = self._build_arg_type_map() - self._invocable_cache: dict[ - tuple[str, tuple[tuple[str, str, str], ...]], object - ] = {} - self._prebundled_specializations: dict[ - tuple[str, tuple[tuple[str, str, str], ...]], - Algorithm, - ] = {} - self._deferred_launch_dim_inference = False - - def _lookup_block_definition(self, name: str) -> object: - """Look up a block definition only for an unambiguous name. - - Before SSA reconstruction, branches and loops may assign the same name - more than once. The block map stores only the last assignment, which - can follow the use being analyzed. Refuse that shortcut when the - function's definition table contains multiple assignments. This avoids - mistaking a branch or later rebinding for a constant. - - Parameters - ---------- - name : str - IR variable name to find in the current block's assignment map. - - Returns - ------- - object or None - Local definition, or None for a missing or multiply assigned name. - """ - - # Before SSA, a name can be rebound in a loop or branch. The block map - # retains only its last assignment, which may follow the current use. - if len(self._func_ir._definitions.get(name, ())) > 1: - return None - return self._block_defs.get(name) - - def _infer_constant(self, value): - """Try local and scalar facts before the compiler's constant resolver. - - A unique block definition can supply a constant, global, or closure - value. Otherwise use static scalar provenance, then let compiler - inference return a value or raise its normal inference error. - """ - - if isinstance(value, ir.Var): - definition = self._lookup_block_definition(value.name) - if isinstance(definition, (ir.Const, ir.Global, ir.FreeVar)): - return definition.value - scalar = self._resolve_static_scalar_provenance(value) - if isinstance(scalar, StaticScalarProvenance): - return scalar.value - return self._func_ir.infer_constant(value) - - def _resolve_static_scalar_value( - self, - value: object, - ) -> object: - """Resolve a scalar only when its IR provenance is explicitly static. - - Use the shared scalar-provenance resolver with this function's - definitions and specialized argument types. Callers use the result to - choose between compile-time scalar bindings and runtime operands; - ordinary constant inference does not establish this binding contract. - - Parameters - ---------- - value : object - Scalar value or IR reference for the provenance resolver. - - Returns - ------- - object - Resolved scalar, including None when it is explicitly static, - or the ``_UNRESOLVED`` sentinel when static provenance cannot - be established. Compare the sentinel by identity; None can - mean an omitted control. - """ - - arg_types = tuple(getattr(self._state, "args", ()) or ()) - resolved, scalar = try_resolve_static_scalar( - value, - definitions=self._lookup_definitions, - argument_type=lambda index: ( - arg_types[index] if 0 <= index < len(arg_types) else None - ), - ) - return scalar if resolved else _UNRESOLVED - - def _resolve_static_scalar_provenance(self, value): - """Resolve a scalar together with its established source dtype. - - Return ``_UNRESOLVED`` when the shared resolver cannot prove static - origin. Keeping dtype provenance lets scalar coercion distinguish a - Python literal from a value with an explicit numeric width. - """ - - arg_types = tuple(getattr(self._state, "args", ()) or ()) - resolved, scalar = try_resolve_static_scalar_provenance( - value, - definitions=self._lookup_definitions, - argument_type=lambda index: ( - arg_types[index] if 0 <= index < len(arg_types) else None - ), - ) - return scalar if resolved else _UNRESOLVED - - def _build_arg_type_map(self) -> dict[str, object]: - """Pair argument names with compiler types only when counts agree.""" - - arg_names = tuple(getattr(self._func_ir, "arg_names", ()) or ()) - arg_types = tuple(getattr(self._state, "args", ()) or ()) - if len(arg_names) != len(arg_types): - return {} - return dict(zip(arg_names, arg_types)) - - def _lookup_definition(self, value): - """Resolve one unambiguous assignment for a provenance query. - - Queries for a callee or constructor need one source value before they - can inspect it. Before static single assignment (SSA) reconstruction, a - name can have several assignments in branches or loops. Choosing the - last one would lose the other possibilities, so the guarded block lookup - and function lookup only accept an unambiguous definition. Use - ``_lookup_definitions`` when the caller can combine several - alternatives. - - Parameters - ---------- - value : ir.Var, str, or object - Variable or variable name to look up. An already resolved definition - is returned unchanged so recursive queries can accept either form. - - Returns - ------- - object or None - The defining value, or ``None`` when the name has no unique - definition. - """ - - if isinstance(value, ir.Var): - definition = self._lookup_block_definition(value.name) - if definition is not None: - return definition - try: - return self._func_ir.get_definition(value) - except KeyError: - return None - if isinstance(value, str): - definition = self._lookup_block_definition(value) - if definition is not None: - return definition - try: - return self._func_ir.get_definition(value) - except KeyError: - return None - return value - - def _lookup_definitions(self, value) -> list[object]: - """Collect recorded sources without discarding branch inputs. - - Payload and descriptor queries need to compare every possible - assignment: selecting just one branch could accept incompatible dtypes - or storage settings. This helper supplies the definitions recorded for a - name, plus an unambiguous block-local definition when available. It - deduplicates by object identity and leaves the caller to follow aliases - and check agreement. This is a conservative collection, not a proof that - each assignment can reach the current statement. - - Parameters - ---------- - value : ir.Var, str, or object - Variable or name whose definitions are needed. A resolved definition - is returned as a one-element list for the recursive callers. - - Returns - ------- - list of object - Candidate defining values. An unknown variable produces an empty - list. - """ - - defs: list[object] = [] - seen_ids: set[int] = set() - - def add(candidate) -> None: - if candidate is None: - return - cid = id(candidate) - if cid in seen_ids: - return - seen_ids.add(cid) - defs.append(candidate) - - if isinstance(value, ir.Var): - add(self._lookup_block_definition(value.name)) - for definition in ( - getattr(self._func_ir, "_definitions", {}) or {} - ).get(value.name, ()): - add(definition) - return defs - if isinstance(value, str): - add(self._lookup_block_definition(value)) - for definition in ( - getattr(self._func_ir, "_definitions", {}) or {} - ).get(value, ()): - add(definition) - return defs - return [value] - - def _resolve_attribute_chain(self, func_var): - """Trace a Python object and attribute path through untyped IR. - - Return ``(root, attributes)`` in lookup order. Return ``None`` for a - cycle or a root that is not a constant, global, or closure value. - """ - - attrs: list[str] = [] - current = func_var - seen: set[str] = set() - while True: - if isinstance(current, ir.Var): - if current.name in seen: - return None - seen.add(current.name) - current = self._lookup_definition(current) - elif isinstance(current, ir.Expr) and current.op == "getattr": - attrs.append(current.attr) - current = current.value - else: - break - if isinstance(current, (ir.Global, ir.FreeVar, ir.Const)): - root = current.value - else: - return None - attrs.reverse() - return (root, attrs) - - def _resolve_python_value(self, value): - """Read a resolved attribute chain to identify a Python callable. - - Missing attributes and failed imports return ``None``. Reading an - attribute can trigger the backend's lazy registration. - """ - - chain = self._resolve_attribute_chain(value) - if chain is None: - return None - root, attrs = chain - obj = root - try: - for attr in attrs: - obj = getattr(obj, attr) - except (AttributeError, ImportError): - return None - return obj - - def _is_common_root_member(self, value, name: str) -> bool: - """Require both common API identity and its backend-member marker.""" - - member = getattr(_common_api, name) - return ( - self._resolve_python_value(value) is member - and getattr(member, "__cuda_coop_backend_member__", None) == name - ) - - def _is_supported_factory(self, obj) -> bool: - """Require a registered factory in an accepted provider namespace.""" - - if not callable(obj): - return False - metadata = factory_operation(obj) - if metadata is None: - return False - from ._operations import rewrite_operation - - specification = rewrite_operation(metadata.operation) - if specification is None: - return False - return metadata.namespace in specification.factory_namespaces - - def _resolve_factory_from_var(self, func_var): - """Resolve a callee through the operation registry. - - Constants, globals, closure values, and directly supplied callables - can identify it. A function's spelling alone does not identify a - cooperative provider. - """ - - direct = None - direct_def = self._lookup_definition(func_var) - if isinstance(direct_def, (ir.Global, ir.FreeVar, ir.Const)): - direct = direct_def.value - elif callable(direct_def): - direct = direct_def - elif direct_def is None: - try: - direct = self._infer_constant(func_var) - except _INFERENCE_EXCEPTIONS: - direct = None - return direct if self._is_supported_factory(direct) else None - - def _extract_1d_extent_literal(self, value_ref): - """Read a native array's element count from its shape argument. - - Array provenance queries call this before provider specialization needs - ``items_per_thread``. The extent is the number of elements in the - array's single dimension, not a byte count. For example, both ``4`` and - ``(4,)`` describe four elements. Native-array typing later checks the - constructor; this query does not apply ``ThreadData``'s positivity and - boolean rules. - - Parameters - ---------- - value_ref : ir.Var or object - Shape argument to resolve with the rewrite's constant inference. - - Returns - ------- - int or None - Scalar shape or the sole integer in a tuple or list. ``None`` means - the shape is unresolved or is not a recognized one-dimensional - shape. - """ - - try: - value = self._infer_constant(value_ref) - except _INFERENCE_EXCEPTIONS: - return None - if isinstance(value, int): - return value - if ( - isinstance(value, tuple) - and len(value) == 1 - and isinstance(value[0], int) - ): - return int(value[0]) - if ( - isinstance(value, list) - and len(value) == 1 - and isinstance(value[0], int) - ): - return int(value[0]) - return None - - @staticmethod - def _is_jitted_dispatcher(obj: object) -> bool: - """Recognize the dispatcher interface of a resolved helper callee. - - Parameters - ---------- - obj : object - Python value resolved from a call's callee. - - Returns - ------- - bool - Whether the value is callable and has ``py_func`` and a - ``targetoptions`` dictionary. This checks the interface; it does - not require the device-function option to be enabled. - """ - - return ( - obj is not None - and callable(obj) - and hasattr(obj, "py_func") - and isinstance(getattr(obj, "targetoptions", None), dict) - ) - - def _is_temp_storage_ctor_call(self, call: ir.Expr) -> bool: - if self._is_common_root_member(call.func, "TempStorage"): - return True - return self._resolve_python_value(call.func) is TempStorage - - def _is_thread_data_ctor_call(self, call: ir.Expr) -> bool: - if self._is_common_root_member(call.func, "ThreadData"): - return True - return self._resolve_python_value(call.func) is ThreadData - - def _extract_thread_data_specification( - self, call: ir.Expr - ) -> _ThreadDataSpecification: - """Recover the compile-time ``ThreadData`` constructor contract. - - Require a positive integral item count, rejecting booleans, and - normalize an explicit alignment. A directly referenced function - argument used for the extent requests literal specialization before - continuing. An omitted or explicitly None dtype remains unresolved so - operation consumers or typed writes can supply it later. - - Record whether the constructor came from the common API; ``apply`` - uses that flag when validating the inferred numeric dtype before - lowering the descriptor to a local array. This method does not - allocate that array or mutate the constructor expression. - - Parameters - ---------- - call : ir.Expr - Recognized common or Numba-CUDA-MLIR ``ThreadData`` - constructor call. - - Returns - ------- - _ThreadDataSpecification - Static extent, optional dtype and alignment, and whether the - constructor came from the common API. - - Raises - ------ - ForceLiteralArg - The extent is a function argument awaiting literal specialization. - CoopSinglePhaseRewriteError - Constructor arguments are invalid or required compile-time - values cannot be resolved. - """ - - kw_map = {name: value for name, value in call.kws} - is_common_root = self._is_common_root_member(call.func, "ThreadData") - allowed_keywords = {"items_per_thread", "dtype", "alignment"} - unexpected_keywords = sorted(set(kw_map) - allowed_keywords) - if unexpected_keywords: - names = ", ".join(unexpected_keywords) - scope = "cuda.coop" if is_common_root else "cuda.coop.numba_mlir" - raise CoopSinglePhaseRewriteError( - f"{scope}.ThreadData got unexpected keyword(s): {names}" - ) - extent_refs = [] - if call.args: - extent_refs.append(("positional items_per_thread", call.args[0])) - if "items_per_thread" in kw_map: - extent_refs.append(("items_per_thread", kw_map["items_per_thread"])) - if len(call.args) > 2: - raise CoopSinglePhaseRewriteError( - "coop.ThreadData accepts at most items_per_thread and dtype " - "positional arguments." - ) - if len(extent_refs) > 1: - names = " and ".join((name for name, _ in extent_refs)) - raise CoopSinglePhaseRewriteError( - f"coop.ThreadData received both {names}; specify only one." - ) - if not extent_refs: - raise CoopSinglePhaseRewriteError( - "coop.ThreadData requires items_per_thread." - ) - items_ref = extent_refs[0][1] - try: - definition = self._func_ir.get_definition(items_ref) - except KeyError: - definition = None - if isinstance(definition, ir.Arg) and not isinstance( - self._state.args[definition.index], numba_types.Literal - ): - raise ForceLiteralArg({definition.index}) - dtype_ref = None - if len(call.args) == 2: - dtype_ref = call.args[1] - if "dtype" in kw_map: - if dtype_ref is not None: - raise CoopSinglePhaseRewriteError( - "coop.ThreadData received dtype " - "both positionally and by keyword." - ) - dtype_ref = kw_map["dtype"] - alignment = None - alignment_ref = kw_map.get("alignment") - if alignment_ref is not None: - try: - raw_alignment = self._infer_constant(alignment_ref) - except _INFERENCE_EXCEPTIONS as exc: - raise CoopSinglePhaseRewriteError( - "coop.ThreadData alignment must be a compile-time " - "integer or None" - ) from exc - try: - alignment = _normalize_thread_data_alignment(raw_alignment) - except (TypeError, ValueError) as exc: - raise CoopSinglePhaseRewriteError( - f"coop.ThreadData {exc}" - ) from exc - try: - raw_items_per_thread = self._infer_constant(items_ref) - except _INFERENCE_EXCEPTIONS as exc: - raise CoopSinglePhaseRewriteError( - "items_per_thread must be a compile-time integer" - ) from exc - if isinstance(raw_items_per_thread, bool): - raise CoopSinglePhaseRewriteError( - "items_per_thread must be an integer" - ) - try: - items_per_thread = operator.index(raw_items_per_thread) - except TypeError as exc: - raise CoopSinglePhaseRewriteError( - "items_per_thread must be an integer" - ) from exc - if items_per_thread <= 0: - raise CoopSinglePhaseRewriteError( - "items_per_thread must be a positive integer" - ) - dtype = None - if dtype_ref is not None: - dtype = self._resolve_dtype_ref(dtype_ref) - if dtype is None: - try: - self._infer_constant(dtype_ref) - except _INFERENCE_EXCEPTIONS as exc: - raise CoopSinglePhaseRewriteError( - "coop.ThreadData dtype must resolve to a " - "compile-time dtype or None" - ) from exc - if any(dtype is alias for alias in (bool, int, float, complex)): - dtype = normalize_dtype_param(dtype) - return _ThreadDataSpecification( - items_per_thread=items_per_thread, - dtype=dtype, - common_root=is_common_root, - alignment=alignment, - ) - - @staticmethod - def _merge_thread_data_specifications( - existing: _ThreadDataSpecification | None, - observed: _ThreadDataSpecification, - ) -> _ThreadDataSpecification: - """Combine payload facts that must describe one compatible allocation. - - Provenance queries call this when aliases or branches provide more than - one description of a payload. A later provider needs one element count - and dtype, so unequal known values raise a rewrite error. Unknown fields - can be filled from the other description. Keep the larger alignment and - retain common API origin if either input has it. - - Parameters - ---------- - existing : _ThreadDataSpecification or None - Facts accumulated so far, or ``None`` for the first description. - observed : _ThreadDataSpecification - Facts from another source of the payload. ``items_per_thread`` - counts elements; ``alignment`` is in bytes. Either may still be - unknown. - - Returns - ------- - _ThreadDataSpecification - Combined facts, or ``observed`` unchanged when it is the first - input. - """ - - if existing is None: - return observed - if ( - existing.items_per_thread is not None - and observed.items_per_thread is not None - and (existing.items_per_thread != observed.items_per_thread) - ): - raise CoopSinglePhaseRewriteError( - "Inconsistent items_per_thread " - "across merged coop.ThreadData aliases." - ) - if ( - existing.dtype is not None - and observed.dtype is not None - and (existing.dtype != observed.dtype) - ): - raise CoopSinglePhaseRewriteError( - "Inconsistent dtype across merged coop.ThreadData aliases." - ) - items_per_thread = existing.items_per_thread - if items_per_thread is None: - items_per_thread = observed.items_per_thread - dtype = existing.dtype - if dtype is None: - dtype = observed.dtype - alignment = existing.alignment - if alignment is None: - alignment = observed.alignment - elif observed.alignment is not None: - alignment = max(alignment, observed.alignment) - return _ThreadDataSpecification( - items_per_thread=items_per_thread, - dtype=dtype, - common_root=existing.common_root or observed.common_root, - alignment=alignment, - ) - - @staticmethod - def _merge_temp_storage_ctor_specifications( - existing: _TempStorageCtorSpecification | None, - observed: _TempStorageCtorSpecification, - ) -> _TempStorageCtorSpecification: - """Require identical settings when storage aliases merge.""" - - if existing is None: - return observed - if existing != observed: - raise CoopSinglePhaseRewriteError( - "Inconsistent TempStorage constructor metadata " - "across merged aliases." - ) - return existing - - def _record_temp_storage_ctor(self, inst: ir.Assign, call: ir.Expr) -> None: - """Record one constructor site for the variable it defines. - - The rewrite runs before SSA reconstruction, so a name rebound inside a - branch or loop body keeps one name for several constructor sites. - Every site must agree on the effective contract, and the sites are - counted so that auto_sync=False storage cannot be silently collapsed. - """ - - name = inst.target.name - specification = self._extract_temp_storage_ctor_specification(call) - existing = self._temp_storage_ctor_specifications.get(name) - if existing is not None and ( - self._temp_storage_contract(existing) - != self._temp_storage_contract(specification) - ): - raise CoopSinglePhaseRewriteError( - f"TempStorage aliases have inconsistent contracts across " - f"constructor instances ({name})." - ) - self._temp_storage_ctor_specifications.setdefault(name, specification) - self._temp_storage_ctor_sites.setdefault(name, set()).add(id(inst)) - - @staticmethod - def _collapsed_manual_sync_error( - names: str, sites: int - ) -> CoopSinglePhaseRewriteError: - return CoopSinglePhaseRewriteError( - f"TempStorage with auto_sync=False must be constructed at " - f"exactly one site; {names} reaches {sites} constructor sites. " - f"The compiler cannot verify caller synchronization when it " - f"merges these regions. Construct the descriptor once or set " - f"auto_sync=True." - ) - - def _validate_temp_storage_ctor_sites(self) -> None: - """Reject repeated constructor sites when the caller manages barriers. - - Equal settings do not prove that independently constructed regions can - share storage safely after the rewrite combines them. - """ - - for name, sites in sorted(self._temp_storage_ctor_sites.items()): - if len(sites) < 2: - continue - specification = self._temp_storage_ctor_specifications[name] - if not self._temp_storage_contract(specification)[2]: - raise self._collapsed_manual_sync_error(repr(name), len(sites)) - - def _record_inferred_thread_data_dtype( - self, value: ir.Var, dtype, seen: set[str] | None = None - ) -> None: - """Record an inferred dtype on payload constructors and aliases. - - Update partial descriptions and reject a conflicting known dtype. - Follow copies, casts, merged paths, and static tuple items. ``seen`` - prevents revisiting names in loops or alias cycles. - """ - - if not isinstance(value, ir.Var) or dtype is None: - return - if seen is None: - seen = set() - if value.name in seen: - return - seen.add(value.name) - specification = self._thread_data_specifications.get(value.name) - if specification is None: - specification = self._resolve_thread_data_specification(value) - if specification is not None: - self._thread_data_specifications[value.name] = specification - if specification is not None: - if specification.dtype is None: - self._thread_data_specifications[value.name] = ( - _ThreadDataSpecification( - items_per_thread=specification.items_per_thread, - dtype=dtype, - common_root=specification.common_root, - alignment=specification.alignment, - ) - ) - elif specification.dtype != dtype: - raise CoopSinglePhaseRewriteError( - "Inconsistent inferred dtype for coop.ThreadData usage." - ) - for definition in self._lookup_definitions(value): - if isinstance(definition, ir.Var): - self._record_inferred_thread_data_dtype(definition, dtype, seen) - continue - if not isinstance(definition, ir.Expr): - continue - if definition.op == "cast": - cast_value = getattr(definition, "value", None) - if isinstance(cast_value, ir.Var): - self._record_inferred_thread_data_dtype( - cast_value, dtype, seen - ) - elif definition.op == "phi": - for incoming in _phi_incoming_values(definition): - if isinstance(incoming, ir.Var): - self._record_inferred_thread_data_dtype( - incoming, dtype, seen - ) - elif definition.op == "static_getitem": - for item in self._resolve_static_tuple_item_vars(definition): - self._record_inferred_thread_data_dtype(item, dtype, seen) - - def _extract_temp_storage_ctor_specification( - self, call: ir.Expr - ) -> _TempStorageCtorSpecification: - """Resolve storage constructor settings before layout planning. - - Delegate argument rules to the shared descriptor parser, require - compile-time values, and apply the backend's pointer-alignment - minimum. Convert parser failures into rewrite diagnostics. - """ - - def constant(value, *, name): - try: - return self._infer_constant(value) - except _INFERENCE_EXCEPTIONS as exc: - raise CoopSinglePhaseRewriteError( - f"TempStorage {name} must be a compile-time literal." - ) from exc - - try: - descriptor = temp_storage_constructor(call, constant) - except (TypeError, ValueError) as exc: - message = str(exc) - if not message.startswith("TempStorage"): - message = f"TempStorage {message}" - raise CoopSinglePhaseRewriteError(message) from exc - alignment = descriptor.alignment - if alignment is not None: - alignment = _normalize_temp_storage_alignment(alignment) - return _TempStorageCtorSpecification( - size_in_bytes=descriptor.size_in_bytes, - alignment=alignment, - auto_sync=descriptor.auto_sync, - sharing=descriptor.sharing, - ) - - @staticmethod - def _mixed_temp_storage_binding_error( - name: str, - ) -> CoopSinglePhaseRewriteError: - return CoopSinglePhaseRewriteError( - f"TempStorage variables must be bound to a TempStorage " - f"descriptor on every path; {name!r} is also bound to a " - f"non-descriptor value such as None. Remove the None initializer " - f"or construct the descriptor unconditionally." - ) - - def _collect_temp_storage_ctor_keys( - self, - value: ir.Var, - seen: set[str], - *, - display_name: str | None = None, - ) -> set[str]: - """Collect reachable storage constructor owners. - - Follow the shared descriptor-provenance traversal through aliases and - control-flow joins, recording constructor specifications as they are - found. Once backing emission has begun, previously validated - constructor owners still count even though their calls have been - replaced by slices. A join containing both a descriptor and a - non-descriptor such as None is invalid; returning only the descriptor - branch would hide an unsafe path. - - Parameters - ---------- - value : ir.Var - Variable whose possible storage origins are needed. - seen : set of str - Visited variable names used to stop cycles in the provenance walk. - display_name : str or None, optional - User-facing variable name for diagnostics. When absent, prefer - the current name unless it is a compiler temporary. - - Returns - ------- - set of str - Constructor owner names before canonicalization. Empty means - no descriptor origin was found, including when traversal stops - at a cycle. - - Raises - ------ - CoopSinglePhaseRewriteError - A constructor contract is invalid or inconsistent, or one - variable merges descriptor and non-descriptor definitions. - """ - - if not isinstance(value, ir.Var): - return set() - if value.name in seen: - return set() - # Diagnostics name the user's variable, not a compiler temporary that - # a conditional expression or an alias chain introduced. - if display_name is None and not value.name.startswith("$"): - display_name = value.name - shown = display_name or value.name - keys: set[str] = set() - non_descriptor = False - for owner, definition in descriptor_definitions( - value, self._lookup_definitions, seen=seen - ): - if ( - isinstance(definition, ir.Expr) - and definition.op == "call" - and self._is_temp_storage_ctor_call(definition) - ): - assert owner is not None - specification = self._extract_temp_storage_ctor_specification( - definition - ) - self._temp_storage_ctor_specifications[owner] = ( - self._merge_temp_storage_ctor_specifications( - self._temp_storage_ctor_specifications.get(owner), - specification, - ) - ) - keys.add(owner) - elif ( - self._temp_storage_backing_emitted - and owner in self._temp_storage_ctor_specifications - ): - # Previously validated constructors are replaced by backing - # slices as individual blocks are rewritten. - keys.add(owner) - else: - non_descriptor = True - if keys and non_descriptor: - raise self._mixed_temp_storage_binding_error(shown) - return keys - - @staticmethod - def _temp_storage_contract( - specification: _TempStorageCtorSpecification, - ) -> tuple[int | None, int | None, bool, str]: - """Return effective settings used to compare storage descriptors. - - An omitted ``auto_sync`` counts as false. - """ - - auto_sync = ( - False - if specification.auto_sync is None - else specification.auto_sync - ) - return ( - specification.size_in_bytes, - specification.alignment, - auto_sync, - specification.sharing, - ) - - def _canonical_temp_storage_ctor_key(self, key: str) -> str: - """Return the canonical owner of a merged descriptor. - - Point ``key`` directly at that owner to shorten later lookups. - """ - - roots = getattr(self, "_temp_storage_ctor_roots", {}) - root = roots.get(key, key) - while roots.get(root, root) != root: - root = roots[root] - if key != root: - roots[key] = root - return root - - def _merge_temp_storage_ctor_keys(self, keys: set[str]) -> str: - """Unify compatible storage constructor owners. - - All reachable roots must have the same effective size, alignment, - sharing, and synchronization contract. Distinct roots may merge only - with automatic synchronization: the compiler cannot prove that - caller-managed barriers still protect reuse after independently - constructed regions collapse into one allocation. - - Choose the earliest recorded constructor, breaking ties by variable - name, and redirect existing root mappings to it. Requirement - collection later uses this canonical identity to combine every alias's - primitive uses. This method checks distinct roots; repeated - constructor sites under one pre-SSA name are checked separately by - constructor-site validation. - - Parameters - ---------- - keys : set of str - Nonempty set of known constructor owner names to unify. - - Returns - ------- - str - Canonical owner name. ``_temp_storage_ctor_roots`` is updated - in place. - - Raises - ------ - CoopSinglePhaseRewriteError - Effective contracts differ, or multiple roots disable - automatic sync. - """ - - roots = {self._canonical_temp_storage_ctor_key(key) for key in keys} - contracts = { - self._temp_storage_contract( - self._temp_storage_ctor_specifications[key] - ) - for key in roots - } - if len(contracts) != 1: - ordered_keys = sorted( - roots, - key=lambda key: ( - self._temp_storage_ctor_order.get(key, 1 << 30), - key, - ), - ) - names = ", ".join(ordered_keys) - raise CoopSinglePhaseRewriteError( - f"TempStorage aliases have inconsistent contracts across " - f"constructor instances ({names})." - ) - if len(roots) > 1 and not next(iter(contracts))[2]: - # Caller synchronization may suffice, but this planner cannot - # prove that it protects reuse after the regions merge. - ordered = sorted( - roots, - key=lambda key: ( - self._temp_storage_ctor_order.get(key, 1 << 30), - key, - ), - ) - raise self._collapsed_manual_sync_error( - ", ".join(repr(key) for key in ordered), len(roots) - ) - canonical = min( - roots, - key=lambda key: ( - self._temp_storage_ctor_order.get(key, 1 << 30), - key, - ), - ) - for key, root in tuple(self._temp_storage_ctor_roots.items()): - if self._canonical_temp_storage_ctor_key(root) in roots: - self._temp_storage_ctor_roots[key] = canonical - for key in roots | keys: - self._temp_storage_ctor_roots[key] = canonical - return canonical - - def _resolve_temp_storage_ctor_key(self, value: ir.Var) -> str | None: - """Find and unify the constructors that can supply a storage value. - - Return ``None`` when no owner is known. Conflicting constructor - contracts or an unsafe merge propagate a rewrite diagnostic. - """ - - if not isinstance(value, ir.Var): - return None - keys = self._collect_temp_storage_ctor_keys(value, seen=set()) - if not keys: - return None - return self._merge_temp_storage_ctor_keys(keys) - - def _resolve_temp_storage_plan( - self, value: ir.Var - ) -> _TempStoragePlan | None: - """Find a descriptor's layout, assigning global offsets when needed. - - Canonicalize aliases first. Return ``None`` for an unknown descriptor; - otherwise reuse or finalize its plan against the known requirements. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - key = self._resolve_temp_storage_ctor_key(value) - if key is None: - return None - if ( - self._temp_storage_global_plan is None - and self._temp_storage_ctor_specifications - ): - rewrite._ensure_temp_storage_global_plan() - return self._finalize_temp_storage_plan_for_var(key) - - @staticmethod - def _temp_storage_domain_key( - entry: _TempStorageUseRequirement, - ) -> tuple[object, ...]: - """Identify storage uses that may reuse one region. - - Legacy block providers share one domain. Caller-owned storage also - uses one domain, preserving its explicit reuse contract; its - block-only restriction is checked by storage-plan validation. - Implementation-owned storage instead partitions uses by group topology - and reuse-barrier scope so incompatible group instances cannot alias. - - When a group has an execution scope but no reuse barrier, append the - call order to its key. Such calls receive distinct domains even if - their topologies match: completion before scratch reuse is uncertain. - - Parameters - ---------- - entry : _TempStorageUseRequirement - Primitive use with its lowering plan and scan-order position. - - Returns - ------- - tuple of object - Placement key for ``_layout_temp_storage_uses``. - - Raises - ------ - CoopSinglePhaseRewriteError - The lowering plan is unsupported, lacks storage contracts, or - its storage instances disagree with its topology. - """ - - lowering_plan = entry.lowering_plan - if lowering_plan is None: - return ("legacy-provider",) - if lowering_plan.unsupported is not None: - raise CoopSinglePhaseRewriteError( - "cooperative provider storage received an unsupported group " - "lowering plan." - ) - topology = lowering_plan.topology - synchronization = lowering_plan.synchronization - storage = lowering_plan.temp_storage - if topology is None or synchronization is None or storage is None: - raise CoopSinglePhaseRewriteError( - "cooperative provider storage requires complete group " - "topology, synchronization, and storage contracts." - ) - if storage.ownership is StorageOwnership.NONE: - raise CoopSinglePhaseRewriteError( - "a storage-bearing cooperative provider received a " - "storage-free lowering plan." - ) - if storage.instances != topology.instances or ( - storage.instance_index != topology.instance_index - ): - raise CoopSinglePhaseRewriteError( - "cooperative provider storage layout disagrees with its " - "group topology." - ) - if storage.ownership is StorageOwnership.CALLER: - # Explicit single-block storage keeps the caller's reuse layout. - # The backend separates topology domains only when it owns the - # storage and controls every slice. - return ("caller-storage",) - domain = ( - topology.group_kind, - topology.execution_scope.value, - topology.logical_width, - topology.instances, - topology.instance_index, - topology.thread_rank, - synchronization.storage_reuse_barrier.value, - ) - if ( - synchronization.storage_reuse_barrier is SynchronizationScope.NONE - and topology.execution_scope is not SynchronizationScope.NONE - ): - return (*domain, "non-reusable", entry.order) - return domain - - def _layout_temp_storage_uses( - self, - uses: list[_TempStorageUseRequirement], - *, - sharing: str, - ) -> tuple[int, int, dict[int, _TempStorageSlice]]: - """Lay out scratch for each call and group instance. - - Exclusive sharing gives each use a distinct domain. Shared placement - reuses a domain only when ``_temp_storage_domain_key`` permits it. - Within a domain, reserve the largest per-instance requirement and - align its stride for every consumer. Multiple group instances receive - separate strides; compatible calls reuse those same instance slots. - - Offsets are relative to the region, before the function-wide backing's - base offset is assigned. Views retain each call's actual byte count - even when another call determines the larger shared stride. No - lifetime or control-flow overlap analysis is performed here. - - Parameters - ---------- - uses : list of _TempStorageUseRequirement - Nonempty list of validated requirements and lowering plans. - sharing : str - Validated sharing policy: ``"exclusive"`` separates every - call; ``"shared"`` permits reuse within compatible domains. - - Returns - ------- - required_size : int - Region size in bytes, including alignment gaps and all instances. - required_alignment : int - Maximum required alignment, at least that of the storage pointer. - slices_by_call_id : dict of int to _TempStorageSlice - Region-relative offset, byte count, instance stride, and lowering - plan keyed by the object identity of each call assignment. - """ - - ordered_uses = sorted(uses, key=lambda entry: entry.order) - required_alignment = max( - _MIN_TEMP_STORAGE_ALIGNMENT, - *(max(1, int(entry.alignment)) for entry in ordered_uses), - ) - domains: dict[tuple[object, ...], list[_TempStorageUseRequirement]] = {} - for entry in ordered_uses: - domain_key = ( - ("exclusive", entry.order) - if sharing == "exclusive" - else self._temp_storage_domain_key(entry) - ) - domains.setdefault(domain_key, []).append(entry) - - required_size = 0 - slices_by_call_id: dict[int, _TempStorageSlice] = {} - for domain_uses in domains.values(): - domain_alignment = max( - _MIN_TEMP_STORAGE_ALIGNMENT, - *(max(1, int(entry.alignment)) for entry in domain_uses), - ) - per_instance_size = max( - int(entry.size_in_bytes) for entry in domain_uses - ) - instance_stride = _align_up(per_instance_size, domain_alignment) - # Topology tells us how many independent groups need scratch: four - # 32-thread warps in a 128-thread block need four separate slots. - # Calls in a reuse domain agree on that count, so the first plan - # is sufficient; an exclusive domain contains only one call. - # Legacy block providers have no topology and need one slot. The - # aligned instance_stride spaces groups' slots apart, while - # compatible calls reuse the same set of slots. - first_plan = domain_uses[0].lowering_plan - instances = ( - 1 - if first_plan is None or first_plan.topology is None - else first_plan.topology.instances - ) - required_size = _align_up(required_size, domain_alignment) - domain_offset = required_size - for entry in domain_uses: - slices_by_call_id[id(entry.call_assign)] = _TempStorageSlice( - offset=domain_offset, - size_in_bytes=int(entry.size_in_bytes), - stride=instance_stride, - instances=instances, - lowering_plan=entry.lowering_plan, - ) - required_size += ( - per_instance_size - if instances == 1 - else instance_stride * instances - ) - return required_size, required_alignment, slices_by_call_id - - def _finalize_temp_storage_plan_for_var( - self, var_name: str - ) -> _TempStoragePlan: - """Finalize the storage layout for one ``TempStorage`` variable. - - Requirement collection has already recorded the constructor options - and the bytes needed by each primitive using this variable. Turn - those requirements into a plan for its shared memory: the total size, - alignment, and slice assigned to each call. The constructor's - ``sharing`` setting controls which calls may reuse the same bytes. - - Infer an omitted ``size_in_bytes`` from the layout, or check that an - explicit size is large enough. Choose an alignment that satisfies the - constructor and all calls, with a pointer-sized minimum. Record an - omitted ``auto_sync`` as false, leaving synchronization to the caller. - - Cache the plan under the constructor's canonical variable name. New - plans describe offsets within this variable's storage region; - ``_ensure_temp_storage_global_plan`` later assigns the region's base - offset within the function's shared-memory allocation. - - Parameters - ---------- - var_name : str - IR variable name identifying a recorded ``TempStorage`` - constructor. Merged constructor names share one cached plan. - - Returns - ------- - _TempStoragePlan - Existing cached plan or the newly computed layout and - sharing/synchronization settings. - - Raises - ------ - CoopSinglePhaseRewriteError - Constructor metadata is missing, size cannot be inferred or is - invalid or insufficient, or alignment is invalid. With no - recorded uses, this helper requires an explicit size, - ``sharing="shared"``, and ``auto_sync=False`` (or omitted). - - Notes - ----- - Whole-function validation separately rejects constructors with no - primitive consumer. An explicit size alone does not make an unused - descriptor valid. - """ - - var_name = self._canonical_temp_storage_ctor_key(var_name) - cached = self._temp_storage_plans.get(var_name) - if cached is not None: - return cached - ctor_specification = self._temp_storage_ctor_specifications.get( - var_name - ) - if ctor_specification is None: - raise CoopSinglePhaseRewriteError( - f"Missing TempStorage constructor metadata for " - f"variable '{var_name}'." - ) - requirements = self._func_temp_storage_requirements.get(var_name) - uses = list(requirements.uses) if requirements is not None else [] - uses.sort(key=lambda entry: entry.order) - if uses: - ( - required_size, - required_alignment, - slices_by_call_id, - ) = self._layout_temp_storage_uses( - uses, - sharing=ctor_specification.sharing, - ) - else: - required_size = 0 - required_alignment = max( - _MIN_TEMP_STORAGE_ALIGNMENT, - int(ctor_specification.alignment or 1), - ) - slices_by_call_id = {} - if ctor_specification.size_in_bytes is None: - if required_size <= 0: - raise CoopSinglePhaseRewriteError( - "TempStorage size_in_bytes must be specified until a " - "cooperative primitive provides a storage requirement." - ) - size_in_bytes = required_size - else: - size_in_bytes = int(ctor_specification.size_in_bytes) - if size_in_bytes <= 0: - raise CoopSinglePhaseRewriteError( - "TempStorage size_in_bytes must be a positive integer." - ) - if required_size > 0 and size_in_bytes < required_size: - raise CoopSinglePhaseRewriteError( - f"TempStorage size_in_bytes is smaller than required by " - f"primitive uses ({size_in_bytes} < {required_size})." - ) - if ctor_specification.alignment is None: - alignment = _default_temp_storage_alignment(required_alignment) - else: - requested_alignment = _normalize_temp_storage_alignment( - ctor_specification.alignment - ) - alignment = max( - requested_alignment, - _default_temp_storage_alignment(required_alignment), - ) - _validate_temp_storage_alignment(alignment) - auto_sync = ( - False - if ctor_specification.auto_sync is None - else ctor_specification.auto_sync - ) - if not uses and (ctor_specification.sharing != "shared" or auto_sync): - raise CoopSinglePhaseRewriteError( - "TempStorage non-default sharing or auto_sync requires a " - "cooperative primitive to consume the storage descriptor." - ) - plan = _TempStoragePlan( - size_in_bytes=size_in_bytes, - alignment=alignment, - sharing=ctor_specification.sharing, - auto_sync=auto_sync, - slices_by_call_id=slices_by_call_id, - ) - self._temp_storage_plans[var_name] = plan - return plan - - def _is_local_array_ctor_call(self, call: ir.Expr) -> bool: - return self._resolve_python_value(call.func) is _cuda_local_array - - def _extract_local_array_specification( - self, call: ir.Expr - ) -> _ThreadDataSpecification: - """Recover a native local array's one-dimensional extent and dtype. - - Read positional or keyword constructor inputs. Leave unavailable facts - as ``None``; normal compiler typing still validates the native call. - """ - - kw_map = {name: value for name, value in call.kws} - items_ref = None - if call.args: - items_ref = call.args[0] - elif "shape" in kw_map: - items_ref = kw_map["shape"] - dtype_ref = None - if len(call.args) >= 2: - dtype_ref = call.args[1] - elif "dtype" in kw_map: - dtype_ref = kw_map["dtype"] - items_per_thread = None - if items_ref is not None: - items_per_thread = self._extract_1d_extent_literal(items_ref) - dtype = None - if dtype_ref is not None: - dtype = self._resolve_dtype_ref(dtype_ref) - return _ThreadDataSpecification( - items_per_thread=items_per_thread, dtype=dtype - ) - - def _is_shared_array_ctor_call(self, call: ir.Expr) -> bool: - return self._resolve_python_value(call.func) is _cuda_shared_array - - def _extract_shared_array_specification( - self, call: ir.Expr - ) -> _ThreadDataSpecification: - """Recover a native shared array's extent and dtype when known. - - Use the payload-description record to carry these facts without - claiming that the array is a public ``ThreadData`` value. - """ - - kw_map = {name: value for name, value in call.kws} - shape_ref = None - if call.args: - shape_ref = call.args[0] - elif "shape" in kw_map: - shape_ref = kw_map["shape"] - dtype_ref = None - if len(call.args) >= 2: - dtype_ref = call.args[1] - elif "dtype" in kw_map: - dtype_ref = kw_map["dtype"] - extent = None - if shape_ref is not None: - extent = self._extract_1d_extent_literal(shape_ref) - dtype = None - if dtype_ref is not None: - dtype = self._resolve_dtype_ref(dtype_ref) - return _ThreadDataSpecification(items_per_thread=extent, dtype=dtype) - - def _resolve_array_specification_from_var( - self, value: ir.Var, seen: set[str] - ) -> _ThreadDataSpecification | None: - """Merge array facts across constructors, aliases, and branches. - - Recognize public payloads and native local or shared arrays. ``seen`` - stops recursive paths; unknown paths contribute no facts. The result - is partial inference, not proof that every path supplies an array. - """ - - if not isinstance(value, ir.Var): - return None - if value.name in seen: - return None - seen.add(value.name) - merged: _ThreadDataSpecification | None = None - for definition in self._lookup_definitions(value): - candidate: _ThreadDataSpecification | None = None - if isinstance(definition, ir.Expr): - if definition.op == "call": - if self._is_thread_data_ctor_call(definition): - candidate = self._extract_thread_data_specification( - definition - ) - elif self._is_local_array_ctor_call(definition): - candidate = self._extract_local_array_specification( - definition - ) - elif self._is_shared_array_ctor_call(definition): - candidate = self._extract_shared_array_specification( - definition - ) - elif definition.op == "cast": - cast_value = getattr(definition, "value", None) - if isinstance(cast_value, ir.Var): - candidate = self._resolve_array_specification_from_var( - cast_value, seen - ) - elif definition.op == "static_getitem": - for item in self._resolve_static_tuple_item_vars( - definition - ): - item_specification = ( - self._resolve_array_specification_from_var( - item, set(seen) - ) - ) - if item_specification is None: - continue - merged = self._merge_thread_data_specifications( - merged, item_specification - ) - continue - elif definition.op == "phi": - for incoming in _phi_incoming_values(definition): - if not isinstance(incoming, ir.Var): - continue - incoming_specification = ( - self._resolve_array_specification_from_var( - incoming, set(seen) - ) - ) - if incoming_specification is None: - continue - merged = self._merge_thread_data_specifications( - merged, incoming_specification - ) - continue - elif isinstance(definition, ir.Var): - candidate = self._resolve_array_specification_from_var( - definition, seen - ) - if candidate is not None: - merged = self._merge_thread_data_specifications( - merged, candidate - ) - return merged - - def _resolve_thread_data_specification_from_var( - self, value: ir.Var, seen: set[str] - ) -> _ThreadDataSpecification | None: - """Infer payload shape and dtype through variable origins. - - Follow aliases, casts, static tuple selections, and phi inputs to - ``ThreadData`` or local-array constructors. Reuse complete cached - specifications; otherwise merge discovered facts with any partial - cached information and cache the result. Conflicting known extents or - dtypes are errors, while alignment constraints merge by taking the - larger minimum. - - Unrecognized or cyclic paths contribute no facts. A returned - specification is therefore partial inference, not proof that every - path is a public payload. ``_is_thread_data_like_var`` supplies that - stricter origin check when rewriting the public ``items_per_thread`` - attribute. Shared arrays are handled by the separate - array-specification resolver. - - Parameters - ---------- - value : ir.Var - Variable whose payload facts are needed. - seen : set of str - Active traversal names, extended in place. Branches receive - copies so each incoming path can contribute its facts. - - Returns - ------- - _ThreadDataSpecification or None - Merged known facts, possibly with an unresolved dtype, or None - when no payload specification can be recovered. - - Raises - ------ - CoopSinglePhaseRewriteError - A constructor is invalid, or incoming payload facts conflict. - """ - - if not isinstance(value, ir.Var): - return None - cached = self._thread_data_specifications.get(value.name) - if ( - cached is not None - and cached.items_per_thread is not None - and (cached.dtype is not None) - ): - return cached - if value.name in seen: - return None - seen.add(value.name) - merged: _ThreadDataSpecification | None = cached - for definition in self._lookup_definitions(value): - candidate: _ThreadDataSpecification | None = None - if isinstance(definition, ir.Expr): - if definition.op == "call": - if self._is_thread_data_ctor_call(definition): - candidate = self._extract_thread_data_specification( - definition - ) - elif self._is_local_array_ctor_call(definition): - candidate = self._extract_local_array_specification( - definition - ) - elif definition.op == "cast": - cast_value = getattr(definition, "value", None) - if isinstance(cast_value, ir.Var): - candidate = ( - self._resolve_thread_data_specification_from_var( - cast_value, seen - ) - ) - elif definition.op == "static_getitem": - for item in self._resolve_static_tuple_item_vars( - definition - ): - item_specification = ( - self._resolve_thread_data_specification_from_var( - item, set(seen) - ) - ) - if item_specification is None: - continue - merged = self._merge_thread_data_specifications( - merged, item_specification - ) - continue - elif definition.op == "phi": - for incoming in _phi_incoming_values(definition): - if not isinstance(incoming, ir.Var): - continue - incoming_specification = ( - self._resolve_thread_data_specification_from_var( - incoming, set(seen) - ) - ) - if incoming_specification is None: - continue - merged = self._merge_thread_data_specifications( - merged, incoming_specification - ) - continue - elif isinstance(definition, ir.Var): - candidate = self._resolve_thread_data_specification_from_var( - definition, seen - ) - if candidate is not None: - merged = self._merge_thread_data_specifications( - merged, candidate - ) - if merged is not None: - self._thread_data_specifications[value.name] = merged - return merged - - def _resolve_static_tuple_item_vars( - self, definition: ir.Expr - ) -> list[ir.Var]: - """Recover candidate variables for a static tuple item.""" - - index = getattr(definition, "index", None) - tuple_value = getattr(definition, "value", None) - if not isinstance(index, int) or not isinstance(tuple_value, ir.Var): - return [] - return self._resolve_tuple_item_vars(tuple_value, index, seen=set()) - - def _resolve_tuple_item_vars( - self, tuple_value: ir.Var, index: int, seen: set[str] - ) -> list[ir.Var]: - """Trace a selected tuple element back to its possible source variables. - - A payload or descriptor may pass through a tuple before a cooperative - call uses it. Provenance queries use this walk to recover the selected - item without treating the whole tuple as a payload. Keep alternatives - from merged control-flow paths so later checks can compare their facts. - - Parameters - ---------- - tuple_value : ir.Var - Variable holding the tuple or an alias of it. - index : int - Compile-time element index. Valid negative indices count from the - end. - seen : set of str - Names already visited on this recursive path. The method adds the - current name; recursive alternatives receive copies to detect cycles - without suppressing a separate branch. - - Returns - ------- - list of ir.Var - Recognized source variables. Unknown, cyclic, and out-of-range paths - contribute no entries, so this result may be incomplete. - """ - - if tuple_value.name in seen: - return [] - seen.add(tuple_value.name) - items: list[ir.Var] = [] - for tuple_definition in self._lookup_definitions(tuple_value): - if isinstance(tuple_definition, ir.Var): - items.extend( - self._resolve_tuple_item_vars( - tuple_definition, index, seen=set(seen) - ) - ) - continue - if not isinstance(tuple_definition, ir.Expr): - continue - if tuple_definition.op == "build_tuple": - tuple_items = tuple(getattr(tuple_definition, "items", ())) - if -len(tuple_items) <= index < len(tuple_items): - item = tuple_items[index] - if isinstance(item, ir.Var): - items.append(item) - continue - if tuple_definition.op in {"cast", "exhaust_iter"}: - source = getattr(tuple_definition, "value", None) - if isinstance(source, ir.Var): - items.extend( - self._resolve_tuple_item_vars( - source, index, seen=set(seen) - ) - ) - continue - if tuple_definition.op == "phi": - for incoming in _phi_incoming_values(tuple_definition): - if isinstance(incoming, ir.Var): - items.extend( - self._resolve_tuple_item_vars( - incoming, index, seen=set(seen) - ) - ) - return items - - def _resolve_thread_data_specification( - self, value: ir.Var - ) -> _ThreadDataSpecification | None: - if not isinstance(value, ir.Var): - return None - return self._resolve_thread_data_specification_from_var( - value, seen=set() - ) - - def _is_thread_data_like_var( - self, value: ir.Var, seen: set[str] | None = None - ) -> bool: - """Check that known origins identify a public ``ThreadData`` payload. - - Unlike shape inference, this check rejects an incoming non-payload - origin, including a native local array, even when another branch is - ``ThreadData``. Trace aliases, casts, iterator exhaustion, tuple - items, and phi inputs; cycle-only paths remain unknown rather than - proving or disproving origin. At least one positive origin and no - negative origin are needed to cache a positive result. This prevents - rewriting ``items_per_thread`` on a mixed or unrelated object just - because some shape information is available. - - Parameters - ---------- - value : ir.Var - Candidate receiver of the public payload attribute. - seen : set of str or None, optional - Initial traversal guard. A copy is used, leaving the supplied - set intact. - - Returns - ------- - bool - Whether public payload provenance is established. Positive - results are cached in ``_thread_data_like_vars``; unknown - results return False. - """ - - def resolve(candidate: ir.Var, active: set[str]) -> bool | None: - if not isinstance(candidate, ir.Var): - return False - if candidate.name in self._thread_data_like_vars: - return True - if candidate.name in active: - return None - active.add(candidate.name) - states: list[bool | None] = [] - for definition in self._lookup_definitions(candidate): - if isinstance(definition, ir.Var): - states.append(resolve(definition, set(active))) - continue - if not isinstance(definition, ir.Expr): - states.append(False) - continue - if definition.op == "call": - states.append(self._is_thread_data_ctor_call(definition)) - continue - sources: list[ir.Var] = [] - if definition.op in {"cast", "exhaust_iter"}: - source = getattr(definition, "value", None) - if isinstance(source, ir.Var): - sources.append(source) - elif definition.op == "phi": - sources.extend( - incoming - for incoming in _phi_incoming_values(definition) - if isinstance(incoming, ir.Var) - ) - elif definition.op == "static_getitem": - sources.extend( - self._resolve_static_tuple_item_vars(definition) - ) - if not sources: - states.append(False) - continue - states.extend( - resolve(source, set(active)) for source in sources - ) - if False in states: - return False - if any(state is True for state in states): - self._thread_data_like_vars.add(candidate.name) - return True - return None - - return resolve(value, set() if seen is None else set(seen)) is True - - def _resolve_var_numba_type(self, value: ir.Var): - """Read an existing compiler type without starting a new typing pass. - - Prefer the state's type map, then known argument names and argument - definitions. Return ``None`` when those sources have no type. - """ - - typemap = getattr(self._state, "typemap", None) - if isinstance(typemap, dict): - mapped = typemap.get(value.name) - if mapped is not None: - return mapped - if value.name in self._arg_type_map: - return self._arg_type_map[value.name] - definition = self._lookup_definition(value) - if isinstance(definition, ir.Arg): - arg_types = tuple(getattr(self._state, "args", ()) or ()) - if 0 <= definition.index < len(arg_types): - return arg_types[definition.index] - return None - - def _resolve_call_result_dtype(self, definition: ir.Expr): - """Probe a recognized scalar callee's result type from operand dtypes. - - Resolve the Python callee first, then use the shared scalar call - rules. Unknown callees or unsupported calls yield no dtype. - """ - - func_obj = None - func_ref = getattr(definition, "func", None) - if isinstance(func_ref, ir.Var): - try: - func_obj = self._infer_constant(func_ref) - except (*_INFERENCE_EXCEPTIONS, ImportError): - func_obj = None - if func_obj is None: - func_def = self._lookup_definition(func_ref) - if isinstance(func_def, (ir.Global, ir.FreeVar, ir.Const)): - func_obj = func_def.value - elif isinstance(func_ref, (ir.Global, ir.FreeVar, ir.Const)): - func_obj = func_ref.value - if func_obj is None: - try: - func_obj = self._resolve_python_value(func_ref) - except _INFERENCE_EXCEPTIONS: - func_obj = None - if func_obj is None: - return None - return scalar_call_dtype( - func_obj, definition.args, self._resolve_var_dtype - ) - - @staticmethod - def _merge_scalar_dtypes(dtypes): - """Keep a scalar dtype only when every candidate is known and agrees. - - Normalize equivalent spellings before comparison. This does not - promote conflicting types or infer a dtype from only one known branch. - """ - - candidates = list(dtypes) - if not candidates or any(dtype is None for dtype in candidates): - return None - resolved = [] - for dtype in candidates: - try: - dtype = normalize_dtype_param(dtype) - except (TypeError, ValueError): - pass - if any( - _dtype_values_match(dtype, existing) for existing in resolved - ): - continue - resolved.append(dtype) - return resolved[0] if len(resolved) == 1 else None - - def _cuda_index_dtype(self, definition: ir.Expr): - return cuda_index_dtype( - definition, self._resolve_attribute_chain, _cuda_module - ) - - def _resolve_definition_dtype(self, definition, *, seen: set[str]): - """Infer an element or scalar dtype from one recognized IR definition. - - Trace argument types, scalar values, array accesses, operators, and - calls. Copy ``seen`` into recursive paths and require agreement at - control-flow joins. Unsupported expressions return ``None``. - """ - - from ._parameters import _python_scalar_dtype - - if isinstance(definition, ir.Arg): - arg_types = tuple(getattr(self._state, "args", ()) or ()) - if 0 <= definition.index < len(arg_types): - arg_type = arg_types[definition.index] - return getattr(arg_type, "dtype", arg_type) - return None - if isinstance(definition, (ir.Const, ir.Global, ir.FreeVar)): - return _python_scalar_dtype(definition.value) - if isinstance(definition, ir.Var): - return self._resolve_var_dtype(definition, seen=set(seen)) - if not isinstance(definition, ir.Expr): - return None - if definition.op in {"cast", "exhaust_iter"}: - source = getattr(definition, "value", None) - if isinstance(source, ir.Var): - return self._resolve_var_dtype(source, seen=set(seen)) - return None - if definition.op == "getattr": - cuda_dtype = self._cuda_index_dtype(definition) - if cuda_dtype is not None: - return cuda_dtype - if definition.attr == "dtype" and isinstance( - definition.value, ir.Var - ): - return self._resolve_var_dtype(definition.value, seen=set(seen)) - return None - if definition.op in {"getitem", "static_getitem"}: - base_value = getattr(definition, "value", None) - if isinstance(base_value, ir.Var): - return self._resolve_var_dtype(base_value, seen=set(seen)) - return None - if definition.op in {"binop", "inplace_binop", "unary"}: - return scalar_expression_dtype( - definition, - lambda value: self._resolve_var_dtype(value, seen=set(seen)), - ) - if definition.op == "phi": - return self._merge_scalar_dtypes( - self._resolve_var_dtype(incoming, seen=set(seen)) - for incoming in _phi_incoming_values(definition) - if isinstance(incoming, ir.Var) - ) - if definition.op == "call": - return self._resolve_call_result_dtype(definition) - return None - - def _infer_thread_data_dtype_from_writes(self, value: ir.Var): - """Infer one payload dtype from its known element assignments. - - Inspect writes only for a recognized payload. Reject differing write - dtypes, record a consistent result on its aliases, and return ``None`` - when the write analysis supplies no dtype. - """ - - if self._resolve_thread_data_specification(value) is None: - return None - inferred = None - for value_dtype in payload_write_dtypes( - self._func_ir, value, self._resolve_var_dtype - ): - if inferred is None: - inferred = value_dtype - elif inferred != value_dtype: - raise CoopSinglePhaseRewriteError( - "Failed to infer a consistent dtype " - "from coop.ThreadData writes." - ) - if inferred is not None: - self._record_inferred_thread_data_dtype(value, inferred) - return inferred - - def _resolve_var_dtype(self, value: ir.Var, seen: set[str] | None = None): - """Infer a dtype from payload facts, types, or definitions. - - Prefer an established payload dtype, then an array element or scalar - compiler type. Otherwise require all recovered definition dtypes to - agree. ``seen`` guards recursion; unknown paths remain unresolved. - """ - - if seen is None: - seen = set() - if value.name in seen: - return None - seen.add(value.name) - specification = self._resolve_thread_data_specification(value) - if specification is not None and specification.dtype is not None: - return specification.dtype - var_type = self._resolve_var_numba_type(value) - dtype = getattr(var_type, "dtype", None) - if dtype is not None: - return dtype - if var_type is not None and hasattr(var_type, "bitwidth"): - return var_type - return self._merge_scalar_dtypes( - self._resolve_definition_dtype(definition, seen=set(seen)) - for definition in self._lookup_definitions(value) - ) - - def _resolve_dtype_ref(self, value_ref): - """Resolve a requested dtype from a constant or an operand's type. - - If ordinary constant inference fails, recognize an array's ``dtype`` - attribute or ask for the variable's known dtype. Return ``None`` when - neither path establishes one. - """ - - try: - return self._infer_constant(value_ref) - except _INFERENCE_EXCEPTIONS: - pass - if isinstance(value_ref, ir.Var): - definition = self._lookup_definition(value_ref) - if ( - isinstance(definition, ir.Expr) - and definition.op == "getattr" - and definition.attr == "dtype" - and isinstance(definition.value, ir.Var) - ): - return self._resolve_var_dtype(definition.value) - return self._resolve_var_dtype(value_ref) - return None - - def _resolve_factory_kwarg_value(self, op_name: str, name: str, value_ref): - """Resolve a factory input without losing an explicit ``None`` value. - - Use the operation's dtype-keyword rules before constant inference. - Compiler argument types can also establish an omitted or explicit - ``None``. Return ``_UNRESOLVED`` for other unknown inputs so argument - validation can distinguish missing facts from an omitted option. - """ - - from ._operations import rewrite_operation - - specification = rewrite_operation(op_name) - if ( - specification is not None - and name in specification.dtype_factory_kwargs - ): - dtype = self._resolve_dtype_ref(value_ref) - if dtype is not None: - return dtype - try: - return self._infer_constant(value_ref) - except _INFERENCE_EXCEPTIONS: - pass - if isinstance(value_ref, ir.Var): - import numba_cuda_mlir.numba_cuda.types as numba_mlir_types - - value_type = self._arg_type_map.get(value_ref.name) - definition = self._lookup_definition(value_ref) - if isinstance(definition, ir.Arg): - value_type = self._state.args[definition.index] - if isinstance(value_type, numba_mlir_types.NoneType) or ( - isinstance(value_type, numba_mlir_types.Omitted) - and value_type.value is None - ): - return None - return _UNRESOLVED - - def _resolve_call_target(self, call: ir.Expr) -> _ResolvedCallTarget | None: - """Resolve a registered factory and any subscripted storage operand. - - Recognize direct provider calls and the ``factory[storage](...)`` IR - form by callable identity. Descriptor-use validation later decides - whether the storage syntax is allowed; this lookup only records it. - - Parameters - ---------- - call : ir.Expr - Call whose callee may resolve to a registered factory. - - Returns - ------- - _ResolvedCallTarget or None - Factory registration, callee names, and optional storage operand. - Return ``None`` when the callee is not a registered provider. - - Raises - ------ - CoopSinglePhaseRewriteError - A subscripted provider lacks an IR variable for its storage. - """ - - factory = self._resolve_factory_from_var(call.func) - if factory is not None: - metadata = factory_operation(factory) - assert metadata is not None - return _ResolvedCallTarget( - factory=factory, - factory_metadata=metadata, - func_var_name=call.func.name, - func_var_name_extra=None, - getitem_temp_storage=None, - ) - func_def = self._lookup_definition(call.func) - if not ( - isinstance(func_def, ir.Expr) - and func_def.op in {"getitem", "static_getitem"} - ): - return None - factory = self._resolve_factory_from_var(func_def.value) - if factory is None: - return None - metadata = factory_operation(factory) - assert metadata is not None - getitem_temp_storage = getattr(func_def, "index", None) - if not isinstance(getitem_temp_storage, ir.Var): - getitem_temp_storage = getattr(func_def, "index_var", None) - if not isinstance(getitem_temp_storage, ir.Var): - raise CoopSinglePhaseRewriteError( - f"coop single-phase getitem syntax expects a runtime " - f"temp-storage variable: '{factory.__name__}" - f"[temp_storage](...)'." - ) - return _ResolvedCallTarget( - factory=factory, - factory_metadata=metadata, - func_var_name=call.func.name, - func_var_name_extra=func_def.value.name, - getitem_temp_storage=getitem_temp_storage, - ) - - -__all__ = ["_ProvenanceRewrite"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_storage.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_storage.py deleted file mode 100644 index 02c798a42332..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_storage.py +++ /dev/null @@ -1,1830 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Allocate shared scratch and insert reuse barriers for cooperative calls. - -The call rewrite scans every block before replacing descriptors so all -provider scratch requirements contribute to one function-wide allocation. -Explicit ``TempStorage`` descriptors and implementation-owned scratch receive -aligned regions and per-call views, partitioned where separate group instances -require independent storage. Emission places the backing allocation before its -consumers and adds the synchronization required by each plan's reuse contract. -""" - -from __future__ import annotations - -import operator -from dataclasses import replace -from typing import TYPE_CHECKING, cast - -from numba_cuda_mlir import cuda as _cuda_module -from numba_cuda_mlir.extending import set_required_dynamic_shared_memory -from numba_cuda_mlir.numba_cuda.types import uint8 - -from cuda.coop._core import ( - GroupLoweringPlan, - StorageOwnership, - SynchronizationScope, -) - -from ._operations import ( - StorageABI, - expected_storage_reuse_barrier, - provider_synchronization_matches, -) -from ._rewrite_support import ( - _DEFAULT_STATIC_SHARED_MEMORY_BYTES, - _DYNAMIC_SHARED_MEMORY_ALIGNMENT, - _GLOBAL_NAME_COUNTER, - _INFERENCE_EXCEPTIONS, - CoopSinglePhaseRewriteError, - _align_up, - _next_global_name, - _normalize_temp_storage_alignment, - _phi_incoming_values, - _query_device_shared_memory_limits, - _RewriteMatch, - _TempStorageGlobalPlan, - _TempStoragePlan, - _TempStorageRequirementSummary, - _TempStorageSlice, - _TempStorageUseRequirement, - ir, -) - -if TYPE_CHECKING: - from cuda.coop._core import GroupTopologyRequirements - - from ._rewrite import CoopSinglePhaseRewrite - - -class _StorageRewrite: - """Plan scratch across the function and emit each call's storage view. - - The provenance mixin supplies descriptor owners and region layouts. This - mixin joins those regions, checks placement limits, and writes allocation, - slice, and synchronization statements into the IR. - """ - - def _validate_storage_match_plan( - self, match: _RewriteMatch, *, ctor_key: str | None = None - ) -> None: - """Check storage contracts before allocating slices and barriers. - - Validate the provider registry contract against the group planner's - execution topology, storage ownership, shared address space, and reuse - barrier. Explicit caller-owned storage is supported only for a single - block instance. With manual synchronization, that caller-owned case - may retain the provider's execution-scope synchronization declaration - even though this rewrite emits no automatic reuse barrier. - - The group planner and this rewrite parse descriptors independently. - When a constructor key is available, compare their effective contracts - so a parser disagreement cannot silently suppress synchronization. - Calls without a lowering plan retain only the legacy block execution - and block synchronization contract. Requirement collection runs these - checks before materializing storage-bearing invocables. - - Parameters - ---------- - match : _RewriteMatch - Validated call with a leading storage pointer in its provider ABI. - ctor_key : str or None, optional - Explicit descriptor owner used to cross-check constructor - metadata. None applies to implicit storage or when no owner - was resolved. - - Returns - ------- - None - Successful return permits subsequent requirement collection. - - Raises - ------ - CoopSinglePhaseRewriteError - The plan cannot be emitted or any provider, ownership, - synchronization, or constructor contract disagrees. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - cls = type(self) - lowering_plan = match.lowering_plan - if lowering_plan is None: - if ( - match.factory_metadata.execution_scope - is not SynchronizationScope.BLOCK - or match.factory_metadata.synchronization_scope - is not SynchronizationScope.BLOCK - ): - raise CoopSinglePhaseRewriteError( - "storage-bearing providers without a group lowering plan " - "require block execution and block synchronization scopes" - ) - return - if lowering_plan.unsupported is not None: - raise CoopSinglePhaseRewriteError( - "cooperative provider storage received an unsupported group " - "lowering plan." - ) - topology = cls._validate_emittable_topology(lowering_plan) - synchronization = lowering_plan.synchronization - storage = lowering_plan.temp_storage - if topology is None or synchronization is None or storage is None: - raise CoopSinglePhaseRewriteError( - "cooperative provider storage requires complete group " - "topology, synchronization, and storage contracts." - ) - if ( - match.factory_metadata.execution_scope - is not topology.execution_scope - ): - raise CoopSinglePhaseRewriteError( - "cooperative provider execution scope disagrees with its " - "group topology." - ) - if storage.ownership is StorageOwnership.NONE: - raise CoopSinglePhaseRewriteError( - "a storage-bearing cooperative provider received a " - "storage-free lowering plan." - ) - if storage.address_space != "shared": - raise CoopSinglePhaseRewriteError( - "storage-bearing cooperative providers require " - "shared-address-space TempStorage." - ) - if storage.instances != topology.instances or ( - storage.instance_index != topology.instance_index - ): - raise CoopSinglePhaseRewriteError( - "cooperative provider storage layout disagrees with its " - "group topology." - ) - caller_owned = storage.ownership is StorageOwnership.CALLER - if caller_owned != (match.runtime_temp_storage_var is not None): - raise CoopSinglePhaseRewriteError( - "cooperative provider TempStorage ownership disagrees with " - "its runtime arguments." - ) - if caller_owned and ( - topology.execution_scope is not SynchronizationScope.BLOCK - or topology.instances != 1 - ): - if topology.execution_scope is SynchronizationScope.WARP: - raise CoopSinglePhaseRewriteError( - "cuda.coop.numba_mlir caller-owned TempStorage is not " - "supported for warp-scoped cooperative primitives; omit " - "temp_storage so the implementation can provide one " - "aligned slice per group instance" - ) - raise CoopSinglePhaseRewriteError( - "cuda.coop.numba_mlir caller-owned TempStorage is supported " - "only for single-instance block-scoped cooperative primitives" - ) - expected_reuse_barrier = expected_storage_reuse_barrier( - topology, storage - ) - planned_reuse_barrier = synchronization.storage_reuse_barrier - if planned_reuse_barrier is not expected_reuse_barrier: - raise CoopSinglePhaseRewriteError( - "cooperative provider TempStorage automatic synchronization " - "disagrees with its planned storage-reuse barrier." - ) - if not provider_synchronization_matches( - match.factory_metadata, topology, synchronization, storage - ): - raise CoopSinglePhaseRewriteError( - "cooperative provider synchronization scope disagrees with " - "its group lowering plan." - ) - if caller_owned and ctor_key is not None: - # The group planner and this rewrite parse the same constructor - # independently. The barrier is emitted only when both agree, so - # any drift must fail loudly rather than drop the barrier. - specification = rewrite._temp_storage_ctor_specifications.get( - rewrite._canonical_temp_storage_ctor_key(ctor_key) - ) - if specification is not None: - size, alignment, auto_sync, sharing = ( - rewrite._temp_storage_contract(specification) - ) - planned_alignment = storage.requested_alignment - if planned_alignment is not None: - planned_alignment = _normalize_temp_storage_alignment( - planned_alignment - ) - planned = ( - storage.requested_size_in_bytes, - planned_alignment, - storage.auto_sync, - storage.sharing, - ) - if (size, alignment, auto_sync, sharing) != planned: - raise CoopSinglePhaseRewriteError( - f"cooperative provider TempStorage contract " - f"disagrees between the group lowering plan " - f"{planned!r} and the descriptor " - f"{(size, alignment, auto_sync, sharing)!r}." - ) - - @staticmethod - def _emit_integer_constant( - block: ir.Block, - *, - scope: ir.Scope | None, - loc: ir.Loc, - stem: str, - value: int, - ) -> ir.Var: - result = ir.Var( - scope, - f"__coop_{stem}_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - block.append(ir.Assign(ir.Const(int(value), loc), result, loc)) - return result - - @staticmethod - def _emit_integer_binop( - block: ir.Block, - *, - scope: ir.Scope | None, - loc: ir.Loc, - stem: str, - fn, - lhs: ir.Var, - rhs: ir.Var, - ) -> ir.Var: - result = ir.Var( - scope, - f"__coop_{stem}_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - block.append(ir.Assign(ir.Expr.binop(fn, lhs, rhs, loc), result, loc)) - return result - - def _emit_linear_thread_rank( - self, - block: ir.Block, - *, - scope: ir.Scope | None, - loc: ir.Loc, - ) -> ir.Var: - """Append IR for the CUDA thread's linear rank within its block. - - Compute ``threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * - threadIdx.z)`` from runtime CUDA indices. Storage instance selection - and logical-warp barrier masks need this rank rather than - ``threadIdx.x`` for blocks launched with more than one dimension. - - Parameters - ---------- - block : ir.Block - IR block receiving index reads and integer arithmetic. - scope : ir.Scope or None - Scope assigned to generated variables. - loc : ir.Loc - Source location for generated statements and variables. - - Returns - ------- - ir.Var - Variable containing the linear block rank, with the x - dimension varying fastest. - """ - - module_var = ir.Var( - scope, - f"__coop_group_topology_module_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - thread_idx_var = ir.Var( - scope, - f"__coop_group_topology_thread_idx_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - block_dim_var = ir.Var( - scope, - f"__coop_group_topology_block_dim_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - block.append( - ir.Assign( - ir.Global( - _next_global_name("group_topology_module"), - _cuda_module, - loc, - ), - module_var, - loc, - ) - ) - block.append( - ir.Assign( - ir.Expr.getattr(module_var, "threadIdx", loc), - thread_idx_var, - loc, - ) - ) - block.append( - ir.Assign( - ir.Expr.getattr(module_var, "blockDim", loc), block_dim_var, loc - ) - ) - - components = {} - for aggregate_name, aggregate in ( - ("thread_idx", thread_idx_var), - ("block_dim", block_dim_var), - ): - for component in ("x", "y", "z"): - value = ir.Var( - scope, - f"__coop_group_topology_{aggregate_name}_{component}_" - f"{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - block.append( - ir.Assign( - ir.Expr.getattr(aggregate, component, loc), value, loc - ) - ) - components[aggregate_name, component] = value - - y_stride = self._emit_integer_binop( - block, - scope=scope, - loc=loc, - stem="group_topology_y_stride", - fn=operator.mul, - lhs=components["block_dim", "y"], - rhs=components["thread_idx", "z"], - ) - yz_rank = self._emit_integer_binop( - block, - scope=scope, - loc=loc, - stem="group_topology_yz_rank", - fn=operator.add, - lhs=components["thread_idx", "y"], - rhs=y_stride, - ) - x_stride = self._emit_integer_binop( - block, - scope=scope, - loc=loc, - stem="group_topology_x_stride", - fn=operator.mul, - lhs=components["block_dim", "x"], - rhs=yz_rank, - ) - return self._emit_integer_binop( - block, - scope=scope, - loc=loc, - stem="group_topology_linear_thread_rank", - fn=operator.add, - lhs=components["thread_idx", "x"], - rhs=x_stride, - ) - - @staticmethod - def _validate_emittable_topology( - lowering_plan: GroupLoweringPlan | None, - ) -> GroupTopologyRequirements | None: - """Require group rank formulas that the storage emitters support. - - A plan must cover the exact block dimensions with its logical width - and instance count. Accept a single block, contiguous power-of-two - logical warps dividing 32, or individual threads. Check the symbolic - instance and rank expressions against those forms; emitters implement - these specific formulas rather than evaluating arbitrary topology - expression strings. - - Parameters - ---------- - lowering_plan : GroupLoweringPlan or None - Plan whose topology and participation requirements govern - storage. None preserves the legacy block-provider path. - - Returns - ------- - GroupTopologyRequirements or None - Validated topology, or None when no plan was supplied. - - Raises - ------ - CoopSinglePhaseRewriteError - Dimensions or execution requirements are missing, coverage is - inconsistent, or no emitter supports the scope and rank formulas. - """ - - if lowering_plan is None: - return None - topology = lowering_plan.topology - participation = lowering_plan.participation - if topology is None or participation is None: - raise CoopSinglePhaseRewriteError( - "cooperative provider storage requires group topology and " - "participation contracts." - ) - exact_block_dim = participation.exact_block_dim - if exact_block_dim is None: - raise CoopSinglePhaseRewriteError( - "cooperative provider storage requires exact block dimensions." - ) - block_threads = ( - exact_block_dim[0] * exact_block_dim[1] * exact_block_dim[2] - ) - if topology.logical_width * topology.instances != block_threads: - raise CoopSinglePhaseRewriteError( - "cooperative provider topology does not cover the exact " - "block dimensions." - ) - scope = topology.execution_scope - if scope is SynchronizationScope.BLOCK: - if ( - topology.instances != 1 - or topology.logical_width != block_threads - or topology.instance_index != "cta" - or topology.thread_rank != "linear_thread_rank" - ): - raise CoopSinglePhaseRewriteError( - "block-scoped cooperative storage requires canonical " - "single-CTA ranks." - ) - elif scope is SynchronizationScope.WARP: - width = topology.logical_width - if ( - width < 1 - or width > 32 - or width & (width - 1) - or 32 % width != 0 - or topology.instance_index != f"linear_thread_rank / {width}" - or topology.thread_rank != f"linear_thread_rank % {width}" - ): - raise CoopSinglePhaseRewriteError( - "warp-scoped cooperative storage requires a power-of-two " - "logical width dividing 32 " - "and canonical contiguous ranks." - ) - elif scope is SynchronizationScope.NONE: - if ( - topology.logical_width != 1 - or topology.instance_index != "linear_thread_rank" - or topology.thread_rank != "0" - ): - raise CoopSinglePhaseRewriteError( - "thread-scoped cooperative storage requires canonical " - "per-thread ranks." - ) - else: - raise CoopSinglePhaseRewriteError( - f"cuda.coop.numba_mlir provider execution scope " - f"{scope.value!r} has no storage emitter" - ) - return topology - - def _emit_storage_instance_index( - self, - block: ir.Block, - *, - lowering_plan: GroupLoweringPlan | None, - scope: ir.Scope | None, - loc: ir.Loc, - ) -> ir.Var: - """Select the executing group's region within a scratch allocation. - - Scratch slicing calls this when a block contains several independent - groups. Each warp or thread needs its own region so concurrent - operations do not overwrite one another. Use zero for a block-wide - group, the linear thread rank for a thread, or that rank divided by the - logical warp width. - - Parameters - ---------- - block : ir.Block - Replacement block receiving the index calculation. - lowering_plan : GroupLoweringPlan or None - Validated operation topology that defines group size and rank. - ``None`` uses the single-region block convention of private provider - calls. - scope : ir.Scope or None - Scope for the generated index variables. - loc : ir.Loc - Original call location used for emitted IR and diagnostics. - - Returns - ------- - ir.Var - Zero-based group instance index. This is a region index, not a byte - offset; the caller multiplies it by the aligned region stride. - """ - - topology = self._validate_emittable_topology(lowering_plan) - if ( - topology is None - or topology.execution_scope is SynchronizationScope.BLOCK - ): - return self._emit_integer_constant( - block, - scope=scope, - loc=loc, - stem="group_topology_instance_index", - value=0, - ) - linear_rank = self._emit_linear_thread_rank(block, scope=scope, loc=loc) - if topology.execution_scope is SynchronizationScope.NONE: - return linear_rank - logical_width = self._emit_integer_constant( - block, - scope=scope, - loc=loc, - stem="group_topology_logical_width", - value=topology.logical_width, - ) - return self._emit_integer_binop( - block, - scope=scope, - loc=loc, - stem="group_topology_instance_index", - fn=operator.floordiv, - lhs=linear_rank, - rhs=logical_width, - ) - - def _has_temp_storage_requirements(self) -> bool: - rewrite = cast("CoopSinglePhaseRewrite", self) - implicit = getattr(self, "_implicit_temp_storage_requirements", None) - return bool(rewrite._func_temp_storage_requirements) or bool( - implicit is not None and implicit.uses - ) - - def _get_device_shared_memory_limits( - self, required_bytes: int - ) -> tuple[int, int]: - """Get the default and opt-in byte limits needed for this allocation. - - Requests up to the conservative static limit need no device query. - Larger ``required_bytes`` values require valid limits from the active - CUDA device; a failed query becomes a rewrite error. - """ - - conservative_default = _DEFAULT_STATIC_SHARED_MEMORY_BYTES - if required_bytes <= conservative_default: - return (conservative_default, conservative_default) - try: - limits = _query_device_shared_memory_limits() - max_default = int(limits["max_default_shared_memory_per_block"]) - max_optin = int(limits["max_optin_shared_memory_per_block"]) - except ( - AttributeError, - KeyError, - OSError, - RuntimeError, - TypeError, - ValueError, - ) as exc: - raise CoopSinglePhaseRewriteError( - "TempStorage requirements above the conservative 49152-byte " - "static shared-memory limit require an exact current-device " - "shared-memory query" - ) from exc - if max_default <= 0 or max_optin <= 0 or max_optin < max_default: - raise CoopSinglePhaseRewriteError( - f"The current device reported invalid shared-memory limits: " - f"default={max_default}, opt-in={max_optin}." - ) - return (max_default, max_optin) - - def _ensure_temp_storage_global_plan(self) -> _TempStorageGlobalPlan: - """Plan one shared backing for explicit and implicit scratch. - - Finalize canonical explicit descriptors in constructor order and - assign aligned base offsets, then append the implementation-owned - region. Calls within each region have already been assigned reusable - or exclusive slices. Round the total size to the greatest alignment so - the backing can satisfy every region with one allocation. - - Use static shared memory when the total fits the default device limit; - otherwise request the exact dynamic byte count in compiler metadata. - Small allocations use the conservative limit without querying a - device. Dynamic placement must fit the opt-in limit and the compiler's - dynamic window alignment guarantee. Cache the plan and updated region - offsets. Allocation and checks for user shared arrays happen later. - - Returns - ------- - _TempStorageGlobalPlan - Cached or newly computed total size, maximum alignment, - placement, and required dynamic byte count. - - Raises - ------ - CoopSinglePhaseRewriteError - A descriptor cannot be finalized, exact device limits are - required but unavailable, or size/alignment requirements - exceed supported limits. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - cached = self._temp_storage_global_plan - if cached is not None: - return cached - ordered_keys = sorted( - { - rewrite._canonical_temp_storage_ctor_key(key) - for key in rewrite._func_temp_storage_requirements - }, - key=lambda name: ( - rewrite._temp_storage_ctor_order.get(name, 1 << 30), - name, - ), - ) - offset = 0 - max_alignment = 1 - for key in ordered_keys: - plan = rewrite._finalize_temp_storage_plan_for_var(key) - alignment = max(1, int(plan.alignment)) - offset = _align_up(offset, alignment) - rewrite._temp_storage_plans[key] = replace(plan, base_offset=offset) - offset += int(plan.size_in_bytes) - max_alignment = max(max_alignment, alignment) - implicit = getattr( - self, - "_implicit_temp_storage_requirements", - _TempStorageRequirementSummary(), - ) - if implicit.uses: - ( - implicit_size, - implicit_alignment, - implicit_slices, - ) = rewrite._layout_temp_storage_uses( - implicit.uses, - sharing="shared", - ) - offset = _align_up(offset, implicit_alignment) - implicit_base_offset = offset - self._implicit_temp_storage_plan = _TempStoragePlan( - size_in_bytes=implicit_size, - alignment=implicit_alignment, - sharing="shared", - auto_sync=True, - slices_by_call_id=implicit_slices, - base_offset=implicit_base_offset, - ) - offset += implicit_size - max_alignment = max(max_alignment, implicit_alignment) - else: - self._implicit_temp_storage_plan = None - total_size = _align_up(offset, max_alignment) - max_default, max_optin = self._get_device_shared_memory_limits( - total_size - ) - uses_dynamic_smem = total_size > max_default - if ( - uses_dynamic_smem - and max_alignment > _DYNAMIC_SHARED_MEMORY_ALIGNMENT - ): - # The static path honors the requested alignment through the - # shared array declaration; the dynamic window only guarantees - # its declared alignment, and telling the optimizer otherwise - # would be a false assumption. - raise CoopSinglePhaseRewriteError( - f"TempStorage requires {max_alignment}-byte alignment, but " - f"the {total_size}-byte backing exceeds the " - f"{max_default}-byte static shared-memory limit and dynamic " - f"shared memory guarantees only " - f"{_DYNAMIC_SHARED_MEMORY_ALIGNMENT}-byte alignment; reduce " - f"the requested alignment or the storage size." - ) - dynamic_shared_bytes = total_size if uses_dynamic_smem else 0 - if dynamic_shared_bytes > max_optin: - raise CoopSinglePhaseRewriteError( - f"TempStorage requires {dynamic_shared_bytes} bytes dynamic " - f"shared memory, but device max opt-in is {max_optin} bytes." - ) - if dynamic_shared_bytes > 0: - set_required_dynamic_shared_memory( - rewrite._state, dynamic_shared_bytes - ) - plan = _TempStorageGlobalPlan( - total_size=total_size, - max_alignment=max_alignment, - uses_dynamic_smem=uses_dynamic_smem, - dynamic_shared_bytes=dynamic_shared_bytes, - ) - self._temp_storage_global_plan = plan - return plan - - def _stage_temp_storage_backing(self) -> ir.Var: - """Stage the scratch allocation before rewriting consumers. - - Place generated allocation statements in the entry block immediately - after argument assignments. Block rewrite visitation need not follow - control flow; emitting next to the first visited consumer could leave - other consumers without a dominating definition. Reject conflicting - user shared-memory declarations before inserting the unified backing. - - Repeated calls return the existing backing variable without inserting - a second allocation. The function's entry block is mutated directly, - even when ``apply`` is currently rewriting a different block. - - Returns - ------- - ir.Var - Shared byte array available to every rewritten consumer. - - Raises - ------ - CoopSinglePhaseRewriteError - Allocation planning fails, user shared arrays would overlap, - or the emission state claims a backing exists without - recording its variable. - """ - rewrite = cast("CoopSinglePhaseRewrite", self) - - if self._temp_storage_backing_emitted: - if self._temp_storage_backing_var is None: - raise CoopSinglePhaseRewriteError( - "TempStorage backing was marked emitted " - "without an IR value." - ) - return self._temp_storage_backing_var - plan = self._ensure_temp_storage_global_plan() - self._reject_conflicting_user_shared_arrays(plan) - entry_block = rewrite._func_ir.blocks[min(rewrite._func_ir.blocks)] - staged = ir.Block(entry_block.scope, entry_block.loc) - backing = self._emit_temp_storage_backing(staged, plan=plan) - insert_at = 0 - while insert_at < len(entry_block.body): - statement = entry_block.body[insert_at] - if not ( - isinstance(statement, ir.Assign) - and isinstance(statement.value, ir.Arg) - ): - break - insert_at += 1 - entry_block.body[insert_at:insert_at] = staged.body - entry_block.verify() - return backing - - def _reject_conflicting_user_shared_arrays( - self, plan: _TempStorageGlobalPlan - ) -> None: - """Keep user shared arrays separate from cooperative scratch memory. - - User data must survive while a primitive writes its temporary storage. - The supported compiler can place static shared allocations over the - dynamic shared-memory window, so these allocations can overwrite each - other. A zero-sized ``cuda.shared.array`` declares a view of that - window; a runtime-sized array also uses it. - - Allow user shared arrays only when both they and the cooperative - backing use static shared-memory allocations. A compile-time-known - cooperative size can still require dynamic backing. In that case, - reject every user shared allocation. Check constructors throughout - the function after device-helper inlining, regardless of barriers or - apparent lifetimes. Restore the block lookup state even on failure. - - Parameters - ---------- - plan : _TempStorageGlobalPlan - Combined cooperative allocation, including its static or dynamic - placement and total byte count. - - Raises - ------ - CoopSinglePhaseRewriteError - A user shared array can overlap the cooperative backing. - - Examples - -------- - This kernel keeps user data in dynamic shared memory across a - transpose load. For 128 threads and four int32 items per thread, the - load's static scratch allocation can overlap ``user_data``:: - - import numpy as np - from numba_cuda_mlir import cuda - from cuda import coop - - @cuda.jit - def conflicting_storage(source, output, items_per_thread): - block = coop.this_block() - thread = cuda.threadIdx.x - user_data = cuda.shared.array(0, dtype=np.int32) - user_data[thread] = source[thread] - cuda.syncthreads() - - items = coop.ThreadData(items_per_thread) - coop.load(block, source, items, algorithm="transpose") - items[0] += user_data[thread] - coop.store(block, output, items) - - source = cuda.to_device(np.arange(512, dtype=np.int32)) - output = cuda.device_array(512, dtype=np.int32) - conflicting_storage[1, 128, 0, 512](source, output, 4) - - Compilation rejects this launch before the arrays can overwrite each - other. The 512 launch bytes provide room for ``user_data``; they do - not separate it from the primitive's scratch. Replacing its shape - ``0`` with ``128`` makes both allocations static and passes this - check; that version needs no dynamic launch bytes. - """ - rewrite = cast("CoopSinglePhaseRewrite", self) - - saved_block = self._block - saved_block_defs = self._block_defs - conflicts: list[tuple[str, ir.Loc]] = [] - try: - for label in sorted(rewrite._func_ir.blocks): - scan_block = rewrite._func_ir.blocks[label] - self._block = scan_block - self._block_defs = { - inst.target.name: inst.value - for inst in scan_block.body - if isinstance(inst, ir.Assign) - } - for inst in scan_block.body: - if not isinstance(inst, ir.Assign): - continue - call = inst.value - if not isinstance(call, ir.Expr) or call.op != "call": - continue - if not rewrite._is_shared_array_ctor_call(call): - continue - shape_ref = ( - call.args[0] - if call.args - else dict(call.kws).get("shape") - ) - try: - shape = rewrite._infer_constant(shape_ref) - except _INFERENCE_EXCEPTIONS: - shape = None - dimensions = shape if isinstance(shape, tuple) else (shape,) - try: - is_static = bool(dimensions) and all( - extent is not None and operator.index(extent) > 0 - for extent in dimensions - ) - except TypeError: - is_static = False - if plan.uses_dynamic_smem or not is_static: - placement = ( - "static" if is_static else "dynamic/runtime-sized" - ) - conflicts.append((placement, inst.loc)) - finally: - self._block = saved_block - self._block_defs = saved_block_defs - if not conflicts: - return - placement = "dynamic" if plan.uses_dynamic_smem else "static" - where = ", ".join( - f"{kind} cuda.shared.array(...) at {loc}" - for kind, loc in conflicts[:3] - ) - raise CoopSinglePhaseRewriteError( - f"cuda.coop temporary storage requires a {plan.total_size}-byte " - f"{placement} shared-memory backing, but this kernel also " - f"declares {where}. The supported numba-cuda-mlir compiler does " - f"not separate these allocations; they would alias. Use " - f"statically sized user shared arrays and keep the combined " - f"cooperative backing within the static shared-memory limit, or " - f"move the user data out of shared memory." - ) - - def _emit_temp_storage_backing( - self, block: ir.Block, *, plan: _TempStorageGlobalPlan - ) -> ir.Var: - """Append one shared byte-array allocation and remember its variable. - - ``plan`` supplies capacity, alignment, and static or dynamic - placement. Dynamic placement declares a zero-length shared array, - which the compiler maps to the launch's dynamic shared memory. - Planning has already recorded the required byte count. Repeated - calls return the existing variable without appending IR. - """ - - if self._temp_storage_backing_emitted: - if self._temp_storage_backing_var is None: - raise CoopSinglePhaseRewriteError( - "TempStorage backing was marked emitted " - "without an IR value." - ) - return self._temp_storage_backing_var - loc = block.loc - scope = block.scope - module_var = ir.Var( - scope, - f"__coop_temp_storage_module_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - shared_var = ir.Var( - scope, - f"__coop_temp_storage_shared_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - array_fn_var = ir.Var( - scope, - f"__coop_temp_storage_array_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - bytes_var = ir.Var( - scope, - f"__coop_temp_storage_bytes_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - align_var = ir.Var( - scope, - f"__coop_temp_storage_alignment_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - dtype_var = ir.Var( - scope, - f"__coop_temp_storage_dtype_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - backing_var = ir.Var( - scope, - f"__coop_temp_storage_backing_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - block.append( - ir.Assign( - ir.Global( - _next_global_name("temp_storage_module"), - _cuda_module, - loc, - ), - module_var, - loc, - ) - ) - block.append( - ir.Assign( - ir.Expr.getattr(module_var, "shared", loc), shared_var, loc - ) - ) - block.append( - ir.Assign( - ir.Expr.getattr(shared_var, "array", loc), array_fn_var, loc - ) - ) - alloc_size = 0 if plan.uses_dynamic_smem else int(plan.total_size) - block.append(ir.Assign(ir.Const(alloc_size, loc), bytes_var, loc)) - block.append( - ir.Assign(ir.Const(plan.max_alignment, loc), align_var, loc) - ) - block.append( - ir.Assign( - ir.Global( - _next_global_name("temp_storage_dtype"), - uint8, - loc, - ), - dtype_var, - loc, - ) - ) - block.append( - ir.Assign( - ir.Expr.call( - array_fn_var, - [bytes_var, dtype_var], - (("alignment", align_var),), - loc, - ), - backing_var, - loc, - ) - ) - self._temp_storage_backing_var = backing_var - self._temp_storage_backing_emitted = True - return backing_var - - def _emit_array_slice( - self, - block: ir.Block, - *, - source_var: ir.Var, - target_var: ir.Var, - start: int | ir.Var, - stop: int | ir.Var, - loc: ir.Loc, - ) -> None: - """Assign a view of a scratch allocation to a provider operand. - - Storage-view emission calls this after computing the required bounds. - The generated slice shares the backing array, allowing each provider to - receive its assigned region without allocating or copying scratch data. - Append constants as needed, the slice object, and the view assignment. - - Parameters - ---------- - block : ir.Block - Replacement block to which the slice statements are appended. - source_var : ir.Var - Backing array or descriptor view to slice. - target_var : ir.Var - Variable receiving the view; its scope is used for new temporaries. - start, stop : int or ir.Var - Inclusive start and exclusive stop in array elements. Scratch arrays - contain bytes, so their element indices are also byte offsets. - Bounds may be constants or variables computed for the current group. - loc : ir.Loc - Source location attached to generated statements. - """ - - slice_ctor_global_name = _next_global_name("temp_storage_slice_ctor") - slice_ctor_var = ir.Var( - target_var.scope, - f"__coop_temp_storage_slice_ctor_var_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - start_var = start - if not isinstance(start_var, ir.Var): - start_var = self._emit_integer_constant( - block, - scope=target_var.scope, - loc=loc, - stem="temp_storage_slice_start", - value=start_var, - ) - stop_var = stop - if not isinstance(stop_var, ir.Var): - stop_var = self._emit_integer_constant( - block, - scope=target_var.scope, - loc=loc, - stem="temp_storage_slice_stop", - value=stop_var, - ) - slice_obj_var = ir.Var( - target_var.scope, - f"__coop_temp_storage_slice_obj_{next(_GLOBAL_NAME_COUNTER)}__", - loc, - ) - block.append( - ir.Assign( - ir.Global(slice_ctor_global_name, slice, loc), - slice_ctor_var, - loc, - ) - ) - block.append( - ir.Assign( - ir.Expr.call(slice_ctor_var, [start_var, stop_var], (), loc), - slice_obj_var, - loc, - ) - ) - block.append( - ir.Assign( - ir.Expr.getitem(source_var, slice_obj_var, loc), target_var, loc - ) - ) - - def _emit_temp_storage_slice_for_call( - self, - block: ir.Block, - *, - source_var: ir.Var, - target_var: ir.Var, - slice_info: _TempStorageSlice, - base_offset: int, - loc: ir.Loc, - ) -> None: - """Emit a call's scratch view for the executing group. - - A single-instance view has constant bounds. Multiple instances add the - emitted group index times the aligned instance stride to the region's - static offset. The view length remains the call's byte requirement, - even if the stride includes padding or room for another consumer. - - Parameters - ---------- - block : ir.Block - IR block receiving bound calculations and the slice assignment. - source_var : ir.Var - Shared byte array or descriptor view to slice. - target_var : ir.Var - Variable to receive the view and supply its IR scope. - slice_info : _TempStorageSlice - Offset in the region, call size, instance layout, and topology. - base_offset : int - Region origin relative to ``source_var``. Use zero for an - already sliced explicit descriptor and the region base for the - unified backing. - loc : ir.Loc - Source location for generated statements. - - Returns - ------- - None - ``target_var`` is defined by statements appended to ``block``. - - Raises - ------ - CoopSinglePhaseRewriteError - A multi-instance slice lacks a lowering plan or has - unsupported topology. - """ - - static_start = int(base_offset) + int(slice_info.offset) - if slice_info.instances == 1: - start: int | ir.Var = static_start - stop: int | ir.Var = static_start + int(slice_info.size_in_bytes) - else: - if slice_info.lowering_plan is None: - raise CoopSinglePhaseRewriteError( - "multi-instance cooperative storage requires a group " - "lowering plan." - ) - instance_index = self._emit_storage_instance_index( - block, - lowering_plan=slice_info.lowering_plan, - scope=target_var.scope, - loc=loc, - ) - stride = self._emit_integer_constant( - block, - scope=target_var.scope, - loc=loc, - stem="temp_storage_instance_stride", - value=( - int(slice_info.size_in_bytes) - if slice_info.stride is None - else int(slice_info.stride) - ), - ) - instance_offset = self._emit_integer_binop( - block, - scope=target_var.scope, - loc=loc, - stem="temp_storage_instance_offset", - fn=operator.mul, - lhs=instance_index, - rhs=stride, - ) - domain_offset = self._emit_integer_constant( - block, - scope=target_var.scope, - loc=loc, - stem="temp_storage_domain_offset", - value=static_start, - ) - start = self._emit_integer_binop( - block, - scope=target_var.scope, - loc=loc, - stem="temp_storage_slice_start", - fn=operator.add, - lhs=domain_offset, - rhs=instance_offset, - ) - slice_size = self._emit_integer_constant( - block, - scope=target_var.scope, - loc=loc, - stem="temp_storage_slice_size", - value=slice_info.size_in_bytes, - ) - stop = self._emit_integer_binop( - block, - scope=target_var.scope, - loc=loc, - stem="temp_storage_slice_stop", - fn=operator.add, - lhs=start, - rhs=slice_size, - ) - self._emit_array_slice( - block, - source_var=source_var, - target_var=target_var, - start=start, - stop=stop, - loc=loc, - ) - - def _runtime_temp_storage_arg_for_call( - self, block: ir.Block, *, source_var: ir.Var, call_assign: ir.Assign - ) -> tuple[ir.Var, _TempStoragePlan | None]: - """Choose the explicit scratch view passed to one provider call. - - ``apply`` uses the function-wide storage plan to replace a descriptor - operand with the call's assigned region. Shared storage at offset zero - can reuse the descriptor view. Exclusive storage or a nonzero offset - needs a narrower view so the provider receives the planned region. - - Parameters - ---------- - block : ir.Block - Replacement block receiving a slice assignment when one is needed. - source_var : ir.Var - Current storage operand, used to find its descriptor's plan. - call_assign : ir.Assign - Original provider-call assignment. Its object identity selects the - recorded slice, and its target supplies the scope for emitted - variables. - - Returns - ------- - tuple of ir.Var and _TempStoragePlan or None - Provider storage operand and its plan for subsequent - synchronization. If no descriptor plan exists, return ``(source_var, - None)`` unchanged. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - temp_storage_arg = source_var - temp_storage_plan = rewrite._resolve_temp_storage_plan(source_var) - if temp_storage_plan is not None: - slice_info = temp_storage_plan.slices_by_call_id.get( - id(call_assign) - ) - if slice_info is None: - raise CoopSinglePhaseRewriteError( - f"Could not resolve TempStorage slice for call at " - f"{call_assign.loc}." - ) - if ( - temp_storage_plan.sharing == "exclusive" - or slice_info.offset != 0 - ): - sliced_var = ir.Var( - call_assign.target.scope, - f"__coop_temp_storage_slice_{next(_GLOBAL_NAME_COUNTER)}__", - call_assign.loc, - ) - self._emit_temp_storage_slice_for_call( - block, - source_var=source_var, - target_var=sliced_var, - slice_info=slice_info, - base_offset=0, - loc=call_assign.loc, - ) - temp_storage_arg = sliced_var - return (temp_storage_arg, temp_storage_plan) - - def _implicit_temp_storage_arg_for_call( - self, block: ir.Block, *, call_assign: ir.Assign - ) -> tuple[ir.Var, _TempStoragePlan]: - """Build a scratch operand for a call that omitted a storage descriptor. - - ``apply`` calls this after allocating the function's shared backing - array. The implementation-owned region has already been sized across - implicit consumers; this method selects the current call's part and - returns its plan for the later synchronization step. - - Parameters - ---------- - block : ir.Block - Replacement block receiving the call's scratch-view statements. - call_assign : ir.Assign - Original provider-call assignment. Its object identity selects the - recorded slice, and its target supplies the scope for emitted - variables. - - Returns - ------- - tuple of ir.Var and _TempStoragePlan - Generated scratch operand and the implicit region's plan. - - Raises - ------ - CoopSinglePhaseRewriteError - The backing array, implicit region plan, or call's slice is missing. - """ - - plan = self._implicit_temp_storage_plan - backing = self._temp_storage_backing_var - if plan is None or backing is None: - raise CoopSinglePhaseRewriteError( - "Missing implementation-owned TempStorage plan for an " - "implicit call." - ) - slice_info = plan.slices_by_call_id.get(id(call_assign)) - if slice_info is None: - raise CoopSinglePhaseRewriteError( - f"Could not resolve implicit TempStorage slice for call at " - f"{call_assign.loc}." - ) - sliced_var = ir.Var( - call_assign.target.scope, - f"__coop_implicit_temp_storage_slice_{next(_GLOBAL_NAME_COUNTER)}__", - call_assign.loc, - ) - self._emit_temp_storage_slice_for_call( - block, - source_var=backing, - target_var=sliced_var, - slice_info=slice_info, - base_offset=plan.base_offset, - loc=call_assign.loc, - ) - return (sliced_var, plan) - - def _emit_temp_storage_auto_sync( - self, - block: ir.Block, - *, - scope: ir.Scope | None, - loc: ir.Loc, - synchronization_scope: SynchronizationScope, - lowering_plan: GroupLoweringPlan | None = None, - ) -> None: - """Emit the post-call barrier required for automatic scratch reuse. - - Emit ``syncthreads`` for block scope and ``syncwarp`` for warp scope; - ``NONE`` emits nothing. A logical warp smaller than 32 receives a mask - covering only its contiguous lanes within the physical warp, computed - from the linear block rank. Full warps use the default warp mask. - - This is the post-call storage-reuse barrier selected by the lowering - contract. It neither establishes uniform participation nor decides - whether a call needs synchronization; the caller has already checked - the storage policy and calls this emitter for automatic sync. - - Parameters - ---------- - block : ir.Block - Destination receiving mask calculations and the barrier call. - scope : ir.Scope or None - Scope assigned to generated variables. - loc : ir.Loc - Source location for generated statements and variables. - synchronization_scope : SynchronizationScope - Requested reuse-barrier scope, converted to the enum on entry. - lowering_plan : GroupLoweringPlan or None, optional - Group topology for validation and logical-warp mask - construction. None uses the legacy scope-only emission path. - - Returns - ------- - None - Append barrier IR in place unless the requested scope is ``NONE``. - - Raises - ------ - CoopSinglePhaseRewriteError - The topology is unsupported or conflicts with the requested scope. - """ - - synchronization_scope = SynchronizationScope(synchronization_scope) - if synchronization_scope is SynchronizationScope.NONE: - return - topology = self._validate_emittable_topology(lowering_plan) - if topology is not None and ( - synchronization_scope is not topology.execution_scope - ): - raise CoopSinglePhaseRewriteError( - "cooperative provider synchronization scope disagrees with " - "its group topology." - ) - sync_attr = { - SynchronizationScope.WARP: "syncwarp", - SynchronizationScope.BLOCK: "syncthreads", - }.get(synchronization_scope) - if sync_attr is None: - raise CoopSinglePhaseRewriteError( - f"cuda.coop.numba_mlir provider synchronization scope " - f"{SynchronizationScope(synchronization_scope).value!r} has " - f"no emitter" - ) - sync_args = [] - if ( - synchronization_scope is SynchronizationScope.WARP - and topology is not None - and topology.logical_width < 32 - ): - linear_rank = self._emit_linear_thread_rank( - block, - scope=scope, - loc=loc, - ) - lane_mask = self._emit_integer_constant( - block, - scope=scope, - loc=loc, - stem="group_topology_lane_mask", - value=31, - ) - lane = self._emit_integer_binop( - block, - scope=scope, - loc=loc, - stem="group_topology_lane", - fn=operator.and_, - lhs=linear_rank, - rhs=lane_mask, - ) - logical_width = self._emit_integer_constant( - block, - scope=scope, - loc=loc, - stem="group_topology_logical_width", - value=topology.logical_width, - ) - logical_group = self._emit_integer_binop( - block, - scope=scope, - loc=loc, - stem="group_topology_logical_group", - fn=operator.floordiv, - lhs=lane, - rhs=logical_width, - ) - shift = self._emit_integer_binop( - block, - scope=scope, - loc=loc, - stem="group_topology_mask_shift", - fn=operator.mul, - lhs=logical_group, - rhs=logical_width, - ) - base_mask = self._emit_integer_constant( - block, - scope=scope, - loc=loc, - stem="group_topology_base_mask", - value=(1 << topology.logical_width) - 1, - ) - sync_args.append( - self._emit_integer_binop( - block, - scope=scope, - loc=loc, - stem="group_topology_sync_mask", - fn=operator.lshift, - lhs=base_mask, - rhs=shift, - ) - ) - sync_module_global_name = _next_global_name("temp_storage_sync_mod") - sync_module_var = ir.Var( - scope, f"__coop_sync_mod_var_{next(_GLOBAL_NAME_COUNTER)}__", loc - ) - sync_fn_var = ir.Var( - scope, f"__coop_sync_fn_{next(_GLOBAL_NAME_COUNTER)}__", loc - ) - sync_result_var = ir.Var( - scope, f"__coop_sync_result_{next(_GLOBAL_NAME_COUNTER)}__", loc - ) - block.append( - ir.Assign( - ir.Global(sync_module_global_name, _cuda_module, loc), - sync_module_var, - loc, - ) - ) - block.append( - ir.Assign( - ir.Expr.getattr(sync_module_var, sync_attr, loc), - sync_fn_var, - loc, - ) - ) - block.append( - ir.Assign( - ir.Expr.call(sync_fn_var, sync_args, (), loc), - sync_result_var, - loc, - ) - ) - - def _temp_storage_alias_ctor_key(self, inst: ir.Assign) -> str | None: - """Recognize an assignment that only forwards a storage descriptor. - - Accept copies, casts, and merged variables only when every source has - known constructor provenance. Return the target's canonical owner so - descriptor-use validation can allow the alias. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - value = inst.value - if isinstance(value, ir.Var): - sources = (value,) - elif isinstance(value, ir.Expr) and value.op == "cast": - sources = (getattr(value, "value", None),) - elif isinstance(value, ir.Expr) and value.op == "phi": - sources = _phi_incoming_values(value) - else: - return None - if not sources or any( - not isinstance(source, ir.Var) for source in sources - ): - return None - keys = [ - rewrite._resolve_temp_storage_ctor_key(cast(ir.Var, source)) - for source in sources - ] - if any(key is None for key in keys): - return None - return rewrite._resolve_temp_storage_ctor_key(inst.target) - - def _validate_temp_storage_uses( - self, func_ir: ir.FunctionIR, matches: dict[ir.Assign, _RewriteMatch] - ) -> None: - """Require TempStorage descriptors to be used only for storage planning. - - A descriptor describes scratch requirements; it is not a runtime - array. An "escape" is a use such as ``storage[0]``, ``len(storage)``, - or ``return storage`` that would expose the descriptor as runtime - data. The rewrite later replaces it with an internal shared-memory - view or ``None``, depending on its consumers. Rejecting these uses - prevents user code from depending on that implementation detail. - - After recording all constructors and matches, follow aliases to - check every use. Allow simple alias, cast, and phi assignments and - one ``temp_storage=`` keyword on a recognized primitive. - Subscripted-provider syntax and other call arguments are not - accepted descriptor uses. Every canonical constructor must have a - primitive consumer, so an otherwise unused - ``storage = coop.TempStorage()`` is also rejected. - - Validation runs after device-helper inlining. Passing a descriptor - to a helper is supported when inlining exposes its uses in the - kernel; a surviving helper call receives an inlining diagnostic. - - Check the entire function before preparing provider bundles or - compiling provider invocables ("materialization"). Even a descriptor - with a valid primitive consumer may have an invalid later use, as - below. Reject that use before compiling the otherwise valid call. - The scan updates the current block lookup state; its caller restores - it. - - Parameters - ---------- - func_ir : FunctionIR - Entire function whose constructors have been recorded. - matches : dict of ir.Assign to _RewriteMatch - Recognized calls keyed by their original IR assignments. - - Returns - ------- - None - All descriptor uses and constructor consumers are valid. - - Raises - ------ - CoopSinglePhaseRewriteError - A storage argument lacks local constructor provenance, a - descriptor has an unsupported use, or a constructor has no - primitive consumer. - - Examples - -------- - With ``from cuda import coop``, this kernel body contains a valid - primitive consumer followed by an invalid descriptor read:: - - block = coop.this_block() - storage = coop.TempStorage() - items = coop.ThreadData(items_per_thread=1) - coop.load( - block, source, items, - algorithm="transpose", temp_storage=storage, - ) - output[0] = storage[0] # Rejected: storage is a descriptor. - - Replacing the last line with ``output[0] = len(storage)`` or - ``storage[0] = 0`` is also invalid, even though ``load`` consumes - the descriptor correctly. - """ - rewrite = cast("CoopSinglePhaseRewrite", self) - - if not rewrite._temp_storage_ctor_specifications: - for match in matches.values(): - if match.runtime_temp_storage_var is not None: - raise CoopSinglePhaseRewriteError( - "cooperative group temp_storage= must originate from " - "a TempStorage constructor in the compiled function." - ) - return - - consumed_ctor_keys: set[str] = set() - for label in sorted(func_ir.blocks): - scan_block = func_ir.blocks[label] - self._block = scan_block - self._block_defs = { - inst.target.name: inst.value - for inst in scan_block.body - if isinstance(inst, ir.Assign) - } - for inst in scan_block.body: - used_vars = list(inst.list_vars()) - if isinstance(inst, ir.Assign): - used_vars = [ - var for var in used_vars if var.name != inst.target.name - ] - descriptor_vars = [] - for value in used_vars: - if ( - rewrite._resolve_temp_storage_ctor_key(value) - is not None - ): - descriptor_vars.append(value) - if not descriptor_vars: - match = matches.get(inst) - if ( - match is not None - and match.runtime_temp_storage_var is not None - ): - raise CoopSinglePhaseRewriteError( - "cooperative group temp_storage= must originate " - "from a TempStorage constructor " - "in the compiled function." - ) - continue - if isinstance(inst, ir.Assign) and ( - self._temp_storage_alias_ctor_key(inst) is not None - ): - continue - match = matches.get(inst) - if ( - match is not None - and match.runtime_temp_storage_var is not None - ): - storage_var = match.runtime_temp_storage_var - storage_key = rewrite._resolve_temp_storage_ctor_key( - storage_var - ) - assert isinstance(inst.value, ir.Expr) - keyword_storage_vars = [ - value - for name, value in inst.value.kws - if name == "temp_storage" - ] - descriptor_runtime_args = [ - value - for value in match.runtime_args - if rewrite._resolve_temp_storage_ctor_key(value) - is not None - ] - if ( - storage_key is not None - and len(keyword_storage_vars) == 1 - and keyword_storage_vars[0].name == storage_var.name - and not descriptor_runtime_args - and all( - value.name == storage_var.name - for value in descriptor_vars - ) - ): - consumed_ctor_keys.add(storage_key) - continue - names = ", ".join( - sorted({value.name for value in descriptor_vars}) - ) - if ( - isinstance(inst, ir.Assign) - and isinstance(inst.value, ir.Expr) - and inst.value.op == "call" - and ( - helper := rewrite._resolve_python_value(inst.value.func) - ) - is not None - and rewrite._is_jitted_dispatcher(helper) - ): - helper_name = helper.py_func.__qualname__ - raise CoopSinglePhaseRewriteError( - f"TempStorage descriptor {names!r} is passed to a " - f"device function that was not inlined into this " - f"kernel ({helper_name!r}); let Numba-CUDA-MLIR " - f"inline the collective helper (inline='always') or " - f"move its cooperative calls into the kernel." - ) - raise CoopSinglePhaseRewriteError( - f"TempStorage values are opaque compile-time descriptors " - f"and may only be passed as temp_storage= to a " - f"registered cooperative primitive; a use involving " - f"{names!r} would escape to runtime." - ) - - constructor_keys = { - rewrite._canonical_temp_storage_ctor_key(key) - for key in rewrite._temp_storage_ctor_specifications - } - consumed_ctor_keys = { - rewrite._canonical_temp_storage_ctor_key(key) - for key in consumed_ctor_keys - } - unconsumed = constructor_keys - consumed_ctor_keys - if unconsumed: - names = ", ".join(sorted(unconsumed)) - raise CoopSinglePhaseRewriteError( - f"TempStorage values are opaque compile-time descriptors and " - f"must be passed as temp_storage= to a registered " - f"cooperative primitive; constructor(s) {names!r} have no " - f"primitive consumer." - ) - - def _collect_temp_storage_uses( - self, func_ir: ir.FunctionIR, matches: dict[ir.Assign, _RewriteMatch] - ) -> list[tuple[int, ir.Assign, _RewriteMatch, str | None]]: - """Validate scratch ownership before provider compilation. - - Consume the analyzed calls in source order, retaining each storage - use's assignment and descriptor owner. Check all descriptor escapes - before returning, including calls whose ABI has no scratch pointer. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - storage_uses: list[ - tuple[int, ir.Assign, _RewriteMatch, str | None] - ] = [] - saved_block_defs = self._block_defs - saved_block = self._block - try: - source_order = 0 - for label in sorted(func_ir.blocks): - scan_block = func_ir.blocks[label] - self._block = scan_block - self._block_defs = { - inst.target.name: inst.value - for inst in scan_block.body - if isinstance(inst, ir.Assign) - } - for inst in scan_block.body: - current_order = source_order - source_order += 1 - match = matches.get(inst) - if match is None: - continue - ctor_key = ( - None - if match.runtime_temp_storage_var is None - else rewrite._resolve_temp_storage_ctor_key( - match.runtime_temp_storage_var - ) - ) - if ( - match.factory_metadata.storage_abi - is StorageABI.LEADING_POINTER - ): - self._validate_storage_match_plan( - match, ctor_key=ctor_key - ) - storage_uses.append( - (current_order, inst, match, ctor_key) - ) - self._validate_temp_storage_uses(func_ir, matches) - finally: - self._block_defs = saved_block_defs - self._block = saved_block - return storage_uses - - def _compute_func_temp_storage_requirements( - self, - storage_uses: list[tuple[int, ir.Assign, _RewriteMatch, str | None]], - ) -> dict[str, _TempStorageRequirementSummary]: - """Collect sizes and alignments from validated provider specializations. - - Bundle preparation has already supplied layouts where possible. - Materialization reuses those artifacts or compiles individual - providers as a fallback. Explicit descriptors accumulate under their - canonical owners; implicit scratch has a separate summary. Allocation - offsets are assigned later, when emitting the shared backing. - """ - - rewrite = cast("CoopSinglePhaseRewrite", self) - requirements: dict[str, _TempStorageRequirementSummary] = {} - for use_order, inst, match, ctor_key in storage_uses: - if ctor_key is not None: - ctor_key = rewrite._canonical_temp_storage_ctor_key(ctor_key) - invocable, _ = rewrite._materialize_invocable(match) - size_in_bytes = max( - 1, int(getattr(invocable, "temp_storage_bytes", 0) or 0) - ) - alignment = max( - 1, int(getattr(invocable, "temp_storage_alignment", 0) or 0) - ) - summary = ( - rewrite._implicit_temp_storage_requirements - if ctor_key is None - else requirements.setdefault( - ctor_key, _TempStorageRequirementSummary() - ) - ) - summary.max_size_in_bytes = max( - summary.max_size_in_bytes, size_in_bytes - ) - summary.max_alignment = max(summary.max_alignment, alignment) - summary.uses.append( - _TempStorageUseRequirement( - call_assign=inst, - order=use_order, - size_in_bytes=size_in_bytes, - alignment=alignment, - lowering_plan=match.lowering_plan, - ) - ) - return requirements - - -__all__ = ["_StorageRewrite"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_support.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_support.py deleted file mode 100644 index b4b02bd3e651..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_rewrite_support.py +++ /dev/null @@ -1,414 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Share call-rewrite records, alignment rules, and compiler support helpers. - -Match records carry resolved provider inputs; payload and storage records -retain facts needed between the whole-function scan and block replacement. -Storage requirements describe sizes and alignments before plans assign backing -offsets. ``_UNRESOLVED`` distinguishes failed static inference from an -explicit ``None`` value. ``_DeferredCoopRewrite`` asks for a retry once -launch metadata is available; it does not report an invalid user call. -""" - -from __future__ import annotations - -import struct -from collections.abc import Callable -from dataclasses import dataclass, field -from itertools import count -from typing import TYPE_CHECKING, Any - -from numba_cuda_mlir.numba_cuda.core.errors import ConstantInferenceError -from numba_cuda_mlir.numba_cuda.core.rewrites import ( - Rewrite as Rewrite, # noqa: PLC0414 - Re-export for typing. -) - -from cuda.coop._core import GroupLoweringPlan - -from ._operations import FactoryOperation -from ._parameters import normalize_dtype_param - -if TYPE_CHECKING: - from numba_cuda_mlir.numba_cuda.core import ir -else: - from numba_cuda_mlir.numbair_transforms import ir - -_INFERENCE_EXCEPTIONS = ( - KeyError, - ValueError, - TypeError, - AttributeError, - ConstantInferenceError, -) -_GLOBAL_NAME_COUNTER = count() -_UNRESOLVED = object() -_MIN_TEMP_STORAGE_ALIGNMENT = max(1, struct.calcsize("P")) -_DEFAULT_STATIC_SHARED_MEMORY_BYTES = 48 * 1024 -# numba-cuda-mlir declares its dynamic shared-memory window with 16-byte -# alignment; a larger request cannot be honored once the backing goes dynamic. -_DYNAMIC_SHARED_MEMORY_ALIGNMENT = 16 - - -class CoopSinglePhaseRewriteError(Exception): - """Report a matched cooperative call that cannot be rewritten.""" - - -class _DeferredCoopRewrite(Exception): - """Signal that a compiler pass must leave a cooperative call for later. - - ``CoopSinglePhaseRewrite.match`` catches this when launch-dependent work - needs exact launch metadata and records the deferral on the rewrite. - ``_CallRewriting._rewrite_calls`` reads that flag and retries kernel work - with the metadata. Deferral leaves the affected IR unchanged; it is not an - application error. - """ - - -def _next_global_name(stem: str) -> str: - """Give an injected Python object a process-unique IR global name.""" - - return f"__cuda_coop_numba_mlir_{stem}_{next(_GLOBAL_NAME_COUNTER)}__" - - -def _phi_incoming_values(definition): - """Read the alternatives of a control-flow merge expression. - - Check the compiler IR shape at this boundary so unsupported phi forms - produce a rewrite diagnostic instead of losing an incoming value. - """ - - if not hasattr(definition, "incoming_values"): - raise CoopSinglePhaseRewriteError( - "Unsupported Numba phi expression shape: missing incoming_values." - ) - incoming_values = definition.incoming_values - if not isinstance(incoming_values, (list, tuple)): - raise CoopSinglePhaseRewriteError( - "Unsupported Numba phi expression shape: incoming_values is not " - "a sequence." - ) - return tuple(incoming_values) - - -def _align_up(value: int, alignment: int) -> int: - if alignment <= 1: - return value - return (value + alignment - 1) // alignment * alignment - - -def _next_power_of_two(value: int) -> int: - if value <= 1: - return 1 - return 1 << (value - 1).bit_length() - - -def _default_temp_storage_alignment(required_alignment: int) -> int: - """Round alignment up to a power of two, at least pointer size.""" - - return max( - _MIN_TEMP_STORAGE_ALIGNMENT, _next_power_of_two(required_alignment) - ) - - -def _normalize_temp_storage_alignment( - alignment: int, *, context: str = "TempStorage alignment" -) -> int: - """Validate power-of-two alignment and apply the pointer-size minimum. - - ``context`` names the setting in diagnostics. A smaller valid request is - raised to the minimum needed by the generated storage pointer. - """ - - if alignment <= 0: - raise CoopSinglePhaseRewriteError( - f"{context} must be a positive integer." - ) - if alignment & alignment - 1 != 0: - raise CoopSinglePhaseRewriteError(f"{context} must be a power of 2.") - return max(_MIN_TEMP_STORAGE_ALIGNMENT, alignment) - - -def _dtype_values_match(lhs, rhs) -> bool: - """Compare dtype spellings after normalization when possible. - - If normalization fails, use the original or partially normalized values. - Callers remain responsible for validating accepted dtypes. - """ - - try: - lhs = normalize_dtype_param(lhs) - rhs = normalize_dtype_param(rhs) - except (TypeError, ValueError, AttributeError): - pass - return lhs == rhs - - -def _validate_temp_storage_alignment( - alignment: int, *, context: str = "TempStorage alignment" -) -> None: - alignment = _normalize_temp_storage_alignment(alignment, context=context) - if alignment % _MIN_TEMP_STORAGE_ALIGNMENT != 0: - raise CoopSinglePhaseRewriteError( - f"{context} must be a multiple of {_MIN_TEMP_STORAGE_ALIGNMENT}." - ) - - -def _check_driver_error(err, op: str) -> None: - if err.value != 0: - raise RuntimeError(f"{op} failed with CUDA driver error {err}") - - -def _query_device_shared_memory_limits() -> dict[str, int]: - """Read per-block shared-memory limits for the current CUDA device. - - Obtain the active context and initialize the driver before querying - default and opt-in capacities. A nonpositive opt-in value falls back to - the default limit. Driver failures raise ``RuntimeError``. - """ - - from numba_cuda_mlir.numba_cuda.cudadrv import devices - - import cuda.bindings.driver as _driver - - context = devices.get_context() - (err,) = _driver.cuInit(0) - _check_driver_error(err, "cuInit") - err, device = _driver.cuDeviceGet(int(context.device.id)) - _check_driver_error(err, "cuDeviceGet") - err, max_default = _driver.cuDeviceGetAttribute( - _driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK, - device, - ) - _check_driver_error( - err, "cuDeviceGetAttribute(MAX_SHARED_MEMORY_PER_BLOCK)" - ) - err, max_optin = _driver.cuDeviceGetAttribute( - _driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, - device, - ) - _check_driver_error( - err, "cuDeviceGetAttribute(MAX_SHARED_MEMORY_PER_BLOCK_OPTIN)" - ) - if int(max_optin) <= 0: - max_optin = max_default - return { - "device_id": int(context.device.id), - "max_default_shared_memory_per_block": int(max_default), - "max_optin_shared_memory_per_block": int(max_optin), - } - - -@dataclass(frozen=True) -class _RewriteMatch: - """Keep one validated provider call ready for compilation and emission. - - ``factory`` and ``factory_metadata`` identify the implementation and - declared call contract. ``factory_kwargs`` holds the resolved Python - values and binding descriptors passed to the host factory to specialize - the provider. ``runtime_args`` keeps device operands without storage, - which is tracked separately in ``runtime_temp_storage_var``. - - ``factory_kw_value_vars`` retains original IR variables consumed at - compile time so their assignments can be removed if no rewritten block - uses them. These are cleanup candidates, not another set of factory - inputs. Inferred keywords need no source variable; lowering metadata and - omitted controls can leave a cleanup candidate without a factory keyword. - - ``func_var_name`` and its optional extra alias identify callee assignments - that may also become unused after replacement. ``family_metadata`` carries - a hook's analysis, such as scalar boxing. ``lowering_plan`` keeps the shared - group contract after its private keyword is removed from factory inputs. - ``op_name`` and ``loc`` identify the operation and source site. - """ - - op_name: str - factory: Callable[..., Any] - factory_metadata: FactoryOperation - func_var_name: str - func_var_name_extra: str | None - runtime_args: tuple[ir.Var, ...] - runtime_temp_storage_var: ir.Var | None - factory_kwargs: dict[str, object] - factory_kw_value_vars: tuple[ir.Var, ...] - loc: ir.Loc - family_metadata: object = None - lowering_plan: GroupLoweringPlan | None = None - - -@dataclass(frozen=True) -class _ResolvedCallTarget: - """Separate callee recognition from argument validation. - - Keep the registered factory and contract, plus the callee names that - replacement may remove. ``getitem_temp_storage`` records an operand from - subscript syntax; later descriptor-use checks decide whether the call may - use that syntax. - """ - - factory: Callable[..., Any] - factory_metadata: FactoryOperation - func_var_name: str - func_var_name_extra: str | None - getitem_temp_storage: ir.Var | None - - @property - def operation(self) -> str: - return self.factory_metadata.operation - - -@dataclass(frozen=True) -class _ThreadDataSpecification: - """Carry known payload facts while ordinary typing is incomplete. - - ``items_per_thread`` and ``dtype`` may be unknown. Native local/shared - array queries also use this record, so its presence alone does not - establish public ``ThreadData`` origin. ``common_root`` records a common - API constructor for later numeric validation. ``alignment`` is an optional - byte alignment. Construction normalizes recognized dtypes and leaves - unsupported values for later validation. - """ - - items_per_thread: int | None - dtype: object | None - common_root: bool = False - alignment: int | None = None - - def __post_init__(self) -> None: - if self.dtype is None: - return - try: - canonical = normalize_dtype_param(self.dtype) - except (TypeError, ValueError): - return - object.__setattr__(self, "dtype", canonical) - - -@dataclass(frozen=True) -class _TempStorageCtorSpecification: - """Retain a descriptor's requested storage policy before layout. - - Capacity and alignment are in bytes and may be inferred later. - ``auto_sync=None`` has the effective value false. ``sharing`` selects - reuse among compatible calls or separate storage for every call. - """ - - size_in_bytes: int | None - alignment: int | None - auto_sync: bool | None - sharing: str - - -@dataclass(frozen=True) -class _TempStorageUseRequirement: - """Record one call's per-group scratch requirement before placement. - - Size and alignment are in bytes. ``call_assign`` retains the original IR - identity for slice lookup; ``order`` is the whole-function scan order, not - runtime execution order. ``lowering_plan`` supplies group instances and - reuse rules when the shared planner produced the call. - """ - - call_assign: ir.Assign - order: int - size_in_bytes: int - alignment: int - lowering_plan: GroupLoweringPlan | None = None - - -@dataclass -class _TempStorageRequirementSummary: - """Accumulate calls that share one descriptor or implicit region. - - The maxima summarize individual byte and alignment requirements. ``uses`` - retains each call so layout can account for alignment gaps, group - instances, and whether storage can be reused. - """ - - max_size_in_bytes: int = 0 - max_alignment: int = 1 - uses: list[_TempStorageUseRequirement] = field(default_factory=list) - - -@dataclass(frozen=True) -class _TempStorageSlice: - """Describe one call's view before the region's base offset is added. - - ``offset`` and ``size_in_bytes`` locate its bytes within the region. - Multiple group ``instances`` are separated by ``stride`` bytes; a missing - stride falls back to the call's size during emission. The lowering plan - supplies the formula that selects the current instance. - """ - - offset: int - size_in_bytes: int - stride: int | None = None - instances: int = 1 - lowering_plan: GroupLoweringPlan | None = None - - -@dataclass(frozen=True) -class _TempStoragePlan: - """Place one explicit descriptor or the implicit scratch region. - - Capacity and alignment include all of the region's calls and group - instances. ``slices_by_call_id`` maps IR assignment identities to their - views. ``base_offset`` places the region within the global byte array; - individual slice offsets remain relative to this region. ``sharing`` and - ``auto_sync`` retain the reuse policy for call emission. - """ - - size_in_bytes: int - alignment: int - sharing: str - auto_sync: bool - slices_by_call_id: dict[int, _TempStorageSlice] - base_offset: int = 0 - - -@dataclass(frozen=True) -class _TempStorageGlobalPlan: - """Describe the backing allocation shared by all scratch regions. - - ``total_size`` includes padding to ``max_alignment``; both are in bytes. - ``uses_dynamic_smem`` selects dynamic placement, in which case - ``dynamic_shared_bytes`` is the launch requirement. Static placement - leaves that requirement at zero. - """ - - total_size: int - max_alignment: int - uses_dynamic_smem: bool - dynamic_shared_bytes: int - - -__all__ = [ - "_DEFAULT_STATIC_SHARED_MEMORY_BYTES", - "_DYNAMIC_SHARED_MEMORY_ALIGNMENT", - "_GLOBAL_NAME_COUNTER", - "_INFERENCE_EXCEPTIONS", - "_MIN_TEMP_STORAGE_ALIGNMENT", - "_UNRESOLVED", - "CoopSinglePhaseRewriteError", - "Rewrite", - "_DeferredCoopRewrite", - "_ResolvedCallTarget", - "_RewriteMatch", - "_TempStorageCtorSpecification", - "_TempStorageGlobalPlan", - "_TempStoragePlan", - "_TempStorageRequirementSummary", - "_TempStorageSlice", - "_TempStorageUseRequirement", - "_ThreadDataSpecification", - "_align_up", - "_default_temp_storage_alignment", - "_dtype_values_match", - "_next_global_name", - "_normalize_temp_storage_alignment", - "_phi_incoming_values", - "_query_device_shared_memory_limits", - "_validate_temp_storage_alignment", - "ir", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_scalar_provenance.py b/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_scalar_provenance.py deleted file mode 100644 index e22309f2ce5b..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_compiler/_scalar_provenance.py +++ /dev/null @@ -1,336 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Track static scalar values without losing their original numeric type. - -A Python literal may take the element type required by a cooperative -operation, but a value explicitly typed as ``int64`` must not silently become -``int32``. ``StaticScalarProvenance`` keeps the value together with any dtype -already attached to it. Here, "provenance" means the constants, arguments, and -forwarding assignments that establish both facts. - -Group planning and call rewriting share this analysis when deciding whether a -control such as ``offset`` or ``oob_default`` can be embedded in generated -code. The static-value resolver accepts only explicit constants and literal -arguments whose possible definitions agree; it does not evaluate runtime -expressions to make them static. Separate helpers infer result types for -scalar operators, casts, and CUDA indices so runtime values can still be -checked against an operation's numeric requirements. Neither analysis rewrites -the function IR or requests a new compiler specialization. -""" - -from __future__ import annotations - -from collections.abc import Callable, Iterable, Sequence -from dataclasses import dataclass -from types import ModuleType -from typing import TYPE_CHECKING, Any - -import numba_cuda_mlir.numba_cuda.types as _numba_types -import numpy as np - -if TYPE_CHECKING: - from numba_cuda_mlir.numba_cuda.core import ir -else: - from numba_cuda_mlir.numbair_transforms import ir - - -@dataclass(frozen=True) -class StaticScalarProvenance: - """Keep a static value together with any dtype its source established. - - A Python literal can use the type required by an operation. An explicitly - typed scalar must retain its width when planning validates or embeds it. - This record lets callers make that distinction after following IR aliases. - - Parameters - ---------- - value : object - Value recovered from a constant or literal argument. - dtype : object, optional - Compiler or NumPy dtype attached to the source. ``None`` means the - source supplied no dtype; it does not mean the value is unknown. - """ - - value: Any - dtype: Any | None = None - - -def _static_scalar( - value: Any, dtype: Any | None = None -) -> StaticScalarProvenance: - if dtype is None and isinstance(value, np.generic): - dtype = value.dtype - return StaticScalarProvenance(value=value, dtype=dtype) - - -def _typed_static_value(scalar: StaticScalarProvenance) -> Any: - """Unwrap a static value while retaining its known numeric width. - - Use a NumPy scalar for a compiler dtype when conversion is supported. - Return the original value if conversion fails. Callers that must validate - the recorded dtype should keep the provenance object itself. - """ - - if scalar.dtype is None or isinstance(scalar.value, np.generic): - return scalar.value - try: - return np.dtype(str(scalar.dtype)).type(scalar.value) - except (TypeError, ValueError, OverflowError): - return scalar.value - - -def try_resolve_static_scalar_provenance( - value: object, - *, - definitions: Callable[[ir.Var], Iterable[object]], - argument_type: Callable[[int], _numba_types.Type | None], - seen: set[str] | None = None, -) -> tuple[bool, StaticScalarProvenance | None]: - """Resolve an explicitly static value while retaining its compiler dtype. - - Planning must distinguish a literal supplied by the user from a runtime - expression that general constant inference happens to evaluate. Accept - constants, globals, free variables, literal arguments, and arguments known - to be ``None``; follow only aliases, casts, and phi inputs. Every reaching - leaf must agree in Python value type, value, and recorded dtype. Runtime - expressions, unresolved paths, and cycles make the result unresolved. This - traversal neither evaluates scalar operators nor requests dispatcher - specialization. - - Parameters - ---------- - value : object - Value to inspect. Non-variable inputs are accepted directly as static; - callers are responsible for restricting them to the intended scalar - domain. - definitions : callable - Return all reaching definitions for an IR variable. - argument_type : callable - Return the compiler type for a function argument index, or ``None`` - when unavailable. Literal types retain their ``literal_type``. - seen : set of str, optional - Names already visited on this recursion path. The current variable is - added in place; recursive branches receive separate copies. - - Returns - ------- - resolved : bool - Whether all inspected definitions establish the same static value. - scalar : StaticScalarProvenance or None - Resolved value and its known dtype, or ``None`` on failure. A resolved - ``None`` value is represented by a provenance object and is distinct - from failure. NumPy scalars contribute their own dtype; ordinary - Python constants leave the dtype unspecified for contextual coercion. - """ - - if not isinstance(value, ir.Var): - return (True, _static_scalar(value)) - if seen is None: - seen = set() - if value.name in seen: - return (False, None) - seen.add(value.name) - - resolved_values: list[StaticScalarProvenance] = [] - for definition in definitions(value): - if isinstance(definition, ir.Arg): - arg_type = argument_type(definition.index) - if isinstance(arg_type, _numba_types.Literal): - resolved_values.append( - _static_scalar( - arg_type.literal_value, arg_type.literal_type - ) - ) - continue - if isinstance(arg_type, _numba_types.NoneType) or ( - isinstance(arg_type, _numba_types.Omitted) - and arg_type.value is None - ): - resolved_values.append(_static_scalar(None, arg_type)) - continue - return (False, None) - if isinstance(definition, (ir.Global, ir.FreeVar, ir.Const)): - resolved_values.append(_static_scalar(definition.value)) - continue - if isinstance(definition, ir.Var): - resolved, scalar = try_resolve_static_scalar_provenance( - definition, - definitions=definitions, - argument_type=argument_type, - seen=set(seen), - ) - if not resolved or scalar is None: - return (False, None) - resolved_values.append(scalar) - continue - if isinstance(definition, ir.Expr) and definition.op == "cast": - resolved, scalar = try_resolve_static_scalar_provenance( - definition.value, - definitions=definitions, - argument_type=argument_type, - seen=set(seen), - ) - if not resolved or scalar is None: - return (False, None) - resolved_values.append(scalar) - continue - if isinstance(definition, ir.Expr) and definition.op == "phi": - incoming_values = getattr(definition, "incoming_values", ()) - if ( - not isinstance(incoming_values, (list, tuple)) - or not incoming_values - ): - return (False, None) - for incoming in incoming_values: - resolved, scalar = try_resolve_static_scalar_provenance( - incoming, - definitions=definitions, - argument_type=argument_type, - seen=set(seen), - ) - if not resolved or scalar is None: - return (False, None) - resolved_values.append(scalar) - continue - return (False, None) - - if not resolved_values: - return (False, None) - first = resolved_values[0] - if any( - type(candidate.value) is not type(first.value) - or candidate.value != first.value - or candidate.dtype != first.dtype - for candidate in resolved_values[1:] - ): - return (False, None) - return (True, first) - - -def try_resolve_static_scalar( - value: object, - *, - definitions: Callable[[ir.Var], Iterable[object]], - argument_type: Callable[[int], _numba_types.Type | None], - seen: set[str] | None = None, -) -> tuple[bool, Any]: - """Resolve a static value and preserve a known scalar width when possible. - - Use ``try_resolve_static_scalar_provenance`` to require agreement across - all reaching definitions without evaluating runtime expressions. Unwrap - its result, converting a value with a known compiler dtype to the matching - NumPy scalar when possible. If that conversion is unsupported or fails, - return the original value. Use the provenance-returning helper when the - recorded dtype itself is needed for validation. - - Parameters - ---------- - value : object - IR value or already-static Python value to resolve. - definitions : callable - Return all reaching definitions for an IR variable. - argument_type : callable - Return the compiler type for a function argument index, or ``None`` - when unavailable. - seen : set of str, optional - Names on the current recursion path. The provenance resolver adds - the current variable in place and copies the set for each branch. - - Returns - ------- - resolved : bool - Whether the value has consistent, explicitly static provenance. - value : object - Unwrapped static value, or ``None`` on failure. Consult ``resolved`` - to distinguish an unresolved value from a statically known ``None``. - """ - - resolved, scalar = try_resolve_static_scalar_provenance( - value, - definitions=definitions, - argument_type=argument_type, - seen=seen, - ) - return (resolved, None if scalar is None else _typed_static_value(scalar)) - - -def scalar_expression_dtype( - definition: ir.Expr, - dtype: Callable[[ir.Var], object], -) -> _numba_types.Type | None: - """Infer an operator's result from the known scalar operand types. - - Accept a binary, in-place binary, or unary IR expression. Delegate to the - compiler's scalar operator rules; unknown or unsupported typing returns - ``None``. This query determines a type without evaluating the expression. - """ - - from ._parameters import _scalar_operator_result_dtype - - if definition.op in {"binop", "inplace_binop"}: - operands = (definition.lhs, definition.rhs) - else: - operands = (definition.value,) - return _scalar_operator_result_dtype( - definition.fn, - *( - dtype(value) if isinstance(value, ir.Var) else None - for value in operands - ), - ) - - -def scalar_call_dtype( - function: object, - arguments: Sequence[object], - dtype: Callable[[ir.Var], object], -) -> _numba_types.Type | None: - """Infer the result type of a recognized scalar cast. - - For one IR argument, try the compiler's scalar typing rules first. Fall - back to the cast's declared dtype when operand typing gives no result. - Return ``None`` for an unrecognized cast; this does not infer arbitrary - function return types or validate the whole call signature. - """ - - from ._parameters import _scalar_cast_dtype, _scalar_operator_result_dtype - - cast_dtype = _scalar_cast_dtype(function) - if cast_dtype is None: - return None - if len(arguments) == 1 and isinstance(arguments[0], ir.Var): - inferred = _scalar_operator_result_dtype(function, dtype(arguments[0])) - if inferred is not None: - return inferred - return cast_dtype - - -def cuda_index_dtype( - definition: ir.Expr, - attribute_chain: Callable[[ir.Var], tuple[object, Sequence[str]] | None], - cuda_module: ModuleType, -) -> _numba_types.Type | None: - """Recognize CUDA index and dimension components as ``int32``. - - Accept x, y, or z of blockDim, blockIdx, gridDim, or threadIdx only when - the attribute chain starts at the supplied CUDA module object. Unrelated - objects with the same attribute names return ``None``. - """ - - if definition.op != "getattr": - return None - chain = attribute_chain(definition.value) - if chain is not None: - root, attributes = chain - if root is cuda_module and (*attributes, definition.attr) in { - (index, component) - for index in ("blockDim", "blockIdx", "gridDim", "threadIdx") - for component in ("x", "y", "z") - }: - return _numba_types.int32 - return None - - -__all__: tuple[str, ...] = () diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_group/__init__.py b/python/cuda_coop/cuda/coop/numba_mlir/_group/__init__.py deleted file mode 100644 index 8bbe3ce1ab86..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_group/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_group/_load_store.py b/python/cuda_coop/cuda/coop/numba_mlir/_group/_load_store.py deleted file mode 100644 index df438cb65450..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_group/_load_store.py +++ /dev/null @@ -1,153 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Expose qualified Load and Store markers to Numba-CUDA-MLIR kernels. - -The operation decorator identifies each call and its lowering family. The -whole-function planner replaces supported calls before device compilation. -Calling a marker directly in Python raises the ``RuntimeError`` from -``_marker.group_primitive_marker`` because only the planner can replace it. -""" - -from __future__ import annotations - -try: - import numpy -except ModuleNotFoundError as exc: - if exc.name != "numpy": - raise - - -from cuda.coop._typing import _CommonNumericT - -from ..._core.api import ThreadDataLike -from ..._typing import ( - BlockLoadStoreAlgorithm, - CommonThreadDataLike, - IntegerValue, - TempStorageLike, - ValidItems, - WarpLoadStoreAlgorithm, -) -from .._compiler._operations import group_operation -from .._thread_group import BlockGroup, WarpGroup -from ._marker import group_primitive_marker - - -@group_operation( - "load", - family_module="cuda.coop.numba_mlir._compiler._group_load_store", -) -def 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 :func:`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 :ref:`per-thread payloads ` and - :ref:`temporary storage ` for allocation rules. - The executable example in :func:`cuda.coop.load` activates this backend - explicitly and shows a guarded final tile. - - See Also - -------- - :cpp:class:`cub::BlockLoad`, :cpp:class:`cub::WarpLoad` - C++ Load primitives used for these group scopes. - - Examples - -------- - Copy a partial final tile using a Numba local array. - - .. literalinclude:: - ../../python/cuda_coop/tests/backends/numba_mlir/runtime/test_qualified_load_store_examples.py - :language: python - :start-after: # qualified-load-store-example-begin - :end-before: # qualified-load-store-example-end - :dedent: 4 - """ - - group_primitive_marker( - "load", - group, - source, - output, - algorithm=algorithm, - valid_items=valid_items, - oob_default=oob_default, - offset=offset, - temp_storage=temp_storage, - ) - - -@group_operation( - "store", - family_module="cuda.coop.numba_mlir._compiler._group_load_store", -) -def 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 :func:`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 :ref:`per-thread payloads ` and - :ref:`temporary storage ` for allocation rules, and - :func:`cuda.coop.store` for an executable partial-tile example. - - See Also - -------- - :cpp:class:`cub::BlockStore`, :cpp:class:`cub::WarpStore` - C++ Store primitives used for these group scopes. - - Examples - -------- - Store only the valid prefix of the final tile. - - .. literalinclude:: - ../../python/cuda_coop/tests/backends/numba_mlir/runtime/test_qualified_load_store_examples.py - :language: python - :start-after: # qualified-load-store-example-begin - :end-before: # qualified-load-store-example-end - :dedent: 4 - """ - - group_primitive_marker( - "store", - group, - destination, - value, - algorithm=algorithm, - valid_items=valid_items, - offset=offset, - temp_storage=temp_storage, - ) - - -__all__ = ["load", "store"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_group/_load_store.pyi b/python/cuda_coop/cuda/coop/numba_mlir/_group/_load_store.pyi deleted file mode 100644 index b52e656ed94c..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_group/_load_store.pyi +++ /dev/null @@ -1,98 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from typing import overload - -from typing_extensions import TypeVar - -from ..._typing import ( - BlockLoadStoreAlgorithm, - CommonNumericScalar, - CommonThreadDataLike, - IntegerValue, - TempStorageLike, - ThreadDataLike, - ValidItems, - WarpLoadStoreAlgorithm, -) -from .._thread_group import BlockGroup, WarpGroup - -_CommonNumericT = TypeVar("_CommonNumericT", bound=CommonNumericScalar) - -@overload -def load( - group: BlockGroup, - source: object, - output: ThreadDataLike[_CommonNumericT], - /, - *, - algorithm: BlockLoadStoreAlgorithm = "direct", - valid_items: ValidItems | None = None, - oob_default: None = None, - offset: IntegerValue | None = None, - temp_storage: TempStorageLike | None = None, -) -> None: ... -@overload -def load( - group: BlockGroup, - source: object, - output: ThreadDataLike[_CommonNumericT], - /, - *, - algorithm: BlockLoadStoreAlgorithm = "direct", - valid_items: ValidItems, - oob_default: _CommonNumericT | float, - offset: IntegerValue | None = None, - temp_storage: TempStorageLike | None = None, -) -> None: ... -@overload -def load( - group: WarpGroup, - source: object, - output: ThreadDataLike[_CommonNumericT], - /, - *, - algorithm: WarpLoadStoreAlgorithm = "direct", - valid_items: ValidItems | None = None, - oob_default: None = None, - offset: IntegerValue | None = None, - temp_storage: None = None, -) -> None: ... -@overload -def load( - group: WarpGroup, - source: object, - output: ThreadDataLike[_CommonNumericT], - /, - *, - algorithm: WarpLoadStoreAlgorithm = "direct", - valid_items: ValidItems, - oob_default: _CommonNumericT | float, - offset: IntegerValue | None = None, - temp_storage: None = None, -) -> None: ... -@overload -def store( - group: BlockGroup, - destination: object, - value: _CommonNumericT | CommonThreadDataLike[_CommonNumericT], - /, - *, - algorithm: BlockLoadStoreAlgorithm = "direct", - valid_items: ValidItems | None = None, - offset: IntegerValue | None = None, - temp_storage: TempStorageLike | None = None, -) -> None: ... -@overload -def store( - group: WarpGroup, - destination: object, - value: _CommonNumericT | CommonThreadDataLike[_CommonNumericT], - /, - *, - algorithm: WarpLoadStoreAlgorithm = "direct", - valid_items: ValidItems | None = None, - offset: IntegerValue | None = None, - temp_storage: None = None, -) -> None: ... diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_group/_marker.py b/python/cuda_coop/cuda/coop/numba_mlir/_group/_marker.py deleted file mode 100644 index 848bcdd93f2c..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_group/_marker.py +++ /dev/null @@ -1,26 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Report cooperative markers that reach ordinary Python execution. - -Kernel compilation replaces these calls with planned device operations. -Reaching this function means that replacement has not occurred. -""" - -from __future__ import annotations - -from typing import Any - - -def group_primitive_marker(operation: str, *args: Any, **kwargs: Any) -> Any: - """Reject execution of a marker that should have been compiler-lowered.""" - - del args, kwargs - raise RuntimeError( - f"cuda.coop.numba_mlir.{operation} is a compile-time kernel construct " - "and must be lowered by the whole-function planner" - ) - - -__all__ = ["group_primitive_marker"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_lowering/__init__.py b/python/cuda_coop/cuda/coop/numba_mlir/_lowering/__init__.py deleted file mode 100644 index e4ca7ca8555b..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_lowering/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Hold the private Load/Store factories and the Numba core adapter. - -Group planning selects a factory in ``_load_store`` after validating a -public call. The package re-exports the storage-free block ``load`` and -``store``. Calling a factory specializes device code; it does not move data. -""" - -from ._load_store import load as load -from ._load_store import store as store - -__all__: tuple[str, ...] = () diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_lowering/_core.py b/python/cuda_coop/cuda/coop/numba_mlir/_lowering/_core.py deleted file mode 100644 index decc93f087f4..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_lowering/_core.py +++ /dev/null @@ -1,684 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Translate shared algorithm descriptions to Numba descriptors. - -The common core describes C++ algorithms, parameters, and template arguments -without depending on a Python compiler. This adapter supplies the Numba types -and calling conventions needed to turn those descriptions into an -``Algorithm``. For example, an array parameter needs a pointer-based wrapper, -and a runtime valid-item count can require checked integer narrowing. - -The adapter also carries scratch and synchronization requirements into source -generation. It creates a specialized algorithm description; compilation and -creation of a callable happen later in ``make_invocable_from_specialization`` -or after several descriptions have been collected for a shared compilation. -""" - -from __future__ import annotations - -from collections.abc import Mapping -from copy import copy -from dataclasses import dataclass -from types import SimpleNamespace -from typing import Any, ClassVar - -import numba_cuda_mlir.numba_cuda.types as numba_types - -from cuda.coop._core import ( - FLOAT32, - FLOAT64, - INT8, - INT16, - INT32, - INT64, - UINT8, - UINT16, - UINT32, - UINT64, - Algorithm, - ArgumentBinding, - Array, - BuiltinDType, - Constant, - CoreBackendAdapter, - CxxFunction, - Dependency, - Pointer, - PointerOffset, - Reference, - SynchronizationScope, - TempStorageParameter, - Value, - lower_method_parameters, -) - -from .. import _types as backend -from .._compiler._operations import StorageABI - - -@dataclass(frozen=True) -class NumbaMlirArrayInputTransform: - """Describe the element conversion a provider wrapper must generate. - - An input array can have a different element type from the CUB parameter. - The adapter resolves the parameter's target type and fixed extent. - Source generation then emits a local array of converted elements and - passes it to the CUB call. Only input-only arrays support this conversion. - - Parameters - ---------- - source_dtype : object - Element type of the array supplied by the kernel. - cpp_expression : str - C++ expression for one converted element. It must contain ``{value}``, - which source generation replaces with the input element expression. - This record checks the placeholder; the C++ compiler checks the - generated expression. - """ - - source_dtype: Any - cpp_expression: str - - def __post_init__(self) -> None: - if "{value}" not in self.cpp_expression: - raise ValueError("array input transform must reference {value}") - - -def _optional_binding(value: object) -> ArgumentBinding: - """Translate legacy presence markers without making their values static. - - Lowering factories accept both explicit binding descriptors and older - arguments whose mere presence requested a runtime overload. Preserve an - ``ArgumentBinding`` as supplied; otherwise ``None`` means omitted and - every other object means runtime. In particular, a plain integer here is - not a compile-time value. Callers must use ``ArgumentBinding.static`` to - embed it. - - Parameters - ---------- - value : object - Binding descriptor, omitted sentinel, or legacy presence marker. - - Returns - ------- - ArgumentBinding - Existing descriptor or a new omitted/runtime binding. - """ - - if isinstance(value, ArgumentBinding): - return value - if value is None: - return ArgumentBinding.omitted() - return ArgumentBinding.runtime() - - -class NumbaMlirCoreAdapter(CoreBackendAdapter): - """Translate shared algorithm parameters to the Numba calling convention. - - The core describes the C++ operation without Numba types. This adapter - maps its types and parameters to the descriptors used by Numba source - generation. ``materialize`` produces a specialized backend ``Algorithm``; - compilation and linking happen later. - - Named overrides cover cases the C++ signature alone cannot express. - For example, a runtime item count needs a checked integer ABI, and an - input conversion needs a local array with the target element type. - The adapter checks overrides against every matching core parameter when - it materializes an algorithm. - - Parameters - ---------- - input_transforms : Mapping[str, NumbaMlirArrayInputTransform], optional - Conversions keyed by core parameter name. Each target must be an - input-only array with a positive specialized extent. - value_abis : mapping of str to backend.Value, optional - Scalar ABI descriptors keyed by core parameter name. An override must - preserve the provider dtype and output role of each matching scalar. - Construction copies the mappings; materialization copies parameter - descriptors before attaching their names. - """ - - _BUILTIN_DTYPES: ClassVar[dict[BuiltinDType, numba_types.Type]] = { - INT8: numba_types.int8, - UINT8: numba_types.uint8, - INT16: numba_types.int16, - UINT16: numba_types.uint16, - INT32: numba_types.int32, - UINT32: numba_types.uint32, - INT64: numba_types.int64, - UINT64: numba_types.uint64, - FLOAT32: numba_types.float32, - FLOAT64: numba_types.float64, - } - _CORE_DTYPES: ClassVar[dict[numba_types.Type, BuiltinDType]] = { - value: key for key, value in _BUILTIN_DTYPES.items() - } - - def __init__( - self, - *, - input_transforms: Mapping[str, NumbaMlirArrayInputTransform] - | None = None, - value_abis: Mapping[str, backend.Value] | None = None, - ) -> None: - self._input_transforms = dict(input_transforms or {}) - self._value_abis = dict(value_abis or {}) - - def normalize_dtype(self, dtype: Any) -> Any: - """Map a shared builtin dtype to its matching Numba numeric type. - - Keep already backend-specific values unchanged. Reject a shared - builtin that the backend mapping does not support. - """ - - if isinstance(dtype, BuiltinDType): - try: - return self._BUILTIN_DTYPES[dtype] - except KeyError as exc: - raise TypeError( - f"unsupported core dtype {dtype.name!r}" - ) from exc - return dtype - - def core_dtype(self, dtype: Any) -> Any: - """Translate a Numba dtype for a shared core specialization. - - Provider factories call this before specializing the common algorithm. - Return the matching ``BuiltinDType`` token, or pass ``dtype`` through - when it has no builtin mapping. - """ - - return self._CORE_DTYPES.get(dtype, dtype) - - def cpp_type(self, dtype: Any) -> str: - """Spell a specialization's dtype for generated C++ declarations. - - Normalize a shared builtin or accept a backend dtype, then return - its builtin C++ spelling or ``storage_t`` for an opaque payload. - """ - - return backend.numba_type_to_cpp(self.normalize_dtype(dtype)) - - def _resolvable(self, value: Any) -> Any: - """Preserve a parameter dependency until template arguments are known. - - ``lower_parameter`` uses this for element dtypes and array lengths. - ``value`` may be a shared ``Dependency``, a ``Constant``, or a direct - value. Return the corresponding backend descriptor for substitution. - Translate shared dependencies and constants to backend descriptors. - Normalize dtype before wrapping a value as a backend constant. - """ - - if isinstance(value, Dependency): - return backend.Dependency(value.name) - if isinstance(value, Constant): - return backend.Constant(value.value) - return backend.Constant(self.normalize_dtype(value)) - - @staticmethod - def _is_backend_output(parameter: Any) -> bool: - """Choose whether a C++ output becomes a backend return value. - - An explicit ``is_return`` overrides ``is_output``. - """ - - if parameter.is_return is None: - return parameter.is_output - return parameter.is_return - - def lower_parameter( - self, - parameter: Any, - *, - specialization: Algorithm, - ) -> backend.Parameter: - """Describe one core parameter in the Numba provider ABI. - - Preserve template dependencies for arrays, pointers, and references - until provider construction resolves them. Scalar values require a - concrete dtype; a named scalar ABI override can impose stricter - runtime typing or checked narrowing. Output ownership follows - ``is_return`` when explicitly set, otherwise ``is_output`` determines - the backend return value. - - Named input transforms are different from ordinary arrays: resolve - their extent and target dtype now so source generation can emit a - fixed local array and per-element C++ conversions. Only input-only - arrays may use this path. Dependent C++ functors substitute bracketed - type placeholders only, leaving unrelated bare tokens unchanged. - - Parameters - ---------- - parameter : object - Core pointer offset, array, pointer, reference, value, or C++ - functor descriptor. Temporary storage is handled by - ``lower_temp_storage``. - specialization : Algorithm - Core specialization providing template arguments for eager - dependency resolution in transforms and C++ functors. - - Returns - ------- - Parameter - Backend descriptor; a scalar override may be the adapter's - existing descriptor. ``materialize`` copies descriptors before - attaching names. - - Raises - ------ - TypeError - The parameter kind is unsupported or a scalar value has a - dependent dtype. - ValueError - An input transform targets an output/inout array or lacks a - positive specialized integer extent. - """ - - if isinstance(parameter, PointerOffset): - return backend.PointerOffset( - self.normalize_dtype(parameter.dtype), - parameter.pointer_arg_index, - static_value=parameter.static_value, - ) - if isinstance(parameter, Array): - transform = ( - self._input_transforms.get(parameter.name) - if parameter.name is not None - else None - ) - if transform is not None: - if parameter.is_output or parameter.is_inout: - raise ValueError( - f"input transform {parameter.name!r}" - f" targets an output array" - ) - target_dtype = parameter.dtype - if isinstance(target_dtype, Dependency): - target_dtype = target_dtype.resolve( - specialization.template_arguments - ) - elif isinstance(target_dtype, Constant): - target_dtype = target_dtype.value - size = parameter.size - if isinstance(size, Dependency): - size = size.resolve(specialization.template_arguments) - elif isinstance(size, Constant): - size = size.value - if ( - not isinstance(size, int) - or isinstance(size, bool) - or size < 1 - ): - raise ValueError( - "Numba-CUDA-MLIR input transforms require a positive " - "specialized array extent" - ) - return backend.TransformedArray( - self.normalize_dtype(transform.source_dtype), - self.normalize_dtype(target_dtype), - size, - transform.cpp_expression, - ) - if isinstance( - parameter.dtype, (Dependency, Constant) - ) or isinstance(parameter.size, (Dependency, Constant)): - return backend.DependentArray( - self._resolvable(parameter.dtype), - self._resolvable(parameter.size), - is_output=self._is_backend_output(parameter), - ) - return backend.Array( - self.normalize_dtype(parameter.dtype), - parameter.size, - is_output=self._is_backend_output(parameter), - ) - if isinstance(parameter, Pointer): - dtype = parameter.dtype - if isinstance(dtype, Dependency): - pointer_type = ( - backend.DependentPointerReference - if parameter.deref_on_call - else backend.DependentPointer - ) - return pointer_type( - backend.Dependency(dtype.name), - is_output=self._is_backend_output(parameter), - ) - pointer_type = ( - backend.PointerReference - if parameter.deref_on_call - else backend.Pointer - ) - return pointer_type( - self.normalize_dtype(dtype), - is_output=self._is_backend_output(parameter), - ) - if isinstance(parameter, Reference): - dtype = parameter.dtype - if isinstance(dtype, Dependency): - return backend.DependentReference( - backend.Dependency(dtype.name), - is_output=self._is_backend_output(parameter), - ) - return backend.Reference( - self.normalize_dtype(dtype), - is_output=self._is_backend_output(parameter), - ) - if isinstance(parameter, Value): - dtype = parameter.dtype - if isinstance(dtype, Dependency): - raise TypeError( - "Numba-CUDA-MLIR does not support dependent scalar values" - ) - normalized_dtype = self.normalize_dtype(dtype) - value_abi = ( - self._value_abis.get(parameter.name) - if parameter.name is not None - else None - ) - if value_abi is not None: - return value_abi - return backend.Value( - normalized_dtype, - is_output=self._is_backend_output(parameter), - ) - if isinstance(parameter, CxxFunction): - dtype = parameter.dtype - cpp = parameter.cpp - if isinstance(dtype, Dependency): - dependency = dtype - dtype = dependency.resolve(specialization.template_arguments) - # Substitute bracketed type placeholders, leaving bare tokens - # unchanged. - cpp = cpp.replace( - f"<{dependency.name}>", - f"<{self.cpp_type(dtype)}>", - ) - return backend.CxxFunction( - cpp, - self.normalize_dtype(dtype), - ) - raise TypeError( - f"unsupported Numba-CUDA-MLIR core parameter {parameter!r}" - ) - - def lower_temp_storage( - self, - parameter: TempStorageParameter, - *, - specialization: Algorithm, - ) -> backend.Pointer: - """Describe scratch as a pointer in the provider's calling convention. - - The shared adapter traversal calls this for a ``TempStorageParameter`` - when the storage ABI includes a leading pointer. ``parameter.dtype`` - supplies its element type; ``specialization`` is part of the adapter - interface but is unused here. Return a backend pointer descriptor. - Later stages handle allocation, byte layout checks, and reuse - barriers. - """ - - del specialization - return backend.Pointer(self.normalize_dtype(parameter.dtype)) - - def materialize( - self, - specialization: Algorithm, - *, - storage_abi: StorageABI, - execution_scope: SynchronizationScope, - synchronization_scope: SynchronizationScope, - extra_type_definitions: tuple[Any, ...] = (), - **kwargs: Any, - ) -> backend.Algorithm: - """Build a backend algorithm from a core specialization. - - Validate named scalar ABI overrides and array input transforms against - all matching core parameters before lowering. Scalar overrides must - preserve provider dtype and output role; transforms apply only to - input-only arrays. Copy each lowered descriptor before attaching its - parameter name because an override may be reused by several parameters - or specializations. - - Keep checked scalar overloads ahead of runtime pointer-offset - overloads. Numba-CUDA-MLIR selects the first convertible signature, - and the offset integer domain is intentionally broader than an exact - scalar ABI. The stable ordering otherwise preserves the core's method - order. Include the leading scratch parameter only for - ``LEADING_POINTER`` storage, translate type declarations, and finish - template substitution without compiling LTO. - - Parameters - ---------- - specialization : Algorithm - Core specialization with ordered template arguments and overloads. - storage_abi : StorageABI - Whether the backend receives a leading scratch pointer or no - scratch. - execution_scope : SynchronizationScope - Participating scope for scratch allocation by the source emitter. - synchronization_scope : SynchronizationScope - Post-call synchronization for allocating wrappers; must be - ``NONE`` or match ``execution_scope``. - extra_type_definitions : tuple, optional - Backend type definitions prepended to the core's declarations, - including any supporting link images. - **kwargs : dict - Reserved for the adapter interface; extra options are rejected. - - Returns - ------- - Algorithm - Specialized backend provider with copied, named parameter - descriptors and explicit storage/execution contracts. - - Raises - ------ - TypeError - Options or scalar ABI declarations are unsupported, or parameter - lowering encounters an unsupported descriptor. - ValueError - Named overrides/transforms do not match eligible core parameters, - their dtypes or output roles conflict, or scope/ABI values are - invalid. - """ - - if kwargs: - unexpected = ", ".join(sorted(kwargs)) - raise TypeError( - f"unexpected Numba-CUDA-MLIR " - f"materialization options: {unexpected}" - ) - - storage_abi = StorageABI(storage_abi) - execution_scope = SynchronizationScope(execution_scope) - synchronization_scope = SynchronizationScope(synchronization_scope) - include_temp_storage = storage_abi is StorageABI.LEADING_POINTER - - invalid_value_abis = { - name - for name, value_abi in self._value_abis.items() - if ( - not isinstance(name, str) - or not name - or not isinstance(value_abi, backend.Value) - ) - } - if invalid_value_abis: - names = ", ".join(sorted(repr(name) for name in invalid_value_abis)) - raise TypeError( - "Numba-CUDA-MLIR value ABI declarations require non-empty " - f"parameter names and backend Value instances: {names}" - ) - value_parameters_by_name = { - name: tuple( - parameter - for method in specialization.parameters - for parameter in method - if parameter.name == name - ) - for name in self._value_abis - } - unknown_value_abis = { - name - for name, parameters in value_parameters_by_name.items() - if not parameters - } - if unknown_value_abis: - names = ", ".join(sorted(unknown_value_abis)) - raise ValueError(f"unknown Numba-CUDA-MLIR value ABI(s): {names}") - invalid_value_targets = { - name - for name, parameters in value_parameters_by_name.items() - if any( - not isinstance(parameter, Value) - or isinstance(parameter, PointerOffset) - or parameter.is_inout - for parameter in parameters - ) - } - if invalid_value_targets: - names = ", ".join(sorted(invalid_value_targets)) - raise ValueError( - f"Numba-CUDA-MLIR value ABIs require " - f"scalar Value parameters: {names}" - ) - dtype_mismatches = { - name - for name, parameters in value_parameters_by_name.items() - if any( - self.normalize_dtype(parameter.dtype) - != self._value_abis[name].provider_dtype - for parameter in parameters - ) - } - if dtype_mismatches: - names = ", ".join(sorted(dtype_mismatches)) - raise ValueError( - "Numba-CUDA-MLIR value ABI provider dtypes do not match " - f"core Value dtypes: {names}" - ) - output_mismatches = { - name - for name, parameters in value_parameters_by_name.items() - if any( - self._is_backend_output(parameter) - != self._value_abis[name].is_output - for parameter in parameters - ) - } - if output_mismatches: - names = ", ".join(sorted(output_mismatches)) - raise ValueError( - "Numba-CUDA-MLIR value ABI output roles do not match core " - f"Value parameters: {names}" - ) - - parameters_by_name = { - name: tuple( - parameter - for method in specialization.parameters - for parameter in method - if parameter.name == name - ) - for name in self._input_transforms - } - unknown_transforms = { - name - for name, parameters in parameters_by_name.items() - if not parameters - } - if unknown_transforms: - names = ", ".join(sorted(unknown_transforms)) - raise ValueError( - f"unknown Numba-CUDA-MLIR input transform(s): {names}" - ) - invalid_transforms = { - name - for name, parameters in parameters_by_name.items() - if any( - not isinstance(parameter, Array) - or parameter.is_output - or parameter.is_inout - for parameter in parameters - ) - } - if invalid_transforms: - names = ", ".join(sorted(invalid_transforms)) - raise ValueError( - "Numba-CUDA-MLIR input transforms require input-only array " - f"parameters: {names}" - ) - - # Pointer-offset overloads deliberately accept a broadly convertible - # integer in the same position where ABI-checked scalar overloads - # require an exact dtype. Numba-CUDA-MLIR selects the first - # convertible overload, so keep checked forms ahead of offset forms. - # The stable sort otherwise preserves the canonical core ordering. - ordered_parameters = sorted( - specialization.parameters, - key=lambda method: any( - isinstance(parameter, PointerOffset) - and parameter.static_value is None - for parameter in method - ), - ) - methods = [] - for method in ordered_parameters: - core_parameters = [ - parameter - for parameter in method - if include_temp_storage - or not isinstance(parameter, TempStorageParameter) - ] - lowered_parameters = lower_method_parameters( - self, - specialization, - method, - include_temp_storage=include_temp_storage, - ) - named_parameters = [] - for parameter, lowered in zip(core_parameters, lowered_parameters): - # Scalar ABI overrides may be reused by several parameters - # or specializations. - named = copy(lowered) - named.parameter_name = parameter.name - named_parameters.append(named) - methods.append(named_parameters) - - definitions = [*extra_type_definitions] - definitions.extend( - SimpleNamespace(code=item.code, lto_irs=[]) - for item in specialization.type_definitions - ) - template_arguments = { - name: self.normalize_dtype(value) - for name, value in specialization.ordered_specialization_arguments - } - return backend.Algorithm( - specialization.struct_name, - specialization.method_name, - specialization.c_name, - list(specialization.includes), - [ - backend.TemplateParameter(parameter.name) - for parameter in specialization.template_parameters - ], - methods, - template_arguments=template_arguments, - storage_abi=storage_abi, - execution_scope=execution_scope, - synchronization_scope=synchronization_scope, - type_definitions=definitions, - fake_return=specialization.fake_return, - output_by_reference=specialization.output_by_reference, - ) - - -__all__ = [ - "NumbaMlirArrayInputTransform", - "NumbaMlirCoreAdapter", - "_optional_binding", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_lowering/_load_store.py b/python/cuda_coop/cuda/coop/numba_mlir/_lowering/_load_store.py deleted file mode 100644 index 0fee6ac6601f..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_lowering/_load_store.py +++ /dev/null @@ -1,765 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Build the compiled implementations chosen for cooperative load/store calls. - -A factory in this module is a host-side function that specializes a CUB -operation for a payload dtype, group size, per-thread extent, and algorithm. -The group planner selects the registered block or warp factory and describes -which optional controls are omitted, fixed at compile time, or supplied at -runtime. Calling the factory creates a device callable; data movement occurs -later when a kernel invokes that callable. - -Separate factories represent algorithms that need shared scratch and those -that do not, so their argument lists agree with the planner's storage choice. -Both paths lower a common core specialization through -``NumbaMlirCoreAdapter``. During specialization collection they return an -``Algorithm`` description instead of compiling it immediately, allowing -several operations to share an NVRTC compilation. -""" - -import operator -from collections.abc import Callable, Iterable -from typing import SupportsIndex, TypedDict, cast - -import numba_cuda_mlir.numba_cuda.types as numba_types - -from cuda.coop._core import ArgumentBinding, BindingKind, SynchronizationScope -from cuda.coop._core.block import ( - make_block_load_specialization, - make_block_store_specialization, -) -from cuda.coop._core.warp import ( - make_warp_load_specialization, - make_warp_store_specialization, -) - -from .._compiler._load_store_algorithms import ( - _BLOCK_LOAD_STORE_ALGORITHMS, - _STORAGE_FREE_ALGORITHMS, - _WARP_LOAD_STORE_ALGORITHMS, - _resolve_algorithm, -) -from .._compiler._operations import ( - StorageABI, - factory_operation, - register_factory, -) -from .._compiler._parameters import ( - _validate_common_numeric_dtype, - _validate_static_oob_default, - normalize_dim_param, -) -from .._types import ( - Algorithm, - BoundedInteger, - ExactValue, - Invocable, - Value, - _validate_logical_warp_threads, - make_invocable_from_specialization, - numba_type_to_wrapper, -) -from ._core import NumbaMlirCoreAdapter, _optional_binding - - -class _GroupTopology(TypedDict, total=False): - """Supply exactly one group shape to runtime-control ABI planning. - - A block uses ``block_dim``; a warp uses ``threads_in_warp``. That thread - count times the per-thread extent bounds the valid-item count. For runtime - counts, ``_load_store_value_abis`` requires exactly one of these shapes. - """ - - block_dim: Iterable[int] - threads_in_warp: int - - -def _positive_int(value, *, name: str) -> int: - """Normalize a per-thread item count or logical warp width. - - Load/Store factories pass an index-compatible ``value`` and the parameter - ``name`` to use in diagnostics. Return a positive Python integer; reject - booleans and non-integers with ``TypeError``, and nonpositive counts with - ``ValueError``. - """ - - if isinstance(value, bool): - raise TypeError(f"{name} must be an integer") - try: - value = operator.index(value) - except TypeError as exc: - raise TypeError(f"{name} must be an integer") from exc - if value <= 0: - raise ValueError(f"{name} must be a positive integer") - return value - - -def _registered_provider_metadata(factory, algorithm): - """Check that a factory's storage ABI matches the selected algorithm. - - ``_load`` and ``_store`` call this after resolving ``algorithm`` to a - canonical name. ``factory`` is the registered callable that selected the - block or warp provider and its scratch policy. - Algorithms that use shared scratch require a leading pointer. Others - require the storage-free factory. Return the registered metadata, or - reject an unknown factory or mismatched choice before materialization. - """ - - registered = factory_operation(factory) - if registered is None: - raise RuntimeError(f"unregistered cuda.coop provider {factory!r}") - expected_storage_abi = ( - StorageABI.NONE - if algorithm in _STORAGE_FREE_ALGORITHMS - else StorageABI.LEADING_POINTER - ) - if registered.storage_abi is not expected_storage_abi: - raise ValueError( - f"{registered.operation} algorithm {algorithm!r} " - f"requires the {expected_storage_abi.value!r} provider" - ) - return registered - - -def _materialization_metadata(registered): - """Pass the chosen factory's scope and scratch rules to the adapter. - - ``_load`` and ``_store`` supply the registry entry returned by - ``_registered_provider_metadata``. Return keyword arguments for - ``NumbaMlirCoreAdapter.materialize`` so provider construction uses those - same storage, execution, and synchronization contracts. - """ - - return { - "storage_abi": registered.storage_abi, - "execution_scope": registered.execution_scope, - "synchronization_scope": registered.synchronization_scope, - } - - -def _load_store_value_abis( - *, - dtype: numba_types.Type, - items_per_thread: int, - valid_items: ArgumentBinding, - oob_default: ArgumentBinding | None = None, - block_dim: Iterable[int] | None = None, - threads_in_warp: int | None = None, -) -> dict[str, Value]: - """Describe runtime controls whose ABI differs from ordinary core scalars. - - A runtime valid-item count enters through a signed 64-bit value so bounds - can be checked before narrowing to CUB's signed 32-bit count. Its - inclusive limit is the exact group tile size. The load operation's - ``oob_default`` is the padding value for items beyond that count. When - supplied at runtime, it uses ``ExactValue`` so typing requires the payload - dtype without an implicit conversion. Omitted and static controls need no - runtime ABI override. - - Parameters - ---------- - dtype : numba_types.Type - Normalized payload dtype, also required for a runtime default. - items_per_thread : int - Positive specialized payload extent per participating thread. - valid_items : ArgumentBinding - Binding kind for the optional valid-item count. - oob_default : ArgumentBinding or None, optional - Binding for the load operation's padding value; ``None`` supplies no - override. - block_dim : iterable of int, optional - Block dimensions used to compute a block tile's capacity. - threads_in_warp : int, optional - Logical warp width used to compute a warp tile's capacity. - - Returns - ------- - dict of str to Value - Overrides for ``num_valid_items``, ``oob_default``, or both. - - Raises - ------ - ValueError - A runtime count does not have exactly one block/warp topology source, - or its bounds cannot fit the integer ABI. - """ - - value_abis: dict[str, Value] = {} - if valid_items.kind is BindingKind.RUNTIME: - tile_items = items_per_thread - if (block_dim is None) == (threads_in_warp is None): - raise ValueError( - "exactly one of block_dim or threads_in_warp must be provided" - ) - if block_dim is not None: - for dimension in block_dim: - tile_items *= dimension - else: - tile_items *= cast(int, threads_in_warp) - value_abis["num_valid_items"] = BoundedInteger( - numba_types.int32, - minimum=0, - maximum=tile_items, - ) - if oob_default is not None and oob_default.kind is BindingKind.RUNTIME: - value_abis["oob_default"] = ExactValue(dtype) - return value_abis - - -def _warp_threads(value) -> int: - value = _positive_int(value, name="threads_in_warp") - return _validate_logical_warp_threads(value) - - -def _load( - provider_factory: Callable[..., object], - dtype: object, - threads_per_block: SupportsIndex | tuple[SupportsIndex, ...] | None = None, - items_per_thread: SupportsIndex = 1, - algorithm: str = "direct", - num_valid_items: object = None, - oob_default: object = None, - offset: object = None, - threads_in_warp: SupportsIndex | None = None, -) -> Invocable | Algorithm: - """Build the callable implementation for a planned load operation. - - Use the registered factory identity to select block/warp semantics and - verify whether the chosen algorithm requires scratch. Lower the common - load specialization through the Numba adapter, including checked runtime - counts, exact-dtype runtime defaults, and embedded static controls. The - factory describes a compiled callable; it does not load data. - - Explicit ``ArgumentBinding`` objects distinguish omitted, static, and - runtime controls. Legacy non-``None`` values indicate runtime presence, - not a literal to embed. For a legacy count, retain the full-tile overload - as well. Legacy offset arguments request offset overloads even when - ``None``; explicit bindings let the planner select their precise form. - - Parameters - ---------- - provider_factory : callable - Registered load factory whose namespace, storage ABI, and - synchronization metadata govern materialization. - dtype : object - Payload dtype accepted by the common numeric profile. - threads_per_block : SupportsIndex or tuple of SupportsIndex - Required enclosing block dimensions, including for warp providers. - items_per_thread : SupportsIndex, optional - Positive per-thread payload extent; defaults to one. - algorithm : str, optional - Specialized load algorithm, resolved within the factory's namespace. - num_valid_items : ArgumentBinding or object, optional - Valid-count binding or legacy runtime-presence marker. - oob_default : ArgumentBinding or object, optional - Binding or runtime-presence marker for the value used to pad items - beyond the valid count. Requires a valid-count binding; static padding - values are checked against the payload dtype. - offset : ArgumentBinding or object, optional - Pointer-offset binding or legacy overload-presence input. - threads_in_warp : SupportsIndex, optional - Required logical width for warp providers; invalid for blocks. - - Returns - ------- - Invocable or Algorithm - Compiled provider callable, or the specialization recorded when - ``collect_specializations`` is active. - - Raises - ------ - TypeError - A dtype, integer extent, or static scalar is unsupported. - ValueError - Topology, algorithm/storage pairing, or optional bindings are invalid. - RuntimeError - The factory is unregistered or artifact construction fails. - """ - - valid_items_binding = _optional_binding(num_valid_items) - oob_default_binding = _optional_binding(oob_default) - offset_binding = _optional_binding(offset) - if ( - oob_default_binding.kind is not BindingKind.OMITTED - and valid_items_binding.kind is BindingKind.OMITTED - ): - raise ValueError("oob_default requires num_valid_items to be provided") - if threads_per_block is None: - raise ValueError("threads_per_block must be provided") - block_dim = normalize_dim_param(threads_per_block) - dtype = _validate_common_numeric_dtype(dtype, operation="load") - if oob_default_binding.kind is BindingKind.STATIC: - oob_default_binding = ArgumentBinding.static( - _validate_static_oob_default(oob_default_binding.value, dtype) - ) - items_per_thread = _positive_int(items_per_thread, name="items_per_thread") - registered = factory_operation(provider_factory) - if registered is None: - raise RuntimeError( - f"unregistered cuda.coop provider {provider_factory!r}" - ) - group_kwargs: _GroupTopology - if registered.namespace == "block": - if threads_in_warp is not None: - raise ValueError("block Load does not accept threads_in_warp") - algorithm = _resolve_algorithm( - algorithm, - _BLOCK_LOAD_STORE_ALGORITHMS, - "block load", - ) - group_kwargs = {"block_dim": block_dim} - elif registered.namespace == "warp": - threads_in_warp = _warp_threads(threads_in_warp) - algorithm = _resolve_algorithm( - algorithm, - _WARP_LOAD_STORE_ALGORITHMS, - "warp load", - ) - group_kwargs = {"threads_in_warp": threads_in_warp} - else: - raise ValueError( - f"unsupported cuda.coop Load provider namespace " - f"{registered.namespace!r}" - ) - registered = _registered_provider_metadata(provider_factory, algorithm) - adapter = NumbaMlirCoreAdapter( - value_abis=_load_store_value_abis( - dtype=dtype, - items_per_thread=items_per_thread, - valid_items=valid_items_binding, - oob_default=oob_default_binding, - block_dim=group_kwargs.get("block_dim"), - threads_in_warp=group_kwargs.get("threads_in_warp"), - ) - ) - specialization_factory = ( - make_block_load_specialization - if registered.namespace == "block" - else make_warp_load_specialization - ) - core_specialization = specialization_factory( - dtype=adapter.core_dtype(dtype), - items_per_thread=items_per_thread, - algorithm=algorithm, - valid_items=valid_items_binding, - oob_default=oob_default_binding, - include_full_tile=( - not isinstance(num_valid_items, ArgumentBinding) - and num_valid_items is not None - ), - include_pointer_offset=( - offset_binding if isinstance(offset, ArgumentBinding) else True - ), - **( - {"block_dim": tuple(block_dim)} - if registered.namespace == "block" - else {"threads_in_warp": threads_in_warp} - ), - ) - specialization = adapter.materialize( - core_specialization, - **_materialization_metadata(registered), - extra_type_definitions=(numba_type_to_wrapper(dtype),), - ) - invocation_topology = ( - { - "logical_warp_threads": threads_in_warp, - "block_threads": block_dim, - } - if registered.namespace == "warp" - else {} - ) - return make_invocable_from_specialization( - specialization, - **invocation_topology, - ) - - -def load( - dtype, - threads_per_block=None, - items_per_thread=1, - algorithm="direct", - num_valid_items=None, - oob_default=None, - offset=None, -): - """Build a block load callable that needs no shared scratch.""" - - return _load( - load, - dtype, - threads_per_block, - items_per_thread, - algorithm, - num_valid_items, - oob_default, - offset, - ) - - -def _load_with_storage( - dtype, - threads_per_block=None, - items_per_thread=1, - algorithm="transpose", - num_valid_items=None, - oob_default=None, - offset=None, -): - """Build a block transpose load callable with shared scratch.""" - - return _load( - _load_with_storage, - dtype, - threads_per_block, - items_per_thread, - algorithm, - num_valid_items, - oob_default, - offset, - ) - - -def warp_load( - dtype, - threads_per_block=None, - threads_in_warp=32, - items_per_thread=1, - algorithm="direct", - num_valid_items=None, - oob_default=None, - offset=None, -): - """Build a warp load callable that needs no shared scratch.""" - - return _load( - warp_load, - dtype, - threads_per_block=threads_per_block, - threads_in_warp=threads_in_warp, - items_per_thread=items_per_thread, - algorithm=algorithm, - num_valid_items=num_valid_items, - oob_default=oob_default, - offset=offset, - ) - - -def _warp_load_with_storage( - dtype, - threads_per_block=None, - threads_in_warp=32, - items_per_thread=1, - algorithm="transpose", - num_valid_items=None, - oob_default=None, - offset=None, -): - """Build a warp load callable with shared scratch.""" - - return _load( - _warp_load_with_storage, - dtype, - threads_per_block=threads_per_block, - threads_in_warp=threads_in_warp, - items_per_thread=items_per_thread, - algorithm=algorithm, - num_valid_items=num_valid_items, - oob_default=oob_default, - offset=offset, - ) - - -def _store( - provider_factory: Callable[..., object], - dtype: object, - threads_per_block: SupportsIndex | tuple[SupportsIndex, ...] | None = None, - items_per_thread: SupportsIndex = 1, - algorithm: str = "direct", - num_valid_items: object = None, - oob_default: object = None, - offset: object = None, - threads_in_warp: SupportsIndex | None = None, -) -> Invocable | Algorithm: - """Build the callable implementation for a planned store operation. - - Resolve the registered block/warp factory and its storage contract, then - materialize the common store specialization with Numba-specific runtime - count bounds. A warp provider also retains the enclosing block dimensions - for scratch layout. This constructs a compiled callable; it does not write - device memory during host-side factory evaluation. - - As in ``_load``, explicit bindings control whether scalar arguments are - omitted, embedded, or runtime. A legacy non-``None`` count means runtime - presence and retains a full-tile overload. Legacy offset inputs request - offset overloads regardless of their value; an explicit binding selects - the offset form. Store operations do not accept ``oob_default``. - - Parameters - ---------- - provider_factory : callable - Registered store factory supplying namespace and ABI metadata. - dtype : object - Payload dtype accepted by the common numeric profile. - threads_per_block : SupportsIndex or tuple of SupportsIndex - Required enclosing block dimensions, also for warp providers. - items_per_thread : SupportsIndex, optional - Positive per-thread payload extent; defaults to one. - algorithm : str, optional - Specialized store algorithm within the factory's namespace. - num_valid_items : ArgumentBinding or object, optional - Valid-count binding or legacy runtime-presence marker. - oob_default : None, optional - Shared factory-interface slot; any non-``None`` value is rejected. - offset : ArgumentBinding or object, optional - Pointer-offset binding or legacy overload-presence input. - threads_in_warp : SupportsIndex, optional - Required logical width for warp providers; invalid for blocks. - - Returns - ------- - Invocable or Algorithm - Compiled provider callable, or the specialization recorded by an - active ``collect_specializations`` context. - - Raises - ------ - TypeError - A dtype or integer extent is unsupported. - ValueError - Topology, algorithm/storage pairing, or optional bindings are invalid. - RuntimeError - The factory is unregistered or artifact construction fails. - """ - - if oob_default is not None: - raise ValueError("oob_default is only valid for Load") - valid_items_binding = _optional_binding(num_valid_items) - offset_binding = _optional_binding(offset) - if threads_per_block is None: - raise ValueError("threads_per_block must be provided") - block_dim = normalize_dim_param(threads_per_block) - dtype = _validate_common_numeric_dtype(dtype, operation="store") - items_per_thread = _positive_int(items_per_thread, name="items_per_thread") - registered = factory_operation(provider_factory) - if registered is None: - raise RuntimeError( - f"unregistered cuda.coop provider {provider_factory!r}" - ) - group_kwargs: _GroupTopology - if registered.namespace == "block": - if threads_in_warp is not None: - raise ValueError("block Store does not accept threads_in_warp") - algorithm = _resolve_algorithm( - algorithm, - _BLOCK_LOAD_STORE_ALGORITHMS, - "block store", - ) - group_kwargs = {"block_dim": block_dim} - elif registered.namespace == "warp": - threads_in_warp = _warp_threads(threads_in_warp) - algorithm = _resolve_algorithm( - algorithm, - _WARP_LOAD_STORE_ALGORITHMS, - "warp store", - ) - group_kwargs = {"threads_in_warp": threads_in_warp} - else: - raise ValueError( - f"unsupported cuda.coop Store provider namespace " - f"{registered.namespace!r}" - ) - registered = _registered_provider_metadata(provider_factory, algorithm) - adapter = NumbaMlirCoreAdapter( - value_abis=_load_store_value_abis( - dtype=dtype, - items_per_thread=items_per_thread, - valid_items=valid_items_binding, - block_dim=group_kwargs.get("block_dim"), - threads_in_warp=group_kwargs.get("threads_in_warp"), - ) - ) - specialization_factory = ( - make_block_store_specialization - if registered.namespace == "block" - else make_warp_store_specialization - ) - core_specialization = specialization_factory( - dtype=adapter.core_dtype(dtype), - items_per_thread=items_per_thread, - algorithm=algorithm, - valid_items=valid_items_binding, - include_full_tile=( - not isinstance(num_valid_items, ArgumentBinding) - and num_valid_items is not None - ), - include_pointer_offset=( - offset_binding if isinstance(offset, ArgumentBinding) else True - ), - **( - {"block_dim": tuple(block_dim)} - if registered.namespace == "block" - else {"threads_in_warp": threads_in_warp} - ), - ) - specialization = adapter.materialize( - core_specialization, - **_materialization_metadata(registered), - extra_type_definitions=(numba_type_to_wrapper(dtype),), - ) - invocation_topology = ( - { - "logical_warp_threads": threads_in_warp, - "block_threads": block_dim, - } - if registered.namespace == "warp" - else {} - ) - return make_invocable_from_specialization( - specialization, - **invocation_topology, - ) - - -def store( - dtype, - threads_per_block=None, - items_per_thread=1, - algorithm="direct", - num_valid_items=None, - oob_default=None, - offset=None, -): - """Build a block store callable that needs no shared scratch.""" - - return _store( - store, - dtype, - threads_per_block, - items_per_thread, - algorithm, - num_valid_items, - oob_default, - offset, - ) - - -def _store_with_storage( - dtype, - threads_per_block=None, - items_per_thread=1, - algorithm="transpose", - num_valid_items=None, - oob_default=None, - offset=None, -): - """Build a block transpose store callable with shared scratch.""" - - return _store( - _store_with_storage, - dtype, - threads_per_block, - items_per_thread, - algorithm, - num_valid_items, - oob_default, - offset, - ) - - -def warp_store( - dtype, - threads_per_block=None, - threads_in_warp=32, - items_per_thread=1, - algorithm="direct", - num_valid_items=None, - oob_default=None, - offset=None, -): - """Build a warp store callable that needs no shared scratch.""" - - return _store( - warp_store, - dtype, - threads_per_block=threads_per_block, - threads_in_warp=threads_in_warp, - items_per_thread=items_per_thread, - algorithm=algorithm, - num_valid_items=num_valid_items, - oob_default=oob_default, - offset=offset, - ) - - -def _warp_store_with_storage( - dtype, - threads_per_block=None, - threads_in_warp=32, - items_per_thread=1, - algorithm="transpose", - num_valid_items=None, - oob_default=None, - offset=None, -): - """Build a warp store callable with shared scratch.""" - - return _store( - _warp_store_with_storage, - dtype, - threads_per_block=threads_per_block, - threads_in_warp=threads_in_warp, - items_per_thread=items_per_thread, - algorithm=algorithm, - num_valid_items=num_valid_items, - oob_default=oob_default, - offset=offset, - ) - - -for _factory, _operation in ((load, "load"), (store, "store")): - register_factory( - _factory, - operation=_operation, - namespace="block", - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.BLOCK, - synchronization_scope=SynchronizationScope.NONE, - ) -for _factory, _operation in ( - (_load_with_storage, "load"), - (_store_with_storage, "store"), -): - register_factory( - _factory, - operation=_operation, - namespace="block", - storage_abi=StorageABI.LEADING_POINTER, - execution_scope=SynchronizationScope.BLOCK, - synchronization_scope=SynchronizationScope.BLOCK, - ) -for _factory, _operation in ((warp_load, "load"), (warp_store, "store")): - register_factory( - _factory, - operation=_operation, - namespace="warp", - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.WARP, - synchronization_scope=SynchronizationScope.NONE, - ) -for _factory, _operation in ( - (_warp_load_with_storage, "load"), - (_warp_store_with_storage, "store"), -): - register_factory( - _factory, - operation=_operation, - namespace="warp", - storage_abi=StorageABI.LEADING_POINTER, - execution_scope=SynchronizationScope.WARP, - synchronization_scope=SynchronizationScope.WARP, - ) -del _factory, _operation diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_semantic.py b/python/cuda_coop/cuda/coop/numba_mlir/_semantic.py deleted file mode 100644 index 1dddf3f2d725..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_semantic.py +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Give compiler dtypes stable identities for specialization and symbol reuse. - -Two equivalent compiler types should select the same compiled operation even -when their internal bookkeeping differs. Adapt top-level Numba types to a -tagged class-and-name description before using the shared semantic encoder. -All other values follow that encoder's normal rules. The resulting hashable -descriptions feed cache keys and symbol hashes; they do not serialize compiler -objects for later reconstruction. -""" - -from collections.abc import Hashable - -import numba_cuda_mlir.numba_cuda.types as numba_types - -from cuda.coop._core import semantic_token - - -def _numba_semantic_token(value: object) -> Hashable: - """Describe a compiler dtype without serializing its implementation state. - - The common semantic-token encoder can inspect arbitrary object attributes. - For a Numba type, cache and symbol identity instead use its concrete type - class and printed spelling, so compiler bookkeeping does not become part - of the specialization key. Only a top-level Numba type is normalized here; - other values are passed directly to the common encoder. - - Parameters - ---------- - value : object - Compiler dtype or other semantic key component. A Numba ``Type`` is - replaced with a tagged tuple of its class module, class qualified name, - and ``str(value)`` before encoding. - - Returns - ------- - Hashable - Semantic description suitable for compiler cache and symbol - keys. This is a description, not a digest or a reversible encoding. - - Raises - ------ - TypeError - The common encoder encounters an unsupported callable value. - """ - - if isinstance(value, numba_types.Type): - value = ( - "numba-cuda-mlir-type", - type(value).__module__, - type(value).__qualname__, - str(value), - ) - return semantic_token(value) - - -__all__ = [ - "_numba_semantic_token", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_temp_storage.py b/python/cuda_coop/cuda/coop/numba_mlir/_temp_storage.py deleted file mode 100644 index ffbfa4b7409b..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_temp_storage.py +++ /dev/null @@ -1,79 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Describe scratch allocation and reuse for the Numba-CUDA-MLIR planner. - -The descriptor validates host-side options and retains them for compilation. -Constructing it does not allocate shared memory or synchronize threads. -""" - -import operator -from enum import Enum - -from .._core.api._payload import _normalize_alignment - - -class TempStorage: - """Shared-memory requirements for cooperative operations in one kernel. - - Parameters, defaults, synchronization rules, and the executable reuse - example follow :func:`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 :ref:`temporary storage ` - for shared versus exclusive slices and manual reuse synchronization. - """ - - def __init__( - self, - size_in_bytes=None, - *, - alignment=None, - auto_sync=False, - sharing="shared", - ): - if size_in_bytes is not None: - if isinstance(size_in_bytes, bool): - raise TypeError( - "TempStorage size_in_bytes must be an integer or None." - ) - try: - size_in_bytes = operator.index(size_in_bytes) - except TypeError as exc: - raise TypeError( - "TempStorage size_in_bytes must be an integer or None." - ) from exc - if size_in_bytes <= 0: - raise ValueError( - "TempStorage size_in_bytes must be a positive integer." - ) - - alignment = _normalize_alignment(alignment) - - if not isinstance(sharing, str) or isinstance(sharing, Enum): - raise TypeError( - "TempStorage sharing must be a string: 'shared' or 'exclusive'." - ) - sharing_value = sharing.strip().lower() - if sharing_value not in {"shared", "exclusive"}: - raise ValueError( - "TempStorage sharing must be 'shared' or 'exclusive'." - ) - - if auto_sync is not None and not isinstance(auto_sync, bool): - raise TypeError("TempStorage auto_sync must be None/True/False.") - - self.size_in_bytes = size_in_bytes - self.alignment = alignment - self.sharing = sharing_value - # Sharing selects the slice layout; synchronization is independent. - # The caller synchronizes before reuse unless auto_sync=True. - self.auto_sync = False if auto_sync is None else auto_sync - - -__all__ = ["TempStorage"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_temp_storage.pyi b/python/cuda_coop/cuda/coop/numba_mlir/_temp_storage.pyi deleted file mode 100644 index 006e535b101a..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_temp_storage.pyi +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from .._typing import TempStorageSharing - -class TempStorage: - """Explicit opaque byte scratch for planned shared-memory operations.""" - - size_in_bytes: int | None - alignment: int | None - auto_sync: bool - sharing: TempStorageSharing - - def __init__( - self, - size_in_bytes: int | None = None, - *, - alignment: int | None = None, - auto_sync: bool | None = False, - sharing: TempStorageSharing = "shared", - ) -> None: - """Configure scratch size, alignment, synchronization, and sharing.""" - -__all__ = ["TempStorage"] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_thread_data.py b/python/cuda_coop/cuda/coop/numba_mlir/_thread_data.py deleted file mode 100644 index ab0433ca651e..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_thread_data.py +++ /dev/null @@ -1,120 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Expose per-thread payload construction to the Numba-CUDA-MLIR compiler. - -``ThreadData`` is a kernel-language marker: the rewrite replaces a supported -call with a fixed-size local array, inferring its dtype from the surrounding -operations when possible. Calling the Python function directly raises a -compiler-context error. Its alignment helper reconciles the common API's -requested minimum with the compiler's pointer-alignment requirement. - -``local`` and ``shared`` expose the active runtime's array namespaces. They -are loaded on first attribute access so importing this module does not itself -require those runtime namespaces to be initialized. -""" - -from __future__ import annotations - -import struct -from typing import SupportsIndex - -from .._core.api._payload import _normalize_alignment -from .._core.thread_group import CoopCompilerContextRequiredError -from ._compiler._activation import _require_runtime - -# Annotations keep the runtime namespaces lazy while documenting module -# ownership for introspection and static analysis. -local: object -shared: object - - -def _normalize_thread_data_alignment( - alignment: SupportsIndex | None, -) -> int | None: - """Convert a requested payload alignment to the compiler's minimum. - - The common API accepts a positive power-of-two byte alignment, while the - Numba local-array representation also requires pointer alignment. Raise - smaller explicit requests to the host pointer size used by this adapter; - this still satisfies the caller's requested minimum. Leave unspecified - alignment for the compiler to choose. - - Parameters - ---------- - alignment : SupportsIndex or None - Requested minimum in bytes. Integer-index values are normalized by - the common helper; booleans are rejected. ``None`` means unspecified. - - Returns - ------- - int or None - At least ``struct.calcsize("P")`` for an explicit request, otherwise - ``None``. - - Raises - ------ - TypeError - The request is a boolean or cannot be interpreted as an integer. - ValueError - The request is not a positive power of two. - """ - - alignment = _normalize_alignment(alignment) - # The compiler requires pointer-aligned arrays. Stronger alignment also - # satisfies smaller minimum-alignment requests from the common API. - return None if alignment is None else max(struct.calcsize("P"), alignment) - - -def 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 - :func:`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 - :ref:`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 - :ref:`per-thread payloads ` for indexing and lifetime. - """ - - raise CoopCompilerContextRequiredError( - "cuda.coop.numba_mlir.ThreadData must be called from a supported " - "GPU kernel." - ) - - -def __getattr__(name: str): - """Resolve and cache an array namespace from the active runtime.""" - if name in {"local", "shared"}: - value = getattr(_require_runtime(), name) - globals()[name] = value - return value - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -def __dir__() -> list[str]: - return sorted(__all__) - - -__all__ = [ - "ThreadData", - "_normalize_thread_data_alignment", - "local", - "shared", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_thread_data.pyi b/python/cuda_coop/cuda/coop/numba_mlir/_thread_data.pyi deleted file mode 100644 index e8418385631b..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_thread_data.pyi +++ /dev/null @@ -1,68 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from typing import Any, Protocol, TypeAlias, overload - -from typing_extensions import TypeVar - -from .._typing import CommonNumericScalar, ThreadDataLike - -_ItemT = TypeVar("_ItemT", bound=CommonNumericScalar) - -_ArrayShape: TypeAlias = int | tuple[int, ...] - -class _LocalMemory(Protocol): - """Numba-CUDA-MLIR thread-local memory namespace.""" - - def array( - self, - shape: _ArrayShape, - dtype: object, - *, - alignment: int | None = 8, - ) -> Any: - """Allocate thread-local compiler storage.""" - -class _SharedMemory(Protocol): - """Numba-CUDA-MLIR shared-memory namespace.""" - - def array( - self, - shape: _ArrayShape, - dtype: object, - *, - alignment: int | None = 8, - ) -> Any: - """Allocate shared compiler storage.""" - -def _normalize_thread_data_alignment(alignment: int | None) -> int | None: ... - -local: _LocalMemory - -shared: _SharedMemory - -@overload -def ThreadData( - items_per_thread: int, - dtype: type[_ItemT], - *, - alignment: int | None = None, -) -> ThreadDataLike[_ItemT]: - """Construct typed thread-local storage.""" - -@overload -def ThreadData( - items_per_thread: int, - dtype: object | None = None, - *, - alignment: int | None = None, -) -> ThreadDataLike[Any]: - """Construct storage using a compiler dtype token or inferred dtype.""" - -__all__ = [ - "ThreadData", - "_normalize_thread_data_alignment", - "local", - "shared", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_thread_group.py b/python/cuda_coop/cuda/coop/numba_mlir/_thread_group.py deleted file mode 100644 index 74b5772cee30..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_thread_group.py +++ /dev/null @@ -1,100 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Attach the Numba-CUDA-MLIR scope to shared thread-group descriptors. - -The shared descriptor supplies hierarchy and partition rules. This subclass -keeps the qualified API's compiler scope when a group is resolved or split. -Creating a descriptor does not query a running kernel or synchronize threads. -""" - -from __future__ import annotations - -from cuda.coop._core import ThreadHierarchy, make_thread_group -from cuda.coop._core.thread_group import ThreadGroup as CoreThreadGroup - -_ROOT_SCOPE = __name__.rsplit(".", 1)[0] - -Hierarchy = ThreadHierarchy - - -class ThreadGroup(CoreThreadGroup): - """Describe a CUDA thread group for the Numba-CUDA-MLIR planner. - - Membership and launch resolution follow :class:`cuda.coop.ThreadGroup`. - A descriptor can exist even when a particular operation does not support - that group; operation planning checks support separately. - """ - - def group_by( - self, - count: int, - *, - exhaustive: bool = True, - ) -> ThreadGroup: - """Partition the parent and retain this backend's descriptor type. - - ``count`` measures threads within a warp or physical warps within a - block. The validation rules follow ``cuda.coop.ThreadGroup.group_by``. - The returned descriptor does not synchronize or rearrange threads. - """ - return super().group_by(count, exhaustive=exhaustive) - - -# These names support explicit imports used by the adjacent typing stubs. They -# remain private to this module rather than expanding the qualified facade. -ReductionGroup = ThreadGroup -BlockGroup = ThreadGroup -WarpGroup = ThreadGroup - - -def _make_group(kind: str) -> ThreadGroup: - """Create a current-group descriptor with this backend's scope.""" - return make_thread_group( - kind, - group_type=ThreadGroup, - scope=_ROOT_SCOPE, - ) - - -def this_thread() -> ThreadGroup: - """Describe the current thread.""" - - return _make_group("thread") - - -def this_warp() -> ThreadGroup: - """Describe the current physical warp.""" - - return _make_group("warp") - - -def this_block() -> ThreadGroup: - """Describe the current CUDA thread block.""" - - return _make_group("block") - - -def this_cluster() -> ThreadGroup: - """Describe the current cluster where the launch can represent it.""" - - return _make_group("cluster") - - -def this_grid() -> ThreadGroup: - """Describe the current grid where the launch can represent it.""" - - return _make_group("grid") - - -__all__ = [ - "Hierarchy", - "ThreadGroup", - "ThreadHierarchy", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_thread_group.pyi b/python/cuda_coop/cuda/coop/numba_mlir/_thread_group.pyi deleted file mode 100644 index 5cefa811ac41..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_thread_group.pyi +++ /dev/null @@ -1,85 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from typing import Generic, Literal, TypeAlias, overload - -from typing_extensions import TypeVar - -from .. import ThreadHierarchy -from .._core.api.thread_group import ThreadGroup as CommonThreadGroup -from .._typing import ThreadGroupKind - -_GroupKindT_co = TypeVar( - "_GroupKindT_co", - bound=ThreadGroupKind, - covariant=True, - default=ThreadGroupKind, -) - -Hierarchy: TypeAlias = ThreadHierarchy - -class ThreadGroup( - CommonThreadGroup[_GroupKindT_co], - Generic[_GroupKindT_co], -): - """Compile-time CUDA group descriptor for Numba-CUDA-MLIR.""" - - @overload - def group_by( - self: ThreadGroup[Literal["warp"]], - count: int, - *, - exhaustive: bool = True, - ) -> ThreadGroup[Literal["threads_within_warp"]]: - """Partition a physical warp into groups of threads.""" - - @overload - def group_by( - self: ThreadGroup[Literal["block"]], - count: int, - *, - exhaustive: bool = True, - ) -> ThreadGroup[Literal["warps_within_block"]]: - """Partition a block into groups of physical warps.""" - -ReductionGroup: TypeAlias = ThreadGroup[ - Literal[ - "thread", - "warp", - "threads_within_warp", - "warps_within_block", - "block", - "cluster", - ] -] - -BlockGroup: TypeAlias = ThreadGroup[Literal["block"]] - -WarpGroup: TypeAlias = ThreadGroup[Literal["warp", "threads_within_warp"]] - -def this_thread() -> ThreadGroup[Literal["thread"]]: - """Describe the current thread.""" - -def this_warp() -> ThreadGroup[Literal["warp"]]: - """Describe the current complete physical warp.""" - -def this_block() -> ThreadGroup[Literal["block"]]: - """Describe the current CUDA thread block.""" - -def this_cluster() -> ThreadGroup[Literal["cluster"]]: - """Describe the current cluster where the launch can represent it.""" - -def this_grid() -> ThreadGroup[Literal["grid"]]: - """Describe the current grid.""" - -__all__ = [ - "Hierarchy", - "ThreadGroup", - "ThreadHierarchy", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/_types.py b/python/cuda_coop/cuda/coop/numba_mlir/_types.py deleted file mode 100644 index 8daf4c938a1f..000000000000 --- a/python/cuda_coop/cuda/coop/numba_mlir/_types.py +++ /dev/null @@ -1,2978 +0,0 @@ -# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Generate C++ wrappers that Numba kernels can call as device functions. - -A provider is the implementation of a cooperative operation for a particular -configuration, such as a block load with a fixed dtype and tile size. An -``Algorithm`` describes that configuration and emits a C++ wrapper around the -CUB operation. Parameter descriptors are ``Parameter`` objects such as -``Value``, ``Pointer``, and ``Array``. They describe how arguments accepted by -Numba map to the wrapper's C++ signature, including pointers, arrays, scalar -controls, and values embedded at compile time. For example, ``Array`` describes -an array parameter's type and extent; ``ThreadData`` constructs the per-thread -payload passed to that parameter. Group planning separately uses -``ThreadGroup`` and ``ThreadHierarchy`` to describe which threads cooperate. -``TypeWrapper`` supplies matching C++ storage declarations for compiler types -without a builtin C++ spelling. - -Construction specializes an algorithm, which can then be compiled to LTO IR -and exposed through an ``Invocable``. The ``Invocable`` supplies compiler -overloads and keeps the link artifacts while the containing kernel is -compiled. The compiler turns calls to it into device calls. - -Several operation specializations may be collected before compilation. -``prepare_ltoir_bundle`` combines their wrappers into a shared translation unit -and coalesces equivalent providers, reducing repeated NVRTC compilation. The -individual and bundled paths query scratch layouts during the same NVRTC -compilation that emits their provider image. The image and layouts stay -together so the Python storage contract agrees with the compiled C++ types. -""" - -from __future__ import annotations - -import hashlib -import os -import re -import weakref -from collections import OrderedDict -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence -from contextlib import contextmanager -from contextvars import ContextVar -from copy import copy -from io import StringIO -from numbers import Integral -from textwrap import dedent -from types import FunctionType as PyFunctionType -from typing import TYPE_CHECKING, Protocol, cast - -import numba_cuda_mlir.numba_cuda.types as numba_types -from numba_cuda_mlir import cuda, types -from numba_cuda_mlir.compiler import ExternFunction -from numba_cuda_mlir.descriptor import mlir_target -from numba_cuda_mlir.extending import _NumbaCudaMlirOverloadFunctionTemplate -from numba_cuda_mlir.numba_cuda.typing.templates import ( - make_overload_template, - signature, -) - -from cuda.coop._core import SynchronizationScope - -from ._compiler import _nvrtc as nvrtc -from ._compiler._operations import StorageABI -from ._semantic import _numba_semantic_token - -if TYPE_CHECKING: - from numba_cuda_mlir.numba_cuda.typing.templates import ( - _OverloadFunctionTemplate, - ) - -NUMBA_TYPES_TO_CPP = { - numba_types.boolean: "bool", - numba_types.int8: "::cuda::std::int8_t", - numba_types.int16: "::cuda::std::int16_t", - numba_types.int32: "::cuda::std::int32_t", - numba_types.int64: "::cuda::std::int64_t", - numba_types.uint8: "::cuda::std::uint8_t", - numba_types.uint16: "::cuda::std::uint16_t", - numba_types.uint32: "::cuda::std::uint32_t", - numba_types.uint64: "::cuda::std::uint64_t", - numba_types.float16: "__half", - numba_types.float32: "float", - numba_types.float64: "double", -} - -_BlockThreads = int | tuple[int, ...] | list[int] - -_SUPPORTED_LOGICAL_WARP_THREADS = frozenset({1, 2, 4, 8, 16, 32}) - - -_COOP_SPECIALIZATION_COLLECTOR: ContextVar[list[Algorithm] | None] = ContextVar( - "cuda_coop_numba_mlir_specialization_collector", default=None -) - -_KERNEL_COMPILE_TARGET: ContextVar[nvrtc.CompileTarget | None] = ContextVar( - "cuda_coop_numba_mlir_kernel_compile_target", default=None -) - - -@contextmanager -def kernel_compile_target( - target: nvrtc.CompileTarget | None, -) -> Iterator[None]: - """Generate device code for the kernel being compiled in this context. - - The whole-function planner enters this once per compiler attempt with - the kernel's resolved target. Providers created inside then follow the - kernel's ``chip`` and ``lineinfo`` options instead of the current device. - ``None`` selects the current device. Exiting restores the previous target, - so a nested kernel compilation does not leak its target to its caller. - """ - - token = _KERNEL_COMPILE_TARGET.set(target) - try: - yield - finally: - _KERNEL_COMPILE_TARGET.reset(token) - - -def _current_compile_target() -> nvrtc.CompileTarget: - """Return the target for generated device code. - - During a kernel compilation this is the kernel's own target, set by - ``kernel_compile_target``. Outside one, such as when a provider is - specialized directly, fall back to the current device without line - information. - """ - - target = _KERNEL_COMPILE_TARGET.get() - if target is not None: - return target - # The CUDA module reexports this accessor but omits it from its stub. - device = ( - cuda.get_current_device() # pyright: ignore[reportAttributeAccessIssue] - ) - cc_major, cc_minor = device.compute_capability - return nvrtc.CompileTarget(cc=int(cc_major) * 10 + int(cc_minor)) - - -@contextmanager -def collect_specializations() -> Iterator[list[Algorithm]]: - """Collect provider specializations without compiling their wrappers. - - Within this context, ``make_invocable_from_specialization`` records each - qualified ``Algorithm`` and returns it instead of an ``Invocable``. The - rewrite uses the collected algorithms, including their stored thread - dimensions, to compile several providers in one NVRTC translation unit - before creating their callable wrappers. Qualification still resolves - compiler identity; this context only defers artifact creation. - - Batching reduces compiler startup and repeated header processing. For - example, eleven distinct uncached providers can share one NVRTC compilation - when bundling succeeds. Repeated primitive calls can reuse a provider, and - cache hits can avoid compilation entirely. - - Each entry creates a fresh context-local list. Exiting restores the previous - collector even after an exception, so nested collections do not append to - their parent's list. - - Yields - ------ - list of Algorithm - Qualified provider specializations in factory-call order, with their - logical warp widths and enclosing block configurations stored on them. - """ - - collected: list[Algorithm] = [] - token = _COOP_SPECIALIZATION_COLLECTOR.set(collected) - try: - yield collected - finally: - _COOP_SPECIALIZATION_COLLECTOR.reset(token) - - -def numba_type_to_cpp(numba_type): - """Return the builtin C++ spelling or the opaque ``storage_t`` name.""" - cpp_type = NUMBA_TYPES_TO_CPP.get(numba_type) - if cpp_type is not None: - return cpp_type - return "storage_t" - - -def _validate_logical_warp_threads(logical_warp_threads): - """Check the logical warp width before compiling or allocating scratch. - - Warp factories and provider code generation use this to require a Python - integer from the supported widths. ``logical_warp_threads`` counts - participating lanes per logical warp. Return it unchanged, or raise - ``ValueError`` for an unsupported width, boolean, or other type. - """ - - if ( - isinstance(logical_warp_threads, bool) - or not isinstance(logical_warp_threads, int) - or logical_warp_threads not in _SUPPORTED_LOGICAL_WARP_THREADS - ): - supported = ", ".join( - str(value) for value in sorted(_SUPPORTED_LOGICAL_WARP_THREADS) - ) - raise ValueError( - "warp-scoped providers require a logical width in " - f"{{{supported}}}; got {logical_warp_threads!r}" - ) - return logical_warp_threads - - -def _normalize_block_threads(threads_per_block): - """Validate one to three block dimensions and return their thread count. - - Require positive Python integers, reject booleans, and limit the - product to 1024 threads. Allocating warp wrappers use this count to - size scratch. - """ - if isinstance(threads_per_block, bool): - # Keep the established ValueError contract for invalid block dimensions. - raise ValueError( # noqa: TRY004 - "block_threads must be a positive integer or dimension tuple; " - f"got {threads_per_block!r}" - ) - if isinstance(threads_per_block, int): - block_threads = threads_per_block - elif isinstance(threads_per_block, (tuple, list)): - if not 1 <= len(threads_per_block) <= 3: - raise ValueError( - "block_threads dimension tuples require one to three entries; " - f"got {threads_per_block!r}" - ) - block_threads = 1 - for dimension in threads_per_block: - if ( - isinstance(dimension, bool) - or not isinstance(dimension, int) - or dimension <= 0 - ): - raise ValueError( - "block_threads dimension tuple entries must be positive " - f"integers; got {threads_per_block!r}" - ) - block_threads *= dimension - else: - # Keep the established ValueError contract for invalid block dimensions. - raise ValueError( # noqa: TRY004 - "block_threads must be a positive integer or dimension tuple; " - f"got {threads_per_block!r}" - ) - if not 1 <= block_threads <= 1024: - raise ValueError( - "block_threads must describe between 1 and 1024 threads; " - f"got {threads_per_block!r}" - ) - return block_threads - - -def _hash_symbol_value(hasher, value, depth=0): - del depth - hasher.update( - repr(_numba_semantic_token(value)).encode( - "utf-8", errors="backslashreplace" - ) - ) - - -def _align_up(value: int, alignment: int) -> int: - return ((value + alignment - 1) // alignment) * alignment - - -def _lto_ir_digest(lto_ir): - if isinstance(lto_ir, str): - data = lto_ir.encode("utf-8") - else: - data = bytes(lto_ir) - return hashlib.sha1(data).hexdigest() - - -def _struct_size_alignment( - member_types: Iterable[numba_types.Type], -) -> tuple[int, int]: - """Lay out ordered members and return their byte size and alignment. - - Each field starts at its required alignment. Tail padding lets an array of - these structs keep that alignment for every element. - """ - offset = 0 - max_align = 1 - for member_type in member_types: - member_size, member_align = _size_alignment_from_numba_type(member_type) - offset = _align_up(offset, member_align) - offset += member_size - max_align = max(max_align, member_align) - return _align_up(offset, max_align), max_align - - -def _registered_struct_member_types( - numba_type: numba_types.Type, -) -> tuple[numba_types.Type, ...] | None: - """Find a struct layout whose CUDA and MLIR member types agree. - - The storage-layout fallback must describe the same value on both sides of - the compiler boundary. Require a CUDA ``StructModel`` and an MLIR - ``StructModel`` with identical ordered member types before using their - members to compute size and alignment. This compares member descriptions, - not independently measured byte layouts. - - MLIR model lookup can require an active context and location. Supply only - those that are missing, leaving an existing caller context in place. - - Parameters - ---------- - numba_type : numba_types.Type - Compiler type to look up in the CUDA/default and MLIR data managers. - - Returns - ------- - tuple of numba_types.Type or None - Ordered members when both models agree, including an empty tuple for - matching empty models. ``None`` means either lookup has no supported - struct model or the member sequences differ. - """ - - from numba_cuda_mlir import models as mlir_models - - # These classes are provided by MLIR's compiled extension without stubs. - from numba_cuda_mlir._mlir.ir import ( - Context, # pyright: ignore[reportAttributeAccessIssue] - Location, # pyright: ignore[reportAttributeAccessIssue] - ) - from numba_cuda_mlir.numba_cuda.datamodel.models import ( - StructModel as CudaStructModel, - ) - from numba_cuda_mlir.numba_cuda.datamodel.registry import default_manager - from numba_cuda_mlir.numba_cuda.models import cuda_data_manager - - def get_member_types(manager, model_type): - try: - model = manager.lookup(numba_type) - except (KeyError, NotImplementedError, TypeError, ValueError): - return None - if not isinstance(model, model_type): - return None - return tuple( - model.get_type(index) for index in range(model.field_count) - ) - - cuda_manager = cuda_data_manager.chain(default_manager) - cuda_members = get_member_types(cuda_manager, CudaStructModel) - if cuda_members is None: - return None - - try: - current_context = Context.current - except ValueError: - current_context = None - - if current_context is None: - with Context(), Location.unknown(): - mlir_members = get_member_types( - mlir_models.mlir_data_manager, - mlir_models.StructModel, - ) - else: - try: - current_location = Location.current - except ValueError: - current_location = None - if current_location is None: - with Location.unknown(): - mlir_members = get_member_types( - mlir_models.mlir_data_manager, - mlir_models.StructModel, - ) - else: - mlir_members = get_member_types( - mlir_models.mlir_data_manager, - mlir_models.StructModel, - ) - - if mlir_members != cuda_members: - return None - return cuda_members - - -def _size_alignment_from_numba_type( - numba_type: numba_types.Type, -) -> tuple[int, int]: - """Compute a supported value layout when LLVM ABI inspection is unavailable. - - ``TypeWrapper`` needs byte size and alignment to generate the opaque C++ - ``storage_t`` used for non-builtin types. Recursively lay out scalar values, - uniform tuples, native aggregates, and matching registered CUDA/MLIR struct - models. Struct fields receive alignment padding and the final size is - rounded to the largest member alignment. Bitfield aggregates use their - compiler-provided storage type instead of laying out individual bitfields. - - This is a restricted fallback, not a general data-model size estimator. - Unknown types and registered structs with no inspectable members fail rather - than receiving a guessed CUB storage layout. - - Parameters - ---------- - numba_type : numba_types.Type - Type whose by-value representation must be shared with generated C++. - - Returns - ------- - tuple of int - ``(size, alignment)`` in bytes. - - Raises - ------ - TypeError - The type or one of its members has no supported layout description. - """ - - from numba_cuda_mlir.type_defs.aggregate_types import AggregateType - - if isinstance( - numba_type, (numba_types.Boolean, numba_types.BooleanLiteral) - ): - return 1, 1 - if isinstance( - numba_type, (numba_types.Integer, numba_types.IntegerLiteral) - ): - size = max(1, numba_type.bitwidth // 8) - return size, size - if isinstance(numba_type, numba_types.Float): - size = max(1, numba_type.bitwidth // 8) - return size, size - if isinstance(numba_type, numba_types.Complex): - elem_size = max(1, numba_type.underlying_float.bitwidth // 8) - return 2 * elem_size, elem_size - if isinstance(numba_type, numba_types.UniTuple): - elem_size, elem_align = _size_alignment_from_numba_type( - numba_type.dtype - ) - return elem_size * numba_type.count, elem_align - if isinstance(numba_type, AggregateType): - if numba_type.is_bitfield_struct: - return _size_alignment_from_numba_type( - numba_type.get_bitfield_storage_type() - ) - return _struct_size_alignment( - field_type - for field_name, field_type, _bit_width in numba_type.fields - if field_name is not None - ) - - struct_member_types = _registered_struct_member_types(numba_type) - if struct_member_types: - return _struct_size_alignment(struct_member_types) - - dtype_name = getattr(numba_type, "name", str(numba_type)) - raise TypeError( - "cuda.coop.numba_mlir cannot safely materialize CUB storage for dtype " - f"{dtype_name}: exact ABI size and alignment are unavailable; use a " - "compiler-native dtype, Numba-CUDA-MLIR AggregateType, or matching " - "registered CUDA and MLIR StructModels with inspectable member types" - ) - - -def _ltoir_to_ptx(ltoir: bytes, *, name: str, cc: int) -> str: - """Link one LTO-IR image to PTX for explicit artifact inspection. - - Compile tests use this to inspect generated instructions and storage. - Provider creation obtains layouts directly from NVRTC and does not call - this helper. The link uses ``cc`` as its target and ``name`` to identify - the input image; it does not change the provider's retained artifacts. - """ - - # cuda-core chooses cu12/cu13 exports dynamically, outside its stubs. - from cuda.core import ( - Linker, # pyright: ignore[reportAttributeAccessIssue] - LinkerOptions, # pyright: ignore[reportAttributeAccessIssue] - ObjectCode, # pyright: ignore[reportAttributeAccessIssue] - ) - - ltoir_obj = ObjectCode.from_ltoir(ltoir, name=name) - linker_options = LinkerOptions( - arch=f"sm_{cc}", - link_time_optimization=True, - ptx=True, - ) - linked_ptx = Linker(ltoir_obj, options=linker_options).link("ptx") - return linked_ptx.code.decode("utf-8") - - -class TypeWrapper: - """Supply an aligned C++ storage declaration for a compiler value type.""" - - def __init__(self, numba_type: numba_types.Type) -> None: - """Build the C++ storage declaration for a compiler value type. - - Builtin types already have C++ spellings and need no declaration. Other - types use an aligned byte-array ``storage_t`` with the compiler's ABI - size and alignment; no field accessors or value conversion are - generated. First query the target context's LLVM layout. If that - inspection fails, use the restricted native/registered-model fallback - rather than guessing. - - Parameters - ---------- - numba_type : numba_types.Type - Type that will appear in a provider template or parameter. - - Raises - ------ - TypeError - LLVM layout inspection failed and the fallback cannot establish the - type's size and alignment. - - Notes - ----- - The instance stores the declaration in ``code`` and initializes - ``lto_irs`` to an empty list. Constructing this wrapper does not compile - any source. - """ - - self.lto_irs = [] - - if numba_type in NUMBA_TYPES_TO_CPP: - self.code = "" - return - - # TypeWrapper needs LLVM ABI layout info (size/alignment) to materialize - # the storage_t wrapper for user-defined dtypes. - try: - abi_context = mlir_target.target_context - value_type = abi_context.get_value_type(numba_type) - size = value_type.get_abi_size(abi_context.target_data) - alignment = value_type.get_abi_alignment(abi_context.target_data) - except Exception: # noqa: BLE001 - unsupported LLVM layouts use the validated fallback. - size, alignment = _size_alignment_from_numba_type(numba_type) - buf = StringIO() - w = buf.write - w(f"struct __align__({alignment}) storage_t\n") - w("{\n") - w(f" char data[{size}];\n") - w("};\n") - - self.code = buf.getvalue() - - -def numba_type_to_wrapper(numba_type: numba_types.Type) -> TypeWrapper: - return TypeWrapper(numba_type) - - -_CPP_PARAMETER_RESERVED_NAMES = frozenset( - [ - "alignas", - "alignof", - "and", - "and_eq", - "asm", - "atomic_cancel", - "atomic_commit", - "atomic_noexcept", - "auto", - "bitand", - "bitor", - "bool", - "break", - "case", - "catch", - "char", - "char8_t", - "char16_t", - "char32_t", - "class", - "compl", - "concept", - "const", - "consteval", - "constexpr", - "constinit", - "const_cast", - "continue", - "co_await", - "co_return", - "co_yield", - "decltype", - "default", - "delete", - "do", - "double", - "dynamic_cast", - "else", - "enum", - "explicit", - "export", - "extern", - "false", - "float", - "for", - "friend", - "goto", - "if", - "inline", - "int", - "long", - "mutable", - "namespace", - "new", - "noexcept", - "not", - "not_eq", - "nullptr", - "operator", - "or", - "or_eq", - "private", - "protected", - "public", - "reflexpr", - "register", - "reinterpret_cast", - "requires", - "return", - "short", - "signed", - "sizeof", - "static", - "static_assert", - "static_cast", - "struct", - "switch", - "synchronized", - "template", - "this", - "thread_local", - "throw", - "true", - "try", - "typedef", - "typeid", - "typename", - "union", - "unsigned", - "using", - "virtual", - "void", - "volatile", - "wchar_t", - "while", - "xor", - "xor_eq", - "threadIdx", - "blockDim", - "storage_t", - ] -) - - -def _cpp_parameter_names(parameters, *, reserved=()): - """Choose readable identifiers without shadowing generated C++ locals.""" - - used = set(_CPP_PARAMETER_RESERVED_NAMES).union(reserved) - names = [] - for pid, parameter in enumerate(parameters): - raw_name = getattr(parameter, "parameter_name", None) - name = ( - re.sub(r"[^a-zA-Z0-9_]", "_", raw_name).rstrip("_") - if isinstance(raw_name, str) - else "" - ) - if not name or not name[0].isalpha() or "__" in name: - name = f"param_{pid}" - stem = name - suffix = 0 - while True: - # Source and ABI wrappers derive these locals from the formal name. - identifiers = { - name, - f"checked_{name}", - f"transformed_{name}", - f"{name}_state", - f"cast_{name}", - f"ref_{name}", - } - if not used.intersection(identifiers): - break - suffix += 1 - name = f"{stem}_{suffix}" - used.update(identifiers) - names.append(name) - return names - - -class Parameter: - """Describe one argument at the Python-to-C++ call boundary. - - Subclasses define the compiler dtype and C++ representation. An output is - supplied by the wrapper and becomes the call result, so it consumes no - Python input argument. Concrete descriptors return themselves from - ``specialize``; dependent descriptors resolve template arguments first. - """ - - # The formal parameter name is separate from an operator's linker symbol. - parameter_name: str | None = None - - def __init__(self, is_output=False): - self.is_output = is_output - - def __repr__(self) -> str: - return f"Parameter(out={self.is_output})" - - def specialize(self, template_arguments): - return self - - def is_provided_by_user(self): - return not self.is_output - - def dtype(self): - raise NotImplementedError("parameter must be specialized before typing") - - def accepts_actual_type(self, actual_type, typing_context): - return typing_context.can_convert(actual_type, self.dtype()) is not None - - -class Value(Parameter): - """Pass a scalar by value using the descriptor's compiler dtype.""" - - def __init__(self, value_type, is_output=False): - self.value_type = value_type - super().__init__(is_output) - - def __repr__(self) -> str: - return f"Value(dtype={self.value_type}, out={self.is_output})" - - def dtype(self): - return self.value_type - - @property - def provider_dtype(self): - """Return the scalar dtype received by the underlying provider.""" - - return self.value_type - - def cpp_decl(self, name): - return numba_type_to_cpp(self.value_type) + " " + name - - def mangled_name(self): - return f"{self.value_type}" - - -class ExactValue(Value): - """A scalar whose compiler dtype must exactly match the provider ABI.""" - - def __repr__(self) -> str: - return f"ExactValue(dtype={self.value_type}, out={self.is_output})" - - def mangled_name(self): - return f"Exact{self.value_type}" - - def accepts_actual_type(self, actual_type, typing_context): - del typing_context - literal_type = getattr(actual_type, "literal_type", actual_type) - return literal_type == self.value_type - - -def _accepts_runtime_control_integer(actual_type) -> bool: - actual_type = getattr(actual_type, "literal_type", actual_type) - return ( - isinstance(actual_type, numba_types.Integer) - and actual_type.bitwidth <= 64 - and (actual_type.signed or actual_type.bitwidth <= 32) - ) - - -class BoundedInteger(Value): - """A signed 64-bit integer checked before provider-type narrowing.""" - - def __init__(self, provider_dtype, *, minimum, maximum): - if not isinstance(provider_dtype, numba_types.Integer) or isinstance( - provider_dtype, numba_types.IntegerLiteral - ): - raise TypeError( - "provider_dtype must be a non-literal integer dtype" - ) - for name, bound in (("minimum", minimum), ("maximum", maximum)): - if not isinstance(bound, Integral) or isinstance(bound, bool): - raise TypeError(f"{name} must be an integer") - minimum = int(minimum) - maximum = int(maximum) - if minimum > maximum: - raise ValueError("minimum must not exceed maximum") - - if provider_dtype.signed: - provider_minimum = -(1 << (provider_dtype.bitwidth - 1)) - provider_maximum = (1 << (provider_dtype.bitwidth - 1)) - 1 - else: - provider_minimum = 0 - provider_maximum = (1 << provider_dtype.bitwidth) - 1 - if not provider_minimum <= minimum <= maximum <= provider_maximum: - raise ValueError( - "bounded integer limits must fit the provider integer dtype" - ) - if not -(1 << 63) <= minimum <= maximum <= (1 << 63) - 1: - raise ValueError( - "bounded integer limits must fit the signed 64-bit ABI" - ) - - self._provider_dtype = provider_dtype - self.minimum = minimum - self.maximum = maximum - super().__init__(numba_types.int64) - - def __repr__(self) -> str: - return ( - "BoundedInteger(" - f"provider_dtype={self.provider_dtype}, minimum={self.minimum}, " - f"maximum={self.maximum})" - ) - - @property - def provider_dtype(self): - return self._provider_dtype - - def mangled_name(self): - bounds = internal_mangle_cpp(f"{self.minimum}_{self.maximum}") - return f"Bounded{self.value_type}To{self.provider_dtype}_{bounds}" - - def accepts_actual_type(self, actual_type, typing_context): - del typing_context - return _accepts_runtime_control_integer(actual_type) - - -class PointerOffset(Value): - """Adjust an earlier pointer argument by an element count. - - ``pointer_arg_index`` selects the earlier provider argument. A static - value is embedded in C++ source; otherwise the caller supplies the offset - at runtime. The offset itself is not an argument to the CUB method. - """ - - def __init__(self, value_type, pointer_arg_index=0, static_value=None): - if static_value is not None and ( - not isinstance(static_value, Integral) - or isinstance(static_value, bool) - ): - raise TypeError("static pointer offset must be an integer") - self.pointer_arg_index = pointer_arg_index - self.static_value = None if static_value is None else int(static_value) - super().__init__(value_type) - - def __repr__(self) -> str: - return ( - f"PointerOffset(dtype={self.value_type}, " - f"pointer_arg_index={self.pointer_arg_index}, " - f"static_value={self.static_value})" - ) - - def mangled_name(self): - static = ( - "Runtime" - if self.static_value is None - else f"Static{internal_mangle_cpp(str(self.static_value))}" - ) - return f"Offset{self.value_type}{static}" - - def is_provided_by_user(self): - return self.static_value is None - - def accepts_actual_type(self, actual_type, typing_context): - del typing_context - return _accepts_runtime_control_integer(actual_type) - - -class Pointer(Parameter): - """Pass contiguous one-dimensional array storage as a typed C++ pointer. - - Typing requires a C-contiguous array because the wrapper passes only the - address; it cannot retain an arbitrary array stride. - """ - - def __init__(self, value_dtype, is_output=False): - self.value_dtype = value_dtype - super().__init__(is_output) - - def __repr__(self) -> str: - return f"Pointer(dtype={self.value_dtype}, out={self.is_output})" - - def cpp_decl(self, name): - return numba_type_to_cpp(self.value_dtype) + "* " + name - - def dtype(self): - # Provider wrappers receive a raw pointer and therefore cannot preserve - # an arbitrary array stride. Require contiguous one-dimensional memory - # so a sliced device view cannot silently produce adjacent-element - # results. - return numba_types.Array(self.value_dtype, 1, "C") - - def mangled_name(self): - return f"P{self.value_dtype}" - - -class DependentPointer(Parameter): - """Resolve a template dtype before constructing a pointer descriptor.""" - - def __init__(self, value_dtype, is_output=False): - self.value_dtype = value_dtype - super().__init__(is_output) - - def __repr__(self) -> str: - return f"DependentPointer(dep={self.value_dtype}, out={self.is_output})" - - def specialize(self, template_arguments): - return Pointer( - self.value_dtype.resolve(template_arguments), self.is_output - ) - - -class PointerReference(Pointer): - """Pass an array pointer, then dereference it for the CUB method call.""" - - def __init__(self, value_dtype, is_output=False): - self.deref_on_call = True - super().__init__(value_dtype, is_output) - - def __repr__(self) -> str: - return ( - f"PointerReference(dtype={self.value_dtype}, out={self.is_output})" - ) - - -class DependentPointerReference(DependentPointer): - """Resolve the dtype of a pointer that the CUB call will dereference.""" - - def __init__(self, value_dtype, is_output=False): - self.deref_on_call = True - super().__init__(value_dtype, is_output) - - def __repr__(self) -> str: - return ( - f"DependentPointerReference(dep=" - f"{self.value_dtype}, out={self.is_output})" - ) - - def specialize(self, template_arguments): - return PointerReference( - self.value_dtype.resolve(template_arguments), - self.is_output, - ) - - -class Reference(Parameter): - """Represent a C++ reference using a scalar at the external-call boundary. - - The ABI wrapper creates a local scalar and passes it by reference to the - typed wrapper. An output reference becomes the external call's result. - """ - - def __init__(self, value_dtype, is_output=False): - self.value_dtype = value_dtype - super().__init__(is_output) - - def __repr__(self) -> str: - return f"Reference(dtype={self.value_dtype}, out={self.is_output})" - - def cpp_decl(self, name): - return numba_type_to_cpp(self.value_dtype) + "& " + name - - def dtype(self): - return self.value_dtype - - def mangled_name(self): - return f"R{self.value_dtype}" - - -class DependentReference(Parameter): - """Resolve a template dtype before constructing a reference descriptor.""" - - def __init__(self, value_dtype, is_output=False): - self.value_dtype = value_dtype - super().__init__(is_output) - - def __repr__(self) -> str: - return ( - f"DependentReference(dep={self.value_dtype}, out={self.is_output})" - ) - - def specialize(self, template_arguments): - return Reference( - self.value_dtype.resolve(template_arguments), self.is_output - ) - - -class Array(Pointer): - """Pass a fixed-size per-thread array to a CUB array-reference parameter. - - The external ABI receives a pointer. The typed wrapper casts that pointer - to a C++ array reference with ``size`` elements. - """ - - def __init__(self, value_dtype, size, is_output=False): - self.size = size - super().__init__(value_dtype, is_output) - - def __repr__(self) -> str: - return ( - f"Array(dtype={self.value_dtype}, size=" - f"{self.size}, out={self.is_output})" - ) - - def cpp_decl(self, name): - return f"{numba_type_to_cpp(self.value_dtype)} (&{name})[{self.size}]" - - def dtype(self): - return numba_types.Array(self.value_dtype, 1, "C") - - def mangled_name(self): - return f"A{self.size}_{self.value_dtype}" - - -class TransformedArray(Array): - """Array input converted elementwise before the native CUB call.""" - - def __init__( - self, - source_dtype, - target_dtype, - size, - cpp_expression, - ): - if "{value}" not in cpp_expression: - raise ValueError("array input transform must reference {value}") - self.target_dtype = target_dtype - self.cpp_expression = cpp_expression - super().__init__(source_dtype, size, is_output=False) - - def __repr__(self) -> str: - return ( - "TransformedArray(" - f"source_dtype={self.value_dtype}, " - f"target_dtype={self.target_dtype}, " - f"size={self.size}, cpp_expression={self.cpp_expression!r})" - ) - - def mangled_name(self): - expression_digest = hashlib.sha1( - self.cpp_expression.encode("utf-8") - ).hexdigest()[:12] - return ( - f"TA{self.size}_{self.value_dtype}_to_{self.target_dtype}_" - f"{expression_digest}" - ) - - -class SubstitutionFailure(Exception): - """Signal that a missing template value excludes one overload.""" - - def __init__(self, message): - self.message = message - super().__init__(self.message) - - -class Dependency: - """Look up a required template argument during specialization.""" - - def __init__(self, dep): - self.dep = dep - - def resolve(self, template_arguments): - if self.dep not in template_arguments: - raise SubstitutionFailure( - f"Template argument {self.dep} not provided" - ) - if template_arguments[self.dep] is None: - raise SubstitutionFailure(f"Template argument {self.dep} is None") - return template_arguments[self.dep] - - -class Constant: - """Keep one value unchanged while dependent parameters are specialized.""" - - def __init__(self, val): - self.val = val - - def resolve(self, _): - return self.val - - -class CxxFunction(Parameter): - """Embed a C++ expression in the provider call without a runtime argument. - - The expression can name a functor or encode a compile-time scalar value. - Its text contributes to the generated symbol so distinct expressions can - select distinct wrappers. - """ - - def __init__(self, cpp, func_dtype): - super().__init__() - self.cpp = cpp - self.func_dtype = func_dtype - - def __repr__(self) -> str: - return f"CxxFunction(cpp={self.cpp})" - - def mangled_name(self): - return f"cuda_coop_numba_mlir_F{internal_mangle_cpp(self.cpp)}" - - def dtype(self): - return self.func_dtype - - def is_provided_by_user(self): - return False - - -class DependentArray(Parameter): - """Resolve both the dtype and length of a fixed-size array parameter.""" - - def __init__(self, value_dtype, size, is_output=False): - self.value_dtype = value_dtype - self.size = size - super().__init__(is_output) - - def __repr__(self) -> str: - return f"DependentArray(dep={self.value_dtype}, out={self.is_output})" - - def specialize(self, template_arguments): - return Array( - self.value_dtype.resolve(template_arguments), - self.size.resolve(template_arguments), - self.is_output, - ) - - -class TemplateParameter: - """Name a required argument in the CUB class template parameter list.""" - - def __init__(self, name): - self.name = name - - def __repr__(self) -> str: - return f"{self.name}" - - -def internal_mangle_cpp(cpp_name: str) -> str: - """Turn a C++ spelling into a readable fragment of a generated name. - - Replace every character other than an ASCII letter or digit with an - underscore. Callers incorporate the result into private wrapper, type, - and parameter names. Distinct spellings can produce the same fragment, - and a leading digit remains a digit: this helper neither guarantees a - valid standalone C identifier nor performs C++ ABI name mangling. - ``_bind_private_symbol_namespace`` separately binds the provider identity. - - Parameters - ---------- - cpp_name : str - C++ spelling or static parameter text used in a generated name. - - Returns - ------- - str - Text of the same length containing only ASCII letters, digits, and - underscores. - """ - - return re.sub(r"[^a-zA-Z0-9]", "_", cpp_name) - - -def mangle_symbol(name, template_parameters): - return "_".join( - [name] - + [ - template_parameter.mangled_name() - for template_parameter in template_parameters - ] - ) - - -def war_introspection(fn, n): - """Give a variadic typing function an inspectable fixed-arity signature. - - Numba overload registration inspects the Python argument list. Generate a - small forwarding function with ``n`` named arguments and capture ``fn`` - in its globals. The returned function is called later during typing. - """ - arglist = ", ".join(f"param{i}" for i in range(n)) - mod_str = dedent(f""" - def impl({arglist}): - return fn({arglist}) - """) - mod_code = compile(mod_str, "", "exec") - func_code = mod_code.co_consts[0] - return PyFunctionType(func_code, {"fn": fn}) - - -def war_introspection_call_with_transforms( - fn: ExternFunction | Callable[..., object], - transforms: Sequence[str], - returns_value: bool, -) -> PyFunctionType: - """Generate an inspectable Python call adapter for one provider ABI. - - Numba overload implementations need an explicit argument list. Build a - function with one named argument per transform instead of forwarding - ``*args``. Pointer-backed payloads and scratch are converted with - ``types.ptr``; scalars pass through unchanged. This bridges the Python - argument representations selected by typing to ``ExternFunction`` inputs. - The generated function is returned for compilation, not executed here. - - Parameters - ---------- - fn : ExternFunction or callable - External provider function captured in the generated function's globals. - transforms : sequence of {"ptr", "value"} - Conversion for each positional argument, in ABI order. - returns_value : bool - Whether to return the external result. Otherwise call ``fn`` for its - side effects and return ``None`` explicitly. - - Returns - ------- - types.FunctionType - Function with the requested fixed arity and argument conversions. - - Raises - ------ - ValueError - A transform is neither ``"ptr"`` nor ``"value"``. - """ - - n = len(transforms) - arglist = ", ".join(f"param{i}" for i in range(n)) - mod_lines = [f"def impl({arglist}):"] - for i, transform in enumerate(transforms): - src_name = f"param{i}" - dst_name = f"arg{i}" - if transform == "ptr": - mod_lines.append(f" {dst_name} = types.ptr({src_name})") - elif transform == "value": - mod_lines.append(f" {dst_name} = {src_name}") - else: - raise ValueError(f"Unexpected transform kind: {transform}") - - call_args = ", ".join(f"arg{i}" for i in range(n)) - if returns_value: - mod_lines.append(f" return fn({call_args})") - else: - mod_lines.append(f" fn({call_args})") - mod_lines.append(" return") - - mod_str = "\n".join(mod_lines) + "\n" - mod_code = compile(mod_str, "", "exec") - func_code = mod_code.co_consts[0] - return PyFunctionType(func_code, {"fn": fn, "types": types}) - - -class Algorithm: - """Generate and compile wrappers for one concrete CUB specialization. - - Construction resolves ``template_arguments`` and keeps only applicable - parameter overloads. Storage and synchronization settings describe the - wrapper contract. Compilation later binds a device target, qualifies - private symbols, and records the compiled scratch size and alignment. - """ - - def __init__( - self, - struct_name, - method_name, - c_name, - includes, - template_parameters, - parameters, - *, - template_arguments: Mapping[str, object], - storage_abi, - execution_scope, - synchronization_scope, - type_definitions=None, - fake_return=False, - output_by_reference=False, - compile_context=None, - ): - """Resolve template arguments and initialize provider compile state. - - ``parameters`` lists the candidate method signatures. A missing - dependent value excludes that signature. Every declared class-template - parameter requires a value. ``storage_abi`` determines whether scratch - is the first runtime argument. ``synchronization_scope`` is ``NONE`` - or equal to ``execution_scope``; allocating wrappers emit the matching - barrier after the call. Construction does not run NVRTC. - """ - self.struct_name = struct_name - self.method_name = method_name - self.c_name = ( - c_name - if c_name.startswith("cuda_coop_numba_mlir_") - else f"cuda_coop_numba_mlir_{c_name}" - ) - self.includes = includes - self.template_parameters = template_parameters - self.parameters = parameters - self.storage_abi = StorageABI(storage_abi) - self.execution_scope = SynchronizationScope(execution_scope) - self.synchronization_scope = SynchronizationScope(synchronization_scope) - if self.synchronization_scope not in { - SynchronizationScope.NONE, - self.execution_scope, - }: - raise ValueError( - "synchronization_scope must be NONE or match execution_scope" - ) - self.type_definitions = type_definitions - self.fake_return = fake_return - self.output_by_reference = output_by_reference - self._compile_context = compile_context - self._temp_storage_bytes = None - self._temp_storage_alignment = None - self.logical_warp_threads: int | None = None - self.block_threads: _BlockThreads | None = None - self._private_symbol_digest = None - self._private_symbol_key = None - self._provider_compile_identity: nvrtc.CompilerIdentity | None = None - self._specialize(template_arguments) - - def __repr__(self) -> str: - return ( - f"{self.struct_name}::" - f"{self.method_name}{self.template_parameters}: {self.parameters}" - ) - - def _symbol_base_name(self): - namespace = ( - "" - if self._private_symbol_digest is None - else f"_S{self._private_symbol_digest}" - ) - return f"{self.c_name}{namespace}_{self.method_name}" - - def _current_provider_compile_identity(self) -> nvrtc.CompilerIdentity: - """Describe compilation for the current target and this provider. - - ``_bind_provider_compile_identity`` calls this when its caller has - not supplied an identity. Query the target even when this object has - cached artifacts so switching targets cannot silently reuse them. - Return the NVRTC identity for relocatable LTO output, including the - resolved header/toolkit context and the target from - ``_current_compile_target``. - """ - - return nvrtc.compiler_identity( - context=self._resolved_compile_context(), - rdc=True, - code="lto", - **_current_compile_target().nvrtc_kwargs(), - ) - - def _bind_provider_compile_identity( - self, compile_identity: nvrtc.CompilerIdentity | None = None - ) -> nvrtc.CompilerIdentity: - """Bind compilation inputs before naming or reusing provider code. - - Symbol qualification, cache-key construction, and ``get_lto_ir`` call - this so a specialization cannot reuse code for a different target or - compiler setup. ``compile_identity`` is the record returned by - ``nvrtc.compiler_identity``; ``None`` queries the current target and - this provider's resolved context. - - Store and return the first identity. Later calls return the same - identity or raise ``RuntimeError`` if the supplied or observed inputs - differ. This check does not compile code or validate thread counts. - - Binding belongs to this ``Algorithm`` instance. Separate kernel - compilations create fresh specializations; reusing this instance on - another device is valid only with the same compute capability and - compiler inputs. - """ - observed = ( - self._current_provider_compile_identity() - if compile_identity is None - else compile_identity - ) - existing = self._provider_compile_identity - # For example, an Algorithm bound while targeting compute_90 cannot - # later supply its artifacts to a compute_80 compilation. That compile - # needs a fresh Algorithm. A different device with the same target and - # compiler inputs is compatible; device ordinal is not part of identity. - if existing is not None and existing != observed: - raise RuntimeError( - "Provider artifacts were already qualified for a different " - "compute capability, output kind, or compiler-option set." - ) - self._provider_compile_identity = observed - return observed - - def _bind_private_symbol_namespace( - self, - *, - logical_warp_threads: int | None = None, - block_threads: _BlockThreads | None = None, - compile_identity: nvrtc.CompilerIdentity | None = None, - ) -> nvrtc.CompilerIdentity: - """Bind this provider to a deterministic private symbol namespace. - - Source emission and overload registration must agree on exported names, - including when several specializations share one translation unit. Hash - the provider's coalescing key, which includes its ABI, thread - configuration, compiler context, and compilation target. Equivalent - providers therefore share names; providers with different interfaces or - targets remain distinct. - - The first call stores the compile identity, key, and symbol suffix on - this object. Later calls validate the current inputs and reuse that - suffix without hashing again. Collection, bundling, and source emission - each call this because provider descriptors are mutable: once names - are bound, changing the provider requires a fresh specialization. - When no identity is supplied, query the current CUDA device even if a - previous call has already bound the provider. - - Parameters - ---------- - logical_warp_threads : int, optional - Logical warp width override used in the coalescing key. ``None`` - uses ``self.logical_warp_threads``. - block_threads : int, tuple of int, or list of int, optional - Enclosing block configuration override. ``None`` uses - ``self.block_threads``. - compile_identity : nvrtc.CompilerIdentity, optional - Target and options returned by ``nvrtc.compiler_identity``. ``None`` - resolves them from this provider's compiler context and current - device. - - Returns - ------- - nvrtc.CompilerIdentity - Bound compilation identity for subsequent source and artifact - creation. - - Raises - ------ - RuntimeError - This object was already bound to a different compilation identity or - coalescing key. - """ - - compile_identity = self._bind_provider_compile_identity( - compile_identity - ) - key = algo_coalesce_key( - self, - logical_warp_threads=logical_warp_threads, - block_threads=block_threads, - ) - if self._private_symbol_key is not None: - # The target/options check above has already passed. Here we also - # protect the meaning of names that callers may already reference: - # changing this provider's warp width from 16 to 32, for example, - # needs a fresh specialization with its own symbol namespace. - if self._private_symbol_key != key: - raise RuntimeError( - "Provider symbols were already qualified for a different " - "provider definition or thread configuration. " - "Create a new specialization for the changed inputs." - ) - # Reuse the digest, but validate the mutable provider on every call. - return compile_identity - - self._private_symbol_key = key - self._private_symbol_digest = hashlib.sha1( - repr(key).encode("utf-8") - ).hexdigest()[:20] - return compile_identity - - def _resolved_compile_context(self): - """Resolve header and toolkit identity once for this provider.""" - if self._compile_context is None: - self._compile_context = nvrtc.resolve_shared_compile_context() - return self._compile_context - - def mangled_name(self, parameters): - return mangle_symbol(self._symbol_base_name(), parameters) - - def _specialize(self, template_arguments: Mapping[str, object]) -> None: - """Resolve template dependencies during provider construction. - - A substitution failure removes its entire overload while other - overloads remain available. Copy resolved descriptors before retaining - parameter names so shared descriptors remain unchanged. Bind the C++ - struct, symbol prefix, and logical warp width before code generation. - """ - - # Every template parameter requires an argument. - template_list = [] - for template_parameter in self.template_parameters: - if template_parameter.name not in template_arguments: - raise ValueError( - f"Template argument {template_parameter.name} not provided" - ) - template_argument = template_arguments[template_parameter.name] - if isinstance(template_argument, int): - template_list.append(str(template_argument)) - elif isinstance(template_argument, str): - template_list.append(template_argument) - else: - template_list.append(numba_type_to_cpp(template_argument)) - template_list = ", ".join(template_list) - - # '::cuda::std::int32_t, 32' -> __cuda__std__int32_t__32 - mangle = internal_mangle_cpp(template_list) - - specialized_parameters = [] - for method in self.parameters: - specialized_signature = [] - try: - for parameter in method: - specialized_parameter = copy( - parameter.specialize(template_arguments) - ) - specialized_parameter.parameter_name = getattr( - parameter, "parameter_name", None - ) - specialized_signature.append(specialized_parameter) - specialized_parameters.append(specialized_signature) - except SubstitutionFailure: - pass # Substitution failure is not an error - - if self.template_parameters: - self.struct_name = f"{self.struct_name}<{template_list}>" - self.c_name += mangle - self.template_parameters = [] - self.parameters = specialized_parameters - self.logical_warp_threads = cast( - int | None, - template_arguments.get( - "LOGICAL_WARP_THREADS", - template_arguments.get("VIRTUAL_WARP_THREADS"), - ), - ) - - @property - def temp_storage_bytes(self): - if self._temp_storage_bytes is None: - raise RuntimeError( - "Temporary storage bytes not computed yet. " - "Call get_lto_ir() first." - ) - return self._temp_storage_bytes - - @property - def temp_storage_alignment(self): - if self._temp_storage_alignment is None: - raise RuntimeError( - "Temporary storage alignment not computed yet. " - "Call get_lto_ir() first." - ) - return self._temp_storage_alignment - - @staticmethod - def _ignore_codegen_param(param): - return isinstance(param, CxxFunction) or ( - isinstance(param, PointerOffset) and param.static_value is not None - ) - - def _emit_abi_wrapper( - self, - w: Callable[[str], object], - method: Sequence[Parameter], - exported_name: str, - internal_name: str, - temp_storage_type_name: str | None = "temp_storage_t", - temp_storage_param_pid: int | None = 0, - ) -> None: - """Write the external-call ABI shim for one typed C++ wrapper. - - ``ExternFunction`` calls a symbol ending in ``__abi`` with a leading - return slot and an integer status result. Convert untyped pointer - arguments back to their C++ pointee types, materialize scalar references - as local values, and route the single output through the return slot. - Emit status zero after the call; the shim does not allocate scratch or - add synchronization. - - C++ functors and static pointer offsets are already embedded in the - internal wrapper, so neither appears in this runtime ABI. Output - parameters are likewise removed from the input list and represented by - ``__ret``. - - Parameters - ---------- - w : callable - Text writer receiving generated C++ fragments. - method : sequence of Parameter - Specialized parameters in internal-wrapper order, with at most one - output. An allocating wrapper supplies the method without scratch. - exported_name : str - Symbol prefix; the emitted entry point appends ``__abi``. - internal_name : str - Typed C++ wrapper called by the shim. - temp_storage_type_name : str or None, optional - C++ scratch type used when casting the designated byte pointer. - temp_storage_param_pid : int or None, optional - Index of an explicit scratch pointer in ``method``. ``None`` - disables that special cast for allocating or storage-free wrappers. - - Returns - ------- - None - Generated source is appended through ``w``. - - Raises - ------ - ValueError - More than one runtime output parameter is present. - """ - - output_param = None - user_params = [] - for pid, param in enumerate(method): - if self._ignore_codegen_param(param): - continue - if param.is_output: - if output_param is not None: - raise ValueError("Multiple output parameters not supported") - output_param = (pid, param) - else: - user_params.append((pid, param)) - - abi_param_decls = ["void *__ret"] - body_lines = [] - call_args_by_pid = {} - - parameter_names = _cpp_parameter_names( - [param for _, param in user_params], - reserved={"__ret", "out", internal_name, temp_storage_type_name}, - ) - for (pid, param), name in zip(user_params, parameter_names): - if isinstance(param, (Pointer, Array, PointerReference)): - abi_param_decls.append(f"void *{name}") - cast_name = f"cast_{name}" - if ( - pid == temp_storage_param_pid - and isinstance(param, Pointer) - and param.value_dtype == numba_types.uint8 - ): - # Non-alloc wrappers receive raw temporary storage - # explicitly. - body_lines.append( - " " - f"{temp_storage_type_name} *{cast_name} = " - f"reinterpret_cast<{temp_storage_type_name} *>({name});" - ) - else: - pointee_type = numba_type_to_cpp(param.value_dtype) - body_lines.append( - f" {pointee_type} *{cast_name} = reinterpret_cast<" - f"{pointee_type} *>({name});" - ) - call_args_by_pid[pid] = cast_name - elif isinstance(param, Reference): - value_type = numba_type_to_cpp(param.value_dtype) - abi_param_decls.append(f"{value_type} {name}") - local_name = f"ref_{name}" - body_lines.append(f" {value_type} {local_name} = {name};") - call_args_by_pid[pid] = local_name - else: - value_type = numba_type_to_cpp(param.dtype()) - abi_param_decls.append(f"{value_type} {name}") - call_args_by_pid[pid] = name - - output_var = None - output_cpp_type = None - if output_param is not None: - output_cpp_type = numba_type_to_cpp(output_param[1].dtype()) - output_var = "out" - body_lines.append(f" {output_cpp_type} {output_var};") - call_args_by_pid[output_param[0]] = output_var - - call_args = [ - call_args_by_pid[pid] - for pid, param in enumerate(method) - if not self._ignore_codegen_param(param) - ] - args_csv = ", ".join(call_args) - if args_csv: - body_lines.append(f" {internal_name}({args_csv});") - else: - body_lines.append(f" {internal_name}();") - - if output_var is not None: - body_lines.append( - f" *reinterpret_cast<{output_cpp_type}" - f" *>(__ret) = {output_var};" - ) - body_lines.append(" return 0;") - - abi_params_csv = ", ".join(abi_param_decls) - w( - f'extern "C" __device__ int {exported_name}' - f"__abi({abi_params_csv}) {{\n" - ) - for line in body_lines: - w(f"{line}\n") - w("}\n\n") - - def _collect_support_ltoirs_and_udf_declarations(self): - lto_irs = [] - udf_declarations = OrderedDict() - - if self.type_definitions: - for type_definition in self.type_definitions: - lto_irs.extend(type_definition.lto_irs) - - return _dedupe_ltoirs(lto_irs), udf_declarations - - def _source_code( - self, - logical_warp_threads: int | None = None, - block_threads: _BlockThreads | None = None, - *, - compile_identity: nvrtc.CompilerIdentity | None = None, - ) -> tuple[str, list[bytes], tuple[str, ...], OrderedDict[str, str]]: - """Generate C++ wrappers and compile-time storage metadata. - - Emit a typed wrapper plus an ``__abi`` shim for each specialized method. - The typed wrapper adapts array pointers to CUB array references, applies - input transforms, folds pointer offsets into earlier pointer arguments, - and traps on out-of-range ``BoundedInteger`` values before narrowing - them. C++ functors and static offsets are embedded in source instead of - passed as runtime arguments. - - For ``LEADING_POINTER`` storage, emit both explicit-scratch and - ``_alloc`` entry points. Only ``_alloc`` allocates scratch and emits the - declared post-call synchronization. It allocates one shared object per - block or logical warp, or one local object for scope ``NONE``. - Warp scratch uses the linear thread rank in the exact enclosing block; - its width must divide the block size. Explicit-scratch wrappers leave - allocation and reuse synchronization to their caller. Storage-free - providers emit neither scratch metadata nor allocating variants. - - This method qualifies private symbols without running NVRTC. It - returns scratch type names so artifact creation can query their C++ - size and alignment during provider compilation. Each name describes - one ``TempStorage`` instance. The scratch planner uses group topology - to set the instance count for explicit-scratch calls. - - Parameters - ---------- - logical_warp_threads : int, optional - Logical warp width for allocating warp wrappers. ``None`` uses - ``self.logical_warp_threads``. - block_threads : int, tuple of int, or list of int, optional - Exact enclosing block size or dimensions. ``None`` uses - ``self.block_threads``. - compile_identity : nvrtc.CompilerIdentity, optional - Bound target/options identity. ``None`` resolves the current - device's identity through ``_bind_private_symbol_namespace``. - - Returns - ------- - src : str - Complete CUDA C++ translation unit for this specialization. - support_lto_irs : list of bytes - Deduplicated supporting link images from type definitions. - temp_storage_types : tuple of str - C++ scratch type names to query, or an empty tuple without scratch. - udf_declarations : collections.OrderedDict - Declaration table used when constructing a shared source preamble; - currently empty. - - Raises - ------ - ValueError - A pointer offset has no earlier pointer target, multiple outputs are - requested, or allocating warp storage has an invalid width/block - size. - RuntimeError - The provider was already qualified for incompatible compiler inputs. - NotImplementedError - The requested allocating execution or synchronization scope has no - source emitter. - """ - - self._bind_private_symbol_namespace( - logical_warp_threads=logical_warp_threads, - block_threads=block_threads, - compile_identity=compile_identity, - ) - support_lto_irs, udf_declarations = ( - self._collect_support_ltoirs_and_udf_declarations() - ) - - algorithm_name = self.struct_name - includes = self.includes or [] - type_definitions = self.type_definitions or [] - method_name = self.method_name - - alias_suffix = internal_mangle_cpp(self._symbol_base_name()) - algorithm_type_name = f"algorithm_t_{alias_suffix}" - temp_storage_type_name = None - temp_storage_types = () - if self.storage_abi is StorageABI.LEADING_POINTER: - temp_storage_type_name = f"temp_storage_t_{alias_suffix}" - temp_storage_types = (temp_storage_type_name,) - - buf = StringIO() - w = buf.write - - w("#include \n") - for include in includes: - w(f"#include <{include}>\n") - for type_definition in type_definitions: - w(f"{type_definition.code}\n") - w("\n") - for decl in udf_declarations.values(): - w(f"{decl}\n") - w("\n") - - w(f"using {algorithm_type_name} = cub::{algorithm_name};\n") - if temp_storage_type_name is not None: - w( - f"using {temp_storage_type_name} = typename " - f"{algorithm_type_name}::TempStorage;\n" - ) - - src = buf.getvalue() - - for method in self.parameters: - param_decls = [] - func_decls = [] - param_args = [] - out_param = None - - param_arg_positions_by_pid = {} - provider_parameters = ( - method[1:] - if self.storage_abi is StorageABI.LEADING_POINTER - else method - ) - parameter_names = _cpp_parameter_names( - provider_parameters, - reserved={ - "temp_storage", - "temp_storages", - algorithm_type_name, - temp_storage_type_name, - }, - ) - for pid, (param, name) in enumerate( - zip(provider_parameters, parameter_names) - ): - if isinstance(param, CxxFunction): - param_args.append(param.cpp) - else: - if isinstance(param, PointerOffset): - offset_expression = name - if param.static_value is None: - param_decls.append(param.cpp_decl(name)) - else: - offset_expression = str(param.static_value) - pointer_arg_pos = param_arg_positions_by_pid.get( - param.pointer_arg_index - ) - if pointer_arg_pos is None: - raise ValueError( - "PointerOffset must reference an earlier " - "pointer parameter." - ) - param_args[pointer_arg_pos] = ( - f"({param_args[pointer_arg_pos]} + " - f"{offset_expression})" - ) - continue - if isinstance(param, BoundedInteger): - checked_name = f"checked_{name}" - param_decls.append(param.cpp_decl(name)) - func_decls.append( - f"if ({name} < {param.minimum} || " - f"{name} > {param.maximum}) {{" - ) - func_decls.append(' asm volatile("trap;" : : :);') - func_decls.append("}") - provider_cpp = numba_type_to_cpp(param.provider_dtype) - func_decls.append( - f"{provider_cpp} {checked_name} = " - f"static_cast<{provider_cpp}>({name});" - ) - param_arg = checked_name - elif isinstance(param, TransformedArray): - source_type = numba_type_to_cpp(param.value_dtype) - target_type = numba_type_to_cpp(param.target_dtype) - transformed_name = f"transformed_{name}" - param_decls.append(f"{source_type} *{name}") - func_decls.append( - f"{target_type} {transformed_name}[{param.size}];" - ) - for index in range(param.size): - expression = param.cpp_expression.format( - value=f"{name}[{index}]" - ) - func_decls.append( - f"{transformed_name}[{index}] = {expression};" - ) - param_arg = transformed_name - elif isinstance(param, Array): - value_type = numba_type_to_cpp(param.value_dtype) - param_decls.append(f"{value_type} *{name}") - param_arg = ( - f"*reinterpret_cast<{value_type}" - f" (*)[{param.size}]>({name})" - ) - else: - param_decls.append(param.cpp_decl(name)) - if getattr(param, "deref_on_call", False): - param_arg = f"*{name}" - else: - param_arg = name - if not self.fake_return and param.is_output: - if out_param is not None: - raise ValueError( - "Multiple output parameters not supported" - ) - out_param = name - if self.output_by_reference: - param_arg_positions_by_pid[pid] = len(param_args) - param_args.append(param_arg) - else: - param_arg_positions_by_pid[pid] = len(param_args) - param_args.append(param_arg) - - provide_alloc_version = ( - self.storage_abi is StorageABI.LEADING_POINTER - ) - if provide_alloc_version: - assert temp_storage_type_name is not None - logical_width = None - if self.execution_scope is SynchronizationScope.BLOCK: - storage = ( - f"__shared__ {temp_storage_type_name} temp_storage;" - ) - elif self.execution_scope is SynchronizationScope.WARP: - logical_width = _validate_logical_warp_threads( - logical_warp_threads - if logical_warp_threads is not None - else self.logical_warp_threads - ) - resolved_block_threads = _normalize_block_threads( - block_threads - if block_threads is not None - else self.block_threads - ) - if resolved_block_threads % logical_width != 0: - raise ValueError( - "warp-scoped provider width must divide the exact " - f"block size; got width={logical_width} and " - f"block_threads={resolved_block_threads}" - ) - instances = resolved_block_threads // logical_width - storage = ( - "unsigned __coop_thread_rank " - "= threadIdx.x + blockDim.x * " - "(threadIdx.y + blockDim.y * threadIdx.z);\n" - " " - f"__shared__ {temp_storage_type_name} temp_storages" - f"[{instances}];\n" - f" {temp_storage_type_name} " - f"&temp_storage = temp_storages" - f"[__coop_thread_rank / {logical_width}];" - ) - elif self.execution_scope is SynchronizationScope.NONE: - storage = f"{temp_storage_type_name} temp_storage;" - else: - raise NotImplementedError( - "cuda.coop.numba_mlir provider execution scope " - f"{self.execution_scope.value!r} has no emitter" - ) - if self.synchronization_scope is SynchronizationScope.WARP: - assert logical_width is not None - sync = ( - "__syncwarp();" - if logical_width == 32 - else ( - "__syncwarp((((1u << " - f"{logical_width}" - ") - 1u) << (((__coop_thread_rank & 31) / " - f"{logical_width}" - f") * {logical_width})));" - ) - ) - else: - sync = { - SynchronizationScope.NONE: "", - SynchronizationScope.BLOCK: "__syncthreads();", - }.get(self.synchronization_scope) - if sync is None: - raise NotImplementedError( - "cuda.coop.numba_mlir provider synchronization scope " - f"{self.synchronization_scope.value!r} has no emitter" - ) - else: - storage = "" - sync = "" - - buf = StringIO() - w = buf.write - - mangled_name = self.mangled_name(method) - param_decls_csv = ", ".join(param_decls) - param_args_csv = ", ".join(param_args) - - if provide_alloc_version: - w(f'extern "C" __device__ void {mangled_name}_alloc(') - w(f"{param_decls_csv}) {{\n") - w(f" {storage}\n") - for decl in func_decls: - w(f" {decl}\n") - if out_param and not self.output_by_reference: - w(f" {out_param} = ") - else: - w(" ") - w(f"{algorithm_type_name}(temp_storage).") - w(f"{method_name}({param_args_csv});\n") - w(f" {sync}\n") - w("}\n") - self._emit_abi_wrapper( - w, - method[1:], - exported_name=f"{mangled_name}_alloc", - internal_name=f"{mangled_name}_alloc", - temp_storage_type_name=temp_storage_type_name, - temp_storage_param_pid=None, - ) - - w(f'extern "C" __device__ void {mangled_name}(') - if self.storage_abi is StorageABI.LEADING_POINTER: - assert temp_storage_type_name is not None - w(f"{temp_storage_type_name} *temp_storage") - if param_decls_csv: - w(", ") - w(f"{param_decls_csv}) {{\n") - for decl in func_decls: - w(f" {decl}\n") - if out_param and not self.output_by_reference: - w(f" {out_param} = ") - else: - w(" ") - if self.storage_abi is StorageABI.LEADING_POINTER: - w(f"{algorithm_type_name}(*temp_storage).") - else: - w(f"{algorithm_type_name}().") - w(f"{method_name}({param_args_csv});\n") - w("}\n\n") - self._emit_abi_wrapper( - w, - method, - exported_name=mangled_name, - internal_name=mangled_name, - temp_storage_type_name=( - "temp_storage_t" - if temp_storage_type_name is None - else temp_storage_type_name - ), - temp_storage_param_pid=( - 0 - if self.storage_abi is StorageABI.LEADING_POINTER - else None - ), - ) - - src += buf.getvalue() - - return ( - src, - support_lto_irs, - temp_storage_types, - udf_declarations, - ) - - def _make_lto_ir_cache_key( - self, - logical_warp_threads=None, - block_threads=None, - *, - compile_identity=None, - ): - """Validate warp topology and identify reusable link images. - - The key includes storage policy, execution and synchronization scopes, - thread configuration, and compiler identity. Reusing an object with a - different bound target fails before its cached artifact is returned. - """ - resolved_logical_warp_threads = ( - logical_warp_threads - if logical_warp_threads is not None - else self.logical_warp_threads - ) - resolved_block_threads = ( - block_threads if block_threads is not None else self.block_threads - ) - if self.execution_scope is SynchronizationScope.WARP: - resolved_logical_warp_threads = _validate_logical_warp_threads( - resolved_logical_warp_threads - ) - resolved_block_threads = _normalize_block_threads( - resolved_block_threads - ) - if resolved_block_threads % resolved_logical_warp_threads != 0: - raise ValueError( - "warp-scoped provider width must divide the exact block " - f"size; got width={resolved_logical_warp_threads} and " - f"block_threads={resolved_block_threads}" - ) - compile_identity = self._bind_provider_compile_identity( - compile_identity - ) - return ( - self.storage_abi.value, - self.execution_scope.value, - self.synchronization_scope.value, - resolved_logical_warp_threads, - resolved_block_threads, - compile_identity, - ) - - def get_lto_ir( - self, - logical_warp_threads: int | None = None, - block_threads: _BlockThreads | None = None, - *, - compile_identity: nvrtc.CompilerIdentity | None = None, - ) -> list[bytes]: - """Compile or reuse this specialization's matching link images. - - On first use, compile a provider translation unit to LTO IR and query - its C++ scratch size and alignment through NVRTC's lowered template - names. Cache the image and layouts together, then retain the layouts, - link images, and filename suffixes on the algorithm. Storage-free - providers skip layout queries and receive size zero and alignment one. - - Reuse is confined to the same bound target, options, and thread - configuration. Without an explicit identity, re-query the current device - before returning cached artifacts. A specialization cannot silently - reuse artifacts compiled for a previous device. For a prebundled - algorithm, ``lto_irs`` contains only its extra link images; the shared - provider image is retained separately in ``_precompiled_ltoir_files``. - - Parameters - ---------- - logical_warp_threads : int, optional - Logical warp width override; ``None`` uses - ``self.logical_warp_threads``. - block_threads : int, tuple of int, or list of int, optional - Exact enclosing block configuration; ``None`` uses - ``self.block_threads``. - compile_identity : nvrtc.CompilerIdentity, optional - Previously resolved target/options identity, or ``None`` to query - it. - - Returns - ------- - list of bytes - Cached supporting link images followed by the generated provider LTO - image, or only extra images after bundling. The list is returned - directly. - - Raises - ------ - RuntimeError - Cached or qualified artifacts use different compiler inputs, or the - compiler reports an error. - ValueError - Warp topology is invalid, or NVRTC returns an unrecognized layout - name or invalid scratch size/alignment. - """ - - # With no explicit identity, re-query the current device even when an - # artifact is already cached. Reusing one Algorithm across devices must - # fail closed instead of returning LTO compiled for the earlier target. - compile_identity = self._bind_provider_compile_identity( - compile_identity - ) - cache_key = self._make_lto_ir_cache_key( - logical_warp_threads, - block_threads, - compile_identity=compile_identity, - ) - existing = self.__dict__.get("lto_irs") - if existing is not None: - existing_key = self.__dict__.get("_lto_ir_cache_key") - if existing_key != cache_key: - raise RuntimeError( - "LTO IR was already compiled for a different thread " - "configuration or compilation target." - ) - return existing - - src, support_lto_irs, temp_storage_types, _ = self._source_code( - logical_warp_threads=logical_warp_threads, - block_threads=block_threads, - compile_identity=compile_identity, - ) - - _, (ltoir, layouts) = nvrtc.compile_with_layouts( - cpp=src, - layout_types=temp_storage_types, - rdc=True, - code="lto", - context=self._resolved_compile_context(), - **compile_identity.target.nvrtc_kwargs(), - ) - ltoir_blob = bytes(cast(bytes, ltoir)) - - lto_irs = list(support_lto_irs) - lto_irs.append(ltoir_blob) - - if temp_storage_types: - self._temp_storage_bytes, self._temp_storage_alignment = layouts[0] - else: - self._temp_storage_bytes = 0 - self._temp_storage_alignment = 1 - self.__dict__["_lto_ir_cache_key"] = cache_key - self.__dict__["_link_input_suffixes"] = [".ltoir"] * len(lto_irs) - self.__dict__["lto_irs"] = lto_irs - return lto_irs - - def codegen(self, func_to_overload): - """Build callable overloads for each specialized C++ method. - - Providers with a leading scratch pointer also expose an allocating - variant whose Python signature omits that pointer. Return the templates - for this callable without adding them to a global typing registry. - """ - if len(self.template_parameters): - raise ValueError("Cannot generate codegen for a template") - - overloads = [] - for method in self.parameters: - overloads.append( - self.codegen_method( - func_to_overload, - method, - self.mangled_name(method), - ) - ) - if self.storage_abi is StorageABI.LEADING_POINTER: - overloads.append( - self.codegen_method( - func_to_overload, - method[1:], - self.mangled_name(method) + "_alloc", - ) - ) - return tuple(overloads) - - def codegen_method( - self, - func_to_overload: Callable[..., object], - method: Sequence[Parameter], - mangled_name: str, - ) -> type[_OverloadFunctionTemplate]: - """Build a Numba overload template for one specialized provider method. - - Pair the method's Python-facing argument checks with the generated - ``__abi`` symbol. Pointer-backed arguments use an untyped pointer ABI - and ``types.ptr`` conversion; scalar arguments use their descriptor - dtype. An output becomes the external call's return value. Embedded - functors and static pointer offsets consume no runtime arguments. - - The resulting typing implementation returns ``None`` when the arity or - an input descriptor rejects the actual compiler types, allowing overload - selection to continue. On a match, it returns a fixed-arity - implementation and attaches the invocable's link paths. Registration - remains local to the returned template instead of modifying a global - typing registry. - - Parameters - ---------- - func_to_overload : callable - Invocable used as the overload key. Its ``files`` attribute, when - present, supplies paths needed to link the external symbol. - method : sequence of Parameter - Concrete method signature, including at most one output parameter. - mangled_name : str - Generated wrapper name without the ``__abi`` suffix; may include - ``_alloc`` when ``method`` omits an explicit scratch argument. - - Returns - ------- - type - Strict, always-inline overload template for this method. - - Raises - ------ - ValueError - The algorithm still has template parameters or the method has - multiple output parameters. - """ - - if len(self.template_parameters): - raise ValueError("Cannot generate codegen for a template") - - def ignore_param(param): - # C++ functions do not require additional argument handling. - ignore = isinstance(param, CxxFunction) or ( - isinstance(param, PointerOffset) - and param.static_value is not None - ) - return ignore - - num_user_provided_params = sum( - [param.is_provided_by_user() for param in method] - ) - returns_value = any(param.is_output for param in method) - link_files = getattr(func_to_overload, "files", []) - expected_input_parameters = [] - abi_input_types = [] - arg_transforms = [] - ret_type = numba_types.void - - for param in method: - if ignore_param(param) or param.is_output: - if not ignore_param(param) and param.is_output: - if ret_type is not numba_types.void: - raise ValueError( - "Multiple output parameters not supported" - ) - ret_type = param.dtype() - continue - - expected_input_parameters.append(param) - if isinstance(param, (Pointer, Array, PointerReference)): - abi_input_types.append(numba_types.CPointer(numba_types.none)) - arg_transforms.append("ptr") - else: - abi_input_types.append(param.dtype()) - arg_transforms.append("value") - - abi_name = f"{mangled_name}__abi" - extern_fn = ExternFunction( - abi_name, - signature(ret_type, *abi_input_types), - link=link_files, - ) - - def algorithm_impl(*args): - if len(args) != len(expected_input_parameters): - return None - typingctx = mlir_target.typing_context - for actual_type, parameter in zip(args, expected_input_parameters): - if not parameter.accepts_actual_type(actual_type, typingctx): - return None - impl = war_introspection_call_with_transforms( - extern_fn, arg_transforms, returns_value=returns_value - ) - if link_files: - impl.__dict__["__numba_cuda_mlir_link__"] = link_files - return impl - - wrapped_algorithm_impl = war_introspection( - algorithm_impl, num_user_provided_params - ) - if link_files: - wrapped_algorithm_impl.__dict__["__numba_cuda_mlir_link__"] = ( - link_files - ) - return make_overload_template( - func_to_overload, - wrapped_algorithm_impl, - {"no_cpython_wrapper": True, "nopython": True}, - strict=True, - inline="always", - prefer_literal=False, - base=_NumbaCudaMlirOverloadFunctionTemplate, - ) - - -def _dedupe_ltoirs(lto_irs): - """Keep the first occurrence of each supporting link image. - - Provider compilation and bundling collect images from type definitions - and parameter descriptors. ``lto_irs`` is their ordered iterable of byte - images or text; compare content digests and return a list of the original - objects, preserving first-use order. - """ - - seen = set() - deduped = [] - for lto_ir in lto_irs: - key = _lto_ir_digest(lto_ir) - if key in seen: - continue - seen.add(key) - deduped.append(lto_ir) - return deduped - - -class _NamedTempFile(Protocol): - @property - def name(self) -> str: ... - - -class _SharedTempFile: - """Keep a shared bundle file alive until its final owner is released. - - Invocables retain this object with their own link files. Its finalizer - removes the shared path only when no owner retains this object. - """ - - def __init__(self, temp_file: _NamedTempFile): - self._temp_file = temp_file - self._temp_file_finalizer = weakref.finalize( - self, _cleanup_temp_files, (temp_file.name,) - ) - - @property - def name(self): - return self._temp_file.name - - -def _collect_extra_ltoirs(algo): - """Retain a bundled algorithm's supporting images as separate inputs. - - ``prepare_ltoir_bundle`` calls this after compiling the shared provider - source. ``algo`` supplies type definitions and specialized parameter - descriptors whose link images are still needed alongside that bundle. - Return their content-deduplicated images in first-use order, or an empty - list when none are attached. - """ - - extras = [] - if algo.type_definitions: - for type_definition in algo.type_definitions: - extras.extend(getattr(type_definition, "lto_irs", [])) - for method in algo.parameters: - for param in method: - lto_ir = getattr(param, "ltoir", None) - if lto_ir is not None: - extras.append(lto_ir) - return _dedupe_ltoirs(extras) - - -def _param_coalesce_key(param): - """Describe one parameter's generated C++ and runtime ABI shape. - - Parameter display names are omitted. Types, transforms, bounds, and - embedded expressions remain part of the key because they affect code. - """ - if isinstance(param, TransformedArray): - return ( - "TransformedArray", - str(param.value_dtype), - str(param.target_dtype), - param.size, - param.cpp_expression, - ) - if isinstance(param, Array): - return ("Array", str(param.value_dtype), param.size, param.is_output) - if isinstance(param, Pointer): - return ( - "Pointer", - str(param.value_dtype), - param.is_output, - getattr(param, "deref_on_call", False), - ) - if isinstance(param, Reference): - return ("Reference", str(param.value_dtype), param.is_output) - if isinstance(param, PointerOffset): - return ( - "PointerOffset", - str(param.value_type), - param.pointer_arg_index, - param.static_value, - param.is_output, - ) - if isinstance(param, BoundedInteger): - return ( - "BoundedInteger", - str(param.value_type), - str(param.provider_dtype), - param.minimum, - param.maximum, - param.is_output, - ) - if isinstance(param, ExactValue): - return ("ExactValue", str(param.value_type), param.is_output) - if isinstance(param, Value): - return ("Value", str(param.value_type), param.is_output) - if isinstance(param, CxxFunction): - return ("CxxFunction", param.cpp, str(param.func_dtype)) - return (type(param).__name__, repr(param)) - - -def algo_coalesce_key( - algo: Algorithm, - *, - logical_warp_threads: int | None = None, - block_threads: _BlockThreads | None = None, -) -> tuple[object, ...]: - """Describe a provider for source coalescing and symbol qualification. - - Include the C++ operation, ordered ABI descriptors, type-definition source - and link-image digests, topology, storage/synchronization contracts, and - already resolved compiler identity. Human-readable parameter names are - excluded: renaming a wrapper argument does not create a distinct provider. - The bundler uses this key to choose one source representative for equivalent - algorithms; private symbol qualification hashes the same description. - - This function only reads the supplied state. It does not resolve a compiler - context, bind a device target, or validate thread dimensions, so callers - must establish those inputs before comparing compilation-ready providers. - - Parameters - ---------- - algo : Algorithm - Provider whose current source and ABI state is described. - logical_warp_threads : int, optional - Logical warp width override, falling back to - ``algo.logical_warp_threads``. - block_threads : int, tuple of int, or list of int, optional - Block configuration override, falling back to ``algo.block_threads``. - - Returns - ------- - tuple - Equality key for the provider's current compilation-relevant state. - """ - - type_defs = [] - for type_definition in getattr(algo, "type_definitions", None) or []: - code = getattr(type_definition, "code", None) or "" - lto_irs = getattr(type_definition, "lto_irs", []) or [] - lto_ir_hashes = tuple(_lto_ir_digest(lto_ir) for lto_ir in lto_irs) - type_defs.append((code, lto_ir_hashes)) - - params_key = tuple( - tuple(_param_coalesce_key(param) for param in method) - for method in getattr(algo, "parameters", []) - ) - - resolved_logical_warp_threads = logical_warp_threads - if resolved_logical_warp_threads is None: - resolved_logical_warp_threads = getattr( - algo, "logical_warp_threads", None - ) - resolved_block_threads = block_threads - if resolved_block_threads is None: - resolved_block_threads = getattr(algo, "block_threads", None) - - return ( - getattr(algo, "struct_name", None), - getattr(algo, "method_name", None), - getattr(algo, "c_name", None), - tuple(getattr(algo, "includes", None) or []), - tuple(type_defs), - params_key, - resolved_logical_warp_threads, - resolved_block_threads, - getattr(algo, "fake_return", None), - getattr(algo, "output_by_reference", None), - getattr(algo, "storage_abi", None), - getattr(algo, "execution_scope", None), - getattr(algo, "synchronization_scope", None), - getattr(algo, "_compile_context", None), - getattr(algo, "_provider_compile_identity", None), - ) - - -def _strip_source_preamble(src, algo, udf_decls): - """Remove shared includes and declarations before combining wrappers. - - Only the first exact occurrence of each generated fragment is removed. - The bundle emits the shared preamble once for all representatives. - """ - body = src - body = body.replace("#include \n", "", 1) - for include in algo.includes or []: - body = body.replace(f"#include <{include}>\n", "", 1) - for type_definition in algo.type_definitions or []: - code = getattr(type_definition, "code", "") - if code: - body = body.replace(f"{code}\n", "", 1) - for decl in udf_decls.values(): - body = body.replace(f"{decl}\n", "", 1) - return body.lstrip() - - -def prepare_ltoir_bundle( - algorithms: Sequence[Algorithm], - *, - allow_single: bool = False, -) -> bytes | None: - """Compile distinct provider specializations into one shared LTO artifact. - - Deduplicate input objects by identity, then coalesce equivalent providers - with ``algo_coalesce_key``. Emit one body per representative and one shared - preamble for includes and type declarations. All providers must resolve to - the same compiler context and bind to the current device's target. Query - each representative's C++ scratch layout during the NVRTC compile and cache - those layouts together with the LTO image. - - On success, mutate every supplied algorithm with its scratch size/alignment, - cache key, extra link images, and a reference to the same temporary bundle - file. Extra images remain separate link inputs. The shared file is removed - when its final shared owner is released, so one invocable cannot delete a - bundle still used by another. Symbol/context qualification can occur even - when the function decides there are too few providers to compile. - - Parameters - ---------- - algorithms : sequence of Algorithm - Concrete provider specializations to bundle, using each algorithm's - stored logical warp width and enclosing block configuration. - allow_single : bool, optional - Compile even one distinct representative when true. Defaults to false. - - Returns - ------- - bytes or None - Shared provider LTO image. ``None`` means no input algorithms, or fewer - than two representatives with ``allow_single=False``. - - Raises - ------ - RuntimeError - Providers resolve to different compiler contexts, conflict with an - existing qualification, or compilation fails. - ValueError - Provider source/topology is invalid, or NVRTC returns an unrecognized - layout name or invalid scratch size/alignment. - """ - - if not algorithms: - return None - - deduped_by_id = OrderedDict() - for algo in algorithms: - deduped_by_id[id(algo)] = algo - all_algos = list(deduped_by_id.values()) - - compile_contexts = {algo._resolved_compile_context() for algo in all_algos} - if len(compile_contexts) != 1: - raise RuntimeError( - "coalesced providers must use one exact compiler context" - ) - compile_context = next(iter(compile_contexts)) - target = _current_compile_target() - compile_identity = nvrtc.compiler_identity( - context=compile_context, - rdc=True, - code="lto", - **target.nvrtc_kwargs(), - ) - - key_to_rep = OrderedDict() - rep_for_algo_id = {} - for algo in all_algos: - algo._bind_private_symbol_namespace( - compile_identity=compile_identity, - ) - rep = key_to_rep.setdefault(algo._private_symbol_key, algo) - rep_for_algo_id[id(algo)] = rep - - reps = list(key_to_rep.values()) - if len(reps) < 2 and not allow_single: - return None - - rep_src = {} - rep_storage_types = {} - rep_udf_decls = {} - includes = OrderedDict() - type_defs = OrderedDict() - udf_decls = OrderedDict() - - for rep in reps: - src, _support_lto_irs, storage_types, udf = rep._source_code( - compile_identity=compile_identity, - ) - rep_src[id(rep)] = src - rep_storage_types[id(rep)] = storage_types - rep_udf_decls[id(rep)] = udf - - for include in rep.includes or []: - includes.setdefault(include, None) - for type_definition in rep.type_definitions or []: - code = getattr(type_definition, "code", "") - if code: - type_defs.setdefault(code, None) - for decl in udf.values(): - udf_decls.setdefault(decl, None) - - buf = StringIO() - w = buf.write - w("#include \n") - for include in includes: - w(f"#include <{include}>\n") - for code in type_defs: - w(f"{code}\n") - w("\n") - if udf_decls: - for decl in udf_decls: - w(f"{decl}\n") - w("\n") - - bodies = [] - for rep in reps: - body = _strip_source_preamble( - rep_src[id(rep)], - rep, - rep_udf_decls[id(rep)], - ) - bodies.append(body.rstrip() + "\n") - - src = buf.getvalue() + "\n".join(bodies) - - # Query CUB TempStorage types, not kernel argument types. For example, a - # transpose Load -> Store contributes two types, even when both calls share - # one explicit TempStorage; direct/striped providers contribute none. - # Compiler-allocated scratch needs the same layouts as explicit scratch. - layout_types = tuple( - storage_type - for rep in reps - for storage_type in rep_storage_types[id(rep)] - ) - _, (ltoir, layouts) = nvrtc.compile_with_layouts( - cpp=src, - layout_types=layout_types, - rdc=True, - code="lto", - context=compile_context, - **target.nvrtc_kwargs(), - ) - ltoir_blob = bytes(cast(bytes, ltoir)) - layouts_by_type = dict(zip(layout_types, layouts, strict=True)) - - from ._compiler._artifacts import make_binary_tempfile - - bundle_temp_file = _SharedTempFile( - make_binary_tempfile(ltoir_blob, ".ltoir") - ) - - for algo in all_algos: - rep = rep_for_algo_id[id(algo)] - storage_types = rep_storage_types[id(rep)] - extras = _collect_extra_ltoirs(algo) - if storage_types: - algo._temp_storage_bytes, algo._temp_storage_alignment = ( - layouts_by_type[storage_types[0]] - ) - else: - algo._temp_storage_bytes = 0 - algo._temp_storage_alignment = 1 - algo._precompiled_ltoir_files = (bundle_temp_file,) - algo.__dict__["_lto_ir_cache_key"] = algo._make_lto_ir_cache_key( - compile_identity=compile_identity, - ) - algo.__dict__["_link_input_suffixes"] = [".ltoir"] * len(extras) - algo.__dict__["lto_irs"] = extras - - return ltoir_blob - - -def make_invocable_from_specialization( - specialization: Algorithm, - *, - logical_warp_threads: int | None = None, - block_threads: _BlockThreads | None = None, -) -> Invocable | Algorithm: - """Turn a concrete provider into a callable and retain its link artifacts. - - Store explicit topology overrides and qualify the provider's private - symbols. During ``collect_specializations``, collect and return the - algorithm immediately so the caller can bundle providers before compilation. - Otherwise obtain its link images and scratch ABI, then create an - ``Invocable`` exposing the files to compiler overloads. - - Previously bundled files are shared by reference. New files written for this - invocation are owned by the returned invocable and removed by its finalizer. - This separation keeps a shared bundle alive while any invocable or - specialization still retains it. - - Parameters - ---------- - specialization : Algorithm - Concrete provider to qualify and, outside collection, compile or reuse. - logical_warp_threads : int, optional - Logical warp width stored on the algorithm when supplied. - block_threads : int, tuple of int, or list of int, optional - Exact enclosing block configuration stored when supplied. - - Returns - ------- - Invocable or Algorithm - Callable wrapper outside collection; the same ``specialization`` object - inside collection, retaining its thread configuration. - - Raises - ------ - RuntimeError - Qualification/cache inputs conflict, compilation fails, or cached link - images and suffix metadata have inconsistent lengths. - """ - - if logical_warp_threads is not None: - specialization.logical_warp_threads = logical_warp_threads - if block_threads is not None: - specialization.block_threads = block_threads - - compile_identity = specialization._bind_private_symbol_namespace( - logical_warp_threads=logical_warp_threads, - block_threads=block_threads, - ) - - collector = _COOP_SPECIALIZATION_COLLECTOR.get() - if collector is not None: - collector.append(specialization) - return specialization - - from ._compiler._artifacts import make_binary_tempfile - - lto_irs = specialization.get_lto_ir( - logical_warp_threads=logical_warp_threads, - block_threads=block_threads, - compile_identity=compile_identity, - ) - precompiled_temp_files = list( - getattr(specialization, "_precompiled_ltoir_files", ()) - ) - suffixes = specialization.__dict__.get( - "_link_input_suffixes", [".ltoir"] * len(lto_irs) - ) - if len(suffixes) != len(lto_irs): - raise RuntimeError("Cached coop link input metadata is inconsistent.") - owned_temp_files = [ - make_binary_tempfile(link_input, suffix) - for link_input, suffix in zip(lto_irs, suffixes) - ] - return Invocable( - temp_files=precompiled_temp_files + owned_temp_files, - owned_temp_files=owned_temp_files, - temp_storage_bytes=specialization.temp_storage_bytes, - temp_storage_alignment=specialization.temp_storage_alignment, - algorithm=specialization, - ) - - -class Invocable: - """Expose a compiled provider as a callable and retain its link files. - - The compiler obtains overloads from ``_numba_type_`` on first use. The - object also exposes scratch and synchronization requirements to planning. - Direct Python calls raise an error because execution belongs in a kernel. - """ - - def __init__( - self, - temp_files: Sequence[_NamedTempFile], - temp_storage_bytes: int, - temp_storage_alignment: int, - algorithm: Algorithm, - owned_temp_files: Sequence[_NamedTempFile] | None = None, - ) -> None: - """Retain a provider's link inputs and install cleanup for owned files. - - All link files must remain available while the compiler types and links - calls to this object. Keep references to shared bundle owners alongside - ordinary files, but unlink only the explicitly owned paths when this - invocable is finalized. The Numba callable type is constructed lazily on - first ``_numba_type_`` access. - - Parameters - ---------- - temp_files : sequence of _NamedTempFile - Complete ordered link inputs, including any shared bundle owners. - temp_storage_bytes : int - Compiled scratch size required by the specialization. - temp_storage_alignment : int - Compiled scratch alignment in bytes. - algorithm : Algorithm - Specialized provider supplying code generation and execution - metadata. - owned_temp_files : sequence of _NamedTempFile, optional - Files this invocable alone is responsible for unlinking. ``None`` - owns no paths; membership in ``temp_files`` alone does not imply - ownership. - """ - - self._temp_files = temp_files - self._owned_temp_files = ( - () if owned_temp_files is None else owned_temp_files - ) - self._temp_storage_bytes = temp_storage_bytes - self._temp_storage_alignment = temp_storage_alignment - self.specialization = algorithm - self._temp_file_finalizer = weakref.finalize( - self, - _cleanup_temp_files, - tuple(v.name for v in self._owned_temp_files), - ) - self._numba_type = None - - @property - def temp_storage_bytes(self): - return self._temp_storage_bytes - - @property - def temp_storage_alignment(self): - return self._temp_storage_alignment - - @property - def storage_abi(self): - return self.specialization.storage_abi - - @property - def execution_scope(self): - return self.specialization.execution_scope - - @property - def synchronization_scope(self): - return self.specialization.synchronization_scope - - @property - def files(self): - return [v.name for v in self._temp_files] - - @property - def _numba_type_(self): - """Return a compiler-local callable type without global registration.""" - - if self._numba_type is None: - templates = self.specialization.codegen(self) - self._numba_type = numba_types.Function(templates) - return self._numba_type - - def __call__(self, *args): - raise RuntimeError( - "__call__ should not be called directly outside of a " - "numba_cuda_mlir.cuda.jit(...) kernel." - ) - - -def _cleanup_temp_files(paths): - """Remove owned artifact paths, allowing a path to be already absent.""" - for path in paths: - try: - os.unlink(path) - except FileNotFoundError: - pass - - -__all__ = [ - "NUMBA_TYPES_TO_CPP", - "Algorithm", - "Array", - "BoundedInteger", - "Constant", - "CxxFunction", - "Dependency", - "DependentArray", - "DependentPointer", - "DependentPointerReference", - "DependentReference", - "ExactValue", - "Invocable", - "Parameter", - "Pointer", - "PointerOffset", - "PointerReference", - "Reference", - "SubstitutionFailure", - "TemplateParameter", - "TransformedArray", - "TypeWrapper", - "Value", - "_hash_symbol_value", - "_ltoir_to_ptx", - "_validate_logical_warp_threads", - "algo_coalesce_key", - "collect_specializations", - "internal_mangle_cpp", - "kernel_compile_target", - "make_invocable_from_specialization", - "mangle_symbol", - "numba_type_to_cpp", - "numba_type_to_wrapper", - "prepare_ltoir_bundle", - "war_introspection", - "war_introspection_call_with_transforms", -] diff --git a/python/cuda_coop/cuda/coop/numba_mlir/py.typed b/python/cuda_coop/cuda/coop/numba_mlir/py.typed deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/python/cuda_coop/cuda/coop/py.typed b/python/cuda_coop/cuda/coop/py.typed deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/python/cuda_coop/examples/numba_mlir/block_load.py b/python/cuda_coop/examples/numba_mlir/block_load.py deleted file mode 100644 index c0e1f5400e64..000000000000 --- a/python/cuda_coop/examples/numba_mlir/block_load.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Load a partial tile with an offset and a default for out-of-bounds items.""" - -import numpy as np -from numba_cuda_mlir import cuda - -from cuda import coop - -_THREADS = 32 -_SOURCE_OFFSET = 3 - - -@cuda.jit -def block_load(source, observed, valid_items, items_per_thread): - """Load a tile; invalid payload slots receive the default value -1.""" - - thread = cuda.threadIdx.x - payload = coop.ThreadData(items_per_thread) - coop.load( - coop.this_block(), - source, - payload, - algorithm="direct", - valid_items=valid_items, - oob_default=-1, - offset=_SOURCE_OFFSET, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - -def main(items_per_thread: int = 4) -> None: - tile_items = _THREADS * items_per_thread - valid_items = tile_items - 7 - source = np.arange(_SOURCE_OFFSET + tile_items, dtype=np.int32) - observed = np.zeros(tile_items, dtype=np.int32) - - block_load[1, _THREADS]( - source, observed, np.int32(valid_items), items_per_thread - ) - - expected = np.full(tile_items, -1, dtype=np.int32) - expected[:valid_items] = source[ - _SOURCE_OFFSET : _SOURCE_OFFSET + valid_items - ] - np.testing.assert_array_equal(observed, expected) - - -if __name__ == "__main__": - main() diff --git a/python/cuda_coop/examples/numba_mlir/block_store.py b/python/cuda_coop/examples/numba_mlir/block_store.py deleted file mode 100644 index ea49e33b6199..000000000000 --- a/python/cuda_coop/examples/numba_mlir/block_store.py +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Store a partial tile at an offset, preserving the rest of the destination.""" - -import numpy as np -from numba_cuda_mlir import cuda - -import cuda.coop.numba_mlir as coop - -_THREADS = 32 -_DESTINATION_OFFSET = 5 - - -@cuda.jit -def block_store(source, destination, valid_items, items_per_thread): - """Store the valid tile prefix while retaining the destination suffix.""" - - thread = cuda.threadIdx.x - payload = coop.ThreadData(items_per_thread) - for item in range(items_per_thread): - payload[item] = source[thread * items_per_thread + item] - coop.store( - coop.this_block(), - destination, - payload, - algorithm="direct", - valid_items=valid_items, - offset=_DESTINATION_OFFSET, - ) - - -def main(items_per_thread: int = 4) -> None: - tile_items = _THREADS * items_per_thread - valid_items = tile_items - 9 - source = np.arange(tile_items, dtype=np.int32) + 100 - destination = np.full( - _DESTINATION_OFFSET + tile_items, - -1, - dtype=np.int32, - ) - - block_store[1, _THREADS]( - source, destination, np.int32(valid_items), items_per_thread - ) - - expected = np.full_like(destination, -1) - expected[_DESTINATION_OFFSET : _DESTINATION_OFFSET + valid_items] = source[ - :valid_items - ] - np.testing.assert_array_equal(destination, expected) - - -if __name__ == "__main__": - main() diff --git a/python/cuda_coop/examples/numba_mlir/source_dumps.py b/python/cuda_coop/examples/numba_mlir/source_dumps.py deleted file mode 100644 index e0cb67182df1..000000000000 --- a/python/cuda_coop/examples/numba_mlir/source_dumps.py +++ /dev/null @@ -1,86 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Dump generated CUDA for a selected example kernel. - -Set CUDA_COOP_SOURCE_DUMP_DIR to the output directory before running. -""" - -import argparse - -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 - - -# docs: start dump-direct -@cuda.jit -def copy_direct(source, destination, items_per_thread): - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - coop.load( - block, - source, - items, - algorithm="direct", - ) - coop.store( - block, - destination, - items, - algorithm="direct", - ) - - -# docs: end dump-direct - - -# docs: start dump-transpose -@cuda.jit -def copy_transpose(source, destination, items_per_thread): - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - scratch = coop.TempStorage(auto_sync=True) - coop.load( - block, - source, - items, - algorithm="transpose", - temp_storage=scratch, - ) - coop.store( - block, - destination, - items, - algorithm="transpose", - temp_storage=scratch, - ) - - -# docs: end dump-transpose - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("example", choices=("direct", "transpose")) - example = parser.parse_args().example - count = 256 - source = ((np.arange(count) * 17) % 113 - 51).astype(np.int32) - destination = np.empty_like(source) - kernel = { - "direct": copy_direct, - "transpose": copy_transpose, - }[example] - items_per_thread = 4 - kernel[1, count // items_per_thread](source, destination, items_per_thread) - cuda.synchronize() - expected = source - np.testing.assert_array_equal(destination, expected) - print(f"{example}: result verified") - - -if __name__ == "__main__": - main() diff --git a/python/cuda_coop/pyproject.toml b/python/cuda_coop/pyproject.toml deleted file mode 100644 index 5d08407709eb..000000000000 --- a/python/cuda_coop/pyproject.toml +++ /dev/null @@ -1,176 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -[build-system] -requires = ["scikit-build-core>=0.10", "setuptools_scm"] -build-backend = "scikit_build_core.build" - -[project] -name = "cuda-coop" -description = "Cooperative CUDA Block and Warp Load and Store for Python DSLs" -authors = [{ name = "NVIDIA Corporation" }] -classifiers = [ - "Development Status :: 3 - Alpha", - "Intended Audience :: Developers", - "Topic :: Software Development :: Libraries :: Python Modules", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Environment :: GPU :: NVIDIA CUDA", - "License :: OSI Approved :: Apache Software License", - "Operating System :: POSIX :: Linux", - "Operating System :: Microsoft :: Windows", -] -requires-python = ">=3.10" -dependencies = [] -dynamic = ["version"] -readme = { file = "README.md", content-type = "text/markdown" } - -[project.optional-dependencies] -numba-cuda-mlir-cu12 = [ - "cuda-core>=0.5.1,<2", - "cuda-pathfinder>=1.2.3", - "numba-cuda-mlir[cu12]>=0.5.0,<0.6", - "packaging", - "numpy", - "typing_extensions>=4.12.0", -] -numba-cuda-mlir-cu13 = [ - "cuda-core>=0.5.1,<2", - "cuda-pathfinder>=1.2.3", - "numba-cuda-mlir[cu13]>=0.5.0,<0.6", - "packaging", - "numpy", - "typing_extensions>=4.12.0", -] -test = [ - "cuda-pathfinder>=1.2.3", - "mypy", - "numpy", - "packaging", - "pyright==1.1.414", - "pytest", - "pytest-cov", - "tomli>=2.0; python_version < '3.11'", - "typing_extensions>=4.12.0", -] - -[project.urls] -Homepage = "https://github.com/NVIDIA/cccl" -Repository = "https://github.com/NVIDIA/cccl" -Documentation = "https://nvidia.github.io/cccl" -Issues = "https://github.com/NVIDIA/cccl/issues" - -[tool.scikit-build] -minimum-version = "build-system.requires" -build-dir = "build/{wheel_tag}" - -[tool.scikit-build.cmake] -version = ">=3.21" -build-type = "Release" -source-dir = "." - -[tool.scikit-build.ninja] -version = ">=1.11" -make-fallback = true - -[tool.scikit-build.metadata.version] -provider = "scikit_build_core.metadata.setuptools_scm" - -[tool.setuptools_scm] -root = "../.." -git_describe_command = [ - "git", - "describe", - "--dirty", - "--tags", - "--long", - "--match", - "cuda-coop-v[0-9]*", -] -tag_regex = '^cuda-coop-v(?P[0-9]+(?:\.[0-9]+){2}(?:[.-].*)?)$' -fallback_version = "0.1.0.dev0" - -[tool.scikit-build.wheel] -py-api = "py3" -platlib = false -license-files = [] -reproducible = true - -[tool.scikit-build.wheel.force-include] -"../../LICENSE" = "${SKBUILD_METADATA_DIR}/licenses/LICENSE" -"../../cub/LICENSE.TXT" = "${SKBUILD_METADATA_DIR}/licenses/cub/LICENSE.TXT" -"../../libcudacxx/LICENSE.TXT" = "${SKBUILD_METADATA_DIR}/licenses/libcudacxx/LICENSE.TXT" -"../../thrust/LICENSE" = "${SKBUILD_METADATA_DIR}/licenses/thrust/LICENSE" - -[tool.scikit-build.wheel.packages] -"cuda/coop" = "cuda/coop" - -[tool.mypy] -cache_dir = ".cache/mypy" -python_version = "3.10" - -[[tool.mypy.overrides]] -module = [ - "cuda.pathfinder.*", - "numpy", - "numpy.*", -] -ignore_missing_imports = true -follow_imports = "skip" - -[tool.pyright] -include = ["cuda/coop/numba_mlir"] -extraPaths = ["."] -typeCheckingMode = "standard" -pythonVersion = "3.10" -reportUnusedFunction = "error" - -[tool.ruff] -extend = "../../pyproject.toml" -line-length = 80 - -[tool.ruff.lint] -extend-select = ["E501"] - -[tool.ruff.lint.isort] -known-first-party = [ - "cuda.bindings", - "cuda.coop", - "cuda.coop._core", - "cuda.coop._core.api", - "cuda.coop._core.block", - "cuda.coop._core.group", - "cuda.coop.numba_mlir", - "cuda.coop.numba_mlir._compiler", - "cuda.coop.numba_mlir._lowering", -] - -[tool.pytest.ini_options] -testpaths = ["tests"] -addopts = [ - "-ra", - "--import-mode=importlib", - "--strict-config", - "--strict-markers", -] -xfail_strict = true -markers = [ - "backend_numba_mlir: cuda.coop.numba_mlir backend test", - "compile: compiler or lowering test that does not launch a kernel", - "unit: backend-local unit test without GPU execution", - "runtime: representative backend runtime test", - "gpu: test that launches work on an NVIDIA GPU", -] - -[tool.coverage.run] -branch = true -source = ["cuda.coop._core", "cuda.coop.numba_mlir"] - -[tool.coverage.report] -show_missing = true -skip_covered = true diff --git a/python/cuda_coop/tests/__init__.py b/python/cuda_coop/tests/__init__.py deleted file mode 100644 index 8bbe3ce1ab86..000000000000 --- a/python/cuda_coop/tests/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception diff --git a/python/cuda_coop/tests/_group_planning.py b/python/cuda_coop/tests/_group_planning.py deleted file mode 100644 index 78c5c4cede6a..000000000000 --- a/python/cuda_coop/tests/_group_planning.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Build shared Load and Store requests for planner tests. - -Tests can vary one option while keeping the element type, item count, and -launch dimensions fixed. Planning these requests needs no kernel compiler. -""" - -from cuda.coop._core import ( - GroupLoadStoreSemantics, - LaunchFacts, - make_group_primitive_call, - plan_group_primitive, -) - - -def _load_store(kind="load", **overrides): - options = {"dtype": "int", "items_per_thread": 2, "algorithm": "direct"} - options.update(overrides) - return GroupLoadStoreSemantics(kind=kind, **options) - - -def _plan(group, operation, launch=(64, 1, 1)): - facts = launch if isinstance(launch, LaunchFacts) else LaunchFacts(launch) - return plan_group_primitive( - make_group_primitive_call(group, operation), facts - ) diff --git a/python/cuda_coop/tests/backends/numba_mlir/__init__.py b/python/cuda_coop/tests/backends/numba_mlir/__init__.py deleted file mode 100644 index 8bbe3ce1ab86..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception diff --git a/python/cuda_coop/tests/backends/numba_mlir/compile/__init__.py b/python/cuda_coop/tests/backends/numba_mlir/compile/__init__.py deleted file mode 100644 index 8bbe3ce1ab86..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/compile/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception diff --git a/python/cuda_coop/tests/backends/numba_mlir/compile/test_block_load_store_compile.py b/python/cuda_coop/tests/backends/numba_mlir/compile/test_block_load_store_compile.py deleted file mode 100644 index 7dcf349af043..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/compile/test_block_load_store_compile.py +++ /dev/null @@ -1,975 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Compile Block Load/Store providers with real NVRTC and device linking. - -The compile stage hides all GPUs, so tests replace device discovery with a -fixed sm_90 target. Source and PTX checks cover storage, synchronization, -symbols, and reuse; runtime suites check values from launched kernels. -""" - -from __future__ import annotations - -import gc -import json -import os -import re -import subprocess -import sys -import textwrap -import weakref -from dataclasses import replace -from pathlib import Path -from types import SimpleNamespace - -import pytest - -pytest.importorskip("numba_cuda_mlir") - -from numba_cuda_mlir import types - -from cuda.coop._core import ArgumentBinding, SynchronizationScope -from cuda.coop._core.block import ( - make_block_load_specialization, - make_block_store_specialization, -) -from cuda.coop.numba_mlir import _types -from cuda.coop.numba_mlir._compiler import _caching, _nvrtc -from cuda.coop.numba_mlir._compiler._operations import StorageABI -from cuda.coop.numba_mlir._lowering._core import NumbaMlirCoreAdapter -from cuda.coop.numba_mlir._lowering._load_store import _load_store_value_abis - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.compile] - -_FIXED_COMPUTE_CAPABILITY = (9, 0) - - -@pytest.fixture(autouse=True) -def _fixed_current_device( - monkeypatch: pytest.MonkeyPatch, -) -> list[tuple[int, int]]: - """Fix device discovery while retaining real NVRTC and nvJitLink calls.""" - - assert os.environ.get("CUDA_VISIBLE_DEVICES") == "", ( - "the Numba-CUDA-MLIR compile stage must hide all CUDA devices" - ) - queries: list[tuple[int, int]] = [] - - def current_device() -> SimpleNamespace: - queries.append(_FIXED_COMPUTE_CAPABILITY) - return SimpleNamespace(compute_capability=_FIXED_COMPUTE_CAPABILITY) - - monkeypatch.setattr(_types.cuda, "get_current_device", current_device) - return queries - - -@pytest.fixture(scope="module") -def compile_context() -> _nvrtc.CompileContext: - return _nvrtc.resolve_compile_context() - - -def _algorithm( - context: _nvrtc.CompileContext, - *, - operation: str = "load", - algorithm: str = "direct", - block_dim: tuple[int, int, int] = (32, 1, 1), - dtype=types.int32, - items_per_thread: int = 2, - valid_items: ArgumentBinding | None = None, -) -> _types.Algorithm: - """Materialize a Load or Store provider for one CUB algorithm and shape. - - Storage-free algorithms get no scratch pointer or barrier. Others take a - leading scratch pointer and block synchronization. - """ - if valid_items is None: - valid_items = ArgumentBinding.runtime() - adapter = NumbaMlirCoreAdapter( - value_abis=_load_store_value_abis( - dtype=dtype, - block_dim=block_dim, - items_per_thread=items_per_thread, - valid_items=valid_items, - ) - ) - factory = ( - make_block_load_specialization - if operation == "load" - else make_block_store_specialization - ) - specialization = factory( - dtype=adapter.core_dtype(dtype), - block_dim=block_dim, - items_per_thread=items_per_thread, - algorithm=algorithm, - valid_items=valid_items, - ) - storage_free = algorithm in {"direct", "striped", "vectorize"} - algorithm = adapter.materialize( - specialization, - storage_abi=( - StorageABI.NONE if storage_free else StorageABI.LEADING_POINTER - ), - execution_scope=SynchronizationScope.BLOCK, - synchronization_scope=( - SynchronizationScope.NONE - if storage_free - else SynchronizationScope.BLOCK - ), - extra_type_definitions=(_types.numba_type_to_wrapper(dtype),), - ) - algorithm._compile_context = context - return algorithm - - -def _source(algorithm: _types.Algorithm) -> str: - source, _support_lto_irs, _storage_types, _udf_declarations = ( - algorithm._source_code() - ) - return source - - -def _storage_abi_variant(source: str, algorithm: _types.Algorithm) -> str: - """Change only the explicit temporary-storage wrapper ABI.""" - - wrapper = f"{algorithm.mangled_name(algorithm.parameters[0])}__abi" - original = f"{wrapper}(void *__ret, " - changed = ( - f"{wrapper}(void *__ret, unsigned long long __cuda_coop_storage_abi, " - ) - assert source.count(original) == 1 - return source.replace(original, changed, 1) - - -def _cache_files(root: Path) -> tuple[Path, ...]: - return tuple(sorted(path for path in root.rglob("*") if path.is_file())) - - -def _stat_identity(path: Path) -> tuple[int, int, int]: - stat = path.stat() - return stat.st_ino, stat.st_mtime_ns, stat.st_size - - -def test_direct_load_store_compile_without_temp_storage_or_barriers( - compile_context: _nvrtc.CompileContext, - _fixed_current_device: list[tuple[int, int]], - monkeypatch: pytest.MonkeyPatch, -) -> None: - # Keep the real LTO-to-PTX helper for the checks below. Fail if bundling - # or invocable creation calls it to get scratch metadata. - inspect_ptx = _types._ltoir_to_ptx - monkeypatch.setattr( - _types, - "_ltoir_to_ptx", - lambda *args, **kwargs: pytest.fail( - "storage-free providers need no PTX query" - ), - ) - algorithms = [ - _algorithm(compile_context, operation="load"), - _algorithm(compile_context, operation="store"), - ] - clones = [ - _algorithm(compile_context, operation="load"), - _algorithm(compile_context, operation="store"), - ] - - sources = [_source(algorithm) for algorithm in algorithms] - clone_sources = [_source(algorithm) for algorithm in clones] - assert sources == clone_sources - for source in sources: - assert "::cuda::std::int64_t num_valid_items" in source - assert "if (num_valid_items < 0 || num_valid_items > 64)" in source - assert 'asm volatile("trap;" : : :);' in source - assert "static_cast<::cuda::std::int32_t>(num_valid_items)" in source - assert "checked_num_valid_items);" in source - - for algorithm, source in zip(algorithms, sources): - explicit_wrappers = tuple( - f"{algorithm.mangled_name(method)}__abi" - for method in algorithm.parameters - ) - implicit_wrappers = tuple( - f"{algorithm.mangled_name(method)}_alloc__abi" - for method in algorithm.parameters - ) - assert all(symbol in source for symbol in explicit_wrappers) - assert all(symbol not in source for symbol in implicit_wrappers) - assert "TempStorage" not in source - assert "temp_storage" not in source - assert "__syncthreads" not in source - assert "__syncwarp" not in source - assert "().Load(" in sources[0] - assert "().Store(" in sources[1] - - ltoir = _types.prepare_ltoir_bundle( - algorithms, - ) - assert isinstance(ltoir, bytes) - assert ltoir - assert _fixed_current_device - assert set(_fixed_current_device) == {_FIXED_COMPUTE_CAPABILITY} - - cc = 10 * _FIXED_COMPUTE_CAPABILITY[0] + _FIXED_COMPUTE_CAPABILITY[1] - ptx = inspect_ptx(ltoir, name="load_store_metadata", cc=cc) - assert ".version" in ptx - assert ".shared" not in ptx - assert "bar.sync" not in ptx - assert all(algorithm.temp_storage_bytes == 0 for algorithm in algorithms) - assert all( - algorithm.temp_storage_alignment == 1 for algorithm in algorithms - ) - - shared_artifact = algorithms[0]._precompiled_ltoir_files[0] - assert algorithms[1]._precompiled_ltoir_files[0] is shared_artifact - artifact_path = Path(shared_artifact.name) - shared_artifact_ref = weakref.ref(shared_artifact) - assert artifact_path.suffix == ".ltoir" - assert artifact_path.read_bytes() == ltoir - - invocables = [ - _types.make_invocable_from_specialization(algorithm) - for algorithm in algorithms - ] - assert all( - invocable.files == [str(artifact_path)] for invocable in invocables - ) - assert all( - invocable.temp_storage_bytes - == invocable.specialization.temp_storage_bytes - for invocable in invocables - ) - assert all( - invocable.temp_storage_alignment - == invocable.specialization.temp_storage_alignment - for invocable in invocables - ) - - del shared_artifact, algorithms - gc.collect() - assert shared_artifact_ref() is not None - assert artifact_path.is_file() - assert all( - invocable.files == [str(artifact_path)] for invocable in invocables - ) - - -@pytest.mark.parametrize("bundled", (False, True), ids=("single", "bundle")) -def test_storage_layout_queries_match_cpp_without_metadata_linking( - compile_context: _nvrtc.CompileContext, - monkeypatch: pytest.MonkeyPatch, - bundled: bool, -) -> None: - """Compare queried layouts with independently compiled C++ constants. - - Direct PTX compilation supplies the reference sizeof/alignof values. The - provider paths must obtain those same layouts during LTO compilation, - without using the LTO-to-PTX helper. Include a storage-free provider so - bundle layout ordering is checked when one member has no query. - """ - - cases = (("direct", 32), ("transpose", 32), ("transpose", 64)) - algorithms = [ - _algorithm(compile_context, algorithm=name, block_dim=(threads, 1, 1)) - for name, threads in cases - ] - reference_sources = [ - "#include \n#include \n" - ] - expected_symbols = {} - for index, (name, threads) in enumerate(cases): - if name != "direct": - storage_type = ( - "typename cub::BlockLoad<::cuda::std::int32_t, " - f"{threads}, 2, cub::BLOCK_LOAD_TRANSPOSE>::TempStorage" - ) - size_symbol = f"expected_storage_size_{index}" - alignment_symbol = f"expected_storage_alignment_{index}" - expected_symbols[index] = (size_symbol, alignment_symbol) - reference_sources.append( - f"__device__ constexpr unsigned {size_symbol} = " - f"sizeof({storage_type});\n" - f"__device__ constexpr unsigned {alignment_symbol} = " - f"alignof({storage_type});\n" - ) - _, reference_ptx = _nvrtc.compile( - cpp="\n".join(reference_sources), - cc=90, - rdc=True, - code="ptx", - context=compile_context, - ) - assert isinstance(reference_ptx, str) - reference_values = { - symbol: int(value) - for symbol, value in re.findall( - r"\.global \.align 4 \.u32 " - r"(expected_storage_(?:size|alignment)_[0-9]+) = ([0-9]+);", - reference_ptx, - ) - } - expected_layouts = [ - tuple(reference_values[symbol] for symbol in expected_symbols[index]) - if index in expected_symbols - else (0, 1) - for index in range(len(algorithms)) - ] - assert expected_layouts[1][0] < expected_layouts[2][0] - monkeypatch.setattr( - _types, - "_ltoir_to_ptx", - lambda *args, **kwargs: pytest.fail( - "storage metadata must not link to PTX" - ), - ) - - if bundled: - assert _types.prepare_ltoir_bundle(algorithms) - else: - for algorithm in algorithms: - assert algorithm.get_lto_ir() - - assert [ - (algorithm.temp_storage_bytes, algorithm.temp_storage_alignment) - for algorithm in algorithms - ] == expected_layouts - - -def test_all_block_load_store_algorithms_compile_with_declared_storage( - compile_context: _nvrtc.CompileContext, -) -> None: - storage_free_algorithms = {"direct", "striped", "vectorize"} - algorithm_names = ( - "direct", - "striped", - "vectorize", - "transpose", - "warp_transpose", - "warp_transpose_timesliced", - ) - algorithms = { - (operation, name): _algorithm( - compile_context, - operation=operation, - algorithm=name, - ) - for operation in ("load", "store") - for name in algorithm_names - } - - for (operation, name), specialization in algorithms.items(): - source = _source(specialization) - cub_token = f"::cub::BLOCK_{operation.upper()}_{name.upper()}" - assert cub_token in source - if name in storage_free_algorithms: - assert specialization.storage_abi is StorageABI.NONE - assert "TempStorage" not in source - assert "temp_storage" not in source - assert "__syncthreads" not in source - assert f"().{operation.title()}(" in source - else: - assert specialization.storage_abi is StorageABI.LEADING_POINTER - assert "::TempStorage" in source - assert "temp_storage" in source - assert "__syncthreads" in source - assert f"(*temp_storage).{operation.title()}(" in source - - for operation in ("load", "store"): - providers = [algorithms[(operation, name)] for name in algorithm_names] - assert len({provider.c_name for provider in providers}) == len( - providers - ) - assert len( - {provider._private_symbol_key for provider in providers} - ) == len(providers) - assert len({_source(provider) for provider in providers}) == len( - providers - ) - - bundle = _types.prepare_ltoir_bundle( - list(algorithms.values()), - ) - assert isinstance(bundle, bytes) - assert bundle - for (_operation, name), specialization in algorithms.items(): - if name in storage_free_algorithms: - assert specialization.temp_storage_bytes == 0 - assert specialization.temp_storage_alignment == 1 - else: - assert specialization.temp_storage_bytes > 0 - assert specialization.temp_storage_alignment > 0 - - -def test_unguarded_vectorize_compiles_the_full_tile_overload( - compile_context: _nvrtc.CompileContext, -) -> None: - algorithms = [ - _algorithm( - compile_context, - operation=operation, - algorithm="vectorize", - block_dim=(64, 1, 1), - items_per_thread=4, - valid_items=ArgumentBinding.omitted(), - ) - for operation in ("load", "store") - ] - - for operation, specialization in zip(("load", "store"), algorithms): - source = _source(specialization) - assert f"::cub::BLOCK_{operation.upper()}_VECTORIZE" in source - assert f"().{operation.title()}(" in source - assert "checked_num_valid_items" not in source - assert 'asm volatile("trap;" : : :);' not in source - assert specialization.storage_abi is StorageABI.NONE - - bundle = _types.prepare_ltoir_bundle( - algorithms, - ) - assert isinstance(bundle, bytes) - assert bundle - assert all( - specialization.temp_storage_bytes == 0 for specialization in algorithms - ) - - -def test_timesliced_transpose_uses_less_storage_for_multiple_warps( - compile_context: _nvrtc.CompileContext, -) -> None: - algorithms = { - (operation, algorithm): _algorithm( - compile_context, - operation=operation, - algorithm=algorithm, - block_dim=(64, 1, 1), - valid_items=ArgumentBinding.omitted(), - ) - for operation in ("load", "store") - for algorithm in ("warp_transpose", "warp_transpose_timesliced") - } - - bundle = _types.prepare_ltoir_bundle( - list(algorithms.values()), - ) - assert isinstance(bundle, bytes) - assert bundle - for operation in ("load", "store"): - regular = algorithms[(operation, "warp_transpose")] - timesliced = algorithms[(operation, "warp_transpose_timesliced")] - assert 0 < timesliced.temp_storage_bytes < regular.temp_storage_bytes - assert ( - timesliced.temp_storage_alignment == regular.temp_storage_alignment - ) - - -def test_representative_dtypes_compile_for_each_additional_algorithm( - compile_context: _nvrtc.CompileContext, -) -> None: - cases = ( - ("striped", types.int8), - ("vectorize", types.uint64), - ("transpose", types.float64), - ("warp_transpose", types.int32), - ("warp_transpose_timesliced", types.uint16), - ) - algorithms = [ - _algorithm( - compile_context, - operation=operation, - algorithm=algorithm, - block_dim=(64, 1, 1), - dtype=dtype, - ) - for operation in ("load", "store") - for algorithm, dtype in cases - ] - - bundle = _types.prepare_ltoir_bundle( - algorithms, - ) - assert isinstance(bundle, bytes) - assert bundle - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "bundle_failure", [False, True], ids=["bundle", "fallback"] -) -def test_production_routes_compile_storage_free_and_storage_bearing_kernels( - monkeypatch: pytest.MonkeyPatch, items_per_thread, bundle_failure -) -> None: - import numba_cuda_mlir.tools as numba_mlir_tools - from numba_cuda_mlir import cuda as compiler_cuda - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop - from cuda.coop.numba_mlir._compiler import _rewrite_invocables - - if bundle_failure: - - def fail_bundle(algorithms): - raise RuntimeError("exercise individual provider fallback") - - monkeypatch.setattr( - _rewrite_invocables, "prepare_ltoir_bundle", fail_bundle - ) - - identities = [] - identify = _nvrtc.include_dirs_identity - - def record_identity(paths): - identity = identify(paths) - identities.append(identity.digest) - return identity - - monkeypatch.setattr(_nvrtc, "include_dirs_identity", record_identity) - - fixed_device = SimpleNamespace(compute_capability=_FIXED_COMPUTE_CAPABILITY) - - def fixed_compute_capability(as_type=str): - assert as_type in (str, tuple) - return _FIXED_COMPUTE_CAPABILITY if as_type is tuple else "sm_90" - - monkeypatch.setattr( - numba_mlir_tools, - "get_gpu_compute_capability", - fixed_compute_capability, - ) - monkeypatch.setattr( - compiler_cuda, "get_current_device", lambda: fixed_device - ) - - @compiler_cuda.jit(chip="sm_90") - def storage_free(source, destination, items_per_thread): - storage = coop.TempStorage( - 128 * 1024, alignment=16, sharing="exclusive", auto_sync=True - ) - payload = coop.ThreadData(items_per_thread, dtype=types.int32) - coop.load( - coop.this_block(), - source, - payload, - algorithm="vectorize", - temp_storage=storage, - ) - coop.store( - coop.this_block(), - destination, - payload, - algorithm="vectorize", - temp_storage=storage, - ) - - @compiler_cuda.jit(chip="sm_90") - def storage_bearing(source, destination, items_per_thread): - payload = qualified_coop.ThreadData(items_per_thread, dtype=types.int32) - qualified_coop.load( - qualified_coop.this_block(), - source, - payload, - algorithm="warp_transpose_timesliced", - ) - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm="warp_transpose_timesliced", - ) - - signature = types.void( - types.int32[::1], - types.int32[::1], - types.IntegerLiteral(items_per_thread), - ) - launch_config_key = ( - ("grid", (1, 1, 1)), - ("block", (64, 1, 1)), - ("sharedmem", 0), - ("cluster", None), - ) - results = [] - for dispatcher in (storage_free, storage_bearing): - before = len(identities) - results.append( - dispatcher._compile_launch_config_signature( - signature, launch_config_key - ) - ) - # Newly created Load and Store providers share discovery even when - # bundle preparation fails and materialization retries individually. - assert len(identities) == before + 1 - for result in results: - assert isinstance(result.metadata["ltoir"], bytes) - assert result.metadata["ltoir"] - assert isinstance(result.metadata["cubin"], bytes) - assert result.metadata["cubin"] - assert result.metadata["linked_external_link_items"] - - storage_free_ptx = next(iter(storage_free.inspect_lto_ptx().values())) - storage_bearing_ptx = next(iter(storage_bearing.inspect_lto_ptx().values())) - assert ".visible .entry" in storage_free_ptx - assert ".shared" not in storage_free_ptx - assert "bar.sync" not in storage_free_ptx - assert ".visible .entry" in storage_bearing_ptx - assert ".shared" in storage_bearing_ptx - assert "bar.sync" in storage_bearing_ptx - - -def test_header_edits_do_not_poison_the_persistent_cache( - tmp_path: Path, -) -> None: - """Revert headers in a fresh process after compiling the edited bytes.""" - probe = textwrap.dedent( - r""" - import json - import re - import sys - from pathlib import Path - from cuda.coop.numba_mlir._compiler import _nvrtc - - root = Path(sys.argv[1]) - for name in (*_nvrtc._REQUIRED_HEADERS, "cub/version.cuh"): - path = root / name - path.parent.mkdir(parents=True, exist_ok=True) - path.touch() - header = root / "probe.cuh" - header.write_text("#define PROBE_VALUE 17\n") - with _nvrtc.compile_context_scope(): - original = _nvrtc.resolve_shared_compile_context() - if sys.argv[2] == "prime": - header.write_text("#define PROBE_VALUE 29\n") - with _nvrtc.compile_context_scope(): - context = _nvrtc.resolve_shared_compile_context() - _, ptx = _nvrtc.compile( - cpp='#include "probe.cuh"\nextern "C" __global__ ' - 'void probe(int* out) { *out = PROBE_VALUE; }\n', - cc=80, rdc=False, code="ptx", context=context, - ) - assert isinstance(ptx, str) - print(json.dumps({ - "original": original.header_identity, - "compiled": context.header_identity, - "constants": re.findall(r"mov\.u32\s+%r\d+,\s*(\d+);", ptx), - })) - """ - ) - root = tmp_path / "headers" - env = dict( - os.environ, - CUDA_COOP_CCCL_ROOT=str(root), - XDG_CACHE_HOME=str(tmp_path / "cache"), - LOCALAPPDATA=str(tmp_path / "cache"), - ) - results = [] - for mode, enabled, expected in ( - ("prime", "1", "29"), - ("read", "1", "17"), - ("read", "0", "17"), - ): - completed = subprocess.run( - [sys.executable, "-c", probe, str(root), mode], - env=dict(env, CUDA_COOP_ENABLE_CACHE=enabled), - capture_output=True, - text=True, - check=True, - ) - result = json.loads(completed.stdout) - assert result["constants"] == [expected] - results.append(result) - assert results[0]["original"] != results[0]["compiled"] - assert results[1]["compiled"] == results[0]["original"] - assert results[2] == results[1] - assert _cache_files(tmp_path / "cache") - - -def test_real_nvrtc_cache_hits_and_invalidates_every_compile_axis( - compile_context: _nvrtc.CompileContext, - _fixed_current_device: list[tuple[int, int]], - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - """Exercise both cache tiers without replacing either compiler library.""" - - monkeypatch.setattr(_caching, "_CACHE_USABLE", True) - monkeypatch.setattr(_caching, "_CACHE_LOCATION", str(tmp_path)) - _nvrtc.compile_impl.cache_clear() - - base_algorithm = _algorithm(compile_context) - base_source = _source(base_algorithm) - base_kwargs = { - "cpp": base_source, - "cc": 90, - "rdc": True, - "code": "lto", - "context": compile_context, - } - - try: - _version, first = _nvrtc.compile(**base_kwargs) - first_info = _nvrtc.compile_impl.cache_info() - assert first_info.hits == 0 - assert first_info.misses == 1 - files = _cache_files(tmp_path) - assert len(files) == 1 - first_disk_identity = _stat_identity(files[0]) - - _version, memory_cached = _nvrtc.compile(**base_kwargs) - memory_info = _nvrtc.compile_impl.cache_info() - assert memory_cached == first - assert memory_info.hits == first_info.hits + 1 - assert memory_info.misses == first_info.misses - assert _cache_files(tmp_path) == files - assert _stat_identity(files[0]) == first_disk_identity - - _nvrtc.compile_impl.cache_clear() - _version, disk_cached = _nvrtc.compile(**base_kwargs) - disk_info = _nvrtc.compile_impl.cache_info() - assert disk_cached == first - assert disk_info.hits == 0 - assert disk_info.misses == 1 - assert _cache_files(tmp_path) == files - assert _stat_identity(files[0]) == first_disk_identity - - dimension_source = _source( - _algorithm(compile_context, block_dim=(8, 4, 1)) - ) - dtype_source = _source(_algorithm(compile_context, dtype=types.uint16)) - items_source = _source(_algorithm(compile_context, items_per_thread=3)) - static_binding_source = _source( - _algorithm( - compile_context, - valid_items=ArgumentBinding.static(17), - ) - ) - algorithm_source = _source( - _algorithm(compile_context, algorithm="striped") - ) - storage_abi_source = _storage_abi_variant(base_source, base_algorithm) - changed_context = replace( - compile_context, - header_identity=f"{compile_context.header_identity}-changed", - ) - variants = { - "source": { - **base_kwargs, - "cpp": f"{base_source}\nstatic_assert(true);\n", - }, - "block-dimension": {**base_kwargs, "cpp": dimension_source}, - "dtype": {**base_kwargs, "cpp": dtype_source}, - "items-per-thread": {**base_kwargs, "cpp": items_source}, - "algorithm": {**base_kwargs, "cpp": algorithm_source}, - "static-runtime-binding": { - **base_kwargs, - "cpp": static_binding_source, - }, - "storage-abi": {**base_kwargs, "cpp": storage_abi_source}, - "compile-context": {**base_kwargs, "context": changed_context}, - "compute-capability": {**base_kwargs, "cc": 80}, - "relocatable-device-code": {**base_kwargs, "rdc": False}, - "output-kind": {**base_kwargs, "code": "ptx"}, - } - - previous_misses = disk_info.misses - previous_file_count = len(files) - for axis, kwargs in variants.items(): - _version, result = _nvrtc.compile(**kwargs) - info = _nvrtc.compile_impl.cache_info() - assert info.misses == previous_misses + 1, axis - assert info.hits == 0, axis - assert len(_cache_files(tmp_path)) == previous_file_count + 1, axis - if kwargs["code"] == "lto": - assert isinstance(result, bytes) and result, axis - else: - assert isinstance(result, str) and ".version" in result, axis - previous_misses = info.misses - previous_file_count += 1 - - original_compiler_options = _nvrtc._compiler_options - - def changed_compiler_options(**kwargs): - return ( - *original_compiler_options(**kwargs), - b"-DCUDA_COOP_COMPILE_OPTION_IDENTITY_TEST=1", - ) - - monkeypatch.setattr( - _nvrtc, "_compiler_options", changed_compiler_options - ) - _version, option_changed = _nvrtc.compile(**base_kwargs) - option_info = _nvrtc.compile_impl.cache_info() - assert option_info.misses == previous_misses + 1 - assert len(_cache_files(tmp_path)) == previous_file_count + 1 - assert isinstance(option_changed, bytes) and option_changed - - assert _fixed_current_device - assert set(_fixed_current_device) == {_FIXED_COMPUTE_CAPABILITY} - finally: - _nvrtc.compile_impl.cache_clear() - - -def test_provider_symbols_and_cached_lto_are_bound_to_the_compilation_target( - compile_context: _nvrtc.CompileContext, - monkeypatch: pytest.MonkeyPatch, -) -> None: - # Simulate a target switch with all GPUs hidden by the fixture. - # These are compiler identities, not devices required on the test runner. - current_cc = [(9, 0)] - - monkeypatch.setattr( - _types.cuda, - "get_current_device", - lambda: SimpleNamespace(compute_capability=current_cc[0]), - ) - sm90 = _algorithm(compile_context) - sm90_source = _source(sm90) - sm90_symbol = sm90.mangled_name(sm90.parameters[0]) - - current_cc[0] = (8, 0) - sm80 = _algorithm(compile_context) - sm80_source = _source(sm80) - sm80_symbol = sm80.mangled_name(sm80.parameters[0]) - - assert sm90_source != sm80_source - assert sm90_symbol != sm80_symbol - for algorithm, cc in ((sm90, 90), (sm80, 80)): - identity = algorithm._provider_compile_identity - assert identity.cc == cc - assert identity.rdc is True - assert identity.code == "lto" - - with pytest.raises(RuntimeError, match="different compute capability"): - # Reject reuse before NVRTC compilation or any attempt to launch code. - sm90.get_lto_ir() - - -@pytest.mark.parametrize( - "names", - ( - ("class", "template", "return", "int"), - ("items", "items", "items_1", "items"), - ("value_", "value_", "value_1", "value_"), - ("input", "transformed_input", "cast_input", "temp_storage"), - ("out", "ref_out", "cast_out", "out"), - ("1input", "two-items", "__ret", "Δ"), - ), - ids=( - "keywords", - "duplicates", - "trailing-underscores", - "generated-locals", - "return-local", - "invalid-identifiers", - ), -) -def test_parameter_names_compile_with_keywords_and_collisions( - compile_context: _nvrtc.CompileContext, - names: tuple[str, ...], -) -> None: - parameters = [ - _types.TransformedArray(types.uint8, types.int32, 2, "{value} != 0"), - _types.Reference(types.int32), - _types.Pointer(types.int32), - _types.Value(types.int32), - ] - for parameter, name in zip(parameters, names): - parameter.parameter_name = name - algorithm = _types.Algorithm( - struct_name="NamedProvider", - method_name="Run", - c_name="named_provider", - includes=(), - template_parameters=(), - parameters=(parameters,), - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.NONE, - synchronization_scope=SynchronizationScope.NONE, - type_definitions=[ - SimpleNamespace( - code=( - "namespace cub {\n" - "struct NamedProvider {\n" - " template\n" - " __device__ void Run(Args&&...) {}\n" - "};\n" - "}" - ), - lto_irs=[], - ) - ], - compile_context=compile_context, - template_arguments={}, - ) - source = _source(algorithm) - assert source == _source(algorithm) - _, ltoir = _nvrtc.compile( - cpp=source, - cc=90, - rdc=True, - code="lto", - context=compile_context, - ) - assert bytes(ltoir) - - -def test_parameter_names_do_not_split_identical_bundle_providers( - compile_context: _nvrtc.CompileContext, -) -> None: - original = _algorithm(compile_context) - renamed = _algorithm(compile_context) - for method in renamed.parameters: - for position, parameter in enumerate(method): - parameter.parameter_name = f"renamed_{position}" - - assert _source(original) != _source(renamed) - bundle = _types.prepare_ltoir_bundle([original, renamed]) - assert bundle is None - assert original._private_symbol_key == renamed._private_symbol_key - assert original.mangled_name( - original.parameters[0] - ) == renamed.mangled_name(renamed.parameters[0]) - - -def test_wrapper_parameters_do_not_shadow_function_type_or_return_local( - compile_context: _nvrtc.CompileContext, -) -> None: - from io import StringIO - - method = [ - _types.Pointer(types.uint8), - _types.Reference(types.int32), - _types.Value(types.int32), - _types.Reference(types.int32, is_output=True), - ] - for parameter, name in zip( - method, ("invoke", "scratch_t", "out", "result") - ): - parameter.parameter_name = name - source_buffer = StringIO() - source_buffer.write( - "#include \n" - "struct scratch_t {};\n" - "__device__ void invoke(scratch_t *, ::cuda::std::int32_t &input, " - "::cuda::std::int32_t value, ::cuda::std::int32_t &output) {\n" - " output = input + value;\n" - "}\n" - ) - _algorithm(compile_context)._emit_abi_wrapper( - source_buffer.write, - method, - exported_name="named_wrapper", - internal_name="invoke", - temp_storage_type_name="scratch_t", - ) - source = source_buffer.getvalue() - assert "named_wrapper__abi(void *__ret," in source - assert "*reinterpret_cast<::cuda::std::int32_t *>(__ret) = out;" in source - _, ltoir = _nvrtc.compile( - cpp=source, - cc=90, - rdc=True, - code="lto", - context=compile_context, - ) - assert bytes(ltoir) diff --git a/python/cuda_coop/tests/backends/numba_mlir/compile/test_installed_activation_compile.py b/python/cuda_coop/tests/backends/numba_mlir/compile/test_installed_activation_compile.py deleted file mode 100644 index 6b8ce9c11349..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/compile/test_installed_activation_compile.py +++ /dev/null @@ -1,216 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check activation and compilation from the installed wheel in isolation. - -Each probe uses a fresh isolated interpreter with ``PYTHONPATH`` deliberately -set to the checkout. It checks module and header origins before compiling, -so checkout files cannot hide missing wheel contents. -""" - -from __future__ import annotations - -import importlib.metadata -import os -import subprocess -import sys -import textwrap -from pathlib import Path - -import pytest - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.compile] - -_PACKAGE_ROOT = Path(__file__).parents[4].resolve() -_IMPORT_ORDERS = ( - pytest.param("numba-first-common", id="numba-first-common"), - pytest.param( - "numba-first-explicit-qualified-alias", - id="numba-first-explicit-qualified-alias", - ), - pytest.param( - "root-first-explicit-qualified", - id="root-first-explicit-qualified", - ), - pytest.param("root-first-register", id="root-first-register"), -) - -_COMPILE_PROBE = textwrap.dedent( - """ - import importlib.metadata - import os - import sys - from pathlib import Path - from types import SimpleNamespace - - assert os.environ["CUDA_VISIBLE_DEVICES"] == "" - import_order = os.environ["CUDA_COOP_IMPORT_ORDER"] - - if import_order == "numba-first-common": - import numba_cuda_mlir - from numba_cuda_mlir import cuda, types - from cuda import coop - - assert "cuda.coop.numba_mlir" in sys.modules - elif import_order == "numba-first-explicit-qualified-alias": - import numba_cuda_mlir - from numba_cuda_mlir import cuda, types - - compiler_cuda = cuda - import cuda.coop as coop - import cuda.coop.numba_mlir as numba_coop - - assert cuda is compiler_cuda - assert numba_coop is sys.modules["cuda.coop.numba_mlir"] - elif import_order == "root-first-explicit-qualified": - from cuda import coop - - assert "numba_cuda_mlir" not in sys.modules - assert "cuda.coop.numba_mlir" not in sys.modules - - import cuda.coop.numba_mlir # noqa: F401 - from numba_cuda_mlir import cuda, types - elif import_order == "root-first-register": - from cuda import coop - - assert "numba_cuda_mlir" not in sys.modules - assert "cuda.coop.numba_mlir" not in sys.modules - - assert coop.register("numba-cuda-mlir") is None - from numba_cuda_mlir import cuda, types - else: - raise AssertionError(f"unexpected import order: {import_order!r}") - - distribution = importlib.metadata.distribution("cuda-coop") - distribution_root = Path(distribution.locate_file("")).resolve() - expected_module = Path( - distribution.locate_file("cuda/coop/__init__.py") - ).resolve() - source_root = Path(os.environ["CUDA_COOP_SOURCE_ROOT"]).resolve() - module_file = Path(coop.__file__).resolve() - qualified_file = Path(sys.modules["cuda.coop.numba_mlir"].__file__).resolve() - - assert module_file == expected_module - assert module_file.is_relative_to(distribution_root) - assert qualified_file.is_relative_to(distribution_root) - assert not module_file.is_relative_to(source_root) - - from cuda.coop._headers import resolve_include_paths - - include_paths = resolve_include_paths( - start=Path.cwd(), - required_headers=( - "cub/block/block_load.cuh", - "cuda/std/cstdint", - ), - ) - assert include_paths.origin == "cuda-coop wheel header bundle" - assert all( - path.resolve().is_relative_to(distribution_root) - for path in include_paths.cccl - ) - - # Numba-CUDA-MLIR 0.5 has no public GPU-free entry point that carries - # configured launch metadata. Fix only its current-device queries; the - # compiler, NVRTC provider compilation, and nvJitLink remain real. - import numba_cuda_mlir.tools as numba_mlir_tools - - fixed_device = SimpleNamespace(compute_capability=(9, 0)) - - def fixed_compute_capability(as_type=str): - assert as_type in (str, tuple) - return (9, 0) if as_type is tuple else "sm_90" - - numba_mlir_tools.get_gpu_compute_capability = fixed_compute_capability - cuda.get_current_device = lambda: fixed_device - - @cuda.jit(chip="sm_90") - def block_load(source, destination, items_per_thread): - thread = cuda.threadIdx.x - payload = coop.ThreadData(items_per_thread) - coop.load( - coop.this_block(), - source, - payload, - algorithm="direct", - ) - for item in range(items_per_thread): - destination[thread * items_per_thread + item] = payload[item] - - items_per_thread = int(os.environ["CUDA_COOP_TEST_ITEMS_PER_THREAD"]) - signature = types.void( - types.int32[::1], types.int32[::1], types.IntegerLiteral(items_per_thread) - ) - launch_config_key = ( - ("grid", (1, 1, 1)), - ("block", (32, 1, 1)), - ("sharedmem", 0), - ("cluster", None), - ) - result = block_load._compile_launch_config_signature( - signature, - launch_config_key, - ) - - assert isinstance(result.metadata["ltoir"], bytes) - assert result.metadata["ltoir"] - assert isinstance(result.metadata["cubin"], bytes) - assert result.metadata["cubin"] - - linked_ltoir = tuple( - Path(path) for path in result.metadata["linked_external_link_items"] - ) - assert linked_ltoir - assert all(path.suffix == ".ltoir" and path.stat().st_size for path in linked_ltoir) - - ptx_by_specialization = block_load.inspect_lto_ptx() - assert len(ptx_by_specialization) == 1 - ptx = next(iter(ptx_by_specialization.values())) - assert isinstance(ptx, str) - assert ".visible .entry" in ptx - """ # noqa: E501 - Preserve embedded source bytes. -) - - -@pytest.mark.parametrize("import_order", _IMPORT_ORDERS) -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_installed_wheel_import_order_compiles_block_load( - import_order: str, - tmp_path: Path, - items_per_thread, -) -> None: - try: - distribution = importlib.metadata.distribution("cuda-coop") - except importlib.metadata.PackageNotFoundError: - pytest.fail( - "the Numba-CUDA-MLIR compile stage requires an installed wheel" - ) - - expected_module = Path( - distribution.locate_file("cuda/coop/__init__.py") - ).resolve() - assert expected_module.is_file() - - environment = os.environ.copy() - environment["CUDA_VISIBLE_DEVICES"] = "" - environment["CUDA_COOP_ENABLE_CACHE"] = "0" - environment["CUDA_COOP_IMPORT_ORDER"] = import_order - environment["CUDA_COOP_TEST_ITEMS_PER_THREAD"] = str(items_per_thread) - environment["CUDA_COOP_SOURCE_ROOT"] = str(_PACKAGE_ROOT) - # Isolated mode must ignore this deliberate source-tree contamination. - environment["PYTHONPATH"] = str(_PACKAGE_ROOT) - environment.pop("CUDA_COOP_CCCL_ROOT", None) - environment.pop("CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION", None) - environment.pop("CUDA_COOP_SOURCE_DUMP_DIR", None) - - result = subprocess.run( - [sys.executable, "-I", "-c", _COMPILE_PROBE], - cwd=tmp_path, - env=environment, - capture_output=True, - text=True, - check=False, - ) - - assert result.returncode == 0, result.stdout + result.stderr diff --git a/python/cuda_coop/tests/backends/numba_mlir/compile/test_load_store_target_compile.py b/python/cuda_coop/tests/backends/numba_mlir/compile/test_load_store_target_compile.py deleted file mode 100644 index 7d224bddebea..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/compile/test_load_store_target_compile.py +++ /dev/null @@ -1,182 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Compile Load/Store providers for the enclosing kernel's target and options. - -Device discovery reports SM90 while kernels explicitly target other chips. -Record actual NVRTC requests and compile/link the complete kernels to verify -architecture suffixes, lineinfo options, and debug-link compatibility. -The SM80 case also relinks the actual inputs explicitly for SM80 because the -compiler can promote its normal final link to the current device's target. -""" - -from types import SimpleNamespace - -import pytest - -pytest.importorskip("numba_cuda_mlir") - -from numba_cuda_mlir import types - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.compile] - -_BLOCK_THREADS = 64 -_DEVICE_COMPUTE_CAPABILITY = (9, 0) -_LAUNCH_CONFIG_KEY = ( - ("grid", (1, 1, 1)), - ("block", (_BLOCK_THREADS, 1, 1)), - ("sharedmem", 0), - ("cluster", None), -) - - -def _record_compilations(monkeypatch: pytest.MonkeyPatch): - """Fix device discovery and record provider compilation and link inputs.""" - - import numba_cuda_mlir.tools as numba_mlir_tools - from numba_cuda_mlir import cuda - from numba_cuda_mlir.linker import Linker - - from cuda.coop.numba_mlir._compiler import _nvrtc - - def fixed_compute_capability(as_type=str): - assert as_type in (str, tuple) - major, minor = _DEVICE_COMPUTE_CAPABILITY - return (major, minor) if as_type is tuple else f"sm_{major}{minor}" - - monkeypatch.setattr( - numba_mlir_tools, "get_gpu_compute_capability", fixed_compute_capability - ) - monkeypatch.setattr( - cuda, - "get_current_device", - lambda: SimpleNamespace(compute_capability=_DEVICE_COMPUTE_CAPABILITY), - ) - - recorded = SimpleNamespace(nvrtc=[], links=[]) - compile_impl = _nvrtc.compile_impl - - def recording_compile_impl(*args, **kwargs): - recorded.nvrtc.append(kwargs) - return compile_impl(*args, **kwargs) - - monkeypatch.setattr(_nvrtc, "compile_impl", recording_compile_impl) - complete = Linker.complete - - def recording_complete(self): - result = complete(self) - recorded.links.append( - SimpleNamespace( - arch=self.arch, - inputs=tuple(self._object_codes), - cubin=result.code, - ) - ) - return result - - monkeypatch.setattr(Linker, "complete", recording_complete) - return cuda, recorded - - -def _compile_kernel(cuda, items_per_thread, **jit_options): - """Compile Load/Store providers with explicit shared temporary storage.""" - - from cuda import coop - - @cuda.jit(**jit_options) - def kernel(source, output, items_per_thread): - block = coop.this_block() - payload = coop.ThreadData(items_per_thread) - storage = coop.TempStorage(sharing="shared", auto_sync=True) - coop.load( - block, source, payload, algorithm="transpose", temp_storage=storage - ) - coop.store( - block, output, payload, algorithm="transpose", temp_storage=storage - ) - - signature = types.void( - types.int32[::1], - types.int32[::1], - types.IntegerLiteral(items_per_thread), - ) - return kernel._compile_launch_config_signature( - signature, _LAUNCH_CONFIG_KEY - ) - - -@pytest.mark.parametrize( - ("jit_options", "arch", "cc"), - [ - ({"chip": "sm_80"}, "80", 80), - ({"cc": (8, 6)}, "86", 86), - ({"chip": "sm_90a"}, "90a", 90), - ({"chip": "sm_100f"}, "100f", 100), - ], -) -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_load_store_targets_the_kernel_not_the_current_device( - monkeypatch: pytest.MonkeyPatch, items_per_thread, jit_options, arch, cc -) -> None: - cuda, recorded = _record_compilations(monkeypatch) - - result = _compile_kernel(cuda, items_per_thread, **jit_options) - - assert result.metadata["cubin"] - if arch == "80": - from cuda.core import Linker, LinkerOptions - - final_link = recorded.links[-1] - assert final_link.cubin == result.metadata["cubin"] - assert len(final_link.inputs) > 1 - linked = Linker( - *final_link.inputs, - options=LinkerOptions(arch="sm_80", link_time_optimization=True), - ).link("cubin") - assert linked.code - - assert recorded.nvrtc - for request in recorded.nvrtc: - assert request["cc"] == cc - assert ( - f"--gpu-architecture=compute_{arch}".encode() - in request["compiler_options"] - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_lineinfo_reaches_load_store_providers( - monkeypatch: pytest.MonkeyPatch, - items_per_thread, -) -> None: - cuda, recorded = _record_compilations(monkeypatch) - - result = _compile_kernel(cuda, items_per_thread, lineinfo=True) - - assert result.metadata["cubin"] - assert recorded.nvrtc - assert all( - b"-lineinfo" in request["compiler_options"] - for request in recorded.nvrtc - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("lineinfo", [False, True]) -def test_debug_kernel_links_generated_code_without_debug_information( - monkeypatch: pytest.MonkeyPatch, items_per_thread, lineinfo -) -> None: - # A debug kernel is linked with -g, and nvJitLink then rejects more than - # one LTO input that carries debug or line information. - cuda, recorded = _record_compilations(monkeypatch) - - result = _compile_kernel( - cuda, items_per_thread, debug=True, opt=False, lineinfo=lineinfo - ) - - assert result.metadata["cubin"] - assert recorded.nvrtc - for request in recorded.nvrtc: - assert b"-lineinfo" not in request["compiler_options"] - assert b"-G" not in request["compiler_options"] diff --git a/python/cuda_coop/tests/backends/numba_mlir/compile/test_scalar_store_typing.py b/python/cuda_coop/tests/backends/numba_mlir/compile/test_scalar_store_typing.py deleted file mode 100644 index 798fb0327523..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/compile/test_scalar_store_typing.py +++ /dev/null @@ -1,136 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check scalar Store dtypes through real kernel compilation. - -Use a fixed target with devices hidden. Cases distinguish planner diagnostics -from errors that require ordinary compiler type inference. -""" - -import os -from types import SimpleNamespace - -import numpy as np -import pytest - -pytest.importorskip("numba_cuda_mlir") - -import numba_cuda_mlir.tools as numba_mlir_tools -from numba_cuda_mlir import cuda, types -from numba_cuda_mlir.numba_cuda.core.errors import TypingError - -import cuda.coop.numba_mlir as qualified_coop -from cuda import coop as root_coop -from cuda.coop.numba_mlir._compiler._group_errors import ( - MemoryDtypeMismatchError, -) - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.compile] - - -@pytest.fixture(autouse=True) -def _fixed_current_device(monkeypatch): - """Supply target metadata without requiring a CUDA device.""" - assert os.environ.get("CUDA_VISIBLE_DEVICES") == "" - monkeypatch.setattr( - numba_mlir_tools, - "get_gpu_compute_capability", - lambda as_type=str: (9, 0) if as_type is tuple else "sm_90", - ) - monkeypatch.setattr( - cuda, - "get_current_device", - lambda: SimpleNamespace(compute_capability=(9, 0)), - ) - - -def _compile(kernel, *arg_types): - # Configured launches are the runtime's only route to exact block metadata. - result = kernel._compile_launch_config_signature( - types.void(*arg_types), - ( - ("grid", (1, 1, 1)), - ("block", (32, 1, 1)), - ("sharedmem", 0), - ("cluster", None), - ), - ) - assert result.metadata["ltoir"] - return result - - -@pytest.mark.parametrize( - "coop", (root_coop, qualified_coop), ids=("root", "qualified") -) -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_store_compiles_a_runtime_payload_index(coop, items_per_thread): - @cuda.jit(chip="sm_90") - def kernel(source, destination, index, items_per_thread): - payload = coop.ThreadData(items_per_thread) - group = coop.this_block() - coop.load(group, source, payload) - coop.store(group, destination, payload[index]) - - _compile( - kernel, - types.int32[::1], - types.int32[::1], - types.int64, - types.IntegerLiteral(items_per_thread), - ) - - -@pytest.mark.parametrize( - "coop", (root_coop, qualified_coop), ids=("root", "qualified") -) -@pytest.mark.parametrize("expression", ("abs", "min", "loop")) -@pytest.mark.parametrize("matching", (False, True), ids=("mismatch", "exact")) -def test_store_checks_actual_expression_dtype(coop, expression, matching): - @cuda.jit(chip="sm_90") - def kernel(source, destination): - value = source[cuda.threadIdx.x] - if expression == "abs": - result = abs(value) - elif expression == "min": - result = min(value, 3.5) - else: - result = value - for _ in range(2): - result += np.float32(1.0) - coop.store(coop.this_block(), destination, result) - - result_dtype = types.float64 if expression == "min" else types.float32 - destination_dtype = result_dtype if matching else types.int32 - arg_types = (types.float32[::1], destination_dtype[::1]) - if matching: - _compile(kernel, *arg_types) - elif expression == "loop": - # Loop types are available during group planning, before typing. - with pytest.raises(MemoryDtypeMismatchError, match="payload dtype"): - _compile(kernel, *arg_types) - else: - with pytest.raises( - TypingError, match="does not match destination dtype" - ): - _compile(kernel, *arg_types) - - -@pytest.mark.parametrize( - "coop", (root_coop, qualified_coop), ids=("root", "qualified") -) -@pytest.mark.parametrize( - ("value", "dtype"), - ( - (127, types.int8), - (-128, types.int8), - (3.5, types.float32), - (255, types.uint8), - ), -) -def test_store_preserves_contextual_literal_conversion(coop, value, dtype): - @cuda.jit(chip="sm_90") - def kernel(destination): - coop.store(coop.this_block(), destination, value) - - _compile(kernel, dtype[::1]) diff --git a/python/cuda_coop/tests/backends/numba_mlir/compile/test_storage_diagnostics.py b/python/cuda_coop/tests/backends/numba_mlir/compile/test_storage_diagnostics.py deleted file mode 100644 index 08ff9abb4db4..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/compile/test_storage_diagnostics.py +++ /dev/null @@ -1,452 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Compile descriptor and storage edge cases to check their diagnostics. - -The fixed target permits compilation with devices hidden. Source-location, -helper-inlining, and shared-memory checks exercise the real compiler pipeline. -""" - -import inspect -import os -from types import SimpleNamespace - -import numpy as np -import pytest - -pytest.importorskip("numba_cuda_mlir") - -import numba_cuda_mlir.tools as numba_mlir_tools -from numba_cuda_mlir import cuda, types -from numba_cuda_mlir.numba_cuda.core.errors import TypingError -from numba_cuda_mlir.numba_cuda.misc.special import literal_unroll - -import cuda.coop.numba_mlir as coop -from cuda.coop.numba_mlir._compiler._group_planner_support import ( - GroupRewriteError, -) -from cuda.coop.numba_mlir._compiler._rewrite_support import ( - CoopSinglePhaseRewriteError, -) - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.compile] - - -@pytest.fixture(autouse=True) -def _fixed_current_device(monkeypatch): - """Supply target metadata without requiring a CUDA device.""" - assert os.environ.get("CUDA_VISIBLE_DEVICES") == "" - monkeypatch.setattr( - numba_mlir_tools, - "get_gpu_compute_capability", - lambda as_type=str: (9, 0) if as_type is tuple else "sm_90", - ) - monkeypatch.setattr( - cuda, - "get_current_device", - lambda: SimpleNamespace(compute_capability=(9, 0)), - ) - - -def _compile(kernel, *arg_types, block=(32, 1, 1)): - """Compile with exact launch dimensions; do not execute the kernel.""" - return kernel._compile_launch_config_signature( - types.void(*arg_types), - ( - ("grid", (1, 1, 1)), - ("block", block), - ("sharedmem", 0), - ("cluster", None), - ), - ) - - -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("operation", ("load", "store")) -@pytest.mark.parametrize( - "inlined_helper", (False, True), ids=("kernel", "helper") -) -def test_load_store_unexpected_keyword_raises_type_error( - qualified, operation, inlined_helper -): - from cuda import coop as portable_coop - - module = coop if qualified else portable_coop - primitive = module.load if operation == "load" else module.store - - def copy(data, items_per_thread): - block = module.this_block() - items = module.ThreadData(items_per_thread) - primitive(block, data, items, make_fast=True) - - if inlined_helper: - helper = cuda.jit(device=True, inline="always")(copy) - - @cuda.jit(chip="sm_90") - def kernel(data, items_per_thread): - helper(data, items_per_thread) - - else: - kernel = cuda.jit(chip="sm_90")(copy) - - with pytest.raises( - TypeError, match="unexpected keyword argument 'make_fast'" - ) as exc_info: - _compile(kernel, types.int32[::1], types.IntegerLiteral(2)) - - lines, first_line = inspect.getsourcelines(copy) - call_line = first_line + next( - index for index, line in enumerate(lines) if "make_fast=True" in line - ) - message = str(exc_info.value) - assert ( - f"{operation}() got an unexpected keyword argument 'make_fast'" - in message - ) - assert f'File "{os.path.relpath(__file__)}", line {call_line}:' in message - assert "primitive(block, data, items, make_fast=True)" in message - - -@pytest.mark.parametrize( - "qualified", (False, True), ids=("portable", "qualified") -) -@pytest.mark.parametrize( - "dtype", (None, types.int32), ids=("inferred", "explicit") -) -def test_load_return_cannot_be_used_as_a_store_payload(qualified, dtype): - from cuda import coop as common_coop - - module = coop if qualified else common_coop - - @cuda.jit(chip="sm_90") - def kernel(source, destination): - payload = module.ThreadData(items_per_thread=2, dtype=dtype) - result = module.load(module.this_block(), source, payload) - module.store(module.this_block(), destination, result) - - with pytest.raises(TypingError, match="none|None"): - _compile(kernel, types.int32[::1], types.int32[::1]) - - -def test_rebound_explicit_dtype_is_not_replaced_by_write_inference(): - @cuda.jit(chip="sm_90") - def kernel(destination): - if destination[0] > 0: - dtype = types.int32 - first = coop.ThreadData(items_per_thread=2, dtype=dtype) - first[0] = 16777217 - destination[0] = first[0] - dtype = types.float32 - second = coop.ThreadData(items_per_thread=2, dtype=dtype) - second[0] = 1.5 - destination[1] = second[0] - - with pytest.raises( - (CoopSinglePhaseRewriteError, TypingError), - match="dtype must resolve", - ): - _compile(kernel, types.int32[::1]) - - -def test_rebound_shared_shape_is_rejected_with_cooperative_storage(): - @cuda.jit(chip="sm_90") - def kernel(destination): - for _ in range(2): - size = 0 - dynamic = cuda.shared.array(size, types.int32) - size = 32 - static = cuda.shared.array(size, types.int32) - dynamic[cuda.threadIdx.x] = 1 - static[cuda.threadIdx.x] = 2 - payload = coop.ThreadData(items_per_thread=2, dtype=types.int32) - payload[0] = dynamic[cuda.threadIdx.x] - payload[1] = static[cuda.threadIdx.x] - coop.store( - coop.this_block(), destination, payload, algorithm="transpose" - ) - - with pytest.raises( - (CoopSinglePhaseRewriteError, TypingError), match="would alias" - ): - _compile(kernel, types.int32[::1]) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_default_helper_accepts_descriptor_alias_chain(items_per_thread): - @cuda.jit(device=True) - def helper(source, storage, items_per_thread): - payload = coop.ThreadData(items_per_thread, dtype=types.int32) - coop.load(coop.this_block(), source, payload, temp_storage=storage) - return payload[0] - - @cuda.jit(chip="sm_90") - def kernel(source, destination, items_per_thread): - storage = coop.TempStorage() - alias = storage - chained = alias - destination[cuda.threadIdx.x] = helper( - source, chained, items_per_thread - ) - - assert _compile( - kernel, - types.int32[::1], - types.int32[::1], - types.IntegerLiteral(items_per_thread), - ).metadata["ltoir"] - - -def test_none_alias_reports_descriptor_error_before_type_inference(): - @cuda.jit(chip="sm_90") - def kernel(source, destination): - empty = None - storage = coop.TempStorage() - if source[0] > 0: - storage = empty - payload = coop.ThreadData(items_per_thread=1, dtype=types.int32) - coop.load(coop.this_block(), source, payload, temp_storage=storage) - destination[cuda.threadIdx.x] = payload[0] - - with pytest.raises( - (GroupRewriteError, TypingError), match="TempStorage.*on every path" - ): - _compile(kernel, types.int32[::1], types.int32[::1]) - - -def test_standalone_primitive_helper_reports_inline_requirement(): - @cuda.jit(device=True, inline="never") - def standalone(destination, value): - coop.store(coop.this_block(), destination, value) - - @cuda.jit(chip="sm_90") - def kernel(source, destination): - standalone(destination, source[cuda.threadIdx.x]) - - with pytest.raises( - (GroupRewriteError, TypingError), match="standalone.*must be inlined" - ): - _compile(kernel, types.int32[::1], types.int32[::1]) - - -@pytest.mark.parametrize( - "payload_kind", ["thread_data", "local_array", "local_array_keyword"] -) -def test_literal_unroll_cannot_determine_cooperative_payload_shape( - payload_kind, -): - @cuda.jit(chip="sm_90") - def kernel(source, destination): - for count in literal_unroll((1, 2)): - if payload_kind == "thread_data": - payload = coop.ThreadData( - items_per_thread=count, dtype=types.int32 - ) - elif payload_kind == "local_array": - payload = cuda.local.array(count, types.int32) - else: - payload = cuda.local.array(shape=count, dtype=types.int32) - coop.load(coop.this_block(), source, payload) - destination[cuda.threadIdx.x] = payload[0] - - with pytest.raises( - (GroupRewriteError, TypingError), - match="does not support literal_unroll values", - ): - _compile(kernel, types.int32[::1], types.int32[::1]) - - -def test_unrelated_literal_unroll_can_coexist_with_cooperative_calls(): - @cuda.jit(chip="sm_90") - def kernel(destination): - value = 0 - for item in literal_unroll((1, 2)): - value += item - coop.store(coop.this_block(), destination, value) - - assert _compile(kernel, types.int64[::1]).metadata["ltoir"] - - -def test_storage_free_operation_can_coexist_with_user_dynamic_shared_memory(): - @cuda.jit(chip="sm_90") - def kernel(destination): - mine = cuda.shared.array(0, types.int32) - mine[cuda.threadIdx.x] = 42 - coop.store(coop.this_block(), destination, mine[cuda.threadIdx.x]) - - assert _compile(kernel, types.int32[::1]).metadata["ltoir"] - - -def test_literal_unroll_cannot_determine_cooperative_selector(): - @cuda.jit(chip="sm_90") - def kernel(source, destination): - for algorithm in literal_unroll(("direct", "striped")): - coop.store( - coop.this_block(), - destination, - source[cuda.threadIdx.x], - algorithm=algorithm, - ) - - with pytest.raises( - (GroupRewriteError, TypingError), - match="does not support literal_unroll values", - ): - _compile(kernel, types.int32[::1], types.int32[::1]) - - -@pytest.mark.parametrize( - "shape", - [np.int32(128), np.int64(128), (np.int32(128),)], - ids=["int32", "int64", "tuple"], -) -def test_numpy_integer_shared_shapes_coexist_with_cooperative_storage(shape): - @cuda.jit(chip="sm_90") - def kernel(source, destination): - mine = cuda.shared.array(shape, types.int32) - thread = cuda.threadIdx.x - mine[thread] = source[thread] - payload = coop.ThreadData(items_per_thread=1, dtype=types.int32) - payload[0] = mine[thread] - coop.store( - coop.this_block(), destination, payload, algorithm="transpose" - ) - - assert _compile(kernel, types.int32[::1], types.int32[::1]).metadata[ - "ltoir" - ] - - -@pytest.mark.parametrize( - "dynamic_backing", [False, True], ids=["static", "dynamic"] -) -@pytest.mark.parametrize("user_shape", ["static", "zero", "runtime"]) -def test_inlined_user_shared_allocation_is_checked_after_inlining( - monkeypatch, dynamic_backing, user_shape -): - from cuda.coop.numba_mlir._compiler import _rewrite_storage - - monkeypatch.setattr( - _rewrite_storage, - "_query_device_shared_memory_limits", - lambda: { - "max_default_shared_memory_per_block": 48 * 1024, - "max_optin_shared_memory_per_block": 96 * 1024, - }, - ) - size_in_bytes = 64 * 1024 if dynamic_backing else None - shape = 32 if user_shape == "static" else 0 - - if user_shape == "runtime": - - @cuda.jit(device=True) - def user_allocation(count): - allocate = cuda.shared.array - alias = allocate - return alias(count, types.int32) - - else: - - @cuda.jit(device=True) - def user_allocation(count): - allocate = cuda.shared.array - alias = allocate - return alias(shape, types.int32) - - @cuda.jit(chip="sm_90") - def kernel(source, destination): - mine = user_allocation(source[0]) - thread = cuda.threadIdx.x - mine[thread] = source[thread] - storage = coop.TempStorage(size_in_bytes) - payload = coop.ThreadData(items_per_thread=1, dtype=types.int32) - payload[0] = mine[thread] - coop.store( - coop.this_block(), - destination, - payload, - algorithm="transpose", - temp_storage=storage, - ) - - if not dynamic_backing and user_shape == "static": - assert _compile(kernel, types.int32[::1], types.int32[::1]).metadata[ - "ltoir" - ] - else: - with pytest.raises( - CoopSinglePhaseRewriteError, - match=( - r"shared-memory backing.*test_storage_diagnostics.py.*" - r"would alias" - ), - ): - _compile(kernel, types.int32[::1], types.int32[::1]) - - -@pytest.mark.parametrize("argument", ["storage", "group_by"]) -def test_literal_unroll_cannot_determine_cooperative_storage_or_partition( - argument, -): - @cuda.jit(chip="sm_90") - def kernel(source, destination): - for count in literal_unroll((32, 64)): - block = coop.this_block() - if argument == "storage": - storage = coop.TempStorage(count) - payload = coop.ThreadData(items_per_thread=1, dtype=types.int32) - payload[0] = source[cuda.threadIdx.x] - coop.store( - block, - destination, - payload, - algorithm="transpose", - temp_storage=storage, - ) - else: - group = block.group_by(count) - coop.store(group, destination, source[cuda.threadIdx.x]) - - with pytest.raises( - (GroupRewriteError, TypingError), - match="does not support literal_unroll values", - ): - _compile(kernel, types.int32[::1], types.int32[::1]) - - -def test_implicit_oversized_storage_rejects_user_static_shared_allocation( - monkeypatch, -): - from cuda.coop.numba_mlir._compiler import _rewrite_storage - - monkeypatch.setattr( - _rewrite_storage, - "_query_device_shared_memory_limits", - lambda: { - "max_default_shared_memory_per_block": 48 * 1024, - "max_optin_shared_memory_per_block": 96 * 1024, - }, - ) - - @cuda.jit(chip="sm_90") - def kernel(source, destination): - mine = cuda.shared.array(1024, types.int32) - thread = cuda.threadIdx.x - mine[thread] = source[thread] - payload = coop.ThreadData(items_per_thread=16, dtype=types.int32) - for item in range(16): - payload[item] = mine[thread] - coop.store( - coop.this_block(), destination, payload, algorithm="transpose" - ) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match=( - "dynamic shared-memory backing.*static cuda.shared.array.*" - "would alias" - ), - ): - _compile(kernel, types.int32[::1], types.int32[::1], block=(1024, 1, 1)) diff --git a/python/cuda_coop/tests/backends/numba_mlir/compile/test_warp_load_store_compile.py b/python/cuda_coop/tests/backends/numba_mlir/compile/test_warp_load_store_compile.py deleted file mode 100644 index b90f2e21456d..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/compile/test_warp_load_store_compile.py +++ /dev/null @@ -1,637 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Compile physical and logical Warp providers with GPUs hidden. - -The fixed device supplies target metadata; NVRTC and nvJitLink remain real. -Checks cover group storage, synchronization scope, and public-call linking. -""" - -from __future__ import annotations - -import os -from types import SimpleNamespace - -import pytest - -pytest.importorskip("numba_cuda_mlir") - -from numba_cuda_mlir import types - -from cuda.coop._core import ArgumentBinding, SynchronizationScope -from cuda.coop._core.warp import ( - make_warp_load_specialization, - make_warp_store_specialization, -) -from cuda.coop.numba_mlir import _types -from cuda.coop.numba_mlir._compiler import _nvrtc -from cuda.coop.numba_mlir._compiler._operations import StorageABI -from cuda.coop.numba_mlir._lowering._core import NumbaMlirCoreAdapter -from cuda.coop.numba_mlir._lowering._load_store import _load_store_value_abis - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.compile] - -_FIXED_COMPUTE_CAPABILITY = (9, 0) -_BLOCK_THREADS = 64 -_WARP_THREADS = 32 -_ITEMS_PER_THREAD = 2 -_ALGORITHMS = ("direct", "striped", "vectorize", "transpose") -_DTYPES = ( - pytest.param(types.int8, id="int8"), - pytest.param(types.uint8, id="uint8"), - pytest.param(types.int16, id="int16"), - pytest.param(types.uint16, id="uint16"), - pytest.param(types.int32, id="int32"), - pytest.param(types.uint32, id="uint32"), - pytest.param(types.int64, id="int64"), - pytest.param(types.uint64, id="uint64"), - pytest.param(types.float32, id="float32"), - pytest.param(types.float64, id="float64"), -) - - -@pytest.fixture(autouse=True) -def _fixed_current_device( - monkeypatch: pytest.MonkeyPatch, -) -> list[tuple[int, int]]: - """Hide runtime discovery while leaving NVRTC and nvJitLink real.""" - - assert os.environ.get("CUDA_VISIBLE_DEVICES") == "", ( - "the Numba-CUDA-MLIR compile stage must hide all CUDA devices" - ) - queries: list[tuple[int, int]] = [] - - def current_device() -> SimpleNamespace: - queries.append(_FIXED_COMPUTE_CAPABILITY) - return SimpleNamespace(compute_capability=_FIXED_COMPUTE_CAPABILITY) - - monkeypatch.setattr(_types.cuda, "get_current_device", current_device) - return queries - - -@pytest.fixture(scope="module") -def compile_context() -> _nvrtc.CompileContext: - return _nvrtc.resolve_compile_context() - - -def _algorithm( - context: _nvrtc.CompileContext, - *, - operation: str, - algorithm: str, - dtype=types.int32, - threads_in_warp: int = _WARP_THREADS, -) -> _types.Algorithm: - """Build a Warp provider with a runtime valid count and offset. - - Load also takes a runtime out-of-bounds default. Only transpose needs - scratch. Its provider has two entry points. One takes a caller - TempStorage pointer. The ``_alloc`` entry point declares one shared slot - per logical warp. It synchronizes that warp after the call so the slot - can be reused. - """ - valid_items = ArgumentBinding.runtime() - oob_default = ( - ArgumentBinding.runtime() - if operation == "load" - else ArgumentBinding.omitted() - ) - adapter = NumbaMlirCoreAdapter( - value_abis=_load_store_value_abis( - dtype=dtype, - items_per_thread=_ITEMS_PER_THREAD, - threads_in_warp=threads_in_warp, - valid_items=valid_items, - oob_default=oob_default, - ) - ) - factory = ( - make_warp_load_specialization - if operation == "load" - else make_warp_store_specialization - ) - specialization = factory( - dtype=adapter.core_dtype(dtype), - items_per_thread=_ITEMS_PER_THREAD, - threads_in_warp=threads_in_warp, - algorithm=algorithm, - valid_items=valid_items, - oob_default=oob_default, - include_pointer_offset=True, - ) - storage_free = algorithm != "transpose" - specialization = adapter.materialize( - specialization, - storage_abi=( - StorageABI.NONE if storage_free else StorageABI.LEADING_POINTER - ), - execution_scope=SynchronizationScope.WARP, - synchronization_scope=( - SynchronizationScope.NONE - if storage_free - else SynchronizationScope.WARP - ), - extra_type_definitions=(_types.numba_type_to_wrapper(dtype),), - ) - specialization._compile_context = context - specialization.logical_warp_threads = threads_in_warp - specialization.block_threads = _BLOCK_THREADS - return specialization - - -def _source(algorithm: _types.Algorithm) -> str: - source, _support_lto_irs, _storage_symbols, _udf_declarations = ( - algorithm._source_code() - ) - return source - - -def _production_compile_environment(monkeypatch: pytest.MonkeyPatch): - """Fix device queries while retaining real compilation and linking.""" - import numba_cuda_mlir.tools as numba_mlir_tools - from numba_cuda_mlir import cuda as compiler_cuda - - fixed_device = SimpleNamespace(compute_capability=_FIXED_COMPUTE_CAPABILITY) - - def fixed_compute_capability(as_type=str): - assert as_type in (str, tuple) - return _FIXED_COMPUTE_CAPABILITY if as_type is tuple else "sm_90" - - monkeypatch.setattr( - numba_mlir_tools, - "get_gpu_compute_capability", - fixed_compute_capability, - ) - monkeypatch.setattr( - compiler_cuda, "get_current_device", lambda: fixed_device - ) - return compiler_cuda - - -def _production_launch_config_key() -> tuple[tuple[str, object], ...]: - """Supply exact block metadata for two physical warps.""" - return ( - ("grid", (1, 1, 1)), - ("block", (_BLOCK_THREADS, 1, 1)), - ("sharedmem", 0), - ("cluster", None), - ) - - -@pytest.mark.parametrize( - "threads_in_warp", (32, 8), ids=("physical", "logical-8") -) -def test_all_warp_algorithms_compile_with_scope_owned_storage( - compile_context: _nvrtc.CompileContext, - _fixed_current_device: list[tuple[int, int]], - threads_in_warp: int, -) -> None: - algorithms = { - (operation, name): _algorithm( - compile_context, - operation=operation, - algorithm=name, - threads_in_warp=threads_in_warp, - ) - for operation in ("load", "store") - for name in _ALGORITHMS - } - - for (operation, name), specialization in algorithms.items(): - source = _source(specialization) - assert f"::cub::WARP_{operation.upper()}_{name.upper()}" in source - tile_items = threads_in_warp * _ITEMS_PER_THREAD - assert ( - f"if (num_valid_items < 0 || num_valid_items > {tile_items})" - in source - ) - memory_name = "src" if operation == "load" else "dst" - assert f"({memory_name} + offset)" in source - assert "__syncthreads" not in source - assert "bar.sync" not in source - if name == "transpose": - assert specialization.storage_abi is StorageABI.LEADING_POINTER - assert "TempStorage" in source - instances = _BLOCK_THREADS // threads_in_warp - assert f"temp_storages[{instances}]" in source - assert f"__coop_thread_rank / {threads_in_warp}" in source - if threads_in_warp == 32: - assert "__syncwarp();" in source - else: - assert "__syncwarp(" in source - assert "__syncwarp();" not in source - assert f"(*temp_storage).{operation.title()}(" in source - else: - assert specialization.storage_abi is StorageABI.NONE - assert "TempStorage" not in source - assert "temp_storages" not in source - assert "__syncwarp" not in source - assert f"().{operation.title()}(" in source - - bundle = _types.prepare_ltoir_bundle( - list(algorithms.values()), - ) - assert isinstance(bundle, bytes) - assert bundle - ptx = _types._ltoir_to_ptx( - bundle, - name="warp_load_store_synchronization_scope", - cc=10 * _FIXED_COMPUTE_CAPABILITY[0] + _FIXED_COMPUTE_CAPABILITY[1], - ) - assert "bar.sync" not in ptx - assert _fixed_current_device - for (_operation, name), specialization in algorithms.items(): - if name == "transpose": - assert specialization.temp_storage_bytes > 0 - assert specialization.temp_storage_alignment > 0 - else: - assert specialization.temp_storage_bytes == 0 - assert specialization.temp_storage_alignment == 1 - - -@pytest.mark.parametrize("dtype", _DTYPES) -def test_direct_multi_item_load_store_compiles_for_every_supported_dtype( - compile_context: _nvrtc.CompileContext, - dtype, -) -> None: - algorithms = [ - _algorithm( - compile_context, - operation=operation, - algorithm="direct", - dtype=dtype, - ) - for operation in ("load", "store") - ] - - for operation, specialization in zip(("load", "store"), algorithms): - source = _source(specialization) - assert f"::cub::WARP_{operation.upper()}_DIRECT" in source - assert f"().{operation.title()}(" in source - assert specialization.storage_abi is StorageABI.NONE - - bundle = _types.prepare_ltoir_bundle( - algorithms, - ) - assert isinstance(bundle, bytes) - assert bundle - assert all( - specialization.temp_storage_bytes == 0 for specialization in algorithms - ) - assert all( - specialization.temp_storage_alignment == 1 - for specialization in algorithms - ) - - -def test_logical_warp_widths_have_distinct_specializations_and_cache_keys( - compile_context: _nvrtc.CompileContext, -) -> None: - algorithms = [ - _algorithm( - compile_context, - operation="load", - algorithm="transpose", - threads_in_warp=width, - ) - for width in (1, 2, 4, 8, 16, 32) - ] - - assert [algorithm.logical_warp_threads for algorithm in algorithms] == [ - 1, - 2, - 4, - 8, - 16, - 32, - ] - assert ( - len( - { - algorithm._make_lto_ir_cache_key( - logical_warp_threads=algorithm.logical_warp_threads, - block_threads=_BLOCK_THREADS, - ) - for algorithm in algorithms - } - ) - == 6 - ) - for width, algorithm in zip((1, 2, 4, 8, 16, 32), algorithms): - source = _source(algorithm) - assert f", {width}>" in source - assert f"temp_storages[{_BLOCK_THREADS // width}]" in source - - bundle = _types.prepare_ltoir_bundle( - algorithms, - ) - assert bundle - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_production_routes_compile_common_and_qualified_warp_kernels( - monkeypatch: pytest.MonkeyPatch, items_per_thread -) -> None: - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - compiler_cuda = _production_compile_environment(monkeypatch) - - def qualified_kernel(algorithm: str): - @compiler_cuda.jit(chip="sm_90") - def kernel( - load_source, - store_source, - observed, - destination, - valid_items, - offset, - items_per_thread, - ): - thread = compiler_cuda.threadIdx.x - load_payload = qualified_coop.ThreadData( - items_per_thread, dtype=types.int32 - ) - qualified_coop.load( - qualified_coop.this_warp(), - load_source, - load_payload, - algorithm=algorithm, - valid_items=valid_items, - oob_default=types.int32(-17), - offset=offset, - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - observed[index] = load_payload[item] - payload[item] = store_source[index] - qualified_coop.store( - qualified_coop.this_warp(), - destination, - payload, - algorithm=algorithm, - valid_items=valid_items, - offset=offset, - ) - - return kernel - - def common_kernel(algorithm: str): - @compiler_cuda.jit(chip="sm_90") - def kernel( - load_source, - store_source, - observed, - destination, - valid_items, - offset, - items_per_thread, - ): - thread = compiler_cuda.threadIdx.x - load_payload = root_coop.ThreadData( - items_per_thread, dtype=types.int32 - ) - root_coop.load( - root_coop.this_warp(), - load_source, - load_payload, - algorithm=algorithm, - valid_items=valid_items, - oob_default=types.int32(-17), - offset=offset, - ) - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - observed[index] = load_payload[item] - payload[item] = store_source[index] - root_coop.store( - root_coop.this_warp(), - destination, - payload, - algorithm=algorithm, - valid_items=valid_items, - offset=offset, - ) - - return kernel - - dispatchers = [ - (qualified, algorithm, factory(algorithm)) - for qualified, factory in ( - (False, common_kernel), - (True, qualified_kernel), - ) - for algorithm in ("direct", "transpose") - ] - - signature = types.void( - types.int32[::1], - types.int32[::1], - types.int32[::1], - types.int32[::1], - types.int32, - types.int64, - types.IntegerLiteral(items_per_thread), - ) - launch_config_key = _production_launch_config_key() - for qualified, algorithm, dispatcher in dispatchers: - result = dispatcher._compile_launch_config_signature( - signature, - launch_config_key, - ) - assert isinstance(result.metadata["ltoir"], bytes) - assert result.metadata["ltoir"] - assert isinstance(result.metadata["cubin"], bytes) - assert result.metadata["cubin"] - assert result.metadata["linked_external_link_items"] - - ptx = next(iter(dispatcher.inspect_lto_ptx().values())) - assert ".visible .entry" in ptx - if algorithm == "transpose": - assert ".shared" in ptx - else: - assert ".shared" not in ptx - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_production_routes_compile_common_and_qualified_logical_warp_kernels( - monkeypatch: pytest.MonkeyPatch, items_per_thread -) -> None: - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - compiler_cuda = _production_compile_environment(monkeypatch) - - def qualified_kernel(algorithm: str): - @compiler_cuda.jit(chip="sm_90") - def kernel(source, destination, valid_items, offset, items_per_thread): - group = qualified_coop.this_warp().group_by(8) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - group, - source, - payload, - algorithm=algorithm, - valid_items=valid_items, - oob_default=types.int32(-17), - offset=offset, - ) - qualified_coop.store( - group, - destination, - payload, - algorithm=algorithm, - valid_items=valid_items, - offset=offset, - ) - - return kernel - - def common_kernel(algorithm: str): - @compiler_cuda.jit(chip="sm_90") - def kernel(source, destination, valid_items, offset, items_per_thread): - group = root_coop.this_warp().group_by(8) - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - group, - source, - payload, - algorithm=algorithm, - valid_items=valid_items, - oob_default=types.int32(-17), - offset=offset, - ) - root_coop.store( - group, - destination, - payload, - algorithm=algorithm, - valid_items=valid_items, - offset=offset, - ) - - return kernel - - signature = types.void( - types.int32[::1], - types.int32[::1], - types.int32, - types.int64, - types.IntegerLiteral(items_per_thread), - ) - launch_config_key = _production_launch_config_key() - for algorithm in ("direct", "transpose"): - for factory in (common_kernel, qualified_kernel): - dispatcher = factory(algorithm) - result = dispatcher._compile_launch_config_signature( - signature, - launch_config_key, - ) - assert result.metadata["ltoir"] - assert result.metadata["cubin"] - assert result.metadata["linked_external_link_items"] - ptx = next(iter(dispatcher.inspect_lto_ptx().values())) - assert "bar.sync" not in ptx - if algorithm == "transpose": - assert ".shared" in ptx - else: - assert ".shared" not in ptx - - -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_warp_scalar_literal_store_compiles_with_destination_dtype( - monkeypatch: pytest.MonkeyPatch, - qualified: bool, -) -> None: - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - compiler_cuda = _production_compile_environment(monkeypatch) - if qualified: - - @compiler_cuda.jit(chip="sm_90") - def kernel(destination): - qualified_coop.store( - qualified_coop.this_warp(), - destination, - 23, - algorithm="direct", - ) - - else: - - @compiler_cuda.jit(chip="sm_90") - def kernel(destination): - root_coop.store( - root_coop.this_warp(), - destination, - 23, - algorithm="direct", - ) - - result = kernel._compile_launch_config_signature( - types.void(types.int32[::1]), - _production_launch_config_key(), - ) - assert result.metadata["ltoir"] - assert result.metadata["cubin"] - - -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_warp_scalar_runtime_expression_rejects_implicit_narrowing( - monkeypatch: pytest.MonkeyPatch, - qualified: bool, -) -> None: - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - compiler_cuda = _production_compile_environment(monkeypatch) - if qualified: - - @compiler_cuda.jit(chip="sm_90") - def kernel(source, destination): - thread = compiler_cuda.threadIdx.x - qualified_coop.store( - qualified_coop.this_warp(), - destination, - source[thread] + 1, - algorithm="direct", - ) - - else: - - @compiler_cuda.jit(chip="sm_90") - def kernel(source, destination): - thread = compiler_cuda.threadIdx.x - root_coop.store( - root_coop.this_warp(), - destination, - source[thread] + 1, - algorithm="direct", - ) - - signature = types.void(types.int32[::1], types.int32[::1]) - with pytest.raises(TypeError, match="does not match payload dtype"): - kernel._compile_launch_config_signature( - signature, - _production_launch_config_key(), - ) diff --git a/python/cuda_coop/tests/backends/numba_mlir/runtime/__init__.py b/python/cuda_coop/tests/backends/numba_mlir/runtime/__init__.py deleted file mode 100644 index 8bbe3ce1ab86..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/runtime/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception diff --git a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_block_load_store.py b/python/cuda_coop/tests/backends/numba_mlir/runtime/test_block_load_store.py deleted file mode 100644 index b5a3015f2cbd..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_block_load_store.py +++ /dev/null @@ -1,1830 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Exercise Block Load/Store layouts, dtype handling, and storage on GPUs. - -Independent host references check each direction. They expose layout errors -that matching Load and Store mistakes could hide in a round trip. Cases cover -common and qualified APIs, scalar controls, payload aliases, and -multidimensional block shapes. Trap probes run in child processes so a device -fault leaves the pytest worker's context usable. -""" - -from __future__ import annotations - -import subprocess -import sys -import textwrap -from functools import cache -from pathlib import Path - -import numpy as np -import pytest - -cuda = pytest.importorskip("numba_cuda_mlir.cuda") -if not cuda.is_available(): - pytest.skip("requires a CUDA-capable runtime", allow_module_level=True) - -from numba_cuda_mlir import types - -import cuda.coop.numba_mlir as qualified_coop -from cuda import coop as root_coop - -assert qualified_coop.__file__ is not None -_QUALIFIED_COOP_ORIGIN = Path(qualified_coop.__file__).resolve() -_SAFE_PATH_FLAG = "-P" if sys.version_info >= (3, 11) else "-I" - -pytestmark = [ - pytest.mark.backend_numba_mlir, - pytest.mark.runtime, - pytest.mark.gpu, - pytest.mark.filterwarnings( - "ignore::numba_cuda_mlir.numba_cuda.core.errors.NumbaPerformanceWarning" - ), -] - -_THREADS = 32 -_ITEMS_PER_THREAD = 2 - - -_TILE_ITEMS = _THREADS * _ITEMS_PER_THREAD -_WIDE_THREADS = 64 -_WIDE_ITEMS_PER_THREAD = 4 -_WIDE_TILE_ITEMS = _WIDE_THREADS * _WIDE_ITEMS_PER_THREAD -_WIDE_PARTIAL_TILE_ITEMS = _WIDE_THREADS * _ITEMS_PER_THREAD -_BLOCK_SHAPES = (_THREADS, (8, 4), (4, 4, 2)) -_LOAD_OFFSET = 5 -_STORE_OFFSET = 7 -_GRID_STRIDE_BLOCKS = 3 -_GRID_STRIDE_ITEMS = 7 * _TILE_ITEMS + 11 -_ALGORITHMS = ( - "direct", - "striped", - "vectorize", - "transpose", - "warp_transpose", - "warp_transpose_timesliced", -) - -_DTYPES = ( - pytest.param(np.dtype(np.int8), types.int8, id="int8"), - pytest.param(np.dtype(np.uint8), types.uint8, id="uint8"), - pytest.param(np.dtype(np.int16), types.int16, id="int16"), - pytest.param(np.dtype(np.uint16), types.uint16, id="uint16"), - pytest.param(np.dtype(np.int32), types.int32, id="int32"), - pytest.param(np.dtype(np.uint32), types.uint32, id="uint32"), - pytest.param(np.dtype(np.int64), types.int64, id="int64"), - pytest.param(np.dtype(np.uint64), types.uint64, id="uint64"), - pytest.param(np.dtype(np.float32), types.float32, id="float32"), - pytest.param(np.dtype(np.float64), types.float64, id="float64"), -) - - -def _values(dtype: np.dtype, size: int, *, shift: int = 0) -> np.ndarray: - """Build varied inputs that also expose accidental dtype narrowing. - - Signed values include both signs. Float64 values carry low bits that - float32 would lose; 64-bit integers exceed the 32-bit range. ``shift`` - supplies different patterns for independent Load and Store inputs. - """ - - values = (np.arange(size, dtype=np.int64) * 3 + shift) % 97 - if dtype.kind in {"i", "f"}: - values = values - 48 - values = values.astype(dtype) - if dtype.kind == "f": - values *= dtype.type(0.25) - if dtype.itemsize == 8: - values += dtype.type(2**-30) - elif dtype.itemsize == 8: - values += np.where(values < 0, -(2**40), 2**40).astype(dtype) - return values - - -def _sentinel(dtype: np.dtype) -> object: - return dtype.type(211 if dtype.kind == "u" else -101) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "module", (root_coop, qualified_coop), ids=("common", "qualified") -) -def test_runtime_payload_index_reuses_one_specialization( - module, *, items_per_thread -): - @cuda.jit - def kernel(source, destination, index, items_per_thread): - payload = module.ThreadData(items_per_thread) - group = module.this_block() - module.load(group, source, payload) - module.store(group, destination, payload[index]) - - source = np.arange((_THREADS * items_per_thread), dtype=np.int32) - destination = np.full(_THREADS, -1, dtype=np.int32) - for index in range(items_per_thread): - kernel[1, _THREADS](source, destination, index, items_per_thread) - cuda.synchronize() - np.testing.assert_array_equal( - destination, source[index::items_per_thread] - ) - assert len(kernel._launch_config_overloads) == 1 - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "module", (root_coop, qualified_coop), ids=("portable", "qualified") -) -@pytest.mark.parametrize( - "dtype", (None, types.int32), ids=("inferred", "explicit") -) -@pytest.mark.parametrize("scope", ("block", "warp", "logical-warp")) -def test_load_mutates_original_payload_and_returns_none( - module, dtype, scope, *, items_per_thread -): - @cuda.jit - def kernel(source, destination, returned_none, items_per_thread): - thread = cuda.threadIdx.x - if scope == "block": - group = module.this_block() - elif scope == "warp": - group = module.this_warp() - else: - group = module.this_warp().group_by(8) - payload = module.ThreadData(items_per_thread, dtype) - result = module.load(group, source, payload) - returned_none[thread] = result is None - module.store(group, destination, payload) - - source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=19 - ) - destination = np.full_like(source, -1) - returned_none = np.zeros(_THREADS, dtype=np.int32) - kernel[1, _THREADS](source, destination, returned_none, items_per_thread) - - np.testing.assert_array_equal(destination, source) - np.testing.assert_array_equal(returned_none, np.ones_like(returned_none)) - - -@cache -def _full_load_kernel(numba_dtype, algorithm="direct"): - """Observe each thread's Load results through ordinary indexed stores. - - Keeping Store out of this helper lets the host check Load independently. - The cached dispatcher captures the requested dtype and algorithm. - """ - - @cuda.jit - def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x + cuda.blockDim.x * ( - cuda.threadIdx.y + cuda.blockDim.y * cuda.threadIdx.z - ) - payload = root_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - root_coop.load( - root_coop.this_block(), - source, - payload, - algorithm=algorithm, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@cache -def _full_store_kernel(numba_dtype, algorithm="direct"): - """Fill each thread's payload with ordinary reads before testing Store. - - Load cannot mask a Store layout error in this path. The cached dispatcher - captures the requested dtype and algorithm. - """ - - @cuda.jit - def kernel(source, destination, items_per_thread): - thread = cuda.threadIdx.x + cuda.blockDim.x * ( - cuda.threadIdx.y + cuda.blockDim.y * cuda.threadIdx.z - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - for item in range(items_per_thread): - payload[item] = source[thread * items_per_thread + item] - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm=algorithm, - ) - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize(("numpy_dtype", "numba_dtype"), _DTYPES) -def test_direct_load_matches_an_independent_oracle_for_every_dtype( - numpy_dtype, numba_dtype, *, items_per_thread -): - source = _values(numpy_dtype, (_THREADS * items_per_thread)) - observed = np.full( - (_THREADS * items_per_thread), _sentinel(numpy_dtype), dtype=numpy_dtype - ) - - _full_load_kernel(numba_dtype)[1, _THREADS]( - source, observed, items_per_thread - ) - - np.testing.assert_array_equal(observed, source) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize(("numpy_dtype", "numba_dtype"), _DTYPES) -def test_direct_multi_item_store_matches_an_independent_oracle_for_every_dtype( - numpy_dtype, numba_dtype, *, items_per_thread -): - source = _values(numpy_dtype, (_THREADS * items_per_thread), shift=11) - destination = np.full( - (_THREADS * items_per_thread), - _sentinel(numpy_dtype), - dtype=numpy_dtype, - ) - - _full_store_kernel(numba_dtype)[1, _THREADS]( - source, destination, items_per_thread - ) - - np.testing.assert_array_equal(destination, source) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("block_shape", _BLOCK_SHAPES, ids=("1d", "2d", "3d")) -def test_direct_load_uses_x_major_thread_order_for_exact_block_shape( - block_shape, *, items_per_thread -): - source = _values(np.dtype(np.int32), (_THREADS * items_per_thread)) - observed = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - - _full_load_kernel(types.int32)[1, block_shape]( - source, observed, items_per_thread - ) - - np.testing.assert_array_equal(observed, source) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("block_shape", _BLOCK_SHAPES, ids=("1d", "2d", "3d")) -def test_direct_store_uses_x_major_thread_order_for_exact_block_shape( - block_shape, *, items_per_thread -): - source = _values(np.dtype(np.int32), (_THREADS * items_per_thread), shift=7) - destination = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - - _full_store_kernel(types.int32)[1, block_shape]( - source, destination, items_per_thread - ) - - np.testing.assert_array_equal(destination, source) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_transpose_load_store_use_x_major_order_for_multidimensional_block( - *, items_per_thread -): - block_shape = (8, 4) - load_source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=5 - ) - store_source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=13 - ) - observed = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - destination = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - - _full_load_kernel(types.int32, "transpose")[1, block_shape]( - load_source, observed, items_per_thread - ) - _full_store_kernel(types.int32, "transpose")[1, block_shape]( - store_source, destination, items_per_thread - ) - - np.testing.assert_array_equal(observed, load_source) - np.testing.assert_array_equal(destination, store_source) - - -@cuda.jit -def _load_valid_prefix( - source, observed, valid_items, source_offset, items_per_thread -): - """Observe only valid slots when Load has no out-of-bounds default. - - The kernel writes only valid slots, so other observation entries keep - their host sentinel. Invalid payload slots are unspecified and are not - checked. - """ - - thread = cuda.threadIdx.x - payload = root_coop.ThreadData(items_per_thread, dtype=types.int32) - root_coop.load( - root_coop.this_block(), - source, - payload, - algorithm="direct", - valid_items=valid_items, - offset=source_offset, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - if index < valid_items: - observed[index] = payload[item] - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("valid_case", ("zero", "partial", "full")) -def test_load_reads_valid_prefix_at_an_independent_offset( - valid_case, *, items_per_thread -): - tile = _THREADS * items_per_thread - valid_items = {"zero": 0, "partial": tile - 9, "full": tile}[valid_case] - source = _values( - np.dtype(np.int32), _LOAD_OFFSET + (_THREADS * items_per_thread) - ) - observed = np.full((_THREADS * items_per_thread), 37, dtype=np.int32) - expected = np.full((_THREADS * items_per_thread), 37, dtype=np.int32) - expected[:valid_items] = source[_LOAD_OFFSET : _LOAD_OFFSET + valid_items] - - _load_valid_prefix[1, _THREADS]( - source, - observed, - np.int32(valid_items), - np.int64(_LOAD_OFFSET), - items_per_thread, - ) - - np.testing.assert_array_equal(observed, expected) - - -@cuda.jit -def _load_defaulting_invalid( - source, - observed, - valid_items, - oob_default, - source_offset, - items_per_thread, -): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - for item in range(items_per_thread): - payload[item] = 19 - qualified_coop.load( - qualified_coop.this_block(), - source, - payload, - algorithm="direct", - valid_items=valid_items, - oob_default=oob_default, - offset=source_offset, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("valid_case", ("zero", "partial")) -def test_load_defaults_invalid_slots(valid_case, *, items_per_thread): - tile = _THREADS * items_per_thread - valid_items = {"zero": 0, "partial": tile - 9}[valid_case] - source = _values( - np.dtype(np.int32), _LOAD_OFFSET + (_THREADS * items_per_thread) - ) - observed = np.full((_THREADS * items_per_thread), 37, dtype=np.int32) - expected = np.full((_THREADS * items_per_thread), -29, dtype=np.int32) - expected[:valid_items] = source[_LOAD_OFFSET : _LOAD_OFFSET + valid_items] - - _load_defaulting_invalid[1, _THREADS]( - source, - observed, - np.int32(valid_items), - np.int32(-29), - np.int64(_LOAD_OFFSET), - items_per_thread, - ) - - np.testing.assert_array_equal(observed, expected) - - -@cuda.jit -def _store_valid_prefix( - source, destination, valid_items, destination_offset, items_per_thread -): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - for item in range(items_per_thread): - payload[item] = source[thread * items_per_thread + item] - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm="direct", - valid_items=valid_items, - offset=destination_offset, - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("valid_case", ("zero", "partial", "full")) -def test_store_writes_only_the_valid_prefix_at_an_independent_offset( - valid_case, *, items_per_thread -): - tile = _THREADS * items_per_thread - valid_items = {"zero": 0, "partial": tile - 9, "full": tile}[valid_case] - source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=17 - ) - destination = np.full( - _STORE_OFFSET + (_THREADS * items_per_thread) + 3, - -41, - dtype=np.int32, - ) - expected = destination.copy() - expected[_STORE_OFFSET : _STORE_OFFSET + valid_items] = source[:valid_items] - - _store_valid_prefix[1, _THREADS]( - source, - destination, - np.int32(valid_items), - np.int64(_STORE_OFFSET), - items_per_thread, - ) - - np.testing.assert_array_equal(destination, expected) - - -@cache -def _algorithm_load_kernel( - algorithm: str, qualified: bool, numba_dtype=types.int32 -): - if qualified: - - @cuda.jit - def kernel( - source, - observed, - valid_items, - source_offset, - oob_default, - items_per_thread, - ): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - qualified_coop.load( - qualified_coop.this_block(), - source, - payload, - algorithm=algorithm, - valid_items=valid_items, - oob_default=oob_default, - offset=source_offset, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - else: - - @cuda.jit - def kernel( - source, - observed, - valid_items, - source_offset, - oob_default, - items_per_thread, - ): - thread = cuda.threadIdx.x - payload = root_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - root_coop.load( - root_coop.this_block(), - source, - payload, - algorithm=algorithm, - valid_items=valid_items, - oob_default=oob_default, - offset=source_offset, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@cache -def _algorithm_store_kernel( - algorithm: str, qualified: bool, numba_dtype=types.int32 -): - if qualified: - - @cuda.jit - def kernel( - source, - destination, - valid_items, - destination_offset, - items_per_thread, - ): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - for item in range(items_per_thread): - payload[item] = source[thread * items_per_thread + item] - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm=algorithm, - valid_items=valid_items, - offset=destination_offset, - ) - - else: - - @cuda.jit - def kernel( - source, - destination, - valid_items, - destination_offset, - items_per_thread, - ): - thread = cuda.threadIdx.x - payload = root_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - for item in range(items_per_thread): - payload[item] = source[thread * items_per_thread + item] - root_coop.store( - root_coop.this_block(), - destination, - payload, - algorithm=algorithm, - valid_items=valid_items, - offset=destination_offset, - ) - - return kernel - - -def _expected_loaded_payload( - source: np.ndarray, - *, - algorithm: str, - valid_items: int, - offset: int, - oob_default: int, - items_per_thread, -) -> np.ndarray: - """Compute a block's expected Load payload in thread-major slot order. - - Striped loads map successive items to successive thread-width stripes; the - other selectors expose blocked payload order. Read only valid tile - positions at ``offset`` and fill the remaining slots with ``oob_default``. - """ - - expected = np.full( - (_THREADS * items_per_thread), oob_default, dtype=source.dtype - ) - for thread in range(_THREADS): - for item in range(items_per_thread): - payload_index = thread * items_per_thread + item - tile_index = ( - thread + item * _THREADS - if algorithm == "striped" - else payload_index - ) - if tile_index < valid_items: - expected[payload_index] = source[offset + tile_index] - return expected - - -def _expected_stored_tile( - source: np.ndarray, - destination: np.ndarray, - *, - algorithm: str, - valid_items: int, - offset: int, - items_per_thread, -) -> np.ndarray: - """Apply the Store layout to a copy of the initial destination. - - ``source`` contains thread-major payload slots. Map them to blocked or - striped tile positions, then add ``offset``. Leave the invalid suffix and - surrounding guard entries unchanged. - """ - - expected = destination.copy() - for thread in range(_THREADS): - for item in range(items_per_thread): - payload_index = thread * items_per_thread + item - tile_index = ( - thread + item * _THREADS - if algorithm == "striped" - else payload_index - ) - if tile_index < valid_items: - expected[offset + tile_index] = source[payload_index] - return expected - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("algorithm", _ALGORITHMS) -@pytest.mark.parametrize("valid_case", ("zero", "partial", "full")) -def test_each_block_load_algorithm_matches_its_layout_oracle( - qualified, algorithm, valid_case, *, items_per_thread -): - tile = _THREADS * items_per_thread - valid_items = {"zero": 0, "partial": tile - 9, "full": tile}[valid_case] - source = _values( - np.dtype(np.int32), - _LOAD_OFFSET + (_THREADS * items_per_thread), - shift=31, - ) - observed = np.full((_THREADS * items_per_thread), 71, dtype=np.int32) - expected = _expected_loaded_payload( - source, - algorithm=algorithm, - valid_items=valid_items, - offset=_LOAD_OFFSET, - oob_default=-29, - items_per_thread=items_per_thread, - ) - - _algorithm_load_kernel(algorithm, qualified)[1, _THREADS]( - source, - observed, - np.int32(valid_items), - np.int64(_LOAD_OFFSET), - np.int32(-29), - items_per_thread, - ) - - np.testing.assert_array_equal(observed, expected) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("algorithm", _ALGORITHMS) -@pytest.mark.parametrize("valid_case", ("zero", "partial", "full")) -def test_each_block_store_algorithm_matches_its_layout_oracle( - qualified, algorithm, valid_case, *, items_per_thread -): - tile = _THREADS * items_per_thread - valid_items = {"zero": 0, "partial": tile - 9, "full": tile}[valid_case] - source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=37 - ) - destination = np.full( - _STORE_OFFSET + (_THREADS * items_per_thread) + 3, -41, dtype=np.int32 - ) - expected = _expected_stored_tile( - source, - destination, - algorithm=algorithm, - valid_items=valid_items, - offset=_STORE_OFFSET, - items_per_thread=items_per_thread, - ) - - _algorithm_store_kernel(algorithm, qualified)[1, _THREADS]( - source, - destination, - np.int32(valid_items), - np.int64(_STORE_OFFSET), - items_per_thread, - ) - - np.testing.assert_array_equal(destination, expected) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize(("numpy_dtype", "numba_dtype"), _DTYPES) -@pytest.mark.parametrize("algorithm", ("striped", "transpose")) -def test_non_direct_partial_load_store_matches_dtype_oracles( - numpy_dtype, numba_dtype, algorithm, *, items_per_thread -): - valid_items = (_THREADS * items_per_thread) - 9 - load_source = _values( - numpy_dtype, _LOAD_OFFSET + (_THREADS * items_per_thread), shift=31 - ) - store_source = _values(numpy_dtype, (_THREADS * items_per_thread), shift=37) - sentinel = _sentinel(numpy_dtype) - default = numpy_dtype.type(3.25 if numpy_dtype.kind == "f" else 103) - observed = np.full( - (_THREADS * items_per_thread), sentinel, dtype=numpy_dtype - ) - destination = np.full( - _STORE_OFFSET + (_THREADS * items_per_thread) + 3, - sentinel, - dtype=numpy_dtype, - ) - expected_load = _expected_loaded_payload( - load_source, - algorithm=algorithm, - valid_items=valid_items, - offset=_LOAD_OFFSET, - oob_default=default, - items_per_thread=items_per_thread, - ) - expected_store = _expected_stored_tile( - store_source, - destination, - algorithm=algorithm, - valid_items=valid_items, - offset=_STORE_OFFSET, - items_per_thread=items_per_thread, - ) - - _algorithm_load_kernel(algorithm, False, numba_dtype)[1, _THREADS]( - load_source, - observed, - np.int32(valid_items), - np.int64(_LOAD_OFFSET), - default, - items_per_thread, - ) - _algorithm_store_kernel(algorithm, True, numba_dtype)[1, _THREADS]( - store_source, - destination, - np.int32(valid_items), - np.int64(_STORE_OFFSET), - items_per_thread, - ) - - np.testing.assert_array_equal(observed, expected_load) - np.testing.assert_array_equal(destination, expected_store) - - -@cache -def _partial_transpose_load_kernel(algorithm: str, qualified: bool): - """Check valid transpose results without observing unspecified slots. - - These calls omit a default value. Only the valid prefix reaches the - observation array; its untouched host sentinel checks that boundary. - """ - - if qualified: - - @cuda.jit - def kernel(source, observed, valid_items, items_per_thread): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_block(), - source, - payload, - algorithm=algorithm, - valid_items=valid_items, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - if index < valid_items: - observed[index] = payload[item] - - else: - - @cuda.jit - def kernel(source, observed, valid_items, items_per_thread): - thread = cuda.threadIdx.x - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - root_coop.this_block(), - source, - payload, - algorithm=algorithm, - valid_items=valid_items, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - if index < valid_items: - observed[index] = payload[item] - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize( - "algorithm", - ("transpose", "warp_transpose", "warp_transpose_timesliced"), -) -def test_partial_transpose_load_reads_the_valid_prefix( - qualified, algorithm, *, items_per_thread -): - valid_items = (_WIDE_THREADS * items_per_thread) - 19 - source = _values( - np.dtype(np.int32), (_WIDE_THREADS * items_per_thread), shift=53 - ) - observed = np.full((_WIDE_THREADS * items_per_thread), 71, dtype=np.int32) - expected = observed.copy() - expected[:valid_items] = source[:valid_items] - - _partial_transpose_load_kernel(algorithm, qualified)[1, _WIDE_THREADS]( - source, observed, np.int32(valid_items), items_per_thread - ) - - np.testing.assert_array_equal(observed, expected) - - -@cache -def _unguarded_wide_load_store_kernel(algorithm: str, qualified: bool): - if qualified: - - @cuda.jit - def kernel( - load_source, store_source, observed, destination, items_per_thread - ): - thread = cuda.threadIdx.x - load_payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_block(), - load_source, - load_payload, - algorithm=algorithm, - ) - store_payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - observed[index] = load_payload[item] - store_payload[item] = store_source[index] - qualified_coop.store( - qualified_coop.this_block(), - destination, - store_payload, - algorithm=algorithm, - ) - - else: - - @cuda.jit - def kernel( - load_source, store_source, observed, destination, items_per_thread - ): - thread = cuda.threadIdx.x - load_payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - root_coop.this_block(), - load_source, - load_payload, - algorithm=algorithm, - ) - store_payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - observed[index] = load_payload[item] - store_payload[item] = store_source[index] - root_coop.store( - root_coop.this_block(), - destination, - store_payload, - algorithm=algorithm, - ) - - return kernel - - -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize( - "algorithm", ("vectorize", "warp_transpose_timesliced") -) -def test_unguarded_wide_load_store_executes_full_tile_path( - qualified, algorithm -): - load_source = _values(np.dtype(np.int32), _WIDE_TILE_ITEMS, shift=59) - store_source = _values(np.dtype(np.int32), _WIDE_TILE_ITEMS, shift=67) - observed = np.full(_WIDE_TILE_ITEMS, -1, dtype=np.int32) - destination = np.full(_WIDE_TILE_ITEMS, -1, dtype=np.int32) - - _unguarded_wide_load_store_kernel(algorithm, qualified)[1, _WIDE_THREADS]( - load_source, store_source, observed, destination, _WIDE_ITEMS_PER_THREAD - ) - - np.testing.assert_array_equal(observed, load_source) - np.testing.assert_array_equal(destination, store_source) - - -@cache -def _transpose_reuse_kernel(algorithm: str, dynamic: bool): - """Reuse one caller storage descriptor for Load and Store. - - Both variants request automatic synchronization. The dynamic case requests - 64 KiB of backing storage; the other infers capacity from the transpose - primitives. The test also checks each loaded payload directly. - """ - - if dynamic: - - @cuda.jit - def kernel(source, destination, observed, items_per_thread): - thread = cuda.threadIdx.x - storage = qualified_coop.TempStorage( - 64 * 1024, alignment=16, auto_sync=True - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_block(), - source, - payload, - algorithm=algorithm, - temp_storage=storage, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm=algorithm, - temp_storage=storage, - ) - - else: - - @cuda.jit - def kernel(source, destination, observed, items_per_thread): - thread = cuda.threadIdx.x - storage = qualified_coop.TempStorage( - sharing="shared", auto_sync=True - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_block(), - source, - payload, - algorithm=algorithm, - temp_storage=storage, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm=algorithm, - temp_storage=storage, - ) - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - ("algorithm", "dynamic"), - ( - ("transpose", False), - ("warp_transpose", False), - ("warp_transpose_timesliced", False), - ("transpose", True), - ), - ids=("transpose", "warp-transpose", "timesliced", "dynamic"), -) -def test_transpose_algorithms_reuse_caller_storage( - algorithm, dynamic, *, items_per_thread -): - source = _values( - np.dtype(np.int32), (_WIDE_THREADS * items_per_thread), shift=43 - ) - destination = np.full( - (_WIDE_THREADS * items_per_thread), -1, dtype=np.int32 - ) - observed = np.full((_WIDE_THREADS * items_per_thread), -1, dtype=np.int32) - - dispatcher = _transpose_reuse_kernel(algorithm, dynamic) - dispatcher[1, _WIDE_THREADS]( - source, destination, observed, items_per_thread - ) - - np.testing.assert_array_equal(destination, source) - np.testing.assert_array_equal(observed, source) - if dynamic: - compiled = next(iter(dispatcher._launch_config_overloads.values())) - assert compiled.metadata["required_dynamic_shared_memory"] == 64 * 1024 - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "module", (root_coop, qualified_coop), ids=("common", "qualified") -) -@pytest.mark.parametrize("alignment", (1, np.int64(32))) -@pytest.mark.parametrize("sharing", ("shared", "exclusive")) -def test_transpose_storage_honors_minimum_alignment( - module, alignment, sharing, *, items_per_thread -): - @cuda.jit - def kernel(source, destination, items_per_thread): - storage = module.TempStorage( - alignment=alignment, sharing=sharing, auto_sync=True - ) - payload = module.ThreadData(items_per_thread, dtype=types.int32) - module.load( - module.this_block(), - source, - payload, - algorithm="transpose", - temp_storage=storage, - ) - module.store( - module.this_block(), - destination, - payload, - algorithm="transpose", - temp_storage=storage, - ) - - source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=47 - ) - destination = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - kernel[1, _THREADS](source, destination, items_per_thread) - np.testing.assert_array_equal(destination, source) - - -@cuda.jit -def _common_grid_stride_load_store(source, destination, items_per_thread): - tile_offset = cuda.blockIdx.x * (_THREADS * items_per_thread) - grid_stride = cuda.gridDim.x * (_THREADS * items_per_thread) - while tile_offset < source.size: - valid_items = source.size - tile_offset - # Runtime counts trap rather than saturate. Clamp each grid-stride - # remainder to the exact tile accepted by this block. - valid_items = min(valid_items, (_THREADS * items_per_thread)) - payload = root_coop.ThreadData(items_per_thread, dtype=types.int32) - root_coop.load( - root_coop.this_block(), - source, - payload, - valid_items=valid_items, - offset=tile_offset, - ) - root_coop.store( - root_coop.this_block(), - destination, - payload, - valid_items=valid_items, - offset=tile_offset, - ) - tile_offset += grid_stride - - -@cuda.jit -def _qualified_grid_stride_load_store(source, destination, items_per_thread): - tile_offset = cuda.blockIdx.x * (_THREADS * items_per_thread) - grid_stride = cuda.gridDim.x * (_THREADS * items_per_thread) - while tile_offset < source.size: - valid_items = source.size - tile_offset - # Runtime counts trap rather than saturate. Clamp each grid-stride - # remainder to the exact tile accepted by this block. - valid_items = min(valid_items, (_THREADS * items_per_thread)) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_block(), - source, - payload, - valid_items=valid_items, - offset=tile_offset, - ) - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - valid_items=valid_items, - offset=tile_offset, - ) - tile_offset += grid_stride - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "kernel", - ( - pytest.param(_common_grid_stride_load_store, id="common"), - pytest.param(_qualified_grid_stride_load_store, id="qualified"), - ), -) -def test_multi_block_grid_stride_load_store_handles_a_partial_tail( - kernel, *, items_per_thread -): - source = _values( - np.dtype(np.int32), (7 * (_THREADS * items_per_thread) + 11), shift=19 - ) - destination = np.full( - (7 * (_THREADS * items_per_thread) + 11), -1, dtype=np.int32 - ) - - kernel[_GRID_STRIDE_BLOCKS, _THREADS](source, destination, items_per_thread) - - np.testing.assert_array_equal(destination, source) - - -def _run_invalid_runtime_valid_items_probe( - operation: str, valid_items: int, *, items_per_thread -) -> subprocess.CompletedProcess[str]: - """Launch an invalid count in a disposable process. - - The child checks that it imports the same ``cuda.coop`` package as the - test process, then compiles the requested operation and payload extent. - A device trap poisons its CUDA context, so the child contains that failure. - Return its status and output for the test's device-trap assertion. The - timeout bounds a stalled child; no kernel runs in the pytest worker here. - """ - - operation_body = textwrap.indent( - textwrap.dedent( - { - "load": """ - payload = root_coop.ThreadData( - items_per_thread=_ITEMS_PER_THREAD, dtype=types.int32 - ) - root_coop.load( - root_coop.this_block(), - source, - payload, - valid_items=valid_items, - ) - """, - "store": """ - payload = root_coop.ThreadData( - items_per_thread=_ITEMS_PER_THREAD, dtype=types.int32 - ) - for item in range(_ITEMS_PER_THREAD): - payload[item] = source[thread * _ITEMS_PER_THREAD + item] - root_coop.store( - root_coop.this_block(), - destination, - payload, - valid_items=valid_items, - ) - """, - }[operation] - ).strip(), - " " * 4, - ) - script = f"""\ -import numpy as np -import numba_cuda_mlir.cuda as cuda -from numba_cuda_mlir import types -from pathlib import Path - -import cuda.coop.numba_mlir as numba_coop -from cuda import coop as root_coop - -expected_origin = Path({str(_QUALIFIED_COOP_ORIGIN)!r}) -actual_origin = Path(numba_coop.__file__).resolve() -if actual_origin != expected_origin: - raise RuntimeError( - f"trap probe imported cuda.coop from {{actual_origin}}, " - f"expected {{expected_origin}}" - ) - -_THREADS = {_THREADS} -_ITEMS_PER_THREAD = {items_per_thread} - -@cuda.jit -def kernel(source, destination, valid_items): - thread = cuda.threadIdx.x -{operation_body} - -source = np.arange( - _THREADS * _ITEMS_PER_THREAD, - dtype=np.int32, -) -destination = np.full_like(source, -1) -kernel[1, _THREADS](source, destination, np.int64({valid_items})) -cuda.synchronize() -raise AssertionError("invalid valid_items did not trap") -""" - return subprocess.run( - [sys.executable, _SAFE_PATH_FLAG, "-B", "-c", script], - check=False, - capture_output=True, - text=True, - timeout=180, - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("operation", ("load", "store")) -@pytest.mark.parametrize( - "valid_items", - ( - pytest.param(-1, id="negative"), - pytest.param(_TILE_ITEMS + 1, id="beyond-tile"), - ), -) -def test_runtime_valid_items_out_of_range_traps_in_an_isolated_context( - operation, valid_items, *, items_per_thread -): - valid_items = ( - valid_items if valid_items < 0 else _THREADS * items_per_thread + 1 - ) - result = _run_invalid_runtime_valid_items_probe( - operation, valid_items, items_per_thread=items_per_thread - ) - output = result.stdout + result.stderr - - assert result.returncode != 0, output - # Drivers report PTX ``trap;`` as either an illegal instruction or the - # more general launch failure. Both are deterministic device-side faults. - assert any( - error in output - for error in ( - "CUDA_ERROR_ILLEGAL_INSTRUCTION", - "CUDA_ERROR_LAUNCH_FAILED", - ) - ), output - - source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=23 - ) - observed = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - _full_load_kernel(types.int32)[1, _THREADS]( - source, observed, items_per_thread - ) - np.testing.assert_array_equal(observed, source) - - -@cuda.jit -def _common_scalar_transpose_store(source, destination): - thread = cuda.threadIdx.x - root_coop.store( - root_coop.this_block(), - destination, - source[thread], - algorithm="transpose", - ) - - -def test_common_scalar_transpose_store_matches_an_independent_oracle(): - source = _values(np.dtype(np.int32), _THREADS, shift=23) - destination = np.full(_THREADS, -1, dtype=np.int32) - - _common_scalar_transpose_store[1, _THREADS](source, destination) - - np.testing.assert_array_equal(destination, source) - - -@pytest.mark.parametrize( - "module", (root_coop, qualified_coop), ids=("root", "qualified") -) -def test_thread_data_constructor_alias_across_branch(module, monkeypatch): - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - - original_apply = CoopSinglePhaseRewrite.apply - applications = 0 - - def bounded_apply(rewrite): - nonlocal applications - applications += 1 - assert applications < 10, ( - "constructor alias repeatedly matched the rewrite" - ) - return original_apply(rewrite) - - monkeypatch.setattr(CoopSinglePhaseRewrite, "apply", bounded_apply) - - @cuda.jit - def kernel(source, destination): - constructor = module.ThreadData - alias = constructor - if cuda.threadIdx.x < 16: - first = alias(1, types.int32) - module.load(module.this_block(), source, first) - destination[cuda.threadIdx.x] = first[0] - - source = _values(np.dtype(np.int32), _THREADS) - destination = np.full_like(source, -1) - kernel[1, _THREADS](source, destination) - - np.testing.assert_array_equal(destination[:16], source[:16]) - np.testing.assert_array_equal( - destination[16:], np.full_like(source[16:], -1) - ) - - -@cuda.jit -def _qualified_local_array_transpose_store( - source, destination, items_per_thread -): - thread = cuda.threadIdx.x - payload = cuda.local.array(shape=items_per_thread, dtype=types.int32) - for item in range(items_per_thread): - payload[item] = source[thread * items_per_thread + item] - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm="transpose", - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_qualified_transpose_store_accepts_local_array_payload( - *, items_per_thread -): - source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=29 - ) - destination = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - - _qualified_local_array_transpose_store[1, _THREADS]( - source, destination, items_per_thread - ) - - np.testing.assert_array_equal(destination, source) - - -@cuda.jit -def _qualified_untyped_load(source, observed, items_per_thread): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData(items_per_thread) - qualified_coop.load( - qualified_coop.this_block(), - source, - payload, - algorithm="direct", - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "module", (root_coop, qualified_coop), ids=("root", "qualified") -) -@pytest.mark.parametrize("projection", (0, -1), ids=("first", "last")) -@pytest.mark.parametrize("operation", ("load", "store")) -def test_untyped_thread_data_infers_dtype_through_tuple_aliases( - module, projection, operation, *, items_per_thread -): - @cuda.jit - def kernel(source, destination, items_per_thread): - payload = module.ThreadData(items_per_thread) - packed = (payload,) - alias = packed - projected = alias[projection] - thread = cuda.threadIdx.x - if operation == "load": - module.load(module.this_block(), source, projected) - for item in range(items_per_thread): - destination[thread * items_per_thread + item] = projected[item] - else: - for item in range(items_per_thread): - projected[item] = source[thread * items_per_thread + item] - module.store(module.this_block(), destination, projected) - - source = _values(np.dtype(np.int32), (_THREADS * items_per_thread)) - destination = np.full_like(source, -1) - kernel[1, _THREADS](source, destination, items_per_thread) - - np.testing.assert_array_equal(destination, source) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "module", (root_coop, qualified_coop), ids=("root", "qualified") -) -@pytest.mark.parametrize( - "dtype", (None, types.int32), ids=("inferred", "explicit") -) -def test_thread_data_item_extent_drives_a_kernel_loop( - module, dtype, *, items_per_thread -): - @cuda.jit - def kernel(source, destination, extents, items_per_thread): - payload = module.ThreadData(items_per_thread, dtype) - module.load(module.this_block(), source, payload) - alias = payload - thread = cuda.threadIdx.x - extents[thread] = alias.items_per_thread - for item in range(alias.items_per_thread): - destination[thread * items_per_thread + item] = alias[item] - - source = _values(np.dtype(np.int32), (_THREADS * items_per_thread)) - destination = np.full_like(source, -1) - extents = np.zeros(_THREADS, dtype=np.int32) - kernel[1, _THREADS](source, destination, extents, items_per_thread) - - np.testing.assert_array_equal(destination, source) - np.testing.assert_array_equal( - extents, np.full_like(extents, items_per_thread) - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_qualified_load_infers_an_untyped_payload(*, items_per_thread): - source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=29 - ) - observed = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - - _qualified_untyped_load[1, _THREADS](source, observed, items_per_thread) - - np.testing.assert_array_equal(observed, source) - - -@cuda.jit -def _qualified_untyped_store(source, destination, items_per_thread): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData(items_per_thread) - for item in range(items_per_thread): - payload[item] = source[thread * items_per_thread + item] - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm="direct", - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_qualified_store_infers_an_untyped_payload(*, items_per_thread): - source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=31 - ) - destination = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - - _qualified_untyped_store[1, _THREADS](source, destination, items_per_thread) - - np.testing.assert_array_equal(destination, source) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "module", (root_coop, qualified_coop), ids=("root", "qualified") -) -def test_direct_load_store_ignores_unused_temp_storage( - module, *, items_per_thread -): - @cuda.jit - def kernel(source, destination, items_per_thread): - storage = module.TempStorage( - 128 * 1024, alignment=16, sharing="exclusive", auto_sync=True - ) - payload = module.ThreadData(items_per_thread, dtype=types.int32) - module.load(module.this_block(), source, payload, temp_storage=storage) - module.store( - module.this_block(), destination, payload, temp_storage=storage - ) - - source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=31 - ) - destination = np.full_like(source, -1) - kernel[1, _THREADS](source, destination, items_per_thread) - - np.testing.assert_array_equal(destination, source) - - -@cache -def _divergent_storage_free_load_store(algorithm: str): - @cuda.jit - def kernel(source, destination, observed, items_per_thread): - thread = cuda.threadIdx.x - if thread & 1: - payload = root_coop.ThreadData(items_per_thread, dtype=types.int32) - root_coop.load( - root_coop.this_block(), - source, - payload, - algorithm=algorithm, - ) - root_coop.store( - root_coop.this_block(), - destination, - payload, - algorithm=algorithm, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("algorithm", ("direct", "striped", "vectorize")) -def test_storage_free_load_store_are_safe_in_divergent_control_flow( - algorithm, items_per_thread -): - source = _values(np.dtype(np.int32), _THREADS * items_per_thread, shift=43) - destination = np.full(_THREADS * items_per_thread, -1, dtype=np.int32) - observed = np.full(_THREADS * items_per_thread, -1, dtype=np.int32) - expected = np.full(_THREADS * items_per_thread, -1, dtype=np.int32) - if algorithm == "striped": - expected.reshape(items_per_thread, _THREADS)[:, 1::2] = source.reshape( - items_per_thread, _THREADS - )[:, 1::2] - else: - expected.reshape(_THREADS, items_per_thread)[1::2] = source.reshape( - _THREADS, items_per_thread - )[1::2] - - _divergent_storage_free_load_store(algorithm)[1, _THREADS]( - source, - destination, - observed, - items_per_thread, - ) - - np.testing.assert_array_equal(destination, expected) - if algorithm == "striped": - expected = expected.reshape(items_per_thread, _THREADS).T.ravel() - np.testing.assert_array_equal(observed, expected) - - -@cuda.jit -def _looped_shared_store(source, destination, observed, items_per_thread): - thread = cuda.threadIdx.x - storage = qualified_coop.TempStorage(sharing="shared") - for iteration in range(2): - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - tile_offset = iteration * (_THREADS * items_per_thread) - for item in range(items_per_thread): - payload[item] = source[ - tile_offset + thread * items_per_thread + item - ] - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm="direct", - temp_storage=storage, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - observed[tile_offset + index] = destination[index] - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_direct_store_accepts_one_shared_descriptor_in_a_loop( - *, items_per_thread -): - source = _values( - np.dtype(np.int32), 2 * (_THREADS * items_per_thread), shift=59 - ) - destination = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - observed = np.full(2 * (_THREADS * items_per_thread), -1, dtype=np.int32) - - _looped_shared_store[1, _THREADS]( - source, destination, observed, items_per_thread - ) - - np.testing.assert_array_equal(observed, source) - np.testing.assert_array_equal( - destination, source[(_THREADS * items_per_thread) :] - ) - - -@cuda.jit(device=True, inline=True) -def _inlined_common_load(source, observed, items_per_thread): - thread = cuda.threadIdx.x - payload = root_coop.ThreadData(items_per_thread, dtype=types.int32) - root_coop.load( - root_coop.this_block(), - source, - payload, - algorithm="direct", - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - -@cuda.jit -def _load_through_inlined_device_helper(source, observed, items_per_thread): - _inlined_common_load(source, observed, items_per_thread) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_common_load_is_planned_after_device_helper_inlining( - *, items_per_thread -): - source = _values( - np.dtype(np.int32), (_THREADS * items_per_thread), shift=67 - ) - observed = np.full((_THREADS * items_per_thread), -1, dtype=np.int32) - - _load_through_inlined_device_helper[1, _THREADS]( - source, observed, items_per_thread - ) - - np.testing.assert_array_equal(observed, source) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "module", (root_coop, qualified_coop), ids=("root", "qualified") -) -@pytest.mark.parametrize("alignment", [None, 1, 16, 32]) -def test_thread_data_alignment_with_inferred_load_store( - module, alignment, *, items_per_thread -): - @cuda.jit - def kernel(source, output, items_per_thread): - data = module.ThreadData(items_per_thread, alignment=alignment) - block = module.this_block() - module.load(block, source, data) - module.store(block, output, data) - - source = np.arange((_THREADS * items_per_thread), dtype=np.int32) - device_source = cuda.to_device(source) - output = cuda.device_array_like(device_source) - kernel[1, _THREADS](device_source, output, items_per_thread) - np.testing.assert_array_equal(output.copy_to_host(), source) - - -_LOOPED_THREADS = 1024 -_LOOPED_ITEMS_PER_THREAD = 4 -_LOOPED_TILE_ITEMS = _LOOPED_THREADS * _LOOPED_ITEMS_PER_THREAD -_LOOPED_TILES = 8 - - -@cache -def _looped_exclusive_store_kernel(manual_sync: bool): - # One transpose store call site reuses its exclusive slice on every loop - # iteration. Correctness requires a reuse barrier between iterations. - if manual_sync: - - @cuda.jit - def kernel(destination, items_per_thread): - thread = cuda.threadIdx.x - storage = qualified_coop.TempStorage( - sharing="exclusive", - auto_sync=False, - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - for tile in range(_LOOPED_TILES): - tile_offset = tile * _LOOPED_TILE_ITEMS - for item in range(items_per_thread): - payload[item] = ( - tile_offset + thread * items_per_thread + item - ) - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm="transpose", - offset=tile_offset, - temp_storage=storage, - ) - cuda.syncthreads() - - else: - - @cuda.jit - def kernel(destination, items_per_thread): - thread = cuda.threadIdx.x - storage = qualified_coop.TempStorage( - sharing="exclusive", auto_sync=True - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - for tile in range(_LOOPED_TILES): - tile_offset = tile * _LOOPED_TILE_ITEMS - for item in range(items_per_thread): - payload[item] = ( - tile_offset + thread * items_per_thread + item - ) - qualified_coop.store( - qualified_coop.this_block(), - destination, - payload, - algorithm="transpose", - offset=tile_offset, - temp_storage=storage, - ) - - return kernel - - -@pytest.mark.parametrize( - "manual_sync", - (False, True), - ids=("exclusive-auto-sync", "exclusive-manual-sync"), -) -def test_exclusive_storage_reused_by_a_looped_call_site_stays_ordered( - manual_sync, -): - destination = np.full( - _LOOPED_TILES * _LOOPED_TILE_ITEMS, -1, dtype=np.int32 - ) - - dispatcher = _looped_exclusive_store_kernel(manual_sync) - dispatcher[1, _LOOPED_THREADS](destination, _LOOPED_ITEMS_PER_THREAD) - - np.testing.assert_array_equal( - destination, - np.arange(_LOOPED_TILES * _LOOPED_TILE_ITEMS, dtype=np.int32), - ) - compiled = next(iter(dispatcher._launch_config_overloads.values())) - # Exactly one block barrier per iteration: the compiler's trailing reuse - # barrier with auto_sync=True, or the caller's explicit one otherwise. - assert compiled.metadata["mlir_module_str"].count("gpu.barrier") == 1 - - -@pytest.mark.parametrize("dtype", [np.int32, np.float64]) -@pytest.mark.parametrize("module", [root_coop, qualified_coop]) -def test_kernel_items_per_thread_specializes_repeated_launches(dtype, module): - @cuda.jit - def kernel(source, destination, items_per_thread): - payload = module.ThreadData(items_per_thread) - block = module.this_block() - module.load(block, source, payload) - for item in range(items_per_thread): - payload[item] += 7 - module.store(block, destination, payload) - - for items_per_thread in (1, 4, 1): - source = np.arange(_THREADS * items_per_thread, dtype=dtype) - destination = np.full_like(source, -1) - kernel[1, _THREADS](source, destination, items_per_thread) - np.testing.assert_array_equal(destination, source + 7) - - -@pytest.mark.parametrize("module", [root_coop, qualified_coop]) -def test_kernel_items_per_thread_without_cooperative_producer(module): - @cuda.jit - def kernel(destination, items_per_thread): - payload = module.ThreadData(items_per_thread) - offset = cuda.threadIdx.x * items_per_thread - for item in range(items_per_thread): - payload[item] = np.int32(offset + item) - for item in range(items_per_thread): - destination[offset + item] = payload[item] - - for items_per_thread in (1, 4, 1): - destination = np.full(_THREADS * items_per_thread, -1, dtype=np.int32) - kernel[1, _THREADS](destination, items_per_thread) - np.testing.assert_array_equal( - destination, np.arange(destination.size, dtype=np.int32) - ) diff --git a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_device_helpers_example.py b/python/cuda_coop/tests/backends/numba_mlir/runtime/test_device_helpers_example.py deleted file mode 100644 index 0faa7d3ca9e6..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_device_helpers_example.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import pytest - -cuda = pytest.importorskip("numba_cuda_mlir.cuda") -if not cuda.is_available(): - pytest.skip("requires a CUDA-capable runtime", allow_module_level=True) - -pytestmark = [ - pytest.mark.backend_numba_mlir, - pytest.mark.runtime, - pytest.mark.gpu, - pytest.mark.filterwarnings( - "ignore::numba_cuda_mlir.numba_cuda.core.errors.NumbaPerformanceWarning" - ), -] - - -def test_device_helpers_example(): - # coop-pg-device-helpers-begin - 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) - # coop-pg-device-helpers-end diff --git a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_load_example.py b/python/cuda_coop/tests/backends/numba_mlir/runtime/test_load_example.py deleted file mode 100644 index 2704319499f2..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_load_example.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import pytest - -cuda = pytest.importorskip("numba_cuda_mlir.cuda") -if not cuda.is_available(): - pytest.skip("requires a CUDA-capable runtime", allow_module_level=True) - -pytestmark = [ - pytest.mark.backend_numba_mlir, - pytest.mark.runtime, - pytest.mark.gpu, - pytest.mark.filterwarnings( - "ignore::numba_cuda_mlir.numba_cuda.core.errors.NumbaPerformanceWarning" - ), -] - - -def test_load_example(): - # example-begin - 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) - # example-end diff --git a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_programming_guide_examples.py b/python/cuda_coop/tests/backends/numba_mlir/runtime/test_programming_guide_examples.py deleted file mode 100644 index f43f7bb05595..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_programming_guide_examples.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import numpy as np -import pytest - -cuda = pytest.importorskip("numba_cuda_mlir.cuda") -if not cuda.is_available(): - pytest.skip("requires a CUDA-capable runtime", allow_module_level=True) - -import cuda.coop.numba_mlir as numba_coop # noqa: F401 -- activate the backend -from cuda import coop - -pytestmark = [ - pytest.mark.backend_numba_mlir, - pytest.mark.runtime, - pytest.mark.gpu, - pytest.mark.filterwarnings( - "ignore::numba_cuda_mlir.numba_cuda.core.errors.NumbaPerformanceWarning" - ), -] - - -def test_warp_copy(): - # coop-pg-warp-copy-begin - @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) - # coop-pg-warp-copy-end diff --git a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_qualified_load_store_examples.py b/python/cuda_coop/tests/backends/numba_mlir/runtime/test_qualified_load_store_examples.py deleted file mode 100644 index 1907cd4a5c41..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_qualified_load_store_examples.py +++ /dev/null @@ -1,54 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import pytest - -cuda = pytest.importorskip("numba_cuda_mlir.cuda") -if not cuda.is_available(): - pytest.skip("requires a CUDA-capable runtime", allow_module_level=True) - -pytestmark = [ - pytest.mark.backend_numba_mlir, - pytest.mark.runtime, - pytest.mark.gpu, - pytest.mark.filterwarnings( - "ignore::numba_cuda_mlir.numba_cuda.core.errors.NumbaPerformanceWarning" - ), -] - - -def test_qualified_load_store_example(): - # qualified-load-store-example-begin - 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) - # qualified-load-store-example-end diff --git a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_storage_examples.py b/python/cuda_coop/tests/backends/numba_mlir/runtime/test_storage_examples.py deleted file mode 100644 index a7811805c581..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_storage_examples.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import pytest - -cuda = pytest.importorskip("numba_cuda_mlir.cuda") -if not cuda.is_available(): - pytest.skip("requires a CUDA-capable runtime", allow_module_level=True) - -pytestmark = [ - pytest.mark.backend_numba_mlir, - pytest.mark.runtime, - pytest.mark.gpu, - pytest.mark.filterwarnings( - "ignore::numba_cuda_mlir.numba_cuda.core.errors.NumbaPerformanceWarning" - ), -] - - -def test_thread_data_example(): - # thread-data-example-begin - 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) - # thread-data-example-end - - -def test_temp_storage_example(): - # temp-storage-example-begin - 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) - # temp-storage-example-end diff --git a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_store_example.py b/python/cuda_coop/tests/backends/numba_mlir/runtime/test_store_example.py deleted file mode 100644 index d336aa65ed10..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_store_example.py +++ /dev/null @@ -1,55 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import pytest - -cuda = pytest.importorskip("numba_cuda_mlir.cuda") -if not cuda.is_available(): - pytest.skip("requires a CUDA-capable runtime", allow_module_level=True) - -pytestmark = [ - pytest.mark.backend_numba_mlir, - pytest.mark.runtime, - pytest.mark.gpu, - pytest.mark.filterwarnings( - "ignore::numba_cuda_mlir.numba_cuda.core.errors.NumbaPerformanceWarning" - ), -] - - -def test_store_example(): - # example-begin - 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) - # example-end diff --git a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_warp_load_store.py b/python/cuda_coop/tests/backends/numba_mlir/runtime/test_warp_load_store.py deleted file mode 100644 index 7dc3d276361d..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/runtime/test_warp_load_store.py +++ /dev/null @@ -1,2011 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Exercise independent physical and logical warp tiles on GPUs. - -Two physical warps and smaller logical groups expose incorrect group origins -and barriers that include other groups. Host references add each group's tile -origin separately from the explicit offset and check Load and Store with -independent input patterns. Guarded cases cover counts that differ between -groups while remaining uniform within each group. Trap and divergence probes -run in bounded child processes. -""" - -from __future__ import annotations - -import subprocess -import sys -import textwrap -from functools import cache -from pathlib import Path - -import numpy as np -import pytest - -cuda = pytest.importorskip("numba_cuda_mlir.cuda") -if not cuda.is_available(): - pytest.skip("requires a CUDA-capable runtime", allow_module_level=True) - -from numba_cuda_mlir import types - -import cuda.coop.numba_mlir as qualified_coop -from cuda import coop as root_coop - -pytestmark = [ - pytest.mark.backend_numba_mlir, - pytest.mark.runtime, - pytest.mark.gpu, - pytest.mark.filterwarnings( - "ignore::numba_cuda_mlir.numba_cuda.core.errors.NumbaPerformanceWarning" - ), -] - -_WARP_THREADS = 32 -_LOGICAL_WARP_THREADS = 8 -_BLOCK_THREADS = 2 * _WARP_THREADS -_ITEMS_PER_THREAD = 2 -_WARP_TILE_ITEMS = _WARP_THREADS * _ITEMS_PER_THREAD -_LOGICAL_TILE_ITEMS = _LOGICAL_WARP_THREADS * _ITEMS_PER_THREAD -_LOGICAL_GROUPS = _BLOCK_THREADS // _LOGICAL_WARP_THREADS -_BLOCK_ITEMS = _BLOCK_THREADS * _ITEMS_PER_THREAD -_LOAD_OFFSET = 5 -_STORE_OFFSET = 7 -_ALGORITHMS = ("direct", "striped", "vectorize", "transpose") -_DTYPES = ( - pytest.param(np.dtype(np.int8), types.int8, id="int8"), - pytest.param(np.dtype(np.uint8), types.uint8, id="uint8"), - pytest.param(np.dtype(np.int16), types.int16, id="int16"), - pytest.param(np.dtype(np.uint16), types.uint16, id="uint16"), - pytest.param(np.dtype(np.int32), types.int32, id="int32"), - pytest.param(np.dtype(np.uint32), types.uint32, id="uint32"), - pytest.param(np.dtype(np.int64), types.int64, id="int64"), - pytest.param(np.dtype(np.uint64), types.uint64, id="uint64"), - pytest.param(np.dtype(np.float32), types.float32, id="float32"), - pytest.param(np.dtype(np.float64), types.float64, id="float64"), -) -_GRID_STRIDE_ITEMS = 3 * _BLOCK_ITEMS - 17 -_GRID_STRIDE_BLOCKS = (_GRID_STRIDE_ITEMS + _BLOCK_ITEMS - 1) // _BLOCK_ITEMS -_QUALIFIED_COOP_ORIGIN = Path(qualified_coop.__file__).resolve() -_SAFE_PATH_FLAG = "-P" if sys.version_info >= (3, 11) else "-I" - - -def _values(size: int, *, shift: int = 0) -> np.ndarray: - return ((np.arange(size, dtype=np.int64) * 3 + shift) % 211 - 101).astype( - np.int32 - ) - - -def _dtype_values(dtype: np.dtype, size: int, *, shift: int = 0) -> np.ndarray: - """Build dtype-sensitive inputs that expose accidental narrowing. - - Float64 values retain bits below float32 precision, and 64-bit integers - exceed the 32-bit range. ``shift`` separates Load and Store patterns so - each direction can be checked independently. - """ - - values = (np.arange(size, dtype=np.int64) * 3 + shift) % 97 - if dtype.kind in {"i", "f"}: - values = values - 48 - values = values.astype(dtype) - if dtype.kind == "f": - values *= dtype.type(0.25) - if dtype.itemsize == 8: - values += dtype.type(2**-30) - elif dtype.itemsize == 8: - values += np.where(values < 0, -(2**40), 2**40).astype(dtype) - return values - - -def _dtype_sentinel(dtype: np.dtype) -> object: - return dtype.type(211 if dtype.kind == "u" else -101) - - -@cache -def _load_kernel(algorithm: str, qualified: bool, numba_dtype=types.int32): - if qualified: - selector = algorithm - - @cuda.jit - def kernel( - source, - observed, - valid_items, - source_offset, - oob_default, - items_per_thread, - ): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - qualified_coop.load( - qualified_coop.this_warp(), - source, - payload, - algorithm=selector, - valid_items=valid_items, - oob_default=oob_default, - offset=source_offset, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - else: - - @cuda.jit - def kernel( - source, - observed, - valid_items, - source_offset, - oob_default, - items_per_thread, - ): - thread = cuda.threadIdx.x - payload = root_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - root_coop.load( - root_coop.this_warp(), - source, - payload, - algorithm=algorithm, - valid_items=valid_items, - oob_default=oob_default, - offset=source_offset, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@cache -def _store_kernel(algorithm: str, qualified: bool, numba_dtype=types.int32): - if qualified: - selector = algorithm - - @cuda.jit - def kernel( - source, - destination, - valid_items, - destination_offset, - items_per_thread, - ): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - for item in range(items_per_thread): - payload[item] = source[thread * items_per_thread + item] - qualified_coop.store( - qualified_coop.this_warp(), - destination, - payload, - algorithm=selector, - valid_items=valid_items, - offset=destination_offset, - ) - - else: - - @cuda.jit - def kernel( - source, - destination, - valid_items, - destination_offset, - items_per_thread, - ): - thread = cuda.threadIdx.x - payload = root_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - for item in range(items_per_thread): - payload[item] = source[thread * items_per_thread + item] - root_coop.store( - root_coop.this_warp(), - destination, - payload, - algorithm=algorithm, - valid_items=valid_items, - offset=destination_offset, - ) - - return kernel - - -@cache -def _direct_dtype_load_store_kernel(numba_dtype, qualified: bool): - """Check both direct operations with independent input patterns. - - Observe Load with ordinary stores, then build the Store payload from - different input data. A matching error in the two primitives cannot cancel - through a Load-to-Store round trip. - """ - - if qualified: - - @cuda.jit - def kernel( - load_source, store_source, observed, destination, items_per_thread - ): - thread = cuda.threadIdx.x - load_payload = qualified_coop.ThreadData( - items_per_thread, dtype=numba_dtype - ) - qualified_coop.load( - qualified_coop.this_warp(), - load_source, - load_payload, - algorithm="direct", - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - observed[index] = load_payload[item] - payload[item] = store_source[index] - qualified_coop.store( - qualified_coop.this_warp(), - destination, - payload, - algorithm="direct", - ) - - else: - - @cuda.jit - def kernel( - load_source, store_source, observed, destination, items_per_thread - ): - thread = cuda.threadIdx.x - load_payload = root_coop.ThreadData( - items_per_thread, dtype=numba_dtype - ) - root_coop.load( - root_coop.this_warp(), - load_source, - load_payload, - algorithm="direct", - ) - payload = root_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - observed[index] = load_payload[item] - payload[item] = store_source[index] - root_coop.store( - root_coop.this_warp(), - destination, - payload, - algorithm="direct", - ) - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize(("numpy_dtype", "numba_dtype"), _DTYPES) -def test_direct_multi_item_load_store_matches_oracles_for_every_dtype( - qualified: bool, numpy_dtype: np.dtype, numba_dtype, *, items_per_thread -) -> None: - load_source = _dtype_values( - numpy_dtype, (_BLOCK_THREADS * items_per_thread), shift=7 - ) - store_source = _dtype_values( - numpy_dtype, (_BLOCK_THREADS * items_per_thread), shift=29 - ) - sentinel = _dtype_sentinel(numpy_dtype) - observed = np.full( - (_BLOCK_THREADS * items_per_thread), sentinel, dtype=numpy_dtype - ) - destination = np.full( - (_BLOCK_THREADS * items_per_thread), sentinel, dtype=numpy_dtype - ) - - _direct_dtype_load_store_kernel(numba_dtype, qualified)[1, _BLOCK_THREADS]( - load_source, store_source, observed, destination, items_per_thread - ) - - np.testing.assert_array_equal(observed, load_source) - np.testing.assert_array_equal(destination, store_source) - - -def _tile_index( - algorithm: str, lane: int, item: int, *, items_per_thread -) -> int: - """Map a lane's item to a blocked or striped index within its warp.""" - - if algorithm == "striped": - return lane + item * _WARP_THREADS - return lane * items_per_thread + item - - -def _expected_loaded_payload( - source: np.ndarray, - *, - algorithm: str, - valid_items: int, - offset: int, - oob_default: int, - items_per_thread, -) -> np.ndarray: - """Compute separate Load results for both physical warps. - - Keep observed payload slots in thread-major order. For each valid item, - add its warp's tile origin and the caller's offset to the layout index. - Each warp applies the count independently; other slots use the default. - """ - - expected = np.full( - (_BLOCK_THREADS * items_per_thread), oob_default, dtype=source.dtype - ) - for thread in range(_BLOCK_THREADS): - warp = thread // _WARP_THREADS - lane = thread % _WARP_THREADS - for item in range(items_per_thread): - payload_index = thread * items_per_thread + item - tile_index = _tile_index( - algorithm, lane, item, items_per_thread=items_per_thread - ) - if tile_index < valid_items: - expected[payload_index] = source[ - offset - + warp * (_WARP_THREADS * items_per_thread) - + tile_index - ] - return expected - - -def _expected_stored_payload( - source: np.ndarray, - destination: np.ndarray, - *, - algorithm: str, - valid_items: int, - offset: int, - items_per_thread, -) -> np.ndarray: - """Map per-thread Store inputs to each physical warp's destination tile. - - Add the warp origin and explicit offset after applying the layout. Begin - with the original destination so invalid items and guard regions remain - part of the reference check. - """ - - expected = destination.copy() - for thread in range(_BLOCK_THREADS): - warp = thread // _WARP_THREADS - lane = thread % _WARP_THREADS - for item in range(items_per_thread): - payload_index = thread * items_per_thread + item - tile_index = _tile_index( - algorithm, lane, item, items_per_thread=items_per_thread - ) - if tile_index < valid_items: - expected[ - offset - + warp * (_WARP_THREADS * items_per_thread) - + tile_index - ] = source[payload_index] - return expected - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("algorithm", _ALGORITHMS) -@pytest.mark.parametrize("valid_case", ("zero", "partial", "full")) -def test_each_warp_load_algorithm_matches_an_independent_two_warp_oracle( - qualified: bool, algorithm: str, valid_case: str, *, items_per_thread -) -> None: - tile = _WARP_THREADS * items_per_thread - valid_items = {"zero": 0, "partial": tile - 9, "full": tile}[valid_case] - source = _values( - _LOAD_OFFSET + (_BLOCK_THREADS * items_per_thread) + 3, shift=31 - ) - observed = np.full((_BLOCK_THREADS * items_per_thread), 71, dtype=np.int32) - expected = _expected_loaded_payload( - source, - algorithm=algorithm, - valid_items=valid_items, - offset=_LOAD_OFFSET, - oob_default=-29, - items_per_thread=items_per_thread, - ) - - _load_kernel(algorithm, qualified)[1, _BLOCK_THREADS]( - source, - observed, - np.int32(valid_items), - np.int64(_LOAD_OFFSET), - np.int32(-29), - items_per_thread, - ) - - np.testing.assert_array_equal(observed, expected) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("algorithm", _ALGORITHMS) -@pytest.mark.parametrize("valid_case", ("zero", "partial", "full")) -def test_each_warp_store_algorithm_masks_each_warp( - qualified: bool, algorithm: str, valid_case: str, *, items_per_thread -) -> None: - tile = _WARP_THREADS * items_per_thread - valid_items = {"zero": 0, "partial": tile - 9, "full": tile}[valid_case] - source = _values((_BLOCK_THREADS * items_per_thread), shift=43) - destination = np.full( - _STORE_OFFSET + (_BLOCK_THREADS * items_per_thread) + 3, - -41, - dtype=np.int32, - ) - expected = _expected_stored_payload( - source, - destination, - algorithm=algorithm, - valid_items=valid_items, - offset=_STORE_OFFSET, - items_per_thread=items_per_thread, - ) - - _store_kernel(algorithm, qualified)[1, _BLOCK_THREADS]( - source, - destination, - np.int32(valid_items), - np.int64(_STORE_OFFSET), - items_per_thread, - ) - - np.testing.assert_array_equal(destination, expected) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize(("numpy_dtype", "numba_dtype"), _DTYPES) -@pytest.mark.parametrize("algorithm", ("striped", "transpose")) -def test_non_direct_partial_load_store_matches_dtype_oracles( - numpy_dtype, numba_dtype, algorithm, *, items_per_thread -): - valid_items = (_WARP_THREADS * items_per_thread) - 9 - load_source = _dtype_values( - numpy_dtype, - _LOAD_OFFSET + (_BLOCK_THREADS * items_per_thread) + 3, - shift=31, - ) - store_source = _dtype_values( - numpy_dtype, (_BLOCK_THREADS * items_per_thread), shift=43 - ) - sentinel = _dtype_sentinel(numpy_dtype) - default = numpy_dtype.type(3.25 if numpy_dtype.kind == "f" else 103) - observed = np.full( - (_BLOCK_THREADS * items_per_thread), sentinel, dtype=numpy_dtype - ) - destination = np.full( - _STORE_OFFSET + (_BLOCK_THREADS * items_per_thread) + 3, - sentinel, - dtype=numpy_dtype, - ) - expected_load = _expected_loaded_payload( - load_source, - algorithm=algorithm, - valid_items=valid_items, - offset=_LOAD_OFFSET, - oob_default=default, - items_per_thread=items_per_thread, - ) - expected_store = _expected_stored_payload( - store_source, - destination, - algorithm=algorithm, - valid_items=valid_items, - offset=_STORE_OFFSET, - items_per_thread=items_per_thread, - ) - - _load_kernel(algorithm, False, numba_dtype)[1, _BLOCK_THREADS]( - load_source, - observed, - np.int32(valid_items), - np.int64(_LOAD_OFFSET), - default, - items_per_thread, - ) - _store_kernel(algorithm, True, numba_dtype)[1, _BLOCK_THREADS]( - store_source, - destination, - np.int32(valid_items), - np.int64(_STORE_OFFSET), - items_per_thread, - ) - - np.testing.assert_array_equal(observed, expected_load) - np.testing.assert_array_equal(destination, expected_store) - - -def _logical_tile_index( - algorithm: str, lane: int, item: int, width: int, *, items_per_thread -) -> int: - if algorithm == "striped": - return lane + item * width - return lane * items_per_thread + item - - -@cache -def _logical_load_store_kernel(algorithm: str, qualified: bool): - """Give each logical warp its own count and offset in a shared block. - - Each control value is uniform within its group. Load and Store use - different source data so host references check both directions. - """ - - if qualified: - load_algorithm = algorithm - store_algorithm = algorithm - - @cuda.jit - def kernel( - load_source, - store_source, - observed, - destination, - valid_by_group, - offset_by_group, - items_per_thread, - ): - thread = cuda.threadIdx.x - group_index = thread // _LOGICAL_WARP_THREADS - offset = offset_by_group[group_index] - group = qualified_coop.this_warp().group_by(_LOGICAL_WARP_THREADS) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - group, - load_source, - payload, - algorithm=load_algorithm, - valid_items=valid_by_group[group_index], - oob_default=types.int32(-127), - offset=offset, - ) - for item in range(items_per_thread): - payload_index = thread * items_per_thread + item - observed[payload_index] = payload[item] - payload[item] = store_source[payload_index] - qualified_coop.store( - group, - destination, - payload, - algorithm=store_algorithm, - valid_items=valid_by_group[group_index], - offset=offset, - ) - - else: - - @cuda.jit - def kernel( - load_source, - store_source, - observed, - destination, - valid_by_group, - offset_by_group, - items_per_thread, - ): - thread = cuda.threadIdx.x - group_index = thread // _LOGICAL_WARP_THREADS - offset = offset_by_group[group_index] - group = root_coop.this_warp().group_by(_LOGICAL_WARP_THREADS) - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - group, - load_source, - payload, - algorithm=algorithm, - valid_items=valid_by_group[group_index], - oob_default=types.int32(-127), - offset=offset, - ) - for item in range(items_per_thread): - payload_index = thread * items_per_thread + item - observed[payload_index] = payload[item] - payload[item] = store_source[payload_index] - root_coop.store( - group, - destination, - payload, - algorithm=algorithm, - valid_items=valid_by_group[group_index], - offset=offset, - ) - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("algorithm", _ALGORITHMS) -def test_logical_warp_algorithms_use_independent_group_tiles( - qualified: bool, algorithm: str, *, items_per_thread -) -> None: - offsets_by_group = np.array([0, 1, 3, 4, 6, 7, 9, 10], dtype=np.int64) - allocation_items = ( - int(offsets_by_group.max()) + (_BLOCK_THREADS * items_per_thread) + 3 - ) - load_source = _values(allocation_items, shift=47) - store_source = _values((_BLOCK_THREADS * items_per_thread), shift=53) - observed = np.full((_BLOCK_THREADS * items_per_thread), 19, dtype=np.int32) - destination = np.full(allocation_items, -31, dtype=np.int32) - valid_by_group = np.array( - [0, 1, 5, (_LOGICAL_WARP_THREADS * items_per_thread), 3, 11, 7, 15], - dtype=np.int32, - ) - np.minimum( - valid_by_group, - _LOGICAL_WARP_THREADS * items_per_thread, - out=valid_by_group, - ) - expected_observed = np.full( - (_BLOCK_THREADS * items_per_thread), -127, dtype=np.int32 - ) - expected_destination = destination.copy() - for thread in range(_BLOCK_THREADS): - group_index = thread // _LOGICAL_WARP_THREADS - lane = thread % _LOGICAL_WARP_THREADS - valid_items = valid_by_group[group_index] - offset = offsets_by_group[group_index] - for item in range(items_per_thread): - payload_index = thread * items_per_thread + item - tile_index = _logical_tile_index( - algorithm, - lane, - item, - _LOGICAL_WARP_THREADS, - items_per_thread=items_per_thread, - ) - if tile_index < valid_items: - memory_index = ( - offset - + group_index * (_LOGICAL_WARP_THREADS * items_per_thread) - + tile_index - ) - expected_observed[payload_index] = load_source[memory_index] - expected_destination[memory_index] = store_source[payload_index] - - _logical_load_store_kernel(algorithm, qualified)[1, _BLOCK_THREADS]( - load_source, - store_source, - observed, - destination, - valid_by_group, - offsets_by_group, - items_per_thread, - ) - - np.testing.assert_array_equal(observed, expected_observed) - np.testing.assert_array_equal(destination, expected_destination) - - -@cache -def _logical_partial_transpose_load_kernel(qualified: bool): - """Observe only valid items from each logical warp's transpose load. - - Omit a default and leave invalid payload slots unobserved. The host - sentinel belongs to the observation array, not the payload contract. - """ - - if qualified: - - @cuda.jit - def kernel(source, observed, valid_by_group, items_per_thread): - thread = cuda.threadIdx.x - group_index = thread // _LOGICAL_WARP_THREADS - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_warp().group_by(_LOGICAL_WARP_THREADS), - source, - payload, - algorithm="transpose", - valid_items=valid_by_group[group_index], - ) - for item in range(items_per_thread): - lane = thread % _LOGICAL_WARP_THREADS - if lane * items_per_thread + item < valid_by_group[group_index]: - observed[thread * items_per_thread + item] = payload[item] - - else: - - @cuda.jit - def kernel(source, observed, valid_by_group, items_per_thread): - thread = cuda.threadIdx.x - group_index = thread // _LOGICAL_WARP_THREADS - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - root_coop.this_warp().group_by(_LOGICAL_WARP_THREADS), - source, - payload, - algorithm="transpose", - valid_items=valid_by_group[group_index], - ) - for item in range(items_per_thread): - lane = thread % _LOGICAL_WARP_THREADS - if lane * items_per_thread + item < valid_by_group[group_index]: - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_logical_transpose_load_reads_valid_items_in_nonzero_groups( - qualified: bool, *, items_per_thread -) -> None: - source = _values((_BLOCK_THREADS * items_per_thread), shift=61) - observed = np.full((_BLOCK_THREADS * items_per_thread), 23, dtype=np.int32) - valid_by_group = np.array( - [0, 1, 5, (_LOGICAL_WARP_THREADS * items_per_thread), 3, 11, 7, 15], - dtype=np.int32, - ) - np.minimum( - valid_by_group, - _LOGICAL_WARP_THREADS * items_per_thread, - out=valid_by_group, - ) - expected = observed.copy() - for thread in range(_BLOCK_THREADS): - group_index = thread // _LOGICAL_WARP_THREADS - lane = thread % _LOGICAL_WARP_THREADS - valid_items = valid_by_group[group_index] - for item in range(items_per_thread): - payload_index = thread * items_per_thread + item - tile_index = lane * items_per_thread + item - if tile_index < valid_items: - expected[payload_index] = source[ - group_index * (_LOGICAL_WARP_THREADS * items_per_thread) - + tile_index - ] - - _logical_partial_transpose_load_kernel(qualified)[1, _BLOCK_THREADS]( - source, observed, valid_by_group, items_per_thread - ) - - np.testing.assert_array_equal(observed, expected) - - -@cache -def _logical_width_direct_kernel(width: int, qualified: bool): - if qualified: - - @cuda.jit - def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_warp().group_by(width), - source, - payload, - algorithm="direct", - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - else: - - @cuda.jit - def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - root_coop.this_warp().group_by(width), - source, - payload, - algorithm="direct", - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("width", (1, 2, 4, 8, 16, 32)) -def test_every_logical_warp_width_addresses_consecutive_tiles( - qualified: bool, width: int, *, items_per_thread -) -> None: - source = _values((_BLOCK_THREADS * items_per_thread), shift=67) - observed = np.full((_BLOCK_THREADS * items_per_thread), -1, dtype=np.int32) - - _logical_width_direct_kernel(width, qualified)[1, _BLOCK_THREADS]( - source, observed, items_per_thread - ) - - np.testing.assert_array_equal(observed, source) - - -@cache -def _logical_direct_dtype_load_store_kernel(numba_dtype, qualified: bool): - if qualified: - - @cuda.jit - def kernel( - load_source, store_source, observed, destination, items_per_thread - ): - thread = cuda.threadIdx.x - group = qualified_coop.this_warp().group_by(_LOGICAL_WARP_THREADS) - load_payload = qualified_coop.ThreadData( - items_per_thread, dtype=numba_dtype - ) - qualified_coop.load( - group, - load_source, - load_payload, - algorithm="direct", - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - observed[index] = load_payload[item] - payload[item] = store_source[index] - qualified_coop.store( - group, - destination, - payload, - algorithm="direct", - ) - - else: - - @cuda.jit - def kernel( - load_source, store_source, observed, destination, items_per_thread - ): - thread = cuda.threadIdx.x - group = root_coop.this_warp().group_by(_LOGICAL_WARP_THREADS) - load_payload = root_coop.ThreadData( - items_per_thread, dtype=numba_dtype - ) - root_coop.load( - group, - load_source, - load_payload, - algorithm="direct", - ) - payload = root_coop.ThreadData( - items_per_thread, - dtype=numba_dtype, - ) - for item in range(items_per_thread): - index = thread * items_per_thread + item - observed[index] = load_payload[item] - payload[item] = store_source[index] - root_coop.store( - group, - destination, - payload, - algorithm="direct", - ) - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize(("numpy_dtype", "numba_dtype"), _DTYPES) -def test_logical_direct_load_store_matches_every_dtype_oracle( - qualified: bool, numpy_dtype: np.dtype, numba_dtype, *, items_per_thread -) -> None: - load_source = _dtype_values( - numpy_dtype, (_BLOCK_THREADS * items_per_thread), shift=13 - ) - store_source = _dtype_values( - numpy_dtype, (_BLOCK_THREADS * items_per_thread), shift=37 - ) - sentinel = _dtype_sentinel(numpy_dtype) - observed = np.full( - (_BLOCK_THREADS * items_per_thread), sentinel, dtype=numpy_dtype - ) - destination = np.full( - (_BLOCK_THREADS * items_per_thread), sentinel, dtype=numpy_dtype - ) - - _logical_direct_dtype_load_store_kernel(numba_dtype, qualified)[ - 1, _BLOCK_THREADS - ](load_source, store_source, observed, destination, items_per_thread) - - np.testing.assert_array_equal(observed, load_source) - np.testing.assert_array_equal(destination, store_source) - - -@cache -def _partial_load_kernel(algorithm: str, qualified: bool): - """Observe only valid items according to each warp's selected layout. - - Without a default value, the test checks the valid prefix and leaves - invalid payload slots unobserved. The host reference separately applies - the same layout definition to each warp's tile origin. - """ - - striped = algorithm == "striped" - if qualified: - selector = algorithm - - @cuda.jit - def kernel(source, observed, valid_items, items_per_thread): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_warp(), - source, - payload, - algorithm=selector, - valid_items=valid_items, - ) - for item in range(items_per_thread): - lane = thread % _WARP_THREADS - tile_index = ( - lane + item * _WARP_THREADS - if striped - else lane * items_per_thread + item - ) - if tile_index < valid_items: - observed[thread * items_per_thread + item] = payload[item] - - else: - - @cuda.jit - def kernel(source, observed, valid_items, items_per_thread): - thread = cuda.threadIdx.x - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - root_coop.this_warp(), - source, - payload, - algorithm=algorithm, - valid_items=valid_items, - ) - for item in range(items_per_thread): - lane = thread % _WARP_THREADS - tile_index = ( - lane + item * _WARP_THREADS - if striped - else lane * items_per_thread + item - ) - if tile_index < valid_items: - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("algorithm", _ALGORITHMS) -def test_partial_load_reads_valid_items_for_each_layout_and_warp( - qualified: bool, algorithm: str, *, items_per_thread -) -> None: - valid_items = (_WARP_THREADS * items_per_thread) - 11 - source = _values((_BLOCK_THREADS * items_per_thread), shift=59) - observed = np.full((_BLOCK_THREADS * items_per_thread), 71, dtype=np.int32) - expected = observed.copy() - for thread in range(_BLOCK_THREADS): - warp = thread // _WARP_THREADS - lane = thread % _WARP_THREADS - for item in range(items_per_thread): - payload_index = thread * items_per_thread + item - tile_index = _tile_index( - algorithm, lane, item, items_per_thread=items_per_thread - ) - if tile_index < valid_items: - expected[payload_index] = source[ - warp * (_WARP_THREADS * items_per_thread) + tile_index - ] - - _partial_load_kernel(algorithm, qualified)[1, _BLOCK_THREADS]( - source, observed, np.int32(valid_items), items_per_thread - ) - - np.testing.assert_array_equal(observed, expected) - - -@cache -def _per_warp_valid_items_kernel(qualified: bool): - if qualified: - - @cuda.jit - def kernel(source, observed, valid_by_warp, items_per_thread): - thread = cuda.threadIdx.x - warp = thread // _WARP_THREADS - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_warp(), - source, - payload, - algorithm="direct", - valid_items=valid_by_warp[warp], - oob_default=types.int32(-83), - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - else: - - @cuda.jit - def kernel(source, observed, valid_by_warp, items_per_thread): - thread = cuda.threadIdx.x - warp = thread // _WARP_THREADS - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - root_coop.this_warp(), - source, - payload, - algorithm="direct", - valid_items=valid_by_warp[warp], - oob_default=types.int32(-83), - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_runtime_valid_items_can_differ_between_physical_warps( - qualified: bool, *, items_per_thread -) -> None: - source = _values((_BLOCK_THREADS * items_per_thread), shift=71) - observed = np.full((_BLOCK_THREADS * items_per_thread), 17, dtype=np.int32) - valid_by_warp = np.array( - [13, (_WARP_THREADS * items_per_thread) - 5], dtype=np.int32 - ) - expected = np.full((_BLOCK_THREADS * items_per_thread), -83, dtype=np.int32) - for warp, valid_items in enumerate(valid_by_warp): - begin = warp * (_WARP_THREADS * items_per_thread) - expected[begin : begin + valid_items] = source[ - begin : begin + valid_items - ] - - _per_warp_valid_items_kernel(qualified)[1, _BLOCK_THREADS]( - source, observed, valid_by_warp, items_per_thread - ) - - np.testing.assert_array_equal(observed, expected) - - -@cache -def _multidimensional_load_kernel(qualified: bool): - if qualified: - - @cuda.jit - def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x + cuda.blockDim.x * ( - cuda.threadIdx.y + cuda.blockDim.y * cuda.threadIdx.z - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_warp(), - source, - payload, - algorithm="direct", - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - else: - - @cuda.jit - def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x + cuda.blockDim.x * ( - cuda.threadIdx.y + cuda.blockDim.y * cuda.threadIdx.z - ) - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - root_coop.this_warp(), - source, - payload, - algorithm="direct", - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize( - "block_shape", - ((16, 4), (8, 4, 2)), - ids=("2d", "3d"), -) -def test_physical_warp_origin_uses_x_major_multidimensional_rank( - qualified: bool, block_shape: tuple[int, ...], *, items_per_thread -) -> None: - source = _values((_BLOCK_THREADS * items_per_thread), shift=83) - observed = np.full((_BLOCK_THREADS * items_per_thread), -1, dtype=np.int32) - - _multidimensional_load_kernel(qualified)[1, block_shape]( - source, observed, items_per_thread - ) - - np.testing.assert_array_equal(observed, source) - - -@cache -def _logical_multidimensional_load_kernel(qualified: bool): - if qualified: - - @cuda.jit - def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x + cuda.blockDim.x * ( - cuda.threadIdx.y + cuda.blockDim.y * cuda.threadIdx.z - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_warp().group_by(_LOGICAL_WARP_THREADS), - source, - payload, - algorithm="direct", - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - else: - - @cuda.jit - def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x + cuda.blockDim.x * ( - cuda.threadIdx.y + cuda.blockDim.y * cuda.threadIdx.z - ) - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - root_coop.this_warp().group_by(_LOGICAL_WARP_THREADS), - source, - payload, - algorithm="direct", - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_logical_warp_origin_uses_x_major_multidimensional_rank( - qualified: bool, *, items_per_thread -) -> None: - source = _values((_BLOCK_THREADS * items_per_thread), shift=79) - observed = np.full((_BLOCK_THREADS * items_per_thread), -1, dtype=np.int32) - - _logical_multidimensional_load_kernel(qualified)[1, (8, 4, 2)]( - source, observed, items_per_thread - ) - - np.testing.assert_array_equal(observed, source) - - -@cache -def _static_control_load_kernel(qualified: bool): - if qualified: - - @cuda.jit - def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_warp(), - source, - payload, - algorithm="direct", - valid_items=(_WARP_THREADS * items_per_thread) - 7, - oob_default=-113, - offset=_LOAD_OFFSET, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - else: - - @cuda.jit - def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - root_coop.this_warp(), - source, - payload, - algorithm="direct", - valid_items=(_WARP_THREADS * items_per_thread) - 7, - oob_default=-113, - offset=_LOAD_OFFSET, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_physical_warp_static_controls_share_runtime_addressing( - qualified: bool, *, items_per_thread -) -> None: - source = _values( - _LOAD_OFFSET + (_BLOCK_THREADS * items_per_thread), shift=89 - ) - observed = np.full((_BLOCK_THREADS * items_per_thread), -1, dtype=np.int32) - expected = np.full( - (_BLOCK_THREADS * items_per_thread), -113, dtype=np.int32 - ) - valid_items = (_WARP_THREADS * items_per_thread) - 7 - for warp in range(2): - begin = warp * (_WARP_THREADS * items_per_thread) - source_begin = _LOAD_OFFSET + begin - expected[begin : begin + valid_items] = source[ - source_begin : source_begin + valid_items - ] - - _static_control_load_kernel(qualified)[1, _BLOCK_THREADS]( - source, observed, items_per_thread - ) - - np.testing.assert_array_equal(observed, expected) - - -@cache -def _scalar_store_kernel(qualified: bool): - if qualified: - - @cuda.jit - def kernel(source, destination): - thread = cuda.threadIdx.x - value = types.int32(source[thread] + 1) - qualified_coop.store( - qualified_coop.this_warp(), - destination, - value, - algorithm="direct", - offset=_STORE_OFFSET, - ) - - else: - - @cuda.jit - def kernel(source, destination): - thread = cuda.threadIdx.x - value = types.int32(source[thread] + 1) - root_coop.store( - root_coop.this_warp(), - destination, - value, - algorithm="direct", - offset=_STORE_OFFSET, - ) - - return kernel - - -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_physical_warp_scalar_store_uses_destination_dtype( - qualified: bool, -) -> None: - source = _values(_BLOCK_THREADS, shift=91) - destination = np.full( - _STORE_OFFSET + _BLOCK_THREADS + 3, -17, dtype=np.int32 - ) - expected = destination.copy() - expected[_STORE_OFFSET : _STORE_OFFSET + _BLOCK_THREADS] = source + 1 - - _scalar_store_kernel(qualified)[1, _BLOCK_THREADS](source, destination) - - np.testing.assert_array_equal(destination, expected) - - -@cache -def _literal_scalar_store_kernel(qualified: bool): - if qualified: - - @cuda.jit - def kernel(destination): - qualified_coop.store( - qualified_coop.this_warp(), - destination, - 23, - algorithm="direct", - offset=_STORE_OFFSET, - ) - - else: - - @cuda.jit - def kernel(destination): - root_coop.store( - root_coop.this_warp(), - destination, - 23, - algorithm="direct", - offset=_STORE_OFFSET, - ) - - return kernel - - -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_physical_warp_scalar_literal_infers_the_destination_dtype( - qualified: bool, -) -> None: - destination = np.full( - _STORE_OFFSET + _BLOCK_THREADS + 3, -17, dtype=np.int32 - ) - expected = destination.copy() - expected[_STORE_OFFSET : _STORE_OFFSET + _BLOCK_THREADS] = 23 - - _literal_scalar_store_kernel(qualified)[1, _BLOCK_THREADS](destination) - - np.testing.assert_array_equal(destination, expected) - - -@cache -def _grid_stride_transpose_kernel(qualified: bool): - """Clamp a separate valid count for each warp in each block's tile. - - Pass only the block origin as ``offset``. The provider adds the warp's - tile origin, including for the partial last warp of the final block. - """ - - if qualified: - - @cuda.jit - def kernel(source, destination, items_per_thread): - thread = cuda.threadIdx.x - warp = thread // _WARP_THREADS - block_offset = cuda.blockIdx.x * (_BLOCK_THREADS * items_per_thread) - remaining = ( - source.size - - block_offset - - warp * (_WARP_THREADS * items_per_thread) - ) - valid_items = min( - max(remaining, 0), (_WARP_THREADS * items_per_thread) - ) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - qualified_coop.this_warp(), - source, - payload, - algorithm="transpose", - valid_items=valid_items, - offset=block_offset, - ) - qualified_coop.store( - qualified_coop.this_warp(), - destination, - payload, - algorithm="transpose", - valid_items=valid_items, - offset=block_offset, - ) - - else: - - @cuda.jit - def kernel(source, destination, items_per_thread): - thread = cuda.threadIdx.x - warp = thread // _WARP_THREADS - block_offset = cuda.blockIdx.x * (_BLOCK_THREADS * items_per_thread) - remaining = ( - source.size - - block_offset - - warp * (_WARP_THREADS * items_per_thread) - ) - valid_items = min( - max(remaining, 0), (_WARP_THREADS * items_per_thread) - ) - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - root_coop.this_warp(), - source, - payload, - algorithm="transpose", - valid_items=valid_items, - offset=block_offset, - ) - root_coop.store( - root_coop.this_warp(), - destination, - payload, - algorithm="transpose", - valid_items=valid_items, - offset=block_offset, - ) - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_grid_stride_tail_clamps_valid_items_per_physical_warp( - qualified: bool, *, items_per_thread -) -> None: - source = _values((3 * (_BLOCK_THREADS * items_per_thread) - 17), shift=103) - destination = np.full( - (3 * (_BLOCK_THREADS * items_per_thread) - 17), -1, dtype=np.int32 - ) - - _grid_stride_transpose_kernel(qualified)[ - ( - ( - 3 * (_BLOCK_THREADS * items_per_thread) - - 17 - + _BLOCK_THREADS * items_per_thread - - 1 - ) - // (_BLOCK_THREADS * items_per_thread) - ), - _BLOCK_THREADS, - ](source, destination, items_per_thread) - - np.testing.assert_array_equal(destination, source) - - -@cache -def _logical_grid_stride_transpose_kernel(qualified: bool): - """Clamp the tail count for every logical group in each block's tile. - - Pass the block origin as ``offset`` and let the provider add the logical - group's tile origin. Groups beyond the source use a zero valid count. - """ - - if qualified: - - @cuda.jit - def kernel(source, destination, items_per_thread): - thread = cuda.threadIdx.x - group_index = thread // _LOGICAL_WARP_THREADS - block_offset = cuda.blockIdx.x * (_BLOCK_THREADS * items_per_thread) - remaining = ( - source.size - - block_offset - - group_index * (_LOGICAL_WARP_THREADS * items_per_thread) - ) - valid_items = min( - max(remaining, 0), (_LOGICAL_WARP_THREADS * items_per_thread) - ) - group = qualified_coop.this_warp().group_by(_LOGICAL_WARP_THREADS) - payload = qualified_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - qualified_coop.load( - group, - source, - payload, - algorithm="transpose", - valid_items=valid_items, - offset=block_offset, - ) - qualified_coop.store( - group, - destination, - payload, - algorithm="transpose", - valid_items=valid_items, - offset=block_offset, - ) - - else: - - @cuda.jit - def kernel(source, destination, items_per_thread): - thread = cuda.threadIdx.x - group_index = thread // _LOGICAL_WARP_THREADS - block_offset = cuda.blockIdx.x * (_BLOCK_THREADS * items_per_thread) - remaining = ( - source.size - - block_offset - - group_index * (_LOGICAL_WARP_THREADS * items_per_thread) - ) - valid_items = min( - max(remaining, 0), (_LOGICAL_WARP_THREADS * items_per_thread) - ) - group = root_coop.this_warp().group_by(_LOGICAL_WARP_THREADS) - payload = root_coop.ThreadData( - items_per_thread, - dtype=types.int32, - ) - root_coop.load( - group, - source, - payload, - algorithm="transpose", - valid_items=valid_items, - offset=block_offset, - ) - root_coop.store( - group, - destination, - payload, - algorithm="transpose", - valid_items=valid_items, - offset=block_offset, - ) - - return kernel - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_grid_stride_tail_clamps_valid_items_per_logical_warp( - qualified: bool, *, items_per_thread -) -> None: - source = _values((3 * (_BLOCK_THREADS * items_per_thread) - 17), shift=107) - destination = np.full( - (3 * (_BLOCK_THREADS * items_per_thread) - 17), -1, dtype=np.int32 - ) - - _logical_grid_stride_transpose_kernel(qualified)[ - ( - ( - 3 * (_BLOCK_THREADS * items_per_thread) - - 17 - + _BLOCK_THREADS * items_per_thread - - 1 - ) - // (_BLOCK_THREADS * items_per_thread) - ), - _BLOCK_THREADS, - ](source, destination, items_per_thread) - - np.testing.assert_array_equal(destination, source) - - -def _run_divergent_warp_probe( - qualified: bool, *, items_per_thread -) -> subprocess.CompletedProcess[str]: - """Run a one-warp primitive with a child-process timeout. - - A barrier that includes the inactive warp could stall the kernel. The - child contains that failure and checks that it imports the same - ``cuda.coop`` package as the test process before launching. Return its - status and output for the test assertion. - """ - - if qualified: - thread_data = "qualified_coop.ThreadData" - group = "qualified_coop.this_warp()" - load = "qualified_coop.load" - algorithm = repr("transpose") - else: - thread_data = "root_coop.ThreadData" - group = "root_coop.this_warp()" - load = "root_coop.load" - algorithm = repr("transpose") - - script = f"""\ -import numpy as np -import numba_cuda_mlir.cuda as cuda -from numba_cuda_mlir import types -from pathlib import Path - -import cuda.coop.numba_mlir as qualified_coop -from cuda import coop as root_coop - -expected_origin = Path({str(_QUALIFIED_COOP_ORIGIN)!r}) -actual_origin = Path(qualified_coop.__file__).resolve() -if actual_origin != expected_origin: - raise RuntimeError( - f"divergent probe imported cuda.coop from {{actual_origin}}, " - f"expected {{expected_origin}}" - ) - -_WARP_THREADS = {_WARP_THREADS} -_BLOCK_THREADS = {_BLOCK_THREADS} -_ITEMS_PER_THREAD = {items_per_thread} - -@cuda.jit -def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x - warp = thread // _WARP_THREADS - if warp == 0: - payload = {thread_data}(items_per_thread, dtype=types.int32) - {load}( - {group}, - source, - payload, - algorithm={algorithm}, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - -source = ((np.arange(_BLOCK_THREADS * _ITEMS_PER_THREAD, dtype=np.int64) * 3 + 97) % 211 - 101).astype(np.int32) -observed = np.full_like(source, -37) -expected = observed.copy() -expected[: _WARP_THREADS * _ITEMS_PER_THREAD] = source[: _WARP_THREADS * _ITEMS_PER_THREAD] -kernel[1, _BLOCK_THREADS](source, observed, _ITEMS_PER_THREAD) -cuda.synchronize() -np.testing.assert_array_equal(observed, expected) -""" # noqa: E501 - Preserve embedded source bytes. - return subprocess.run( - [sys.executable, _SAFE_PATH_FLAG, "-B", "-c", script], - check=False, - capture_output=True, - text=True, - timeout=180, - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_one_physical_warp_can_take_a_transpose_primitive_path( - qualified: bool, *, items_per_thread -) -> None: - result = _run_divergent_warp_probe( - qualified, items_per_thread=items_per_thread - ) - assert result.returncode == 0, result.stdout + result.stderr - - -def _run_divergent_logical_warp_probe( - qualified: bool, *, items_per_thread -) -> subprocess.CompletedProcess[str]: - """Probe a nonzero logical group in each physical warp with a timeout. - - Only subgroup two participates, testing masks that must exclude other - lanes. A child process bounds an incorrect barrier and checks the - package's import origin. Return its status and captured output. - """ - - if qualified: - thread_data = "qualified_coop.ThreadData" - group = "qualified_coop.this_warp().group_by(_LOGICAL_WARP_THREADS)" - load = "qualified_coop.load" - algorithm = repr("transpose") - else: - thread_data = "root_coop.ThreadData" - group = "root_coop.this_warp().group_by(_LOGICAL_WARP_THREADS)" - load = "root_coop.load" - algorithm = repr("transpose") - - script = f"""\ -import numpy as np -import numba_cuda_mlir.cuda as cuda -from numba_cuda_mlir import types -from pathlib import Path - -import cuda.coop.numba_mlir as qualified_coop -from cuda import coop as root_coop - -expected_origin = Path({str(_QUALIFIED_COOP_ORIGIN)!r}) -actual_origin = Path(qualified_coop.__file__).resolve() -if actual_origin != expected_origin: - raise RuntimeError( - f"divergent logical-warp probe imported cuda.coop from {{actual_origin}}, " - f"expected {{expected_origin}}" - ) - -_WARP_THREADS = {_WARP_THREADS} -_LOGICAL_WARP_THREADS = {_LOGICAL_WARP_THREADS} -_BLOCK_THREADS = {_BLOCK_THREADS} -_ITEMS_PER_THREAD = {items_per_thread} - -@cuda.jit -def kernel(source, observed, items_per_thread): - thread = cuda.threadIdx.x - subgroup = (thread % _WARP_THREADS) // _LOGICAL_WARP_THREADS - if subgroup == 2: - payload = {thread_data}(items_per_thread, dtype=types.int32) - {load}( - {group}, - source, - payload, - algorithm={algorithm}, - ) - for item in range(items_per_thread): - observed[thread * items_per_thread + item] = payload[item] - -source = ((np.arange(_BLOCK_THREADS * _ITEMS_PER_THREAD, dtype=np.int64) * 3 + 113) % 211 - 101).astype(np.int32) -observed = np.full_like(source, -37) -expected = observed.copy() -for warp in range(_BLOCK_THREADS // _WARP_THREADS): - begin = ( - warp * _WARP_THREADS + 2 * _LOGICAL_WARP_THREADS - ) * _ITEMS_PER_THREAD - end = begin + _LOGICAL_WARP_THREADS * _ITEMS_PER_THREAD - expected[begin:end] = source[begin:end] -kernel[1, _BLOCK_THREADS](source, observed, _ITEMS_PER_THREAD) -cuda.synchronize() -np.testing.assert_array_equal(observed, expected) -""" # noqa: E501 - Preserve embedded source bytes. - return subprocess.run( - [sys.executable, _SAFE_PATH_FLAG, "-B", "-c", script], - check=False, - capture_output=True, - text=True, - timeout=180, - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_one_logical_warp_per_physical_warp_can_diverge_at_transpose( - qualified: bool, *, items_per_thread -) -> None: - result = _run_divergent_logical_warp_probe( - qualified, items_per_thread=items_per_thread - ) - assert result.returncode == 0, result.stdout + result.stderr - - -def _run_invalid_runtime_valid_items_probe( - operation: str, - valid_items: int, - *, - logical_width: int | None = None, - items_per_thread, -) -> subprocess.CompletedProcess[str]: - """Run an invalid per-group count in a disposable CUDA context. - - ``logical_width`` selects a logical group; ``None`` uses a full physical - warp. The child checks that its ``cuda.coop`` import origin matches the - test process, then launches with ``items_per_thread`` items per lane. - A device trap poisons its CUDA context. Return captured output and status - so the test can require that fault while keeping its own context usable. - """ - - group = "root_coop.this_warp()" - if logical_width is not None: - group = f"root_coop.this_warp().group_by({logical_width})" - operation_body = textwrap.indent( - textwrap.dedent( - { - "load": f""" - payload = root_coop.ThreadData( - items_per_thread=_ITEMS_PER_THREAD, dtype=types.int32 - ) - root_coop.load( - {group}, - source, - payload, - valid_items=valid_items, - ) - """, - "store": f""" - payload = root_coop.ThreadData( - items_per_thread=_ITEMS_PER_THREAD, dtype=types.int32 - ) - for item in range(_ITEMS_PER_THREAD): - payload[item] = source[thread * _ITEMS_PER_THREAD + item] - root_coop.store( - {group}, - destination, - payload, - valid_items=valid_items, - ) - """, - }[operation] - ).strip(), - " " * 4, - ) - script = f"""\ -import numpy as np -import numba_cuda_mlir.cuda as cuda -from numba_cuda_mlir import types -from pathlib import Path - -import cuda.coop.numba_mlir as numba_coop -from cuda import coop as root_coop - -expected_origin = Path({str(_QUALIFIED_COOP_ORIGIN)!r}) -actual_origin = Path(numba_coop.__file__).resolve() -if actual_origin != expected_origin: - raise RuntimeError( - f"trap probe imported cuda.coop from {{actual_origin}}, " - f"expected {{expected_origin}}" - ) - -_THREADS = {_BLOCK_THREADS} -_ITEMS_PER_THREAD = {items_per_thread} - -@cuda.jit -def kernel(source, destination, valid_items): - thread = cuda.threadIdx.x -{operation_body} - -source = np.arange( - _THREADS * _ITEMS_PER_THREAD, - dtype=np.int32, -) -destination = np.full_like(source, -1) -kernel[1, _THREADS](source, destination, np.int64({valid_items})) -cuda.synchronize() -raise AssertionError("invalid valid_items did not trap") -""" - return subprocess.run( - [sys.executable, _SAFE_PATH_FLAG, "-B", "-c", script], - check=False, - capture_output=True, - text=True, - timeout=180, - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("operation", ("load", "store")) -@pytest.mark.parametrize( - "valid_items", - ( - pytest.param(-1, id="negative"), - pytest.param(_WARP_TILE_ITEMS + 1, id="beyond-warp-tile"), - ), -) -def test_runtime_valid_items_out_of_range_traps_in_an_isolated_context( - operation: str, valid_items: int, *, items_per_thread -) -> None: - valid_items = ( - valid_items if valid_items < 0 else _WARP_THREADS * items_per_thread + 1 - ) - result = _run_invalid_runtime_valid_items_probe( - operation, valid_items, items_per_thread=items_per_thread - ) - output = result.stdout + result.stderr - - assert result.returncode != 0, output - assert any( - error in output - for error in ( - "CUDA_ERROR_ILLEGAL_INSTRUCTION", - "CUDA_ERROR_LAUNCH_FAILED", - ) - ), output - - source = _values((_BLOCK_THREADS * items_per_thread), shift=109) - observed = np.full((_BLOCK_THREADS * items_per_thread), -1, dtype=np.int32) - _load_kernel("direct", False)[1, _BLOCK_THREADS]( - source, - observed, - np.int32(_WARP_THREADS * items_per_thread), - np.int64(0), - np.int32(-1), - items_per_thread, - ) - np.testing.assert_array_equal(observed, source) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("operation", ("load", "store")) -@pytest.mark.parametrize( - "valid_items", - ( - pytest.param(-1, id="negative"), - pytest.param(_LOGICAL_TILE_ITEMS + 1, id="beyond-logical-tile"), - ), -) -def test_logical_runtime_valid_items_out_of_range_traps_in_isolated_context( - operation: str, valid_items: int, *, items_per_thread -) -> None: - valid_items = ( - valid_items - if valid_items < 0 - else _LOGICAL_WARP_THREADS * items_per_thread + 1 - ) - result = _run_invalid_runtime_valid_items_probe( - operation, - valid_items, - logical_width=_LOGICAL_WARP_THREADS, - items_per_thread=items_per_thread, - ) - output = result.stdout + result.stderr - - assert result.returncode != 0, output - assert any( - error in output - for error in ( - "CUDA_ERROR_ILLEGAL_INSTRUCTION", - "CUDA_ERROR_LAUNCH_FAILED", - ) - ), output diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/__init__.py b/python/cuda_coop/tests/backends/numba_mlir/unit/__init__.py deleted file mode 100644 index 8bbe3ce1ab86..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_caching.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_caching.py deleted file mode 100644 index c38c7d418dfd..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_caching.py +++ /dev/null @@ -1,266 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import hashlib -import ntpath -import os -import posixpath -import subprocess -import sys -import textwrap -from types import SimpleNamespace - -import pytest - -from cuda.coop.numba_mlir import _types -from cuda.coop.numba_mlir._compiler import _caching -from cuda.coop.numba_mlir._semantic import _numba_semantic_token - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - - -@pytest.mark.parametrize("payload", ("[]", "3", "null")) -def test_cache_treats_non_object_json_as_a_miss(tmp_path, payload): - path = tmp_path / "cache-entry.json" - path.write_text(payload, encoding="utf-8") - - assert _caching._read_cache(path) is _caching._CACHE_MISS - - -def test_disk_cache_disables_itself_when_its_directory_is_unavailable( - monkeypatch, -): - calls = [] - - def unavailable(_cache_identity): - raise PermissionError("cache directory is read-only") - - monkeypatch.setattr(_caching, "_CACHE_USABLE", True) - monkeypatch.setattr(_caching, "_cache_identity_path", unavailable) - - @_caching.disk_cache - def compute(value): - calls.append(value) - return value + 1 - - assert compute(3) == 4 - assert compute(3) == 4 - assert calls == [3, 3] - assert not _caching._CACHE_USABLE - - -def test_disk_cache_bypasses_unsupported_key_values(monkeypatch, tmp_path): - calls = [] - unsupported = object() - - monkeypatch.setattr(_caching, "_CACHE_USABLE", True) - monkeypatch.setattr( - _caching, - "_cache_identity_path", - lambda _cache_identity: str(tmp_path), - ) - - @_caching.disk_cache - def compute(value): - calls.append(value) - return 4 if value is unsupported else value + 1 - - assert compute(unsupported) == 4 - assert compute(unsupported) == 4 - assert compute(3) == 4 - assert compute(3) == 4 - assert calls == [unsupported, unsupported, 3] - assert _caching._CACHE_USABLE - - -def test_disk_cache_bypasses_unsupported_result_values(monkeypatch, tmp_path): - calls = [] - unsupported = object() - - monkeypatch.setattr(_caching, "_CACHE_USABLE", True) - monkeypatch.setattr( - _caching, - "_cache_identity_path", - lambda _cache_identity: str(tmp_path), - ) - - @_caching.disk_cache - def compute(value): - calls.append(value) - return unsupported if value == "unsupported" else value + 1 - - assert compute("unsupported") is unsupported - assert compute("unsupported") is unsupported - assert not tuple(tmp_path.glob(".*")) - assert compute(3) == 4 - assert compute(3) == 4 - assert calls == ["unsupported", "unsupported", 3] - assert _caching._CACHE_USABLE - - -def test_disk_cache_persists_byte_valued_compiler_options( - monkeypatch, tmp_path -): - calls = [] - - monkeypatch.setattr(_caching, "_CACHE_USABLE", True) - monkeypatch.setattr( - _caching, - "_cache_identity_path", - lambda _cache_identity: str(tmp_path), - ) - - @_caching.disk_cache - def compute(options): - calls.append(options) - return b"ltoir" - - options = (b"--std=c++17", b"--gpu-architecture=compute_90") - assert compute(options) == b"ltoir" - assert compute(options) == b"ltoir" - assert calls == [options] - assert len(tuple(tmp_path.iterdir())) == 1 - - -def test_disk_cache_persists_layout_results_in_a_fresh_process(tmp_path): - """Require a disk hit to restore nested tuples and arbitrary LTO bytes. - - The read process refuses to compute a result. Separate interpreters rule - out reuse through an in-memory cache or an object retained by the writer. - """ - - script = textwrap.dedent( - """ - import sys - from cuda.coop.numba_mlir._compiler import _caching - - _caching._CACHE_LOCATION = sys.argv[1] - expected = (b"ltoir\\x00\\xff", ((24, 8), (64, 16))) - - @_caching.disk_cache - def compile_provider(*, layout_queries): - assert sys.argv[2] == "write", "disk cache missed in fresh process" - return expected - - result = compile_provider(layout_queries=("sizeof(A)", "sizeof(B)")) - assert result == expected - assert isinstance(result, tuple) - assert isinstance(result[0], bytes) - assert isinstance(result[1], tuple) - assert all(isinstance(layout, tuple) for layout in result[1]) - """ - ) - environment = dict(os.environ, CUDA_COOP_ENABLE_CACHE="1") - for phase in ("write", "read"): - result = subprocess.run( - [sys.executable, "-B", "-c", script, str(tmp_path), phase], - env=environment, - capture_output=True, - text=True, - timeout=60, - check=False, - ) - assert result.returncode == 0, result.stdout + result.stderr - assert len(tuple(tmp_path.rglob("*"))) == 2 - - -def test_layout_cache_identity(monkeypatch, tmp_path): - calls = [] - monkeypatch.setattr(_caching, "_CACHE_USABLE", True) - monkeypatch.setattr(_caching, "_CACHE_LOCATION", str(tmp_path)) - - @_caching.disk_cache - def compile_provider(*, layout_queries): - calls.append(layout_queries) - return b"ltoir", ((len(layout_queries) * 16, 16),) - - first = ("sizeof(A)",) - second = ("sizeof(A)", "sizeof(B)") - assert compile_provider(layout_queries=first) == (b"ltoir", ((16, 16),)) - assert compile_provider(layout_queries=second) == (b"ltoir", ((32, 16),)) - assert compile_provider(layout_queries=first) == (b"ltoir", ((16, 16),)) - assert compile_provider(layout_queries=second) == (b"ltoir", ((32, 16),)) - assert calls == [first, second] - - -def test_symbol_hash_uses_numba_type_identity(): - assert _numba_semantic_token(_types.types.int32) == _numba_semantic_token( - ( - "numba-cuda-mlir-type", - type(_types.types.int32).__module__, - type(_types.types.int32).__qualname__, - "int32", - ) - ) - int32 = hashlib.sha1() - _types._hash_symbol_value(int32, _types.types.int32) - another_int32 = hashlib.sha1() - _types._hash_symbol_value(another_int32, _types.types.int32) - float32 = hashlib.sha1() - _types._hash_symbol_value(float32, _types.types.float32) - - assert int32.digest() == another_int32.digest() - assert int32.digest() != float32.digest() - - -def test_numba_type_tokens_remain_distinct(): - values = ( - _types.types.int32, - _types.types.float32, - _types.types.float64, - _types.types.uint32, - ) - tokens = tuple(_numba_semantic_token(value) for value in values) - - assert len(set(tokens)) == len(values) - - -@pytest.mark.parametrize( - ("platform", "environment", "expected"), - [ - ("posix", {"XDG_CACHE_HOME": "/scratch/cache"}, "/scratch/cache/cccl"), - ("posix", {}, "/home/test/.cache/cccl"), - ("posix", {"XDG_CACHE_HOME": ""}, "/home/test/.cache/cccl"), - ("posix", {"XDG_CACHE_HOME": "relative"}, "/home/test/.cache/cccl"), - ("posix", {"LOCALAPPDATA": "/other/cache"}, "/home/test/.cache/cccl"), - ("nt", {"LOCALAPPDATA": r"D:\Cache"}, r"D:\Cache\cccl"), - ("nt", {}, r"C:\Users\test\AppData\Local\cccl"), - ("nt", {"LOCALAPPDATA": ""}, r"C:\Users\test\AppData\Local\cccl"), - ( - "nt", - {"LOCALAPPDATA": "relative"}, - r"C:\Users\test\AppData\Local\cccl", - ), - ( - "nt", - {"LOCALAPPDATA": r"D:relative"}, - r"C:\Users\test\AppData\Local\cccl", - ), - ( - "nt", - {"XDG_CACHE_HOME": r"D:\Cache"}, - r"C:\Users\test\AppData\Local\cccl", - ), - ], -) -def test_cache_location_uses_platform_cache_directory( - monkeypatch, platform, environment, expected -): - path = ntpath if platform == "nt" else posixpath - user_home = r"C:\Users\test" if platform == "nt" else "/home/test" - monkeypatch.setattr( - _caching, - "os", - SimpleNamespace( - name=platform, - environ=environment, - path=SimpleNamespace( - isabs=path.isabs, - join=path.join, - expanduser=lambda value: user_home if value == "~" else value, - ), - ), - ) - - assert _caching._cache_location() == expected diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_data_movement.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_data_movement.py deleted file mode 100644 index 7e23de0d7282..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_data_movement.py +++ /dev/null @@ -1,2919 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check Load/Store planning, argument validation, and emitted call structure. - -Frontend IR and core plans are real. Tests replace provider materialization -where they need to inspect rewriting without compiling or launching kernels. -The compile and runtime suites check those later stages separately. -""" - -import operator -from collections import Counter -from enum import Enum -from types import SimpleNamespace - -import numpy as np -import pytest - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - - -class _StringAlgorithm(str, Enum): - DIRECT = "direct" - - -@pytest.fixture(autouse=True) -def _fixed_provider_compute_capability(monkeypatch): - """Keep provider identity stable without querying a CUDA device.""" - from cuda.coop.numba_mlir import _types - - monkeypatch.setattr( - _types.cuda, - "get_current_device", - lambda: SimpleNamespace(compute_capability=(9, 0)), - ) - - -def _plan(function, *, arg_types, block=(64, 1, 1)): - """Build frontend IR with exact argument and launch facts.""" - from numba_cuda_mlir.numba_cuda.compiler import run_frontend - - from cuda.coop.numba_mlir._compiler._group_planner import _GroupCallPlanner - - func_ir = run_frontend(function) - state = SimpleNamespace(func_ir=func_ir, args=arg_types) - planner = _GroupCallPlanner( - state, - {"block": block, "grid": (1, 1, 1), "cluster": None}, - ) - return func_ir, planner - - -def _planned_factory_calls(func_ir, ir): - """Find calls and their directly assigned global factory objects in IR.""" - globals_by_name = { - inst.target.name: inst.value.value - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) and isinstance(inst.value, ir.Global) - } - return [ - (globals_by_name.get(inst.value.func.name), inst.value) - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) - and isinstance(inst.value, ir.Expr) - and inst.value.op == "call" - ] - - -def _resolved_python_call_targets(func_ir, ir): - """Resolve rewritten call targets to inspect allocations and barriers.""" - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - - resolver = object.__new__(CoopSinglePhaseRewrite) - resolver._func_ir = func_ir - targets = [] - for block in func_ir.blocks.values(): - resolver._block_defs = { - inst.target.name: inst.value - for inst in block.body - if isinstance(inst, ir.Assign) - } - targets.extend( - resolver._resolve_python_value(inst.value.func) - for inst in block.body - if isinstance(inst, ir.Assign) - and isinstance(inst.value, ir.Expr) - and inst.value.op == "call" - ) - return targets - - -class _MovementInvocable: - """Supply scratch requirements without compiling a provider.""" - - files = ("movement-test.ltoir",) - storage_abi = "leading_pointer" - execution_scope = "block" - synchronization_scope = "block" - - def __init__(self, temp_storage_bytes): - self.temp_storage_bytes = temp_storage_bytes - self.temp_storage_alignment = 16 - - def __call__(self, *args): - del args - - -def _rewrite_planned_movement(function, *, arg_types): - """Run planning and rewriting with fixed provider storage requirements. - - Return the rewritten IR, rewrite state, and fake providers so tests can - inspect calls, scratch allocation, and reuse barriers together. - """ - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - - func_ir, planner = _plan(function, arg_types=arg_types) - assert planner.run() - - class TypingContext: - def refresh(self): - pass - - state = SimpleNamespace( - func_ir=func_ir, - args=arg_types, - typingctx=TypingContext(), - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - invocables = { - "load": _MovementInvocable(48), - "store": _MovementInvocable(64), - } - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: None - rewrite._materialize_invocable = lambda match: ( - invocables[match.op_name], - False, - ) - assert rewrite.prepare_calls_and_storage(func_ir) - rewrite.begin_rewrite() - for label in sorted(func_ir.blocks): - block = func_ir.blocks[label] - while rewrite.match(func_ir, block, state.typemap, state.calltypes): - block = rewrite.apply() - func_ir.blocks[label] = block - rewrite.finish_rewrite() - return func_ir, rewrite, tuple(invocables.values()) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_load_store_calls_across_blocks_are_analyzed_once( - monkeypatch, items_per_thread -): - from numba_cuda_mlir import cuda, types - from numba_cuda_mlir.numbair_transforms import ir - - from cuda import coop - from cuda.coop.numba_mlir._compiler._group_rewriting import ( - GroupRewriteContext, - ) - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - from cuda.coop.numba_mlir._compiler._rewrite_load_store import ( - _LoadStoreRewrite, - ) - - def kernel(source, destination, flag, items_per_thread): - block = coop.this_block() - items = coop.ThreadData(items_per_thread) - coop.load(block, source, items, algorithm="transpose") - written = coop.ThreadData(items_per_thread) - for item in range(items_per_thread): - written[item] = source[item] - if flag: - coop.store(block, destination, written, algorithm="transpose") - else: - coop.store(block, destination, items, algorithm="transpose") - - analyzed = Counter() - split_args = CoopSinglePhaseRewrite._validate_and_split_args - - def record_analysis(self, op_name, call, getitem_temp_storage): - analyzed[id(call)] += 1 - return split_args(self, op_name, call, getitem_temp_storage) - - monkeypatch.setattr( - CoopSinglePhaseRewrite, "_validate_and_split_args", record_analysis - ) - infer_payload = _LoadStoreRewrite._infer_load_store_payload - - def infer_planned_payload(context, inference): - assert inference.lowering_plan is not None - with monkeypatch.context() as patch: - for name in ("dtype", "infer_thread_data_write_dtype"): - patch.setattr( - GroupRewriteContext, - name, - lambda *_args: pytest.fail( - "planned Load/Store repeated payload dtype inference" - ), - ) - infer_payload(context, inference) - - monkeypatch.setattr( - _LoadStoreRewrite, "_infer_load_store_payload", infer_planned_payload - ) - array_type = types.Array(types.int32, 1, "C") - func_ir, _, invocables = _rewrite_planned_movement( - kernel, - arg_types=( - array_type, - array_type, - types.boolean, - types.IntegerLiteral(items_per_thread), - ), - ) - targets = _resolved_python_call_targets(func_ir, ir) - assert len(analyzed) == 3 - assert set(analyzed.values()) == {1} - assert targets.count(invocables[0]) == 1 - assert targets.count(invocables[1]) == 2 - assert targets.count(cuda.shared.array) == 1 - assert targets.count(cuda.local.array) == 2 - - -def _run_single_phase_to_provider_boundary( - function, - *, - arg_types, - monkeypatch, - allow_provider_bundling=False, -): - """Check argument validation before a provider can be materialized. - - Allow the bundling boundary only for inputs expected to pass validation. - Provider compilation and materialization are replaced by test guards. - """ - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - - func_ir, planner = _plan(function, arg_types=arg_types, block=(32, 1, 1)) - assert planner.run() - - class TypingContext: - def refresh(self): - pass - - state = SimpleNamespace( - func_ir=func_ir, - args=arg_types, - typingctx=TypingContext(), - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - if allow_provider_bundling: - monkeypatch.setattr( - rewrite, - "_compute_func_temp_storage_requirements", - lambda _storage_uses: {}, - ) - monkeypatch.setattr( - rewrite, - "_prepare_ltoir_bundle_for_matches", - ( - (lambda _matches: None) - if allow_provider_bundling - else lambda _matches: pytest.fail( - "invalid movement arguments reached provider bundling" - ) - ), - ) - monkeypatch.setattr( - rewrite, - "_materialize_invocable", - lambda _match: pytest.fail( - "invalid movement arguments reached provider materialization" - ), - ) - matched_group_call = False - assert rewrite.prepare_calls_and_storage(func_ir) - for label in sorted(func_ir.blocks): - rewrite.match( - func_ir, - func_ir.blocks[label], - state.typemap, - state.calltypes, - ) - matched_group_call |= bool(rewrite._matches) - assert matched_group_call - - -@pytest.mark.parametrize( - "keyword, change_plan", - [("dtype", False), ("items_per_thread", False), ("items_per_thread", True)], - ids=["factory-dtype", "factory-extent", "plan-extent"], -) -def test_planned_payload_rejects_conflicting_specialization( - monkeypatch, keyword, change_plan -): - from numba_cuda_mlir import types - - from cuda import coop - from cuda.coop.numba_mlir._compiler._group_load_store import ( - _LoadStorePlanning, - ) - from cuda.coop.numba_mlir._compiler._group_planning import ( - GroupPlanningContext, - ) - from cuda.coop.numba_mlir._compiler._rewrite_support import ( - CoopSinglePhaseRewriteError, - ) - - def kernel(source): - items = coop.ThreadData(items_per_thread=1) - coop.load(coop.this_block(), source, items) - - rewrite_call = GroupPlanningContext.rewrite_call - - def conflicting_call(self, inst, **options): - options["kwargs"][keyword] = types.float32 if keyword == "dtype" else 4 - return rewrite_call(self, inst, **options) - - if change_plan: - monkeypatch.setattr( - _LoadStorePlanning, "_planning_items_per_thread", lambda *_args: 4 - ) - else: - monkeypatch.setattr( - GroupPlanningContext, "rewrite_call", conflicting_call - ) - with pytest.raises( - CoopSinglePhaseRewriteError, - match=f"factory argument '{keyword}' does not match", - ): - _run_single_phase_to_provider_boundary( - kernel, - arg_types=(types.int32[::1],), - monkeypatch=monkeypatch, - ) - - -class _FailingAttribute: - """Raise an error when the rewrite resolves a call result dtype.""" - - def __init__(self, exception_type): - self._exception_type = exception_type - - def __getattr__(self, name): - raise self._exception_type(name) - - -@pytest.mark.parametrize( - "exception_type", - (AttributeError, ImportError, KeyError, TypeError, ValueError), -) -def test_call_result_dtype_treats_attribute_failure_as_unresolved( - exception_type, -): - from numba_cuda_mlir.numba_cuda.compiler import run_frontend - from numba_cuda_mlir.numbair_transforms import ir - - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - - failing_attribute = _FailingAttribute(exception_type) - - def kernel(value): - return failing_attribute.cast(value) - - func_ir = run_frontend(kernel) - state = SimpleNamespace(func_ir=func_ir, args=(), typemap={}) - rewrite = CoopSinglePhaseRewrite(state) - call = next( - inst.value - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) - and isinstance(inst.value, ir.Expr) - and inst.value.op == "call" - ) - - assert rewrite._resolve_call_result_dtype(call) is None - - -@pytest.mark.parametrize("positional_value", (31, None)) -def test_positional_static_runtime_control_cannot_be_repeated_by_keyword( - positional_value, -): - from numba_cuda_mlir import types - from numba_cuda_mlir.numba_cuda.compiler import run_frontend - - import cuda.coop.numba_mlir as coop - from cuda.coop.numba_mlir._compiler._rewrite import ( - CoopSinglePhaseRewrite, - CoopSinglePhaseRewriteError, - ) - from cuda.coop.numba_mlir._lowering._load_store import load as provider_load - - def kernel(source, dynamic_valid_items): - output = coop.ThreadData(items_per_thread=2, dtype=types.int32) - return provider_load( - source, - output, - positional_value, - num_valid_items=dynamic_valid_items, - dtype=types.int32, - threads_per_block=32, - items_per_thread=2, - ) - - func_ir = run_frontend(kernel) - state = SimpleNamespace( - func_ir=func_ir, - args=(types.Array(types.int32, 1, "C"), types.int32), - typingctx=SimpleNamespace(refresh=lambda: None), - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="duplicate runtime argument 'num_valid_items'", - ): - rewrite.prepare_calls_and_storage(func_ir) - - -def test_provider_memory_parameters_require_contiguous_arrays(): - from numba_cuda_mlir import types - - from cuda.coop.numba_mlir import _types - - assert _types.Pointer(types.int32).dtype() == types.Array( - types.int32, 1, "C" - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("later_use", ["runtime", "provider"]) -def test_static_factory_value_across_blocks_preserves_remaining_uses( - monkeypatch, - items_per_thread, - later_use, -): - from numba_cuda_mlir import types - from numba_cuda_mlir.numba_cuda.compiler import run_frontend - from numba_cuda_mlir.numbair_transforms import ir - - import cuda.coop.numba_mlir as coop - from cuda.coop._core import ArgumentBinding - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - from cuda.coop.numba_mlir._lowering._load_store import load as provider_load - - def runtime_use(source, flag, items_per_thread): - valid_items = 31 - output = coop.ThreadData(items_per_thread, dtype=types.int32) - provider_load( - source, - output, - num_valid_items=valid_items, - dtype=types.int32, - threads_per_block=32, - items_per_thread=items_per_thread, - ) - if flag: - return valid_items - return 0 - - def provider_use(source, flag, items_per_thread): - valid_items = 31 - output = coop.ThreadData(items_per_thread, dtype=types.int32) - provider_load( - source, - output, - num_valid_items=valid_items, - dtype=types.int32, - threads_per_block=32, - items_per_thread=items_per_thread, - ) - if flag: - provider_load( - source, - output, - num_valid_items=valid_items, - dtype=types.int32, - threads_per_block=32, - items_per_thread=items_per_thread, - ) - return 0 - - kernel = runtime_use if later_use == "runtime" else provider_use - func_ir = run_frontend(kernel) - state = SimpleNamespace( - func_ir=func_ir, - args=( - types.Array(types.int32, 1, "C"), - types.boolean, - types.IntegerLiteral(items_per_thread), - ), - typingctx=SimpleNamespace(refresh=lambda: None), - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - invocable = SimpleNamespace(files=(), specialization=None) - valid_item_bindings = [] - - def materialize(match): - valid_item_bindings.append(match.factory_kwargs["num_valid_items"]) - return invocable, False - - monkeypatch.setattr( - rewrite, "_prepare_ltoir_bundle_for_matches", lambda _: None - ) - monkeypatch.setattr(rewrite, "_materialize_invocable", materialize) - assert rewrite.prepare_calls_and_storage(func_ir) - rewrite.begin_rewrite() - for label in sorted(func_ir.blocks): - block = func_ir.blocks[label] - if rewrite.match(func_ir, block, state.typemap, state.calltypes): - func_ir.blocks[label] = rewrite.apply() - rewrite.finish_rewrite() - - definition_blocks = { - label - for label, block in func_ir.blocks.items() - for stmt in block.body - if isinstance(stmt, ir.Assign) and stmt.target.name == "valid_items" - } - use_blocks = { - label - for label, block in func_ir.blocks.items() - for stmt in block.body - if any( - var.name == "valid_items" - for var in stmt.list_vars() - if not isinstance(stmt, ir.Assign) or var.name != stmt.target.name - ) - } - if later_use == "runtime": - assert len(definition_blocks) == 1 - assert use_blocks - definition_blocks - expected_calls = 1 - else: - assert not definition_blocks - assert not use_blocks - expected_calls = 2 - assert valid_item_bindings == [ArgumentBinding.static(31)] * expected_calls - calls = [ - call - for target, call in _planned_factory_calls(func_ir, ir) - if target is invocable - ] - assert len(calls) == expected_calls - assert all(len(call.args) == 2 and not call.kws for call in calls) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_common_direct_block_load_store_lowers_to_private_factories( - items_per_thread, -): - pytest.importorskip("numba_cuda_mlir") - from numba_cuda_mlir import types - from numba_cuda_mlir.numbair_transforms import ir - - import cuda.coop.numba_mlir._lowering as lowering - from cuda import coop - from cuda.coop.numba_mlir._compiler._group_planner import has_group_markers - - def memory(source, destination, items_per_thread): - storage = coop.TempStorage() - output = coop.ThreadData(items_per_thread) - coop.load( - coop.this_block(), - source, - output, - algorithm="direct", - valid_items=31, - oob_default=-1, - offset=3, - temp_storage=storage, - ) - coop.store( - coop.this_block(), - destination, - output, - algorithm="direct", - valid_items=31, - offset=3, - temp_storage=storage, - ) - - array_type = types.Array(types.int32, 1, "C") - func_ir, planner = _plan( - memory, - arg_types=( - array_type, - array_type, - types.IntegerLiteral(items_per_thread), - ), - ) - assert has_group_markers(func_ir) - assert planner.run() - assert not has_group_markers(func_ir) - - expected = {lowering.load, lowering.store} - calls = [ - (factory, call) - for factory, call in _planned_factory_calls(func_ir, ir) - if factory in expected - ] - assert Counter(factory for factory, _ in calls) == Counter(expected) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("operation", ("load", "store")) -@pytest.mark.parametrize( - ("algorithm", "storage_free"), - ( - ("direct", True), - ("striped", True), - ("vectorize", True), - ("transpose", False), - ("warp_transpose", False), - ("warp_transpose_timesliced", False), - ), -) -def test_group_planner_selects_algorithm_storage_provider( - qualified, operation, algorithm, storage_free, items_per_thread -): - from numba_cuda_mlir import types - from numba_cuda_mlir.numbair_transforms import ir - - import cuda.coop.numba_mlir as numba_coop - from cuda import coop - from cuda.coop.numba_mlir._lowering import _load_store - - module = numba_coop if qualified else coop - - if operation == "load": - - def memory(values, items_per_thread): - payload = module.ThreadData(items_per_thread, dtype=types.int32) - module.load( - module.this_block(), - values, - payload, - algorithm=algorithm, - ) - - else: - - def memory(values, items_per_thread): - payload = module.ThreadData(items_per_thread, dtype=types.int32) - module.store( - module.this_block(), - values, - payload, - algorithm=algorithm, - ) - - array_type = types.Array(types.int32, 1, "C") - func_ir, planner = _plan( - memory, - arg_types=(array_type, types.IntegerLiteral(items_per_thread)), - ) - assert planner.run() - - expected = getattr( - _load_store, - operation if storage_free else f"_{operation}_with_storage", - ) - calls = [ - call - for factory, call in _planned_factory_calls(func_ir, ir) - if factory is expected - ] - assert len(calls) == 1 - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("operation", ("load", "store")) -def test_group_planner_normalizes_algorithm_strings( - qualified, operation, items_per_thread -): - from numba_cuda_mlir import types - from numba_cuda_mlir.numbair_transforms import ir - - import cuda.coop.numba_mlir as numba_coop - from cuda import coop - from cuda.coop.numba_mlir._lowering import _load_store - - module = numba_coop if qualified else coop - - if operation == "load": - - def memory(values, items_per_thread): - payload = module.ThreadData(items_per_thread, dtype=types.int32) - module.load( - module.this_block(), - values, - payload, - algorithm=" WARP-TRANSPOSE ", - ) - - else: - - def memory(values, items_per_thread): - payload = module.ThreadData(items_per_thread, dtype=types.int32) - module.store( - module.this_block(), - values, - payload, - algorithm=" WARP-TRANSPOSE ", - ) - - array_type = types.Array(types.int32, 1, "C") - func_ir, planner = _plan( - memory, - arg_types=(array_type, types.IntegerLiteral(items_per_thread)), - ) - assert planner.run() - - calls = [ - call - for factory, call in _planned_factory_calls(func_ir, ir) - if factory is getattr(_load_store, f"_{operation}_with_storage") - ] - assert len(calls) == 1 - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize( - "logical_width", (None, 8), ids=("physical", "logical") -) -@pytest.mark.parametrize("operation", ("load", "store")) -@pytest.mark.parametrize( - ("algorithm", "storage_free"), - ( - ("direct", True), - ("striped", True), - ("vectorize", True), - ("transpose", False), - ), -) -def test_warp_planner_selects_declared_provider( - qualified, - logical_width, - operation, - algorithm, - storage_free, - items_per_thread, -): - from numba_cuda_mlir import types - from numba_cuda_mlir.numbair_transforms import ir - - import cuda.coop.numba_mlir as numba_coop - from cuda import coop - from cuda.coop.numba_mlir._lowering import _load_store - - module = numba_coop if qualified else coop - group = module.this_warp() - if logical_width is not None: - group = group.group_by(logical_width) - - if operation == "load": - - def memory(values, items_per_thread): - payload = module.ThreadData(items_per_thread, dtype=types.int32) - module.load( - group, - values, - payload, - algorithm=algorithm, - ) - - else: - - def memory(values, items_per_thread): - payload = module.ThreadData(items_per_thread, dtype=types.int32) - module.store( - group, - values, - payload, - algorithm=algorithm, - ) - - array_type = types.Array(types.int32, 1, "C") - func_ir, planner = _plan( - memory, - arg_types=(array_type, types.IntegerLiteral(items_per_thread)), - ) - assert planner.run() - - expected = getattr( - _load_store, - f"warp_{operation}" - if storage_free - else f"_warp_{operation}_with_storage", - ) - calls = [ - call - for factory, call in _planned_factory_calls(func_ir, ir) - if factory is expected - ] - assert len(calls) == 1 - constants = { - inst.target.name: inst.value.value - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) and isinstance(inst.value, ir.Const) - } - keywords = dict(calls[0].kws) - assert constants[keywords["threads_in_warp"].name] == ( - 32 if logical_width is None else logical_width - ) - assert constants[keywords["threads_per_block"].name] == (64, 1, 1) - assert constants[keywords["algorithm"].name] == algorithm - assert "warp_tile_effective_offset" in keywords["offset"].name - - -@pytest.mark.parametrize("group_kind", ("block", "warp", "logical_warp")) -@pytest.mark.parametrize("operation", ("load", "store")) -@pytest.mark.parametrize( - "algorithm", - ( - pytest.param(True, id="bool"), - pytest.param(0, id="integer"), - pytest.param(object(), id="object"), - ), -) -def test_qualified_planner_rejects_non_string_algorithm( - group_kind, - operation, - algorithm, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as coop - - group = ( - coop.this_warp().group_by(8) - if group_kind == "logical_warp" - else getattr(coop, f"this_{group_kind}")() - ) - - if operation == "load": - - def memory(values): - payload = coop.ThreadData(items_per_thread=2, dtype=types.int32) - coop.load(group, values, payload, algorithm=algorithm) - - else: - - def memory(values): - payload = coop.ThreadData(items_per_thread=2, dtype=types.int32) - coop.store(group, values, payload, algorithm=algorithm) - - array_type = types.Array(types.int32, 1, "C") - _, planner = _plan(memory, arg_types=(array_type,)) - with pytest.raises(TypeError, match="algorithm must be a string"): - planner.run() - - -@pytest.mark.parametrize( - ("group_kind", "algorithm"), - ( - ("block", "stripd"), - ("warp", "warp_transpose"), - ("logical_warp", "warp_transpose"), - ), -) -@pytest.mark.parametrize("operation", ("load", "store")) -def test_qualified_planner_rejects_unsupported_algorithm( - group_kind, - operation, - algorithm, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as coop - - group = ( - coop.this_warp().group_by(8) - if group_kind == "logical_warp" - else getattr(coop, f"this_{group_kind}")() - ) - - if operation == "load": - - def memory(values): - payload = coop.ThreadData(items_per_thread=2, dtype=types.int32) - coop.load(group, values, payload, algorithm=algorithm) - - else: - - def memory(values): - payload = coop.ThreadData(items_per_thread=2, dtype=types.int32) - coop.store(group, values, payload, algorithm=algorithm) - - array_type = types.Array(types.int32, 1, "C") - _, planner = _plan(memory, arg_types=(array_type,)) - with pytest.raises(ValueError, match=r"(?s)algorithm.*must be one\s+of"): - planner.run() - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "logical_width", [None, 8], ids=["physical", "logical"] -) -def test_warp_effective_offset_uses_x_major_group_index( - logical_width, items_per_thread -): - from numba_cuda_mlir import types - from numba_cuda_mlir.numbair_transforms import ir - - from cuda import coop - from cuda.coop.numba_mlir._lowering import _load_store - - group = coop.this_warp() - if logical_width is not None: - group = group.group_by(logical_width) - - def memory(values, offset, items_per_thread): - payload = coop.ThreadData(items_per_thread, dtype=types.int32) - coop.load(group, values, payload, offset=offset) - - array_type = types.Array(types.int32, 1, "C") - func_ir, planner = _plan( - memory, - arg_types=( - array_type, - types.int64, - types.IntegerLiteral(items_per_thread), - ), - block=(16, 4, 1), - ) - assert planner.run() - - definitions = { - inst.target.name: inst.value - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) - } - calls = [ - call - for factory, call in _planned_factory_calls(func_ir, ir) - if factory is _load_store.warp_load - ] - assert len(calls) == 1 - effective = dict(calls[0].kws)["offset"] - assert "warp_tile_effective_offset" in effective.name - assert definitions[effective.name].op == "binop" - assert definitions[effective.name].rhs.name == "offset" - - group_index = next( - value - for name, value in definitions.items() - if "warp_tile_group_index" in name - ) - origin = next( - value - for name, value in definitions.items() - if "warp_tile_origin" in name - ) - linear_rank = next( - value - for name, value in definitions.items() - if "warp_tile_linear_rank" in name - ) - assert group_index.fn is operator.floordiv - assert origin.fn is operator.mul - assert linear_rank.fn is operator.add - - group_width = next( - value - for name, value in definitions.items() - if "warp_tile_group_width" in name - ) - group_tile_items = next( - value - for name, value in definitions.items() - if "warp_tile_group_tile_items" in name - ) - expected_width = 32 if logical_width is None else logical_width - assert group_width.value == expected_width - assert group_tile_items.value == expected_width * items_per_thread - - -def test_qualified_store_recovers_keyword_local_array_extent(): - from numba_cuda_mlir import cuda, types - from numba_cuda_mlir.numbair_transforms import ir - - import cuda.coop.numba_mlir as numba_coop - from cuda.coop.numba_mlir._lowering import _load_store - - def memory(values): - payload = cuda.local.array(shape=2, dtype=types.int32) - numba_coop.store( - numba_coop.this_block(), - values, - payload, - algorithm="transpose", - ) - - array_type = types.Array(types.int32, 1, "C") - func_ir, planner = _plan(memory, arg_types=(array_type,)) - assert planner.run() - - calls = [ - call - for factory, call in _planned_factory_calls(func_ir, ir) - if factory is _load_store._store_with_storage - ] - assert len(calls) == 1 - assert len(calls[0].args) == 2 - - -@pytest.mark.parametrize( - "qualified", [False, True], ids=["common", "qualified"] -) -@pytest.mark.parametrize( - "group_factory", - [ - lambda module: module.this_thread(), - lambda module: module.this_cluster(), - lambda module: module.this_grid(), - ], - ids=["thread", "cluster", "grid"], -) -def test_unsupported_load_group_returns_typed_plan_before_compile( - monkeypatch, qualified, group_factory -): - pytest.importorskip("numba_cuda_mlir") - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as numba_coop - from cuda import coop - - module = numba_coop if qualified else coop - group = group_factory(module) - - def memory(source): - output = module.ThreadData(items_per_thread=2, dtype=types.int32) - return module.load(group, source, output) - - array_type = types.Array(types.int32, 1, "C") - _, planner = _plan(memory, arg_types=(array_type,)) - from cuda.coop.numba_mlir._compiler import _group_load_store - - monkeypatch.setattr( - _group_load_store._LoadStorePlanning, - "_scope_factory", - lambda *_args, **_kwargs: pytest.fail( - "unsupported group reached provider compilation" - ), - ) - with pytest.raises( - NotImplementedError, - match=r"does\s+not\s+support\s+group\s+kind|supports\s+this_block\(\)", - ): - planner.run() - - -@pytest.mark.parametrize( - "qualified", [False, True], ids=["common", "qualified"] -) -@pytest.mark.parametrize("operation", ["load", "store"]) -@pytest.mark.parametrize( - "logical_width", [None, 8], ids=["physical", "logical"] -) -def test_warp_rejects_explicit_temp_storage_before_provider( - monkeypatch, qualified, operation, logical_width -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as numba_coop - from cuda import coop - from cuda.coop.numba_mlir._compiler import _group_load_store - - module = numba_coop if qualified else coop - group = module.this_warp() - if logical_width is not None: - group = group.group_by(logical_width) - - if operation == "load": - - def memory(values): - storage = module.TempStorage() - payload = module.ThreadData(items_per_thread=2, dtype=types.int32) - module.load( - group, - values, - payload, - temp_storage=storage, - ) - - else: - - def memory(values): - storage = module.TempStorage() - payload = module.ThreadData(items_per_thread=2, dtype=types.int32) - module.store( - group, - values, - payload, - temp_storage=storage, - ) - - monkeypatch.setattr( - _group_load_store._LoadStorePlanning, - "_scope_factory", - lambda *_args, **_kwargs: pytest.fail( - "explicit Warp storage reached provider selection" - ), - ) - array_type = types.Array(types.int32, 1, "C") - _, planner = _plan(memory, arg_types=(array_type,)) - with pytest.raises(ValueError, match="temp_storage.*Warp groups"): - planner.run() - - -@pytest.mark.parametrize( - "qualified", [False, True], ids=["common", "qualified"] -) -def test_incomplete_physical_warp_fails_before_provider_selection( - monkeypatch, qualified -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as numba_coop - from cuda import coop - from cuda.coop.numba_mlir._compiler import _group_load_store - - module = numba_coop if qualified else coop - - def memory(values): - payload = module.ThreadData(items_per_thread=2, dtype=types.int32) - module.load(module.this_warp(), values, payload) - - monkeypatch.setattr( - _group_load_store._LoadStorePlanning, - "_scope_factory", - lambda *_args, **_kwargs: pytest.fail( - "incomplete physical Warp reached provider selection" - ), - ) - array_type = types.Array(types.int32, 1, "C") - _, planner = _plan(memory, arg_types=(array_type,), block=(48, 1, 1)) - with pytest.raises(NotImplementedError, match="complete 32-thread warps"): - planner.run() - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("algorithm", ("direct", "striped", "vectorize")) -def test_storage_free_load_store_accept_temp_storage_without_using_it( - algorithm, - items_per_thread, -): - pytest.importorskip("numba_cuda_mlir") - from numba_cuda_mlir import cuda, types - from numba_cuda_mlir.numbair_transforms import ir - - from cuda import coop - from cuda.coop.numba_mlir._compiler._rewrite import ( - CoopSinglePhaseRewrite, - ) - - class TypingContext: - def refresh(self): - pass - - class FakeInvocable: - files = ("movement-test.ltoir",) - temp_storage_bytes = 0 - temp_storage_alignment = 1 - storage_abi = "none" - execution_scope = "block" - synchronization_scope = "none" - - def __call__(self, *args): - del args - - def memory(source, destination, items_per_thread): - shared_storage = coop.TempStorage(sharing="shared") - exclusive_storage = coop.TempStorage(sharing="exclusive") - oversized_storage = coop.TempStorage(128 * 1024, alignment=16) - output = coop.ThreadData(items_per_thread, dtype=types.int32) - extra_output = coop.ThreadData(items_per_thread, dtype=types.int32) - coop.load( - coop.this_block(), - source, - output, - algorithm=algorithm, - temp_storage=shared_storage, - ) - coop.store( - coop.this_block(), - destination, - output, - algorithm=algorithm, - temp_storage=exclusive_storage, - ) - coop.load( - coop.this_block(), - source, - extra_output, - algorithm=algorithm, - temp_storage=oversized_storage, - ) - - array_type = types.Array(types.int32, 1, "C") - func_ir, planner = _plan( - memory, - arg_types=( - array_type, - array_type, - types.IntegerLiteral(items_per_thread), - ), - ) - assert planner.run() - state = SimpleNamespace( - func_ir=func_ir, - args=( - array_type, - array_type, - types.IntegerLiteral(items_per_thread), - ), - typingctx=TypingContext(), - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - invocable = FakeInvocable() - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: None - rewrite._materialize_invocable = lambda _match: (invocable, False) - assert rewrite.prepare_calls_and_storage(func_ir) - rewrite.begin_rewrite() - for label in sorted(func_ir.blocks): - block = func_ir.blocks[label] - while rewrite.match(func_ir, block, state.typemap, state.calltypes): - block = rewrite.apply() - func_ir.blocks[label] = block - rewrite.finish_rewrite() - - calls = _planned_factory_calls(func_ir, ir) - invocable_calls = [call for factory, call in calls if factory is invocable] - assert len(invocable_calls) == 3 - assert all(len(call.args) == 2 for call in invocable_calls) - resolved_calls = _resolved_python_call_targets(func_ir, ir) - assert cuda.shared.array not in resolved_calls - assert cuda.syncthreads not in resolved_calls - assert cuda.syncwarp not in resolved_calls - assert rewrite._temp_storage_global_plan is None - assert rewrite._temp_storage_backing_var is None - assert rewrite._temp_storage_plans == {} - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -@pytest.mark.parametrize("alignment", (1, 16, 32)) -@pytest.mark.parametrize( - ("storage_kind", "expected_size", "expected_barriers"), - ( - ("implicit", 64, 2), - ("shared", 64, 2), - ("exclusive", 128, 0), - ("unsized", 64, 0), - ("compat-none", 64, 0), - ), -) -def test_transpose_storage_contract_reaches_whole_function_rewrite( - qualified, - alignment, - storage_kind, - expected_size, - expected_barriers, - items_per_thread, -): - from numba_cuda_mlir import cuda, types - from numba_cuda_mlir.numbair_transforms import ir - - import cuda.coop.numba_mlir as numba_coop - from cuda import coop - - module = numba_coop if qualified else coop - - if storage_kind == "implicit": - - def memory(source, destination, items_per_thread): - output = module.ThreadData(items_per_thread, dtype=types.int32) - module.load( - module.this_block(), - source, - output, - algorithm="transpose", - ) - module.store( - module.this_block(), - destination, - output, - algorithm="transpose", - ) - - elif storage_kind == "shared": - - def memory(source, destination, items_per_thread): - storage = module.TempStorage( - 64, - alignment=alignment, - auto_sync=True, - sharing="shared", - ) - output = module.ThreadData(items_per_thread, dtype=types.int32) - module.load( - module.this_block(), - source, - output, - algorithm="transpose", - temp_storage=storage, - ) - module.store( - module.this_block(), - destination, - output, - algorithm="transpose", - temp_storage=storage, - ) - - elif storage_kind == "exclusive": - - def memory(source, destination, items_per_thread): - storage = module.TempStorage( - 128, - alignment=alignment, - auto_sync=False, - sharing="exclusive", - ) - output = module.ThreadData(items_per_thread, dtype=types.int32) - module.load( - module.this_block(), - source, - output, - algorithm="transpose", - temp_storage=storage, - ) - module.store( - module.this_block(), - destination, - output, - algorithm="transpose", - temp_storage=storage, - ) - - elif storage_kind == "compat-none": - - def memory(source, destination, items_per_thread): - storage = module.TempStorage(alignment=alignment, auto_sync=None) - output = module.ThreadData(items_per_thread, dtype=types.int32) - module.load( - module.this_block(), - source, - output, - algorithm="transpose", - temp_storage=storage, - ) - module.store( - module.this_block(), - destination, - output, - algorithm="transpose", - temp_storage=storage, - ) - - else: - - def memory(source, destination, items_per_thread): - storage = module.TempStorage(alignment=alignment) - output = module.ThreadData(items_per_thread, dtype=types.int32) - module.load( - module.this_block(), - source, - output, - algorithm="transpose", - temp_storage=storage, - ) - module.store( - module.this_block(), - destination, - output, - algorithm="transpose", - temp_storage=storage, - ) - - array_type = types.Array(types.int32, 1, "C") - func_ir, rewrite, invocables = _rewrite_planned_movement( - memory, - arg_types=( - array_type, - array_type, - types.IntegerLiteral(items_per_thread), - ), - ) - targets = _resolved_python_call_targets(func_ir, ir) - invocable_calls = [ - call - for factory, call in _planned_factory_calls(func_ir, ir) - if factory in invocables - ] - - assert targets.count(cuda.shared.array) == 1 - assert targets.count(cuda.syncthreads) == expected_barriers - assert targets.count(cuda.syncwarp) == 0 - assert len(invocable_calls) == 2 - assert all(len(call.args) == 3 for call in invocable_calls) - assert rewrite._temp_storage_backing_emitted - assert rewrite._temp_storage_global_plan.total_size == expected_size - assert rewrite._temp_storage_global_plan.max_alignment == ( - 16 if storage_kind == "implicit" else max(16, alignment) - ) - assert {invocable.temp_storage_bytes for invocable in invocables} == { - 48, - 64, - } - - -@pytest.mark.parametrize( - "qualified", (False, True), ids=("common", "qualified") -) -def test_transpose_storage_contract_rejects_insufficient_capacity_before_codegen( # noqa: E501 - Preserve descriptive test name. - qualified, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as numba_coop - from cuda import coop - from cuda.coop.numba_mlir._compiler._rewrite import ( - CoopSinglePhaseRewriteError, - ) - - module = numba_coop if qualified else coop - - def memory(source): - storage = module.TempStorage( - 47, - alignment=1, - sharing="shared", - ) - output = module.ThreadData(items_per_thread=2, dtype=types.int32) - return module.load( - module.this_block(), - source, - output, - algorithm="transpose", - temp_storage=storage, - ) - - array_type = types.Array(types.int32, 1, "C") - with pytest.raises( - CoopSinglePhaseRewriteError, match="smaller than required" - ): - _rewrite_planned_movement(memory, arg_types=(array_type,)) - - -def _parameter_names(specialization): - return [ - [type(parameter).__name__ for parameter in overload] - for overload in specialization.parameters - ] - - -def _block_provider_metadata(): - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir._compiler._operations import StorageABI - - return { - "storage_abi": StorageABI.NONE, - "execution_scope": SynchronizationScope.BLOCK, - "synchronization_scope": SynchronizationScope.NONE, - } - - -def test_block_load_store_adapters_preserve_offset_overloads(): - pytest.importorskip("numba_cuda_mlir") - from numba_cuda_mlir import types - - from cuda.coop._core import ArgumentBinding - from cuda.coop._core.block import ( - make_block_load_specialization, - make_block_store_specialization, - ) - from cuda.coop.numba_mlir._lowering._core import NumbaMlirCoreAdapter - from cuda.coop.numba_mlir._lowering._load_store import ( - _load_store_value_abis, - ) - - load = NumbaMlirCoreAdapter( - value_abis=_load_store_value_abis( - dtype=types.int32, - block_dim=(16, 2, 1), - items_per_thread=3, - valid_items=ArgumentBinding.runtime(), - oob_default=ArgumentBinding.runtime(), - ) - ).materialize( - make_block_load_specialization( - dtype=types.int32, - block_dim=(16, 2, 1), - items_per_thread=3, - algorithm="direct", - valid_items=True, - oob_default=True, - include_full_tile=True, - include_pointer_offset=True, - ), - **_block_provider_metadata(), - ) - assert _parameter_names(load) == [ - ["Pointer", "Array"], - [ - "Pointer", - "Array", - "BoundedInteger", - "ExactValue", - ], - [ - "Pointer", - "Array", - "BoundedInteger", - "ExactValue", - "PointerOffset", - ], - ["Pointer", "Array", "PointerOffset"], - ] - - store = NumbaMlirCoreAdapter().materialize( - make_block_store_specialization( - dtype=types.int32, - block_dim=(32, 1, 1), - items_per_thread=2, - algorithm="direct", - include_pointer_offset=True, - ), - **_block_provider_metadata(), - ) - assert _parameter_names(store) == [ - ["Pointer", "Array"], - ["Pointer", "Array", "PointerOffset"], - ] - - -def test_static_movement_controls_are_absent_from_numba_runtime_abi(): - pytest.importorskip("numba_cuda_mlir") - from numba_cuda_mlir import types - - from cuda.coop._core import ArgumentBinding - from cuda.coop._core.block import ( - make_block_load_specialization, - make_block_store_specialization, - ) - from cuda.coop.numba_mlir._lowering._core import NumbaMlirCoreAdapter - from cuda.coop.numba_mlir._types import PointerOffset, algo_coalesce_key - - adapter = NumbaMlirCoreAdapter() - load = adapter.materialize( - make_block_load_specialization( - dtype=types.int32, - block_dim=(32, 1, 1), - items_per_thread=2, - algorithm="direct", - valid_items=ArgumentBinding.static(17), - oob_default=ArgumentBinding.static(-1), - include_pointer_offset=ArgumentBinding.static(3), - ), - **_block_provider_metadata(), - ) - assert len(load.parameters) == 1 - method = load.parameters[0] - offset = method[-1] - assert isinstance(offset, PointerOffset) - assert offset.static_value == 3 - assert not offset.is_provided_by_user() - assert sum(parameter.is_provided_by_user() for parameter in method) == 2 - - source = load._source_code()[0] - assert ".Load((src + 3)," in source - assert "::cuda::std::int64_t offset" not in source - assert ", 17, -1);" in source - - static_four = adapter.materialize( - make_block_store_specialization( - dtype=types.int32, - block_dim=(32, 1, 1), - items_per_thread=2, - algorithm="direct", - include_pointer_offset=ArgumentBinding.static(4), - ), - **_block_provider_metadata(), - ) - static_five = adapter.materialize( - make_block_store_specialization( - dtype=types.int32, - block_dim=(32, 1, 1), - items_per_thread=2, - algorithm="direct", - include_pointer_offset=ArgumentBinding.static(5), - ), - **_block_provider_metadata(), - ) - assert static_four.mangled_name(static_four.parameters[0]) != ( - static_five.mangled_name(static_five.parameters[0]) - ) - assert algo_coalesce_key(static_four) != algo_coalesce_key(static_five) - - runtime = adapter.materialize( - make_block_store_specialization( - dtype=types.int32, - block_dim=(32, 1, 1), - items_per_thread=2, - algorithm="direct", - include_pointer_offset=ArgumentBinding.runtime(), - ), - **_block_provider_metadata(), - ) - runtime_source = runtime._source_code()[0] - assert "::cuda::std::int64_t offset" in runtime_source - assert ".Store((dst + offset)," in runtime_source - - -def test_runtime_valid_items_is_checked_before_cub_integer_narrowing(): - from numba_cuda_mlir import types - - from cuda.coop._core import ArgumentBinding - from cuda.coop._core.block import make_block_load_specialization - from cuda.coop.numba_mlir._lowering._core import NumbaMlirCoreAdapter - from cuda.coop.numba_mlir._lowering._load_store import ( - _load_store_value_abis, - ) - from cuda.coop.numba_mlir._types import BoundedInteger, algo_coalesce_key - - adapter = NumbaMlirCoreAdapter( - value_abis=_load_store_value_abis( - dtype=types.int32, - block_dim=(16, 2, 1), - items_per_thread=3, - valid_items=ArgumentBinding.runtime(), - ) - ) - load = adapter.materialize( - make_block_load_specialization( - dtype=types.int32, - block_dim=(16, 2, 1), - items_per_thread=3, - algorithm="direct", - valid_items=ArgumentBinding.runtime(), - ), - **_block_provider_metadata(), - ) - valid_items = load.parameters[0][-1] - assert isinstance(valid_items, BoundedInteger) - assert valid_items.dtype() == types.int64 - assert valid_items.provider_dtype == types.int32 - assert (valid_items.minimum, valid_items.maximum) == (0, 96) - assert all( - valid_items.accepts_actual_type(dtype, None) - for dtype in ( - types.int8, - types.int16, - types.int32, - types.int64, - types.uint8, - types.uint16, - types.uint32, - ) - ) - assert all( - not valid_items.accepts_actual_type(dtype, None) - for dtype in (types.boolean, types.uint64, types.float32, types.float64) - ) - - source = load._source_code()[0] - declaration = "::cuda::std::int64_t num_valid_items" - guard = "if (num_valid_items < 0 || num_valid_items > 96)" - trap = 'asm volatile("trap;" : : :);' - narrowing = ( - "::cuda::std::int32_t checked_num_valid_items = " - "static_cast<::cuda::std::int32_t>(num_valid_items);" - ) - invocation = ".Load(src, *reinterpret_cast<" - assert source.index(declaration) < source.index(guard) - assert source.index(guard) < source.index(trap) < source.index(narrowing) - assert source.index(narrowing) < source.index(invocation) - assert "checked_num_valid_items);" in source - - narrower_adapter = NumbaMlirCoreAdapter( - value_abis=_load_store_value_abis( - dtype=types.int32, - block_dim=(16, 1, 1), - items_per_thread=3, - valid_items=ArgumentBinding.runtime(), - ) - ) - narrower_tile = narrower_adapter.materialize( - make_block_load_specialization( - dtype=types.int32, - block_dim=(16, 1, 1), - items_per_thread=3, - algorithm="direct", - valid_items=ArgumentBinding.runtime(), - ), - **_block_provider_metadata(), - ) - assert algo_coalesce_key(load) != algo_coalesce_key(narrower_tile) - assert load.mangled_name(load.parameters[0]) != narrower_tile.mangled_name( - narrower_tile.parameters[0] - ) - - -def test_arbitrary_primitive_declares_bounded_and_exact_value_abis(): - from numba_cuda_mlir import types - - from cuda.coop._core import ( - FLOAT32, - INT32, - Algorithm, - SynchronizationScope, - Value, - ) - from cuda.coop.numba_mlir._compiler._operations import StorageABI - from cuda.coop.numba_mlir._lowering._core import NumbaMlirCoreAdapter - from cuda.coop.numba_mlir._types import ( - BoundedInteger, - ExactValue, - algo_coalesce_key, - ) - from cuda.coop.numba_mlir._types import ( - Value as BackendValue, - ) - - specialization = Algorithm( - struct_name="FuturePrimitive", - method_name="Run", - c_name="test_future_primitive", - includes=(), - template_parameters=(), - parameters=( - ( - Value(INT32, name="remaining"), - Value(FLOAT32, name="fill"), - ), - ), - template_arguments={}, - metadata={"primitive": "not_load_or_store"}, - ) - bounded = BoundedInteger(types.int32, minimum=-2, maximum=7) - exact = ExactValue(types.float32) - algorithm = NumbaMlirCoreAdapter( - value_abis={"remaining": bounded, "fill": exact} - ).materialize( - specialization, - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.NONE, - synchronization_scope=SynchronizationScope.NONE, - ) - - lowered_bounded, lowered_exact = algorithm.parameters[0] - assert lowered_bounded is not bounded - assert lowered_exact is not exact - assert lowered_bounded.parameter_name == "remaining" - assert lowered_exact.parameter_name == "fill" - assert bounded.provider_dtype == types.int32 - assert bounded.dtype() == types.int64 - assert bounded.accepts_actual_type(types.uint32, None) - assert not bounded.accepts_actual_type(types.uint64, None) - assert exact.accepts_actual_type(types.float32, None) - assert not exact.accepts_actual_type(types.float64, None) - - source = algorithm._source_code()[0] - guard = "if (remaining < -2 || remaining > 7)" - trap = 'asm volatile("trap;" : : :);' - narrowing = ( - "::cuda::std::int32_t checked_remaining = " - "static_cast<::cuda::std::int32_t>(remaining);" - ) - assert source.index(guard) < source.index(trap) < source.index(narrowing) - assert "checked_remaining, fill);" in source - - wider = NumbaMlirCoreAdapter( - value_abis={ - "remaining": BoundedInteger(types.int32, minimum=-2, maximum=8), - "fill": exact, - } - ).materialize( - specialization, - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.NONE, - synchronization_scope=SynchronizationScope.NONE, - ) - assert algo_coalesce_key(algorithm) != algo_coalesce_key(wider) - assert algorithm.mangled_name( - algorithm.parameters[0] - ) != wider.mangled_name(wider.parameters[0]) - - load_like = Algorithm( - struct_name="LegacyLikePrimitive", - method_name="Run", - c_name="test_load_like_primitive", - includes=(), - template_parameters=(), - parameters=( - ( - Value(INT32, name="num_valid_items"), - Value(FLOAT32, name="oob_default"), - ), - ), - template_arguments={}, - metadata={"primitive": "load"}, - ) - without_declarations = NumbaMlirCoreAdapter().materialize( - load_like, - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.NONE, - synchronization_scope=SynchronizationScope.NONE, - ) - assert all( - type(parameter) is BackendValue - for parameter in without_declarations.parameters[0] - ) - - -def test_value_abi_declarations_fail_closed_during_materialization(): - from numba_cuda_mlir import types - - from cuda.coop._core import ( - INT32, - Algorithm, - Pointer, - SynchronizationScope, - Value, - ) - from cuda.coop.numba_mlir._compiler._operations import StorageABI - from cuda.coop.numba_mlir._lowering._core import NumbaMlirCoreAdapter - from cuda.coop.numba_mlir._types import ExactValue - - def specialization(parameter): - return Algorithm( - struct_name="FuturePrimitive", - method_name="Run", - c_name="test_value_abi_validation", - includes=(), - template_parameters=(), - parameters=((parameter,),), - template_arguments={}, - ) - - def materialize(parameter, value_abis): - return NumbaMlirCoreAdapter(value_abis=value_abis).materialize( - specialization(parameter), - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.NONE, - synchronization_scope=SynchronizationScope.NONE, - ) - - with pytest.raises(ValueError, match="unknown .* value ABI"): - materialize( - Value(INT32, name="value"), {"missing": ExactValue(types.int32)} - ) - with pytest.raises(ValueError, match="require scalar Value parameters"): - materialize( - Pointer(INT32, name="value"), - {"value": ExactValue(types.int32)}, - ) - with pytest.raises(ValueError, match="provider dtypes do not match"): - materialize( - Value(INT32, name="value"), - {"value": ExactValue(types.float32)}, - ) - with pytest.raises(ValueError, match="output roles do not match"): - materialize( - Value(INT32, name="value", is_output=True), - {"value": ExactValue(types.int32)}, - ) - with pytest.raises(TypeError, match="backend Value instances"): - materialize(Value(INT32, name="value"), {"value": object()}) - - -def test_runtime_offsets_reject_bool_float_and_uint64_types(): - from numba_cuda_mlir import types - - from cuda.coop.numba_mlir._types import PointerOffset - - offset = PointerOffset(types.int64) - assert all( - offset.accepts_actual_type(dtype, None) - for dtype in ( - types.int8, - types.int16, - types.int32, - types.int64, - types.uint8, - types.uint16, - types.uint32, - ) - ) - assert all( - not offset.accepts_actual_type(dtype, None) - for dtype in (types.boolean, types.uint64, types.float32, types.float64) - ) - - -_INTEGER_LITERAL_DTYPES = ( - pytest.param("int8", np.int8, id="int8"), - pytest.param("uint8", np.uint8, id="uint8"), - pytest.param("int16", np.int16, id="int16"), - pytest.param("uint16", np.uint16, id="uint16"), - pytest.param("int32", np.int32, id="int32"), - pytest.param("uint32", np.uint32, id="uint32"), - pytest.param("int64", np.int64, id="int64"), - pytest.param("uint64", np.uint64, id="uint64"), -) - - -@pytest.mark.parametrize( - ("operation", "parameter"), - [("store", "value"), ("load", "oob_default")], -) -@pytest.mark.parametrize(("dtype_name", "numpy_dtype"), _INTEGER_LITERAL_DTYPES) -def test_python_integer_literal_boundaries_are_checked_against_payload_dtype( - operation, - parameter, - dtype_name, - numpy_dtype, -): - from numba_cuda_mlir import types - - from cuda.coop.numba_mlir._compiler._parameters import coerce_static_scalar - - dtype = getattr(types, dtype_name) - bounds = np.iinfo(numpy_dtype) - for value in (int(bounds.min), int(bounds.max)): - result = coerce_static_scalar( - value, - dtype, - operation=operation, - parameter=parameter, - ) - assert isinstance(result, numpy_dtype) - assert int(result) == value - - for value in (int(bounds.min) - 1, int(bounds.max) + 1): - with pytest.raises(ValueError, match="outside the range"): - coerce_static_scalar( - value, - dtype, - operation=operation, - parameter=parameter, - ) - - -@pytest.mark.parametrize( - ("operation", "parameter"), - [("store", "value"), ("load", "oob_default")], -) -@pytest.mark.parametrize( - ("dtype_name", "numpy_dtype"), - [("float32", np.float32), ("float64", np.float64)], -) -def test_python_float_literal_boundaries_are_checked_against_payload_dtype( - operation, - parameter, - dtype_name, - numpy_dtype, -): - from numba_cuda_mlir import types - - from cuda.coop.numba_mlir._compiler._parameters import coerce_static_scalar - - dtype = getattr(types, dtype_name) - maximum = float(np.finfo(numpy_dtype).max) - for value in (-maximum, 0, 1, 1.25, maximum): - result = coerce_static_scalar( - value, - dtype, - operation=operation, - parameter=parameter, - ) - assert isinstance(result, numpy_dtype) - assert np.isfinite(result) - - for value in (float("-inf"), float("inf"), float("nan"), 1 << 1024): - with pytest.raises(ValueError, match="finite"): - coerce_static_scalar( - value, - dtype, - operation=operation, - parameter=parameter, - ) - - -@pytest.mark.parametrize( - ("operation", "parameter"), - [("store", "value"), ("load", "oob_default")], -) -def test_contextual_scalar_conversion_rejects_bool_float_to_int_and_typed_casts( - operation, - parameter, -): - from numba_cuda_mlir import types - - from cuda.coop.numba_mlir._compiler._parameters import coerce_static_scalar - - with pytest.raises(TypeError, match="must not be bool"): - coerce_static_scalar( - True, - types.int32, - operation=operation, - parameter=parameter, - ) - with pytest.raises(TypeError, match="float-to-integer"): - coerce_static_scalar( - 1.0, - types.int32, - operation=operation, - parameter=parameter, - ) - assert coerce_static_scalar( - np.int32(1), - types.int32, - operation=operation, - parameter=parameter, - ).dtype == np.dtype(np.int32) - with pytest.raises(TypeError, match="does not match payload dtype"): - coerce_static_scalar( - np.int64(1), - types.int32, - operation=operation, - parameter=parameter, - ) - assert coerce_static_scalar( - 1, - types.int32, - operation=operation, - parameter=parameter, - source_dtype=types.int32, - ).dtype == np.dtype(np.int32) - with pytest.raises(TypeError, match="does not match payload dtype"): - coerce_static_scalar( - 1, - types.int32, - operation=operation, - parameter=parameter, - source_dtype=types.int64, - ) - - -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize( - ("dtype_name", "value"), - [ - pytest.param("uint8", 255, id="uint8-upper"), - pytest.param("float32", 1, id="int-to-float"), - pytest.param("float32", 1.25, id="float-rounding"), - pytest.param("int32", np.int32(-7), id="typed-exact"), - ], -) -def test_scalar_store_literals_are_typed_from_the_destination( - qualified, - dtype_name, - value, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - module = qualified_coop if qualified else root_coop - - def memory(destination): - module.store(module.this_block(), destination, value) - - array_type = types.Array(getattr(types, dtype_name), 1, "C") - _, planner = _plan(memory, arg_types=(array_type,)) - assert planner.run() - - -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize( - ("dtype_name", "value", "error"), - [ - pytest.param("uint8", 256, "outside the range", id="out-of-range"), - pytest.param("int32", 1.0, "float-to-integer", id="float-to-int"), - pytest.param("int32", True, "bool", id="bool"), - pytest.param( - "int32", - np.int64(1), - "does not match payload dtype", - id="typed-mismatch", - ), - ], -) -def test_scalar_store_literals_fail_before_provider_selection( - monkeypatch, - qualified, - dtype_name, - value, - error, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - from cuda.coop.numba_mlir._compiler import _group_load_store - - module = qualified_coop if qualified else root_coop - - def memory(destination): - module.store(module.this_block(), destination, value) - - monkeypatch.setattr( - _group_load_store._LoadStorePlanning, - "_scope_factory", - lambda *_args, **_kwargs: pytest.fail( - "invalid scalar Store reached provider selection" - ), - ) - array_type = types.Array(getattr(types, dtype_name), 1, "C") - with pytest.raises((TypeError, ValueError), match=error): - _plan(memory, arg_types=(array_type,))[1].run() - - -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize( - ("value_type_name", "accepted"), - [("int32", True), ("int64", False)], -) -def test_runtime_scalar_store_requires_exact_destination_dtype( - monkeypatch, - qualified, - value_type_name, - accepted, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - module = qualified_coop if qualified else root_coop - - def memory(destination, value): - module.store(module.this_block(), destination, value) - - array_type = types.Array(types.int32, 1, "C") - arg_types = (array_type, getattr(types, value_type_name)) - if accepted: - _run_single_phase_to_provider_boundary( - memory, - arg_types=arg_types, - monkeypatch=monkeypatch, - allow_provider_bundling=True, - ) - else: - with pytest.raises(TypeError, match="does not match payload dtype"): - _plan(memory, arg_types=arg_types)[1].run() - - -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize( - "source_kind", - ["index", "element", "numpy-cast", "compiler-cast"], -) -def test_cuda_and_array_scalars_keep_compiler_dtypes_for_store( - qualified, - source_kind, -): - from numba_cuda_mlir import cuda, types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - module = qualified_coop if qualified else root_coop - - def index_source(source, destination): - module.store(module.this_block(), destination, cuda.threadIdx.x) - - def element_source(source, destination): - module.store(module.this_block(), destination, source[cuda.threadIdx.x]) - - def numpy_cast_source(source, destination): - module.store( - module.this_block(), destination, np.int32(cuda.threadIdx.x + 1) - ) - - def compiler_cast_source(source, destination): - module.store( - module.this_block(), destination, types.int32(cuda.threadIdx.x + 1) - ) - - memory = { - "index": index_source, - "element": element_source, - "numpy-cast": numpy_cast_source, - "compiler-cast": compiler_cast_source, - }[source_kind] - - array_type = types.Array(types.int32, 1, "C") - _, planner = _plan(memory, arg_types=(array_type, array_type)) - assert planner.run() - - -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -def test_runtime_scalar_expression_cannot_narrow_into_store(qualified): - from numba_cuda_mlir import cuda, types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - module = qualified_coop if qualified else root_coop - - def memory(destination): - value = cuda.threadIdx.x + 1 - module.store(module.this_block(), destination, value) - - array_type = types.Array(types.int32, 1, "C") - with pytest.raises(TypeError, match="does not match payload dtype"): - _plan(memory, arg_types=(array_type,))[1].run() - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize("parameter", ["valid_items", "offset"]) -@pytest.mark.parametrize( - "dtype_name", - ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32"], -) -def test_runtime_control_integer_domain_is_accepted_before_materialization( - monkeypatch, - qualified, - parameter, - dtype_name, - items_per_thread, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - module = qualified_coop if qualified else root_coop - - if parameter == "valid_items": - - def memory(source, control, items_per_thread): - output = module.ThreadData(items_per_thread, dtype=types.int32) - return module.load( - module.this_block(), source, output, valid_items=control - ) - - else: - - def memory(source, control, items_per_thread): - output = module.ThreadData(items_per_thread, dtype=types.int32) - return module.load( - module.this_block(), source, output, offset=control - ) - - array_type = types.Array(types.int32, 1, "C") - _run_single_phase_to_provider_boundary( - memory, - arg_types=( - array_type, - getattr(types, dtype_name), - types.IntegerLiteral(items_per_thread), - ), - monkeypatch=monkeypatch, - allow_provider_bundling=True, - ) - - -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize("parameter", ["valid_items", "offset"]) -@pytest.mark.parametrize( - "dtype_name", - ["boolean", "uint64", "float32", "float64"], -) -def test_runtime_control_invalid_types_fail_before_materialization( - monkeypatch, - qualified, - parameter, - dtype_name, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - from cuda.coop.numba_mlir._compiler._rewrite import ( - CoopSinglePhaseRewriteError, - ) - - module = qualified_coop if qualified else root_coop - - if parameter == "valid_items": - - def memory(source, control): - output = module.ThreadData(items_per_thread=2, dtype=types.int32) - return module.load( - module.this_block(), source, output, valid_items=control - ) - - else: - - def memory(source, control): - output = module.ThreadData(items_per_thread=2, dtype=types.int32) - return module.load( - module.this_block(), source, output, offset=control - ) - - array_type = types.Array(types.int32, 1, "C") - with pytest.raises( - CoopSinglePhaseRewriteError, - match=r"must be (?:an integer|a signed integer)", - ): - _run_single_phase_to_provider_boundary( - memory, - arg_types=(array_type, getattr(types, dtype_name)), - monkeypatch=monkeypatch, - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_single_phase_rewrite_preserves_static_block_movement_bindings( - items_per_thread, -): - pytest.importorskip("numba_cuda_mlir") - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as coop - from cuda.coop._core import ArgumentBinding, BindingKind - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - - def memory(source, destination, dynamic_offset, items_per_thread): - output = coop.ThreadData(items_per_thread, dtype=types.int32) - coop.load( - coop.this_block(), - source, - output, - valid_items=31, - oob_default=-1, - offset=3, - ) - coop.store( - coop.this_block(), - destination, - output, - offset=dynamic_offset, - ) - - array_type = types.Array(types.int32, 1, "C") - arg_types = (array_type, array_type, types.int64) - func_ir, planner = _plan( - memory, arg_types=(*arg_types, types.IntegerLiteral(items_per_thread)) - ) - assert planner.run() - - class TypingContext: - def refresh(self): - pass - - state = SimpleNamespace( - func_ir=func_ir, - args=(*arg_types, types.IntegerLiteral(items_per_thread)), - typingctx=TypingContext(), - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - matches = [] - assert rewrite.prepare_calls_and_storage(func_ir) - for label in sorted(func_ir.blocks): - block = func_ir.blocks[label] - if rewrite.match(func_ir, block, state.typemap, state.calltypes): - matches.extend(rewrite._matches.values()) - - load_match = next(match for match in matches if match.op_name == "load") - assert len(load_match.runtime_args) == 2 - assert load_match.factory_kwargs[ - "num_valid_items" - ] == ArgumentBinding.static(31) - assert load_match.factory_kwargs["oob_default"] == ArgumentBinding.static( - np.int32(-1) - ) - assert load_match.factory_kwargs["offset"] == ArgumentBinding.static(3) - - store_match = next(match for match in matches if match.op_name == "store") - assert len(store_match.runtime_args) == 3 - assert "offset" not in store_match.factory_kwargs - assert all( - not isinstance(value, ArgumentBinding) - or value.kind is not BindingKind.RUNTIME - for value in load_match.factory_kwargs.values() - ) - - -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize( - ("default_type_name", "error"), - [ - ("boolean", "supports oob_default dtypes"), - ("float16", "supports oob_default dtypes"), - ("complex64", "supports oob_default dtypes"), - ("optional", "supports oob_default dtypes"), - ("int64", "does not match payload dtype"), - ], - ids=["bool", "float16", "complex", "optional", "mismatched"], -) -def test_runtime_oob_default_rejects_before_provider_materialization( - monkeypatch, - qualified, - default_type_name, - error, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - from cuda.coop.numba_mlir._compiler._rewrite import ( - CoopSinglePhaseRewriteError, - ) - - module = qualified_coop if qualified else root_coop - - def memory(source, valid_items, oob_default): - output = module.ThreadData(items_per_thread=2, dtype=types.int32) - return module.load( - module.this_block(), - source, - output, - valid_items=valid_items, - oob_default=oob_default, - ) - - default_type = ( - types.Optional(types.int32) - if default_type_name == "optional" - else getattr(types, default_type_name) - ) - array_type = types.Array(types.int32, 1, "C") - with pytest.raises( - (CoopSinglePhaseRewriteError, TypeError), match=error - ) as exc: - _run_single_phase_to_provider_boundary( - memory, - arg_types=(array_type, types.int32, default_type), - monkeypatch=monkeypatch, - ) - - if default_type_name == "int64": - from cuda.coop.numba_mlir._compiler._group_errors import ( - DefaultDtypeMismatchError, - ) - - assert isinstance(exc.value, DefaultDtypeMismatchError) - lines = str(exc.value).splitlines() - assert len(lines) == 2 - assert all(len(line) <= 80 for line in lines) - - -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize( - "oob_default", - [ - pytest.param(True, id="bool"), - pytest.param(np.float16(1), id="float16"), - pytest.param(1 + 2j, id="complex"), - pytest.param(object(), id="object"), - pytest.param(float("inf"), id="nonfinite"), - pytest.param(1 << 65, id="outside-64-bit"), - ], -) -def test_static_oob_default_rejects_before_provider_materialization( - monkeypatch, - qualified, - oob_default, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - from cuda.coop.numba_mlir._compiler._rewrite import ( - CoopSinglePhaseRewriteError, - ) - - module = qualified_coop if qualified else root_coop - - def memory(source): - output = module.ThreadData(items_per_thread=2, dtype=types.int32) - return module.load( - module.this_block(), - source, - output, - valid_items=31, - oob_default=oob_default, - ) - - array_type = types.Array(types.int32, 1, "C") - with pytest.raises( - (CoopSinglePhaseRewriteError, TypeError, ValueError), - match="oob_default", - ): - _run_single_phase_to_provider_boundary( - memory, - arg_types=(array_type,), - monkeypatch=monkeypatch, - ) - - -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize("operation", ["load", "store"]) -def test_public_algorithms_require_plain_strings(qualified, operation): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - module = qualified_coop if qualified else root_coop - - def memory(array): - payload = module.ThreadData(items_per_thread=2, dtype=types.int32) - if operation == "load": - return module.load( - module.this_block(), - array, - payload, - algorithm=_StringAlgorithm.DIRECT, - ) - module.store( - module.this_block(), - array, - payload, - algorithm=_StringAlgorithm.DIRECT, - ) - - array_type = types.Array(types.int32, 1, "C") - with pytest.raises(TypeError, match="algorithm must be a string"): - _plan(memory, arg_types=(array_type,))[1].run() - - -@pytest.mark.parametrize("operation", ["load", "store"]) -@pytest.mark.parametrize( - ("algorithm", "storage_free"), - [ - ("direct", True), - ("striped", True), - ("vectorize", True), - ("transpose", False), - ("warp_transpose", False), - ("warp_transpose_timesliced", False), - ], -) -def test_block_algorithms_use_their_declared_provider( - monkeypatch, operation, algorithm, storage_free -): - from numba_cuda_mlir import types - - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir._compiler._operations import StorageABI - from cuda.coop.numba_mlir._lowering import _load_store - - materializations = [] - monkeypatch.setattr( - _load_store.NumbaMlirCoreAdapter, - "materialize", - lambda _self, specialization, **kwargs: ( - materializations.append((specialization, kwargs)) or specialization - ), - ) - monkeypatch.setattr( - _load_store, - "make_invocable_from_specialization", - lambda specialization: specialization, - ) - factory = getattr( - _load_store, - operation if storage_free else f"_{operation}_with_storage", - ) - - specialization = factory( - types.int32, - threads_per_block=32, - algorithm=algorithm, - ) - - assert specialization.metadata["algorithm"] == algorithm - assert len(materializations) == 1 - assert materializations[0][1]["storage_abi"] is ( - StorageABI.NONE if storage_free else StorageABI.LEADING_POINTER - ) - assert ( - materializations[0][1]["execution_scope"] is SynchronizationScope.BLOCK - ) - assert materializations[0][1]["synchronization_scope"] is ( - SynchronizationScope.NONE - if storage_free - else SynchronizationScope.BLOCK - ) - - -@pytest.mark.parametrize("operation", ["load", "store"]) -def test_block_algorithm_string_is_normalized_before_materialization( - monkeypatch, operation -): - from numba_cuda_mlir import types - - from cuda.coop.numba_mlir._lowering import _load_store - - monkeypatch.setattr( - _load_store.NumbaMlirCoreAdapter, - "materialize", - lambda _self, specialization, **_kwargs: specialization, - ) - monkeypatch.setattr( - _load_store, - "make_invocable_from_specialization", - lambda specialization: specialization, - ) - factory = getattr(_load_store, f"_{operation}_with_storage") - - specialization = factory( - types.int32, - threads_per_block=32, - algorithm=" WARP-TRANSPOSE ", - ) - - assert specialization.metadata["algorithm"] == "warp_transpose" - - -@pytest.mark.parametrize("operation", ["load", "store"]) -@pytest.mark.parametrize( - ("algorithm", "storage_free"), - [ - ("direct", True), - ("striped", True), - ("vectorize", True), - ("transpose", False), - ], -) -@pytest.mark.parametrize("threads_in_warp", [1, 2, 4, 8, 16, 32]) -def test_warp_algorithms_use_their_declared_provider( - monkeypatch, operation, algorithm, storage_free, threads_in_warp -): - from numba_cuda_mlir import types - - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir._compiler._operations import StorageABI - from cuda.coop.numba_mlir._lowering import _load_store - - materializations = [] - monkeypatch.setattr( - _load_store.NumbaMlirCoreAdapter, - "materialize", - lambda _self, specialization, **kwargs: ( - materializations.append((specialization, kwargs)) or specialization - ), - ) - monkeypatch.setattr( - _load_store, - "make_invocable_from_specialization", - lambda specialization, **kwargs: (specialization, kwargs), - ) - factory = getattr( - _load_store, - f"warp_{operation}" - if storage_free - else f"_warp_{operation}_with_storage", - ) - - specialization, invocation_kwargs = factory( - types.int32, - threads_in_warp=threads_in_warp, - threads_per_block=(64, 1, 1), - algorithm=algorithm, - ) - - assert specialization.metadata["algorithm"] == algorithm - assert specialization.metadata["scope"] == "warp" - assert len(materializations) == 1 - assert materializations[0][1]["storage_abi"] is ( - StorageABI.NONE if storage_free else StorageABI.LEADING_POINTER - ) - assert ( - materializations[0][1]["execution_scope"] is SynchronizationScope.WARP - ) - assert materializations[0][1]["synchronization_scope"] is ( - SynchronizationScope.NONE if storage_free else SynchronizationScope.WARP - ) - assert invocation_kwargs == { - "logical_warp_threads": threads_in_warp, - "block_threads": (64, 1, 1), - } - - -@pytest.mark.parametrize("operation", ["load", "store"]) -@pytest.mark.parametrize( - "algorithm", - [True, 0, "stripd", "warp_transpose", object()], -) -def test_invalid_physical_warp_algorithm_fails_before_materialization( - monkeypatch, operation, algorithm -): - from numba_cuda_mlir import types - - from cuda.coop.numba_mlir._lowering import _load_store - - monkeypatch.setattr( - _load_store.NumbaMlirCoreAdapter, - "materialize", - lambda *args, **kwargs: pytest.fail( - "invalid Warp algorithm reached provider materialization" - ), - ) - factory = getattr(_load_store, f"warp_{operation}") - error_type = TypeError if not isinstance(algorithm, str) else ValueError - - with pytest.raises(error_type): - factory( - types.int32, - threads_in_warp=32, - threads_per_block=(64, 1, 1), - algorithm=algorithm, - ) - - -@pytest.mark.parametrize("operation", ["load", "store"]) -@pytest.mark.parametrize("threads", [0, 3, 31, 64, True, 32.0]) -def test_warp_provider_rejects_invalid_width_before_materialization( - monkeypatch, operation, threads -): - from numba_cuda_mlir import types - - from cuda.coop.numba_mlir._lowering import _load_store - - monkeypatch.setattr( - _load_store.NumbaMlirCoreAdapter, - "materialize", - lambda *args, **kwargs: pytest.fail( - "invalid Warp width reached provider materialization" - ), - ) - factory = getattr(_load_store, f"warp_{operation}") - - with pytest.raises( - (TypeError, ValueError), - match=r"logical width|integer|positive", - ): - factory( - types.int32, - threads_in_warp=threads, - threads_per_block=(64, 1, 1), - ) - - -@pytest.mark.parametrize("operation", ["load", "store"]) -@pytest.mark.parametrize( - ("algorithm", "error_type", "message"), - [ - (True, TypeError, "algorithm must be a string"), - (0, TypeError, "algorithm must be a string"), - (_StringAlgorithm.DIRECT, TypeError, "algorithm must be a string"), - ("stripd", ValueError, "expected one of"), - (object(), TypeError, "algorithm must be a string"), - ], -) -def test_invalid_block_algorithm_values_fail_before_provider_materialization( - monkeypatch, operation, algorithm, error_type, message -): - from numba_cuda_mlir import types - - from cuda.coop.numba_mlir._lowering import _load_store - - monkeypatch.setattr( - _load_store.NumbaMlirCoreAdapter, - "materialize", - lambda *args, **kwargs: pytest.fail( - "invalid algorithm reached provider materialization" - ), - ) - factory = getattr(_load_store, operation) - with pytest.raises(error_type, match=message): - factory(types.int32, threads_per_block=32, algorithm=algorithm) - - -@pytest.mark.parametrize("operation", ["load", "store"]) -@pytest.mark.parametrize("dtype_name", ["boolean", "float16", "complex64"]) -def test_unsupported_dtypes_fail_before_provider_materialization( - monkeypatch, operation, dtype_name -): - from numba_cuda_mlir import types - - from cuda.coop.numba_mlir._lowering import _load_store - - monkeypatch.setattr( - _load_store.NumbaMlirCoreAdapter, - "materialize", - lambda *args, **kwargs: pytest.fail( - "unsupported dtype reached provider materialization" - ), - ) - factory = getattr(_load_store, operation) - - with pytest.raises(TypeError, match="supports dtypes"): - factory(getattr(types, dtype_name), threads_per_block=32) - - -@pytest.mark.parametrize( - "oob_default", - [ - pytest.param(True, id="bool"), - pytest.param(np.float16(1), id="float16"), - pytest.param(1 + 2j, id="complex"), - pytest.param(object(), id="object"), - pytest.param(float("inf"), id="nonfinite"), - pytest.param(1 << 65, id="outside-64-bit"), - ], -) -def test_direct_provider_rejects_invalid_static_oob_default_before_materialization( # noqa: E501 - Preserve descriptive test name. - monkeypatch, - oob_default, -): - from numba_cuda_mlir import types - - from cuda.coop._core import ArgumentBinding - from cuda.coop.numba_mlir._lowering import _load_store - - monkeypatch.setattr( - _load_store.NumbaMlirCoreAdapter, - "materialize", - lambda *args, **kwargs: pytest.fail( - "invalid oob_default reached provider materialization" - ), - ) - - with pytest.raises((TypeError, ValueError), match="oob_default"): - _load_store.load( - types.int32, - threads_per_block=32, - num_valid_items=ArgumentBinding.static(31), - oob_default=ArgumentBinding.static(oob_default), - ) - - -@pytest.mark.parametrize("operation", ["load", "store"]) -def test_static_valid_items_beyond_the_exact_tile_fail_before_nvrtc( - monkeypatch, operation -): - from numba_cuda_mlir import types - - from cuda.coop._core import ArgumentBinding - from cuda.coop.numba_mlir._lowering import _load_store - - monkeypatch.setattr( - _load_store.NumbaMlirCoreAdapter, - "materialize", - lambda *args, **kwargs: pytest.fail( - "invalid tile reached materialization" - ), - ) - factory = getattr(_load_store, operation) - - with pytest.raises(ValueError, match="block tile size"): - factory( - types.int32, - threads_per_block=32, - items_per_thread=2, - num_valid_items=ArgumentBinding.static(65), - ) - - -@pytest.mark.parametrize("operation", ["load", "store"]) -def test_negative_static_offsets_fail_before_provider_materialization( - monkeypatch, operation -): - from numba_cuda_mlir import types - - from cuda.coop._core import ArgumentBinding - from cuda.coop.numba_mlir._lowering import _load_store - - monkeypatch.setattr( - _load_store.NumbaMlirCoreAdapter, - "materialize", - lambda *args, **kwargs: pytest.fail( - "negative offset reached provider materialization" - ), - ) - factory = getattr(_load_store, operation) - - with pytest.raises(ValueError, match="offset must be nonnegative"): - factory( - types.int32, - threads_per_block=32, - offset=ArgumentBinding.static(-1), - ) diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_family_registration.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_family_registration.py deleted file mode 100644 index e1dae4b75186..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_family_registration.py +++ /dev/null @@ -1,827 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Exercise family registration and rewriting with small test providers. - -The fixtures isolate private registries so test families cannot affect later -operations. Fake invocables expose storage and call metadata without compiling -C++ providers. -""" - -from dataclasses import dataclass -from types import SimpleNamespace - -import pytest - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - - -@pytest.fixture(autouse=True) -def _restore_private_registries(): - """Restore all registries changed by the test, even after a failure.""" - from cuda.coop._core.api import _dispatch as common_dispatch - from cuda.coop._core.group import _dispatch as core_dispatch - from cuda.coop.numba_mlir._compiler import _operations - - registries = ( - core_dispatch._GROUP_OPERATION_FAMILIES, - common_dispatch._COMMON_GROUP_OPERATIONS_BY_NAME, - common_dispatch._COMMON_GROUP_OPERATIONS_BY_FUNCTION, - _operations._GROUP_OPERATIONS, - _operations._GROUP_FAMILY_MODULES, - _operations._FACTORY_OPERATIONS, - _operations._GROUP_PRIMITIVES, - _operations._REWRITE_OPERATIONS, - ) - snapshots = tuple(dict(registry) for registry in registries) - try: - yield - finally: - for registry, snapshot in zip(registries, snapshots): - registry.clear() - registry.update(snapshot) - - -@dataclass(frozen=True) -class _FakeSemantics: - """Provide semantic identity for a temporary operation family.""" - - token: str - - @property - def semantic_key(self): - return ("fake-family", self.token) - - @property - def result_visibility(self): - from cuda.coop._core import ResultVisibility - - return ResultVisibility.PER_MEMBER - - @property - def returns_value(self): - return True - - -@pytest.mark.parametrize( - ("override", "error", "message"), - [ - ( - {"classifications": None}, - TypeError, - "classifications must be callable", - ), - ({"planner": None}, TypeError, "planner must be callable"), - ( - {"group_kinds": frozenset()}, - ValueError, - "group_kinds must not be empty", - ), - ( - {"group_kinds": frozenset({"not_a_group"})}, - ValueError, - "group_kinds contains unsupported values", - ), - ( - {"unsupported_group_message": " "}, - ValueError, - "unsupported_group_message must be a non-empty string", - ), - ], -) -def test_core_family_registration_rejects_invalid_contracts( - override, - error, - message, -): - from cuda.coop._core.group import _dispatch - - arguments = { - "classifications": lambda operation: (), - "planner": lambda call, group, launch, operation: object(), - "group_kinds": frozenset({"block"}), - "unsupported_group_message": "test family requires a block", - } - arguments.update(override) - - with pytest.raises(error, match=message): - _dispatch._register_group_operation_family(_FakeSemantics, **arguments) - - -def test_group_primitive_registration_rejects_noncallable_hooks(): - from cuda.coop.numba_mlir._compiler._operations import ( - GroupPrimitiveRegistration, - ) - - with pytest.raises(TypeError, match="lower must be callable"): - GroupPrimitiveRegistration(lower=None) - with pytest.raises( - TypeError, - match="validate_common_arguments must be callable or None", - ): - GroupPrimitiveRegistration( - lower=lambda: None, - validate_common_arguments=object(), - ) - - -def test_factory_registration_rejects_noncallable_provider(): - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir._compiler._operations import ( - StorageABI, - register_factory, - ) - - with pytest.raises(TypeError, match="lowering factory must be callable"): - register_factory( - object(), - operation="_test_invalid_factory", - namespace="test", - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.NONE, - synchronization_scope=SynchronizationScope.NONE, - ) - - -def _rewrite_specification(**overrides): - """Build a valid baseline so each case can vary one registration rule.""" - from cuda.coop.numba_mlir._compiler._operations import ( - RewriteOperationSpecification, - ) - - arguments = { - "factory_namespaces": frozenset({"test_namespace"}), - "dtype_factory_kwargs": frozenset({"value_type"}), - "runtime_arg_counts": frozenset({1, 2}), - "runtime_factory_kwargs": ("tail",), - "runtime_factory_kw_prerequisites": (), - "allowed_factory_kwargs": frozenset( - {"guard", "offset", "tail", "value_type"} - ), - "required_factory_kwargs": frozenset({"value_type"}), - "accepts_temp_storage": False, - "scalar_binding_kwargs": frozenset({"tail"}), - "runtime_offset_kwarg": "offset", - "infer_payload": lambda context, inference: None, - } - arguments.update(overrides) - return RewriteOperationSpecification(**arguments) - - -@pytest.mark.parametrize( - ("override", "error", "message"), - [ - ( - {"runtime_arg_counts": frozenset()}, - ValueError, - "runtime_arg_counts must not be empty", - ), - ( - {"runtime_arg_counts": frozenset({-1})}, - ValueError, - "runtime_arg_counts must contain non-negative integers", - ), - ( - {"runtime_arg_counts": frozenset({True})}, - ValueError, - "runtime_arg_counts must contain non-negative integers", - ), - ( - {"runtime_arg_counts": frozenset({1, 3})}, - ValueError, - "require more trailing runtime arguments", - ), - ( - {"runtime_factory_kwargs": ("tail", "tail")}, - ValueError, - "runtime_factory_kwargs must be unique", - ), - ( - {"runtime_factory_kwargs": ("unknown",)}, - ValueError, - "runtime_factory_kwargs must be allowed factory kwargs", - ), - ( - {"runtime_factory_kw_prerequisites": (("tail",),)}, - TypeError, - "must contain name pairs", - ), - ( - { - "runtime_factory_kw_prerequisites": ( - ("tail", "guard"), - ("tail", "value_type"), - ) - }, - ValueError, - "prerequisite names must be unique", - ), - ( - {"runtime_factory_kw_prerequisites": (("guard", "tail"),)}, - ValueError, - "targets must be runtime factory kwargs", - ), - ( - {"runtime_factory_kw_prerequisites": (("tail", "unknown"),)}, - ValueError, - "requirements must be known factory kwargs", - ), - ( - {"runtime_factory_kw_prerequisites": (("tail", "tail"),)}, - ValueError, - "cannot require themselves", - ), - ( - {"scalar_binding_kwargs": frozenset({"guard"})}, - ValueError, - "scalar_binding_kwargs must be runtime factory kwargs", - ), - ( - {"runtime_offset_kwarg": ""}, - ValueError, - "runtime_offset_kwarg must be a non-empty string or None", - ), - ( - {"runtime_offset_kwarg": "unknown"}, - ValueError, - "runtime_offset_kwarg must be an allowed factory kwarg", - ), - ( - { - "runtime_arg_counts": frozenset({1, 2, 3}), - "runtime_factory_kwargs": ("tail", "offset"), - }, - ValueError, - "must not also be a runtime factory kwarg", - ), - ( - {"infer_payload": None}, - TypeError, - "infer_payload must be callable", - ), - ( - {"analyze_match": object()}, - TypeError, - "analyze_match must be callable or None", - ), - ], -) -def test_rewrite_operation_specification_rejects_inconsistent_contracts( - override, - error, - message, -): - with pytest.raises(error, match=message): - _rewrite_specification(**override) - - -@pytest.mark.parametrize( - ("struct_name", "scope_name", "sync_token"), - [ - ("WarpNamedButBlockScoped", "block", "__syncthreads();"), - ("BlockNamedButWarpScoped", "warp", "__syncwarp();"), - ("NeutralProvider", "none", None), - ], -) -def test_generated_synchronization_uses_metadata_not_struct_names( - struct_name, - scope_name, - sync_token, -): - from numba_cuda_mlir import types - - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir import _types - from cuda.coop.numba_mlir._compiler._nvrtc import CompilerIdentity - from cuda.coop.numba_mlir._compiler._operations import StorageABI - - scope = SynchronizationScope(scope_name) - algorithm = _types.Algorithm( - struct_name=struct_name, - method_name="Run", - c_name="test_declarative_provider_metadata", - includes=(), - template_parameters=(), - parameters=((_types.Pointer(types.uint8), _types.Value(types.int32)),), - storage_abi=StorageABI.LEADING_POINTER, - execution_scope=scope, - synchronization_scope=scope, - template_arguments={}, - ) - if scope is SynchronizationScope.WARP: - algorithm.logical_warp_threads = 32 - algorithm.block_threads = 64 - - source = algorithm._source_code( - compile_identity=CompilerIdentity( - cc=90, rdc=True, code="lto", compiler_options=() - ) - )[0] - mangled_name = algorithm.mangled_name(algorithm.parameters[0]) - alloc_body = source.split(f"void {mangled_name}_alloc(", 1)[1].split( - "\n}\n", 1 - )[0] - pointer_body = source.split(f"void {mangled_name}(", 1)[1].split( - "\n}\n", 1 - )[0] - - for token in ("__syncthreads();", "__syncwarp();"): - assert (token in source) is (token == sync_token) - assert (token in alloc_body) is (token == sync_token) - assert token not in pointer_body - if scope is SynchronizationScope.WARP: - assert "temp_storages[2]" in source - assert "[__coop_thread_rank / 32]" in source - assert scope.value in repr( - algorithm._make_lto_ir_cache_key( - compile_identity=CompilerIdentity( - cc=90, rdc=True, code="lto", compiler_options=() - ) - ) - ) - - -def test_group_synchronization_scope_fails_with_stable_diagnostic(): - from numba_cuda_mlir import types - - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir import _types - from cuda.coop.numba_mlir._compiler._nvrtc import CompilerIdentity - from cuda.coop.numba_mlir._compiler._operations import StorageABI - - algorithm = _types.Algorithm( - struct_name="Provider", - method_name="Run", - c_name="test_unsupported_group_synchronization", - includes=(), - template_parameters=(), - parameters=((_types.Pointer(types.uint8), _types.Value(types.int32)),), - storage_abi=StorageABI.LEADING_POINTER, - execution_scope=SynchronizationScope.GROUP, - synchronization_scope=SynchronizationScope.GROUP, - template_arguments={}, - ) - - with pytest.raises( - NotImplementedError, match="scope 'group' has no emitter" - ): - algorithm._source_code( - compile_identity=CompilerIdentity( - cc=90, rdc=True, code="lto", compiler_options=() - ) - ) - - -def test_storage_free_provider_uses_default_constructor_and_zero_storage(): - from numba_cuda_mlir import types - - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir import _types - from cuda.coop.numba_mlir._compiler._nvrtc import CompilerIdentity - from cuda.coop.numba_mlir._compiler._operations import StorageABI - - algorithm = _types.Algorithm( - struct_name="StorageFreeProvider", - method_name="Run", - c_name="test_storage_free_provider", - includes=(), - template_parameters=(), - parameters=((_types.Value(types.int32),),), - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.BLOCK, - synchronization_scope=SynchronizationScope.NONE, - template_arguments={}, - ) - source, _, storage_symbols, _ = algorithm._source_code( - compile_identity=CompilerIdentity( - cc=90, rdc=True, code="lto", compiler_options=() - ) - ) - - assert "StorageFreeProvider().Run" not in source - assert "algorithm_t_" in source - assert "().Run(param_0);" in source - assert "TempStorage" not in source - assert "temp_storage" not in source - assert storage_symbols == () - - -class _FakeInvocable: - """Expose scratch and synchronization metadata for one group scope.""" - - files = ("family-registration-test.ltoir",) - specialization = None - temp_storage_bytes = 24 - temp_storage_alignment = 8 - - def __init__(self, scope): - from cuda.coop.numba_mlir._compiler._operations import StorageABI - - self.storage_abi = StorageABI.LEADING_POINTER - self.execution_scope = scope - self.synchronization_scope = scope - - def __call__(self, *args): - del args - - -class _StorageFreeInvocable: - """Model a provider that needs neither scratch nor a reuse barrier.""" - - files = ("storage-free-family-registration-test.ltoir",) - specialization = None - temp_storage_bytes = 0 - temp_storage_alignment = 1 - storage_abi = "none" - execution_scope = "none" - synchronization_scope = "none" - - def __call__(self, *args): - del args - - -class _TypingContext: - """Count refreshes without a complete compiler typing context.""" - - def __init__(self): - self.refresh_count = 0 - - def refresh(self): - self.refresh_count += 1 - - -def _resolved_calls(func_ir): - """Resolve callable objects to inspect rewritten IR.""" - from numba_cuda_mlir.numbair_transforms import ir - - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - - resolver = object.__new__(CoopSinglePhaseRewrite) - resolver._func_ir = func_ir - calls = [] - for block in func_ir.blocks.values(): - resolver._block_defs = { - inst.target.name: inst.value - for inst in block.body - if isinstance(inst, ir.Assign) - } - for inst in block.body: - value = getattr(inst, "value", None) - if isinstance(value, ir.Expr) and value.op == "call": - calls.append( - (resolver._resolve_python_value(value.func), value) - ) - return calls - - -def test_registered_rewrite_callbacks_drive_generic_storage_rewrite(): - import numpy as np - from numba_cuda_mlir import cuda, types - from numba_cuda_mlir.numba_cuda.compiler import run_frontend - from numba_cuda_mlir.numbair_transforms import ir - - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir._compiler import _operations - from cuda.coop.numba_mlir._compiler._group_rewriting import ( - GroupRewriteContext, - ) - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - - operation = "_test_rewrite_family" - synchronization_scope = SynchronizationScope.BLOCK - events = [] - contexts = [] - family_metadata = object() - invocable = _FakeInvocable(synchronization_scope) - - def provider(*runtime_args, **factory_kwargs): - assert not runtime_args - events.append(("factory", dict(factory_kwargs))) - return invocable - - def record_context(context): - assert isinstance(context, GroupRewriteContext) - assert not hasattr(context, "_matches") - assert not hasattr(context, "_temp_storage_global_plan") - contexts.append(context) - - def infer_payload(context, inference): - record_context(context) - events.append(("infer", tuple(inference.runtime_args))) - inference.infer_kwarg("inferred", "from-callback") - inference.infer_kwarg("element_type", np.dtype("int32")) - - def analyze_match( - context, - *, - op_name, - runtime_args, - factory_kwargs, - ): - record_context(context) - events.append( - ( - "analyze", - op_name, - tuple(runtime_args), - dict(factory_kwargs), - ) - ) - return family_metadata - - def prepare_runtime_args( - context, - block, - *, - match, - runtime_args, - scope, - loc, - ): - record_context(context) - assert match.family_metadata is family_metadata - events.append(("prepare", tuple(runtime_args))) - prepared = ir.Var(scope, "__family_prepared_value", loc) - block.append(ir.Assign(ir.Const(29, loc), prepared, loc)) - return [*runtime_args, prepared] - - def validate_runtime_controls( - context, - *, - op_name, - runtime_args, - factory_kwargs, - ): - record_context(context) - events.append( - ( - "validate", - op_name, - tuple(runtime_args), - dict(factory_kwargs), - ) - ) - - _operations.register_factory( - provider, - operation=operation, - namespace="alternate", - storage_abi=_operations.StorageABI.LEADING_POINTER, - execution_scope=synchronization_scope, - synchronization_scope=synchronization_scope, - ) - _operations.register_rewrite_operation( - operation, - _operations.RewriteOperationSpecification( - factory_namespaces=frozenset({"block", "alternate"}), - dtype_factory_kwargs=frozenset({"element_type"}), - runtime_arg_counts=frozenset({1}), - runtime_factory_kwargs=(), - runtime_factory_kw_prerequisites=(), - allowed_factory_kwargs=frozenset( - {"element_type", "inferred", "token"} - ), - required_factory_kwargs=frozenset( - {"element_type", "inferred", "token"} - ), - accepts_temp_storage=False, - scalar_binding_kwargs=frozenset(), - runtime_offset_kwarg=None, - infer_payload=infer_payload, - analyze_match=analyze_match, - prepare_runtime_args=prepare_runtime_args, - validate_runtime_controls=validate_runtime_controls, - ), - ) - - def kernel(value): - return provider(value, token=7, element_type=value.dtype) - - func_ir = run_frontend(kernel) - typingctx = _TypingContext() - state = SimpleNamespace( - func_ir=func_ir, - args=(types.Array(types.int32, 1, "C"),), - typingctx=typingctx, - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - assert rewrite.prepare_calls_and_storage(func_ir) - rewrite.begin_rewrite() - for label in sorted(func_ir.blocks): - block = func_ir.blocks[label] - while rewrite.match(func_ir, block, state.typemap, state.calltypes): - block = rewrite.apply() - func_ir.blocks[label] = block - rewrite.finish_rewrite() - - calls = _resolved_calls(func_ir) - invocable_calls = [call for target, call in calls if target is invocable] - assert len(invocable_calls) == 1 - assert len(invocable_calls[0].args) == 3 - assert sum(target is cuda.shared.array for target, _ in calls) == 1 - sync_targets = { - target - for target, _ in calls - if target in {cuda.syncthreads, cuda.syncwarp} - } - assert sync_targets == {cuda.syncthreads} - assert rewrite._temp_storage_global_plan.total_size == 24 - assert rewrite._temp_storage_global_plan.max_alignment == 8 - assert rewrite._implicit_temp_storage_plan.size_in_bytes == 24 - assert rewrite._implicit_temp_storage_plan.alignment == 8 - assert typingctx.refresh_count == 1 - - factory_events = [event for event in events if event[0] == "factory"] - infer_events = [event for event in events if event[0] == "infer"] - analyze_events = [event for event in events if event[0] == "analyze"] - prepare_events = [event for event in events if event[0] == "prepare"] - validate_events = [event for event in events if event[0] == "validate"] - assert factory_events == [ - ( - "factory", - { - "token": 7, - "inferred": "from-callback", - "element_type": types.int32, - }, - ) - ] - assert len(infer_events) == 1 - assert all(len(event[1]) == 1 for event in infer_events) - assert len(analyze_events) == 1 - assert all(event[1] == operation for event in analyze_events) - assert all( - event[3] - == { - "token": 7, - "inferred": "from-callback", - "element_type": types.int32, - } - for event in analyze_events - ) - assert len(prepare_events) == 1 - assert len(validate_events) == 1 - assert all(event[1] == operation for event in validate_events) - assert len(contexts) == 4 - - resolver = object.__new__(CoopSinglePhaseRewrite) - resolver._func_ir = func_ir - resolver._block_defs = { - inst.target.name: inst.value - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) - } - assert resolver._infer_constant(invocable_calls[0].args[-1]) == 29 - - -def test_storage_free_provider_accepts_unused_temp_storage_descriptor(): - from numba_cuda_mlir import cuda, types - from numba_cuda_mlir.numba_cuda.compiler import run_frontend - - import cuda.coop.numba_mlir as coop - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir._compiler import _operations - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - - operation = "_test_storage_free_family" - invocable = _StorageFreeInvocable() - - def provider(*runtime_args, **factory_kwargs): - assert not runtime_args - assert not factory_kwargs - return invocable - - _operations.register_factory( - provider, - operation=operation, - namespace="alternate", - storage_abi=_operations.StorageABI.NONE, - execution_scope=SynchronizationScope.NONE, - synchronization_scope=SynchronizationScope.NONE, - ) - _operations.register_rewrite_operation( - operation, - _operations.RewriteOperationSpecification( - factory_namespaces=frozenset({"alternate"}), - dtype_factory_kwargs=frozenset(), - runtime_arg_counts=frozenset({1}), - runtime_factory_kwargs=(), - runtime_factory_kw_prerequisites=(), - allowed_factory_kwargs=frozenset(), - required_factory_kwargs=frozenset(), - accepts_temp_storage=True, - scalar_binding_kwargs=frozenset(), - runtime_offset_kwarg=None, - infer_payload=lambda *_args: None, - ), - ) - - def kernel(value): - storage = coop.TempStorage() - return provider(value, temp_storage=storage) - - func_ir = run_frontend(kernel) - typingctx = _TypingContext() - state = SimpleNamespace( - func_ir=func_ir, - args=(types.int32,), - typingctx=typingctx, - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - assert rewrite.prepare_calls_and_storage(func_ir) - rewrite.begin_rewrite() - for label in sorted(func_ir.blocks): - block = func_ir.blocks[label] - while rewrite.match(func_ir, block, state.typemap, state.calltypes): - block = rewrite.apply() - func_ir.blocks[label] = block - rewrite.finish_rewrite() - - calls = _resolved_calls(func_ir) - invocable_calls = [call for target, call in calls if target is invocable] - assert len(invocable_calls) == 1 - assert len(invocable_calls[0].args) == 1 - assert all(target is not cuda.shared.array for target, _ in calls) - assert all( - target not in {cuda.syncthreads, cuda.syncwarp} for target, _ in calls - ) - assert rewrite._temp_storage_global_plan is None - assert rewrite._temp_storage_backing_var is None - - -@pytest.mark.parametrize( - ("attribute", "value"), - [ - ("storage_abi", "leading_pointer"), - ("execution_scope", "block"), - ("synchronization_scope", "block"), - ], -) -def test_provider_metadata_must_match_registered_rewrite_contract( - attribute, - value, -): - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir._compiler import _operations - from cuda.coop.numba_mlir._compiler._rewrite_invocables import ( - _InvocableRewrite, - ) - from cuda.coop.numba_mlir._compiler._rewrite_support import ( - CoopSinglePhaseRewriteError, - ) - - provider_metadata = _operations.FactoryOperation( - operation="_test_provider_contract_mismatch", - namespace="alternate", - storage_abi=_operations.StorageABI.NONE, - execution_scope=SynchronizationScope.NONE, - synchronization_scope=SynchronizationScope.NONE, - ) - invocable = _StorageFreeInvocable() - setattr(invocable, attribute, value) - - with pytest.raises(CoopSinglePhaseRewriteError, match=attribute): - _InvocableRewrite._validate_invocable(invocable, provider_metadata) - - -def test_group_family_loads_lazily_and_registers_additively(monkeypatch): - from cuda.coop.numba_mlir._compiler import _operations - - operation = "_test_lazy_registration" - module_name = "_test_lazy_registration_module" - loads = [] - - def lower(*args): - return [] - - @_operations.group_operation(operation, family_module=module_name) - def marker(group, value): - raise AssertionError("compiler marker must not execute") - - def load_family(name): - loads.append(name) - _operations.register_group_primitive(operation, lower=lower) - - original = dict(_operations._GROUP_PRIMITIVES) - monkeypatch.setattr(_operations, "import_module", load_family) - registration = _operations.group_primitive(operation) - - assert _operations.group_operation_name(marker) == operation - assert registration.lower is lower - assert _operations.group_primitive(operation) is registration - assert loads == [module_name] - assert all( - _operations.group_primitive(name) is existing - for name, existing in original.items() - ) - _operations.register_group_primitive(operation, lower=lower) - with pytest.raises(RuntimeError, match="already registered"): - _operations.register_group_primitive(operation, lower=lambda: None) diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_group_hierarchy.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_group_hierarchy.py deleted file mode 100644 index 38734fb66c1a..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_group_hierarchy.py +++ /dev/null @@ -1,291 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check descriptor recognition and exact launch facts in frontend IR. - -These tests distinguish descriptor resolution from executable support for a -primitive. A known group kind alone does not make its operations available. -""" - -from types import SimpleNamespace - -import pytest -from numba_cuda_mlir.numba_cuda.compiler import run_frontend -from numba_cuda_mlir.numbair_transforms import ir - -import cuda.coop as common_coop -import cuda.coop.numba_mlir as coop -from cuda.coop._core import LaunchFactOrigin, LaunchFacts, resolve_thread_group -from cuda.coop.numba_mlir._compiler._group_planner import ( - GroupRewriteError, - _GroupCallPlanner, - has_group_markers, -) -from cuda.coop.numba_mlir._compiler._planner import CoopWholeFunctionPlanner - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - -_GROUP_KINDS = ( - ("this_thread", "thread"), - ("this_warp", "warp"), - ("this_block", "block"), - ("this_cluster", "cluster"), - ("this_grid", "grid"), -) -_GLOBAL_BLOCK_GROUP = coop.this_block() -_ALTERNATE_GLOBAL_BLOCK_GROUP = coop.this_block() - - -def _state(function, args=()): - return SimpleNamespace(func_ir=run_frontend(function), args=args) - - -def _launch(block=(64, 1, 1), grid=(1, 1, 1), cluster=None): - return {"block": block, "grid": grid, "cluster": cluster} - - -def _assigned_var(state, name): - return next( - inst.target - for block in state.func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) and inst.target.name == name - ) - - -def test_group_marker_detection_distinguishes_group_ir(): - def group_marker_function(): - return coop.this_block() - - def plain_function(value): - return value + 1 - - assert has_group_markers(run_frontend(group_marker_function)) - assert not has_group_markers(run_frontend(plain_function)) - - -def test_group_marker_detection_does_not_semantically_resolve_group_by( - monkeypatch, -): - def grouped(count): - return _GLOBAL_BLOCK_GROUP.group_by(count) - - monkeypatch.setattr( - _GroupCallPlanner, - "_group", - lambda *_args: pytest.fail("marker detection resolved group semantics"), - ) - - assert has_group_markers(run_frontend(grouped)) - - -def test_group_marker_detection_follows_merged_group_by_receivers(monkeypatch): - def grouped(count, choose_alternate): - if choose_alternate: - parent = _ALTERNATE_GLOBAL_BLOCK_GROUP - else: - parent = _GLOBAL_BLOCK_GROUP - return parent.group_by(count) - - monkeypatch.setattr( - _GroupCallPlanner, - "_group", - lambda *_args: pytest.fail("marker detection resolved group semantics"), - ) - - assert has_group_markers(run_frontend(grouped)) - - -@pytest.mark.parametrize( - "api", (common_coop, coop), ids=("common", "qualified") -) -def test_standalone_primitive_helper_is_rejected_without_requesting_launch( - api, monkeypatch -): - from cuda.coop.numba_mlir._compiler import _group_planner - - def device_helper(source): - group = api.this_block() - output = api.ThreadData(items_per_thread=1, dtype="int32") - return api.load(group, source, output) - - state = SimpleNamespace( - func_ir=run_frontend(device_helper), - args=(), - metadata={"targetoptions": {"device": True}}, - ) - before = { - label: tuple(block.body) - for label, block in state.func_ir.blocks.items() - } - - monkeypatch.setattr( - _group_planner, - "require_launch_config", - lambda _state: pytest.fail( - "device-function group planning requested kernel launch metadata" - ), - ) - - with pytest.raises( - GroupRewriteError, match="device_helper.*must be inlined" - ): - CoopWholeFunctionPlanner(state).run() - assert { - label: tuple(block.body) - for label, block in state.func_ir.blocks.items() - } == before - - -@pytest.mark.parametrize( - "api", (common_coop, coop), ids=("common", "qualified") -) -@pytest.mark.parametrize(("constructor_name", "kind"), _GROUP_KINDS) -def test_planner_recognizes_the_full_group_descriptor_vocabulary( - api, - constructor_name, - kind, -): - constructor = getattr(api, constructor_name) - - def describe_group(): - group = constructor() - return group - - state = _state(describe_group) - planner = _GroupCallPlanner(state, _launch()) - planned_group = planner._group(_assigned_var(state, "group")) - - assert planned_group is not None - assert planned_group.kind == kind - expected_source = "common_root" if api is common_coop else "current" - assert planned_group.source == expected_source - - -def test_group_constructor_recognition_uses_exact_callable_identity(): - def impostor_this_block(): - return coop.this_block() - - impostor_this_block.__module__ = coop.this_block.__module__ - impostor_this_block.__name__ = coop.this_block.__name__ - impostor_this_block.__cuda_coop_backend_member__ = "this_block" - - def describe_group(): - group = impostor_this_block() - return group - - state = _state(describe_group) - planner = _GroupCallPlanner(state, _launch()) - - assert planner._group(_assigned_var(state, "group")) is None - assert not has_group_markers(state.func_ir) - - -def test_planner_captures_exact_launch_facts_and_provenance(): - def plain_function(): - return None - - planner = _GroupCallPlanner( - _state(plain_function), - _launch( - block=(8, 4, 2), - grid=(6, 4, 2), - cluster=(2, 2, 1), - ), - ) - - assert planner.launch == LaunchFacts( - exact_block_dim=(8, 4, 2), - exact_grid_dim=(6, 4, 2), - exact_cluster_dim=(2, 2, 1), - cluster_launch=True, - cooperative_launch=False, - provenance=tuple( - LaunchFactOrigin( - fact=fact, - source="numba_cuda_mlir_launch_config", - verified=True, - ) - for fact in ( - "exact_block_dim", - "exact_grid_dim", - "cluster_launch", - "exact_cluster_dim", - ) - ), - ) - assert tuple(origin.fact for origin in planner.launch.provenance) == ( - "exact_block_dim", - "exact_grid_dim", - "cluster_launch", - "exact_cluster_dim", - ) - assert all( - origin.source == "numba_cuda_mlir_launch_config" and origin.verified - for origin in planner.launch.provenance - ) - assert all( - planner.launch.is_verified(fact) - for fact in ( - "exact_block_dim", - "exact_grid_dim", - "cluster_launch", - "exact_cluster_dim", - ) - ) - assert not planner.launch.is_verified("cooperative_launch") - - -@pytest.mark.parametrize(("constructor_name", "kind"), _GROUP_KINDS) -def test_core_resolves_every_physical_group_from_exact_launch_facts( - constructor_name, - kind, -): - launch = LaunchFacts( - exact_block_dim=(8, 4, 2), - exact_grid_dim=(6, 4, 2), - exact_cluster_dim=(2, 2, 1), - cluster_launch=True, - provenance=LaunchFactOrigin( - fact="cluster_launch", - source="test", - verified=True, - ), - ) - group = getattr(coop, constructor_name)() - resolved = resolve_thread_group(group, launch).require_supported() - - assert resolved.kind == kind - if kind == "thread": - assert resolved.hierarchy.implicit - return - assert resolved.hierarchy.block_dim == (8, 4, 2) - if kind in {"cluster", "grid"}: - assert resolved.hierarchy.cluster_dim == (2, 2, 1) - if kind == "grid": - assert resolved.hierarchy.grid_dim == (3, 2, 2) - - -@pytest.mark.parametrize( - "constructor", - (common_coop.this_block, coop.this_block), - ids=("common", "qualified"), -) -def test_descriptor_values_cannot_escape_to_runtime(constructor): - def escapes(): - return constructor() - - from cuda.coop.numba_mlir._compiler._group_errors import ( - EscapingGroupDescriptorError, - ) - - with pytest.raises( - EscapingGroupDescriptorError, match="would escape to runtime" - ) as exc: - _GroupCallPlanner(_state(escapes), _launch()).run() - - assert isinstance(exc.value, GroupRewriteError) - lines = str(exc.value).splitlines() - assert len(lines) > 1 - assert all(len(line) <= 80 for line in lines) diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_group_lowering_plan.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_group_lowering_plan.py deleted file mode 100644 index ca71535feb1c..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_group_lowering_plan.py +++ /dev/null @@ -1,955 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check scalar provenance and the core plans selected from compiler IR. - -Tests inspect plans before provider compilation, including loop dtype -propagation, argument bindings, and storage and synchronization requirements. -""" - -from types import SimpleNamespace - -import numpy as np -import pytest - -from cuda.coop._core import INT32 - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - -_STATIC_PROVENANCE_GLOBAL = np.int32(3) - - -def _planner(function, *, arg_types, block=(64, 1, 1), ssa=False): - """Build a planner with exact launch facts, optionally in SSA form. - - SSA cases exercise merged reaching definitions as well as the frontend - form used during ordinary group planning. - """ - from numba_cuda_mlir.numba_cuda.compiler import run_frontend - - from cuda.coop.numba_mlir._compiler._group_planner import _GroupCallPlanner - - func_ir = run_frontend(function) - if ssa: - from numba_cuda_mlir.numba_cuda.core.ir_utils import build_definitions - from numba_cuda_mlir.numba_cuda.core.ssa import reconstruct_ssa - - func_ir = reconstruct_ssa(func_ir) - func_ir._definitions = build_definitions(func_ir.blocks) - state = SimpleNamespace(func_ir=func_ir, args=arg_types) - return _GroupCallPlanner( - state, - {"block": block, "grid": (1, 1, 1), "cluster": None}, - ) - - -def _assigned_var(func_ir, name): - from numba_cuda_mlir.numbair_transforms import ir - - return next( - statement.target - for block in func_ir.blocks.values() - for statement in block.body - if isinstance(statement, ir.Assign) and statement.target.name == name - ) - - -def test_group_planning_context_exposes_only_declared_operations(): - from cuda.coop.numba_mlir._compiler._group_planning import ( - GroupPlanningContext, - ) - - launch = object() - planner = SimpleNamespace( - launch=launch, - _constant=lambda value: ("constant", value), - _reject_literal_unroll_value=lambda *_args: None, - ) - context = GroupPlanningContext(planner) - - assert context.launch is launch - assert context.constant("value") == ("constant", "value") - assert not hasattr(context, "_planner") - assert not hasattr(context, "__dict__") - with pytest.raises(AttributeError): - context._planner = planner - with pytest.raises(AttributeError): - _ = context.state - - -def test_group_planner_tracks_runtime_scalar_expression_provenance(): - from numba_cuda_mlir import cuda, types - from numba_cuda_mlir.numba_cuda.compiler import run_frontend - - from cuda.coop._core import BindingKind - from cuda.coop.numba_mlir._compiler._group_planner import _GroupCallPlanner - - def provenance(source, scalar): - alias = scalar - index = cuda.threadIdx.x - element = source[index] - binary = alias + index - unary = -alias - cast = np.int32(binary) - if index: - merged = cast - else: - merged = element - return merged, unary - - array_type = types.Array(types.int32, 1, "C") - func_ir = run_frontend(provenance) - planner = _GroupCallPlanner( - SimpleNamespace(func_ir=func_ir, args=(array_type, types.int32)), - {"block": (32, 1, 1), "grid": (1, 1, 1), "cluster": None}, - ) - expected = { - "scalar": types.int32, - "alias": types.int32, - "index": types.int32, - "element": types.int32, - "binary": types.int64, - "unary": types.int64, - "cast": types.int32, - "merged": types.int32, - } - for name, dtype in expected.items(): - value = _assigned_var(func_ir, name) - assert planner.context.dtype(value) == dtype - assert ( - planner.context.planning_binding(value).kind is BindingKind.RUNTIME - ) - - -@pytest.mark.parametrize("ssa", [False, True]) -@pytest.mark.parametrize("step_dtype", ["float32", "float64"]) -def test_loop_dtype_checks_every_reaching_definition(ssa, step_dtype): - from numba_cuda_mlir import types - from numba_cuda_mlir.numbair_transforms import ir - - from cuda.coop.numba_mlir._compiler._group_planner_support import ( - GroupRewriteError, - ) - - def loop(source, step, condition): - value = source[0] - for _ in range(5): - if condition: - value = value + step - return value - - planner = _planner( - loop, - arg_types=( - types.float32[::1], - getattr(types, step_dtype), - types.boolean, - ), - ssa=ssa, - ) - result = next( - statement.value - for block in planner.func_ir.blocks.values() - for statement in block.body - if isinstance(statement, ir.Return) - ) - if step_dtype == "float32": - assert planner.context.dtype(result) == types.float32 - else: - with pytest.raises(GroupRewriteError, match="inconsistent dtypes"): - planner.context.dtype(result) - - -@pytest.mark.parametrize("ssa", [False, True]) -def test_loop_dtype_does_not_ignore_an_unknown_producer(ssa): - from numba_cuda_mlir import types - from numba_cuda_mlir.numbair_transforms import ir - - def opaque(value): - return value - - def loop(source, condition): - value = source[0] - for _ in range(5): - if condition: - value = opaque(value) - return value - - planner = _planner( - loop, arg_types=(types.float32[::1], types.boolean), ssa=ssa - ) - result = next( - statement.value - for block in planner.func_ir.blocks.values() - for statement in block.body - if isinstance(statement, ir.Return) - ) - assert planner.context.dtype(result) is None - # A failed query must not leave speculative types for later operations. - assert planner.context.dtype(result) is None - - -def test_group_planner_marks_only_explicit_static_scalar_provenance_static(): - from numba_cuda_mlir import cuda, types - from numba_cuda_mlir.numba_cuda.compiler import run_frontend - - from cuda.coop._core import BindingKind - from cuda.coop.numba_mlir._compiler._group_planner import _GroupCallPlanner - - free_scalar = np.int32(4) - - def provenance(literal): - constant = 5 - global_value = _STATIC_PROVENANCE_GLOBAL - free_value = free_scalar - alias = constant - if cuda.threadIdx.x: - same_phi = 7 - different_phi = 8 - else: - same_phi = 7 - different_phi = 9 - return ( - literal, - constant, - global_value, - free_value, - alias, - same_phi, - different_phi, - ) - - func_ir = run_frontend(provenance) - planner = _GroupCallPlanner( - SimpleNamespace( - func_ir=func_ir, - args=(types.IntegerLiteral(1),), - ), - {"block": (32, 1, 1), "grid": (1, 1, 1), "cluster": None}, - ) - expected_static = { - "literal": np.int64(1), - "constant": 5, - "global_value": np.int32(3), - "free_value": np.int32(4), - "alias": 5, - "same_phi": 7, - } - for name, expected in expected_static.items(): - binding = planner.context.planning_binding(_assigned_var(func_ir, name)) - assert binding.kind is BindingKind.STATIC - assert binding.value == expected - assert type(binding.value) is type(expected) - - different = planner.context.planning_binding( - _assigned_var(func_ir, "different_phi") - ) - assert different.kind is BindingKind.RUNTIME - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -def test_runtime_arithmetic_controls_share_planning_and_rewrite_paths( - qualified, - items_per_thread, -): - from numba_cuda_mlir import cuda, types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - from cuda.coop._core import BindingKind - from cuda.coop.numba_mlir._compiler._rewrite import CoopSinglePhaseRewrite - - module = qualified_coop if qualified else root_coop - - def memory(source, control, items_per_thread): - index = cuda.threadIdx.x - valid = control + index - default = source[index] - offset = np.int64(control) - output = module.ThreadData(items_per_thread, dtype=types.int32) - return module.load( - module.this_block(), - source, - output, - valid_items=valid, - oob_default=default, - offset=offset, - ) - - array_type = types.Array(types.int32, 1, "C") - arg_types = (array_type, types.int32) - planner = _planner( - memory, - arg_types=(*arg_types, types.IntegerLiteral(items_per_thread)), - block=(32, 1, 1), - ) - assert planner.run() - - state = SimpleNamespace( - func_ir=planner.func_ir, - args=(*arg_types, types.IntegerLiteral(items_per_thread)), - typingctx=SimpleNamespace(refresh=lambda: None), - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - assert rewrite.prepare_calls_and_storage(planner.func_ir) - for label in sorted(planner.func_ir.blocks): - rewrite.match( - planner.func_ir, - planner.func_ir.blocks[label], - state.typemap, - state.calltypes, - ) - assert len(rewrite._matches) == 1 - match = next(iter(rewrite._matches.values())) - assert len(match.runtime_args) == 5 - for name in ("num_valid_items", "oob_default"): - assert match.factory_kwargs[name].kind is BindingKind.RUNTIME - assert rewrite._resolve_var_dtype(match.runtime_args[2]) == types.int64 - assert rewrite._resolve_var_dtype(match.runtime_args[3]) == types.int32 - assert rewrite._resolve_var_dtype(match.runtime_args[4]) == types.int64 - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_direct_load_provider_is_selected_from_complete_core_plan( - monkeypatch, items_per_thread -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as coop - from cuda.coop._core import ( - BindingKind, - GroupLoweringTarget, - StorageOwnership, - SynchronizationScope, - ) - from cuda.coop.numba_mlir._compiler import _group_load_store - - plans = [] - plan_group_primitive = _group_load_store.plan_group_primitive - - def capture_plan(call, launch): - plan = plan_group_primitive(call, launch) - plans.append(plan) - return plan - - monkeypatch.setattr( - _group_load_store, - "plan_group_primitive", - capture_plan, - ) - - def memory_with_storage(source, items_per_thread): - storage = coop.TempStorage( - 256, - alignment=16, - sharing="exclusive", - ) - output = coop.ThreadData(items_per_thread, dtype=types.int32) - return coop.load( - coop.this_block(), - source, - output, - valid_items=31, - oob_default=-1, - temp_storage=storage, - ) - - def memory_without_storage(source, items_per_thread): - output = coop.ThreadData(items_per_thread, dtype=types.int32) - return coop.load( - coop.this_block(), - source, - output, - valid_items=31, - oob_default=-1, - ) - - array_type = types.Array(types.int32, 1, "C") - for memory in (memory_with_storage, memory_without_storage): - planner = _planner( - memory, - arg_types=(array_type, types.IntegerLiteral(items_per_thread)), - ) - assert planner.run() - assert len(plans) == 2 - - plan = plans[0] - implicit_plan = plans[1] - semantic_key = plan.call.operation.semantic_key - # Compiler metadata on a backend dtype does not change the operation. - monkeypatch.setattr( - types.int32, "_coop_test_metadata", {"generation": 1}, raising=False - ) - assert plan.call.operation.semantic_key == semantic_key - assert plan.call.operation == implicit_plan.call.operation - assert plan.target == implicit_plan.target - assert plan.participation == implicit_plan.participation - assert plan.temp_storage == implicit_plan.temp_storage - assert plan.synchronization == implicit_plan.synchronization - assert plan.target is GroupLoweringTarget.CUB_BLOCK - assert plan.unsupported is None - assert plan.artifact_key is not None - assert plan.participation is not None - assert plan.participation.exact_group_size == 64 - assert plan.participation.exact_block_dim == (64, 1, 1) - assert plan.participation.uniform_arguments == ( - "valid_items", - "oob_default", - ) - assert plan.topology is not None - assert plan.topology.thread_rank == "linear_thread_rank" - assert plan.result is None - assert not plan.call.operation.returns_value - assert plan.synchronization is not None - assert ( - plan.synchronization.storage_reuse_barrier is SynchronizationScope.NONE - ) - assert plan.temp_storage is not None - assert plan.temp_storage.ownership is StorageOwnership.NONE - assert plan.temp_storage.address_space is None - assert plan.temp_storage.cpp_type is None - assert plan.temp_storage.instances is None - assert plan.temp_storage.instance_index is None - assert not plan.temp_storage.exact_layout_required - assert plan.temp_storage.sharing is None - assert plan.temp_storage.requested_size_in_bytes is None - assert plan.temp_storage.requested_alignment is None - assert not plan.temp_storage.auto_sync - assert plan.provenance is not None - assert plan.provenance.semantic_key == ( - "CUB", - "cub/block/block_load.cuh", - "cub::BlockLoad", - "Load", - ) - semantics = plan.call.operation - assert semantics.dtype == INT32 - assert semantics.items_per_thread == items_per_thread - assert semantics.valid_items.kind is BindingKind.STATIC - assert semantics.valid_items.value == 31 - assert semantics.oob_default.kind is BindingKind.STATIC - assert semantics.oob_default.value == -1 - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize("operation", ["load", "store"]) -def test_load_store_infer_untyped_payloads_symmetrically( - monkeypatch, - qualified, - operation, - items_per_thread, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - from cuda.coop.numba_mlir._compiler import _group_load_store - - module = qualified_coop if qualified else root_coop - plans = [] - plan_group_primitive = _group_load_store.plan_group_primitive - - def capture_plan(call, launch): - plan = plan_group_primitive(call, launch) - plans.append(plan) - return plan - - monkeypatch.setattr(_group_load_store, "plan_group_primitive", capture_plan) - - if operation == "load": - - def memory(source, destination, items_per_thread): - output = module.ThreadData(items_per_thread) - return module.load( - module.this_block(), - source, - output, - algorithm="direct", - ) - - else: - - def memory(source, destination, items_per_thread): - output = module.ThreadData(items_per_thread) - for item in range(items_per_thread): - output[item] = source[item] - module.store( - module.this_block(), - destination, - output, - algorithm="direct", - ) - - array_type = types.Array(types.int32, 1, "C") - assert _planner( - memory, - arg_types=( - array_type, - array_type, - types.IntegerLiteral(items_per_thread), - ), - ).run() - - assert len(plans) == 1 - assert plans[0].call.operation.dtype == INT32 - - -@pytest.mark.parametrize( - ("metadata_overrides", "diagnostic"), - [ - ({"storage_abi": "leading_pointer"}, "storage_abi"), - ({"execution_scope": "warp"}, "execution_scope"), - ({"synchronization_scope": "block"}, "synchronization_scope"), - ( - {"execution_scope": "warp", "synchronization_scope": "warp"}, - "synchronization_scope", - ), - ], -) -def test_group_plan_rejects_incompatible_provider_metadata( - monkeypatch, - metadata_overrides, - diagnostic, -): - from dataclasses import replace - - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as coop - from cuda.coop.numba_mlir import _lowering - from cuda.coop.numba_mlir._compiler import _group_planning - from cuda.coop.numba_mlir._compiler._group_planner_support import ( - GroupRewriteError, - ) - from cuda.coop.numba_mlir._compiler._operations import factory_operation - - metadata = factory_operation(_lowering.load) - assert metadata is not None - metadata = replace(metadata, **metadata_overrides) - monkeypatch.setattr( - _group_planning, - "factory_operation", - lambda _factory: metadata, - ) - - def memory(source): - output = coop.ThreadData(items_per_thread=2, dtype=types.int32) - return coop.load(coop.this_block(), source, output) - - array_type = types.Array(types.int32, 1, "C") - with pytest.raises(GroupRewriteError, match=diagnostic): - _planner(memory, arg_types=(array_type,)).run() - - -def test_group_plan_allows_declared_sync_scope_when_auto_sync_is_disabled( - monkeypatch, -): - from dataclasses import replace - - from cuda.coop._core import ( - StorageOwnership, - SynchronizationRequirements, - SynchronizationScope, - TempStorageRequirements, - this_block, - ) - from cuda.coop.numba_mlir._compiler import _group_planning - from cuda.coop.numba_mlir._compiler._group_planning import ( - GroupPlanningContext, - ) - from cuda.coop.numba_mlir._compiler._operations import ( - FactoryOperation, - StorageABI, - ) - from tests._group_planning import _load_store, _plan - - plan = _plan(this_block(), _load_store("load")) - plan = replace( - plan, - synchronization=SynchronizationRequirements( - converged_entry=True, - storage_reuse_barrier=SynchronizationScope.NONE, - ), - temp_storage=TempStorageRequirements( - ownership=StorageOwnership.CALLER, - address_space="shared", - cpp_type="TestStorage", - instances=1, - instance_index="cta", - exact_layout_required=True, - sharing="exclusive", - auto_sync=False, - ), - ) - metadata = FactoryOperation( - operation="test", - namespace="test", - storage_abi=StorageABI.LEADING_POINTER, - execution_scope=SynchronizationScope.BLOCK, - synchronization_scope=SynchronizationScope.BLOCK, - ) - monkeypatch.setattr( - _group_planning, - "factory_operation", - lambda _factory: metadata, - ) - - GroupPlanningContext._validate_provider_contract(plan, object()) - - -def test_group_plan_rejects_declared_sync_for_implementation_owned_no_sync( - monkeypatch, -): - from dataclasses import replace - - from cuda.coop._core import ( - StorageOwnership, - SynchronizationRequirements, - SynchronizationScope, - TempStorageRequirements, - this_block, - ) - from cuda.coop.numba_mlir._compiler import _group_planning - from cuda.coop.numba_mlir._compiler._group_planner_support import ( - GroupRewriteError, - ) - from cuda.coop.numba_mlir._compiler._group_planning import ( - GroupPlanningContext, - ) - from cuda.coop.numba_mlir._compiler._operations import ( - FactoryOperation, - StorageABI, - ) - from tests._group_planning import _load_store, _plan - - plan = _plan(this_block(), _load_store("load")) - plan = replace( - plan, - synchronization=SynchronizationRequirements( - converged_entry=True, - storage_reuse_barrier=SynchronizationScope.NONE, - ), - temp_storage=TempStorageRequirements( - ownership=StorageOwnership.IMPLEMENTATION, - address_space="shared", - cpp_type="TestStorage", - instances=1, - instance_index="cta", - exact_layout_required=False, - auto_sync=False, - ), - ) - metadata = FactoryOperation( - operation="test", - namespace="test", - storage_abi=StorageABI.LEADING_POINTER, - execution_scope=SynchronizationScope.BLOCK, - synchronization_scope=SynchronizationScope.BLOCK, - ) - monkeypatch.setattr( - _group_planning, - "factory_operation", - lambda _factory: metadata, - ) - - with pytest.raises(GroupRewriteError, match="synchronization_scope"): - GroupPlanningContext._validate_provider_contract(plan, object()) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_logical_warp_plan_selects_typed_cub_provider( - monkeypatch, items_per_thread -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as coop - from cuda.coop._core import GroupLoweringTarget - from cuda.coop.numba_mlir._compiler import _group_load_store - - plans = [] - plan_group_primitive = _group_load_store.plan_group_primitive - - def capture_plan(call, launch): - plan = plan_group_primitive(call, launch) - plans.append(plan) - return plan - - monkeypatch.setattr( - _group_load_store, - "plan_group_primitive", - capture_plan, - ) - - def memory(source, items_per_thread): - output = coop.ThreadData(items_per_thread, dtype=types.int32) - return coop.load(coop.this_warp().group_by(8), source, output) - - array_type = types.Array(types.int32, 1, "C") - planner = _planner( - memory, - arg_types=(array_type, types.IntegerLiteral(items_per_thread)), - ) - assert planner.run() - - assert len(plans) == 1 - plan = plans[0] - assert plan.target is GroupLoweringTarget.CUB_WARP - assert plan.unsupported is None - assert plan.topology is not None - assert plan.topology.group_kind == "threads_within_warp" - assert plan.topology.logical_width == 8 - assert plan.topology.instances == 8 - assert plan.artifact_key is not None - - -@pytest.mark.parametrize( - "oob_default", - [True, np.float16(1), np.complex64(1 + 2j)], - ids=["bool", "float16", "complex"], -) -def test_static_oob_default_rejects_before_provider_selection( - monkeypatch, oob_default -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as coop - from cuda.coop.numba_mlir._compiler import _group_load_store - - def memory(source): - output = coop.ThreadData(items_per_thread=2, dtype=types.int32) - return coop.load( - coop.this_block(), - source, - output, - valid_items=1, - oob_default=oob_default, - ) - - monkeypatch.setattr( - _group_load_store._LoadStorePlanning, - "_scope_factory", - lambda *_args, **_kwargs: pytest.fail( - "invalid oob_default reached provider selection" - ), - ) - array_type = types.Array(types.int32, 1, "C") - with pytest.raises((TypeError, ValueError), match="oob_default"): - _planner(memory, arg_types=(array_type,)).run() - - -@pytest.mark.parametrize( - "oob_type", - [ - pytest.param("boolean", id="bool"), - pytest.param("float16", id="float16"), - pytest.param("complex64", id="complex"), - pytest.param("optional", id="optional"), - pytest.param("float32", id="mismatched"), - ], -) -def test_runtime_oob_default_rejects_before_provider_selection( - monkeypatch, oob_type -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as coop - from cuda.coop.numba_mlir._compiler import _group_load_store - - def memory(source, oob_default): - output = coop.ThreadData(items_per_thread=2, dtype=types.int32) - return coop.load( - coop.this_block(), - source, - output, - valid_items=1, - oob_default=oob_default, - ) - - monkeypatch.setattr( - _group_load_store._LoadStorePlanning, - "_scope_factory", - lambda *_args, **_kwargs: pytest.fail( - "invalid oob_default reached provider selection" - ), - ) - value_type = ( - types.Optional(types.int32) - if oob_type == "optional" - else getattr(types, oob_type) - ) - array_type = types.Array(types.int32, 1, "C") - with pytest.raises( - (TypeError, ValueError), match="oob_default|supports dtypes" - ): - _planner(memory, arg_types=(array_type, value_type)).run() - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize("operation", ["load", "store"]) -@pytest.mark.parametrize( - "dtype_spelling", - ["builtin", "string", "numpy-type", "numpy-dtype", "backend"], -) -def test_equivalent_dtype_spellings_are_canonicalized_before_planning( - monkeypatch, - qualified, - operation, - dtype_spelling, - items_per_thread, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - from cuda.coop.numba_mlir._compiler import _group_load_store - - spellings = { - "builtin": int, - "string": "int32", - "numpy-type": np.int32, - "numpy-dtype": np.dtype(np.int32), - "backend": types.int32, - } - dtype = spellings[dtype_spelling] - module = qualified_coop if qualified else root_coop - plans = [] - plan_group_primitive = _group_load_store.plan_group_primitive - - def capture_plan(call, launch): - plan = plan_group_primitive(call, launch) - plans.append(plan) - return plan - - monkeypatch.setattr(_group_load_store, "plan_group_primitive", capture_plan) - - if operation == "load": - - def memory(memory, items_per_thread): - payload = module.ThreadData(items_per_thread, dtype=dtype) - return module.load(module.this_block(), memory, payload) - - else: - - def memory(memory, items_per_thread): - payload = module.ThreadData(items_per_thread, dtype=dtype) - module.store(module.this_block(), memory, payload) - - array_type = types.Array(types.int32, 1, "C") - planner = _planner( - memory, - arg_types=(array_type, types.IntegerLiteral(items_per_thread)), - ) - assert planner.run() - assert len(plans) == 1 - assert plans[0].call.operation.dtype == INT32 - - -_STATIC_DEFAULT_DTYPES = ( - pytest.param("int8", np.int8, id="int8"), - pytest.param("uint8", np.uint8, id="uint8"), - pytest.param("int16", np.int16, id="int16"), - pytest.param("uint16", np.uint16, id="uint16"), - pytest.param("int32", np.int32, id="int32"), - pytest.param("uint32", np.uint32, id="uint32"), - pytest.param("int64", np.int64, id="int64"), - pytest.param("uint64", np.uint64, id="uint64"), - pytest.param("float32", np.float32, id="float32"), - pytest.param("float64", np.float64, id="float64"), -) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize(("dtype_name", "numpy_dtype"), _STATIC_DEFAULT_DTYPES) -def test_static_oob_default_boundaries_use_the_load_payload_dtype( - qualified, - dtype_name, - numpy_dtype, - items_per_thread, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - - module = qualified_coop if qualified else root_coop - numpy_kind = np.dtype(numpy_dtype).kind - info = ( - np.iinfo(numpy_dtype) if numpy_kind in "iu" else np.finfo(numpy_dtype) - ) - boundaries = (int(info.min), int(info.max)) - if numpy_kind == "f": - boundaries = (-float(info.max), float(info.max)) - - for oob_default in boundaries: - - def memory(source, items_per_thread): - output = module.ThreadData( - items_per_thread, dtype=getattr(types, dtype_name) - ) - return module.load( - module.this_block(), - source, - output, - valid_items=1, - oob_default=oob_default, # noqa: B023 - the planner consumes each closure in this iteration. - ) - - array_type = types.Array(getattr(types, dtype_name), 1, "C") - assert _planner( - memory, - arg_types=(array_type, types.IntegerLiteral(items_per_thread)), - ).run() - - -@pytest.mark.parametrize("qualified", [False, True], ids=["root", "qualified"]) -@pytest.mark.parametrize( - ("value_type_name", "error"), - [ - ("float32", "does not match payload dtype"), - ("boolean", "supports dtypes"), - ], - ids=["mismatched", "unsupported"], -) -def test_untyped_store_infers_write_dtype_before_destination_fallback( - monkeypatch, - qualified, - value_type_name, - error, -): - from numba_cuda_mlir import types - - import cuda.coop.numba_mlir as qualified_coop - from cuda import coop as root_coop - from cuda.coop.numba_mlir._compiler import _group_load_store - - module = qualified_coop if qualified else root_coop - - def memory(destination, value): - payload = module.ThreadData(items_per_thread=2) - payload[0] = value - payload[1] = value - module.store(module.this_block(), destination, payload) - - array_type = types.Array(types.int32, 1, "C") - planner = _planner( - memory, - arg_types=(array_type, getattr(types, value_type_name)), - ) - monkeypatch.setattr( - _group_load_store._LoadStorePlanning, - "_scope_factory", - lambda *_args, **_kwargs: pytest.fail( - "invalid Store writes reached provider selection" - ), - ) - - with pytest.raises(TypeError, match=error): - planner.run() diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_invocable_identity.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_invocable_identity.py deleted file mode 100644 index 8c71a2afa05e..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_invocable_identity.py +++ /dev/null @@ -1,104 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from types import SimpleNamespace - -import pytest - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - - -class _FakeInvocable: - """Expose a distinct link input for each registered provider factory.""" - - storage_abi = "leading_pointer" - execution_scope = "block" - synchronization_scope = "block" - - def __init__(self, name): - self.name = name - self.files = (f"{name}.ltoir",) - - def __call__(self, *args): - del args - - -def test_invocable_cache_partitions_registered_factory_identities(): - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir._compiler._operations import ( - FactoryOperation, - StorageABI, - ) - from cuda.coop.numba_mlir._compiler._rewrite_invocables import ( - _InvocableRewrite, - ) - from cuda.coop.numba_mlir._compiler._rewrite_support import _RewriteMatch - - def first_factory(**kwargs): - del kwargs - return _FakeInvocable("first") - - def second_factory(**kwargs): - del kwargs - return _FakeInvocable("second") - - common_kwargs = {"threads_per_block": (64, 1, 1), "dtype": "int32"} - provider_metadata = FactoryOperation( - operation="load", - namespace="block", - storage_abi=StorageABI.LEADING_POINTER, - execution_scope=SynchronizationScope.BLOCK, - synchronization_scope=SynchronizationScope.BLOCK, - ) - first_key = _InvocableRewrite._invocable_cache_key( - first_factory, - provider_metadata, - common_kwargs, - ) - second_key = _InvocableRewrite._invocable_cache_key( - second_factory, - provider_metadata, - common_kwargs, - ) - assert first_key != second_key - assert first_key == _InvocableRewrite._invocable_cache_key( - first_factory, - provider_metadata, - dict(common_kwargs), - ) - - rewrite = object.__new__(_InvocableRewrite) - rewrite._invocable_cache = {} - rewrite._prebundled_specializations = {} - rewrite._state = SimpleNamespace(metadata={}) - - def match(factory): - return _RewriteMatch( - op_name="load", - factory=factory, - factory_metadata=provider_metadata, - func_var_name="factory", - func_var_name_extra=None, - runtime_args=(), - runtime_temp_storage_var=None, - factory_kwargs=dict(common_kwargs), - factory_kw_value_vars=(), - loc=None, - ) - - first, first_created = rewrite._materialize_invocable(match(first_factory)) - second, second_created = rewrite._materialize_invocable( - match(second_factory) - ) - assert first_created - assert second_created - assert first.name == "first" - assert second.name == "second" - assert first is not second - - first_again, first_again_created = rewrite._materialize_invocable( - match(first_factory) - ) - assert first_again is first - assert not first_again_created diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_invocable_lifecycle.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_invocable_lifecycle.py deleted file mode 100644 index 458b6b5e1f99..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_invocable_lifecycle.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import gc -import weakref -from pathlib import Path - -import pytest - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - - -def _make_invocable(): - """Build a callable that owns a temporary artifact, without compiling.""" - from numba_cuda_mlir import types - - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir import _types - from cuda.coop.numba_mlir._compiler._artifacts import make_binary_tempfile - from cuda.coop.numba_mlir._compiler._operations import StorageABI - - algorithm = _types.Algorithm( - struct_name="BlockLoad", - method_name="Load", - c_name="cuda_coop_numba_mlir_test_invocable_lifetime", - includes=(), - template_parameters=(), - parameters=((_types.Pointer(types.uint8), _types.Value(types.int32)),), - storage_abi=StorageABI.LEADING_POINTER, - execution_scope=SynchronizationScope.BLOCK, - synchronization_scope=SynchronizationScope.BLOCK, - template_arguments={}, - ) - artifact = make_binary_tempfile(b"test-ltoir", ".ltoir") - return _types.Invocable( - temp_files=(artifact,), - owned_temp_files=(artifact,), - temp_storage_bytes=8, - temp_storage_alignment=8, - algorithm=algorithm, - ) - - -def _registry_sizes(): - """Measure typing registries to detect retained dynamic callables.""" - from numba_cuda_mlir.descriptor import mlir_target - from numba_cuda_mlir.extending import typeof_impl, typing_registry - - return ( - len(typeof_impl.registry), - len(typing_registry.functions), - len(typing_registry.globals), - len(typing_registry.attributes), - len(mlir_target.typing_context._functions), - len(mlir_target.typing_context._globals), - len(mlir_target.typing_context._attributes), - ) - - -def test_invocable_typing_is_local_to_compiler_owners(): - from numba_cuda_mlir.descriptor import mlir_target - from numba_cuda_mlir.extending import ( - _NumbaCudaMlirOverloadFunctionTemplate as overload_template, - ) - - registry_sizes = _registry_sizes() - first = _make_invocable() - second = _make_invocable() - first_path = Path(first.files[0]) - second_path = Path(second.files[0]) - first_ref = weakref.ref(first) - second_ref = weakref.ref(second) - - first_type = mlir_target.typing_context.resolve_value_type(first) - second_type = mlir_target.typing_context.resolve_value_type(second) - - assert len(first_type.templates) == 2 - assert all( - issubclass(template, overload_template) - for template in first_type.templates - ) - assert first_type is mlir_target.typing_context.resolve_value_type(first) - assert second_type is mlir_target.typing_context.resolve_value_type(second) - assert _registry_sizes() == registry_sizes - - # A compiler-owned Function type retains its exact callable and therefore - # the link input until compilation has finished with it. - del first - gc.collect() - assert first_ref() is not None - assert first_path.is_file() - - # No process-global typing registry owns the dynamic template. Releasing - # the compiler's Function type makes the Invocable cycle collectible and - # runs its artifact finalizer. - del first_type - gc.collect() - assert first_ref() is None - assert not first_path.exists() - - del second, second_type - gc.collect() - assert second_ref() is None - assert not second_path.exists() - assert _registry_sizes() == registry_sizes diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_launch_dimension_rewrites.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_launch_dimension_rewrites.py deleted file mode 100644 index fa046ab2cccc..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_launch_dimension_rewrites.py +++ /dev/null @@ -1,359 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check when rewriting can obtain an exact block shape. - -Configured launches can supply the dimensions; launch bounds cannot. Device -helpers defer until inlining supplies their caller's launch context. -""" - -from types import SimpleNamespace - -import pytest - -pytest.importorskip("numba_cuda_mlir") - -from numba_cuda_mlir import types -from numba_cuda_mlir.numba_cuda.compiler import run_frontend - -from cuda.coop.numba_mlir._compiler._planner import CoopWholeFunctionPlanner -from cuda.coop.numba_mlir._compiler._rewrite import ( - CoopSinglePhaseRewrite, - CoopSinglePhaseRewriteError, -) -from cuda.coop.numba_mlir._lowering import load as block_load - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - - -@pytest.fixture(autouse=True) -def _fixed_provider_compute_capability(monkeypatch): - from cuda.coop.numba_mlir import _types - - monkeypatch.setattr( - _types.cuda, - "get_current_device", - lambda: SimpleNamespace(compute_capability=(9, 0)), - ) - - -class _TypingContext: - def __init__(self): - self.refresh_count = 0 - - def refresh(self): - self.refresh_count += 1 - - -def _state(function, *, targetoptions): - """Build frontend IR with the launch metadata needed by each test.""" - array_type = types.Array(types.int32, 1, "C") - return SimpleNamespace( - func_ir=run_frontend(function), - args=(array_type, array_type), - typingctx=_TypingContext(), - typemap={}, - calltypes={}, - metadata={"targetoptions": targetoptions}, - ) - - -def _first_block(state): - return state.func_ir.blocks[min(state.func_ir.blocks)] - - -def _match(state): - """Match a provider call without compiling its source bundle.""" - rewrite = CoopSinglePhaseRewrite(state) - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: None - matched = rewrite.prepare_calls_and_storage( - state.func_ir - ) and rewrite.match( - state.func_ir, - _first_block(state), - state.typemap, - state.calltypes, - ) - return rewrite, matched - - -@pytest.mark.parametrize( - ("block", "expected"), - [ - ((64, 1, 1), 64), - ((8, 4, 1), (8, 4)), - ((8, 4, 2), (8, 4, 2)), - ], -) -def test_launch_block_dimensions_are_canonicalized(block, expected): - rewrite = object.__new__(CoopSinglePhaseRewrite) - rewrite._state = SimpleNamespace( - metadata={"targetoptions": {"__launch_config__": {"block": block}}} - ) - - assert rewrite._infer_threads_per_block_from_launch_config() == expected - - -@pytest.mark.parametrize( - "targetoptions", - [ - {}, - {"launch_bounds": 128}, - {"launch_bounds": (256, 2)}, - {"__launch_config__": {}}, - {"__launch_config__": {"block": (0, 1, 1)}}, - ], -) -def test_launch_bounds_and_inexact_launches_do_not_infer_dimensions( - targetoptions, -): - rewrite = object.__new__(CoopSinglePhaseRewrite) - rewrite._state = SimpleNamespace(metadata={"targetoptions": targetoptions}) - - assert rewrite._infer_threads_per_block_from_launch_config() is None - - -def test_dim_alias_is_rewritten_to_threads_per_block(): - def kernel(source, output): - return block_load(source, output, dtype=types.int32, dim=64) - - state = _state( - kernel, - targetoptions={"__launch_config__": {"block": (64, 1, 1)}}, - ) - rewrite, matched = _match(state) - - assert matched - match = next(iter(rewrite._matches.values())) - assert match.factory_kwargs["threads_per_block"] == 64 - assert "dim" not in match.factory_kwargs - - -def test_explicit_dimension_accepts_an_equivalent_exact_launch_shape(): - def kernel(source, output): - return block_load( - source, - output, - dtype=types.int32, - threads_per_block=64, - ) - - state = _state( - kernel, - targetoptions={"__launch_config__": {"block": (64, 1, 1)}}, - ) - rewrite, matched = _match(state) - - assert matched - match = next(iter(rewrite._matches.values())) - assert match.factory_kwargs["threads_per_block"] == 64 - - -def test_dim_alias_rejects_explicit_threads_per_block(): - def kernel(source, output): - return block_load( - source, - output, - dtype=types.int32, - threads_per_block=32, - dim=64, - ) - - state = _state(kernel, targetoptions={}) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="received both 'threads_per_block' and its 'dim' alias", - ): - _match(state) - - -def test_explicit_dimension_rejects_an_exact_launch_mismatch(): - def kernel(source, output): - return block_load( - source, - output, - dtype=types.int32, - threads_per_block=32, - ) - - state = _state( - kernel, - targetoptions={"__launch_config__": {"block": (64, 1, 1)}}, - ) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match=( - r"factory 'load' received threads_per_block=32, but the exact " - r"kernel launch block is \(64, 1, 1\)" - ), - ): - _match(state) - - -def test_deferred_rewrite_leaves_device_helper_ir_intact(): - def device_function(source, output): - return block_load(source, output, dtype=types.int32) - - state = _state(device_function, targetoptions={"device": True}) - before = { - label: tuple(block.body) - for label, block in state.func_ir.blocks.items() - } - rewrite = CoopSinglePhaseRewrite(state) - - assert not rewrite.prepare_calls_and_storage(state.func_ir) - assert rewrite._deferred_launch_dim_inference - assert { - label: tuple(block.body) - for label, block in state.func_ir.blocks.items() - } == before - - -def test_kernel_planner_retries_with_an_exact_launch(monkeypatch): - from cuda.coop.numba_mlir._compiler import _rewrite as rewrites - - def kernel(source, output): - if source[0]: - return block_load(source, output, dtype=types.int32) - return block_load(source, output, dtype=types.int32) - - class _FakeInvocable: - files = () - temp_storage_bytes = 1 - temp_storage_alignment = 1 - storage_abi = "leading_pointer" - execution_scope = "block" - synchronization_scope = "block" - - def __call__(self, source, output): - del source, output - - state = _state(kernel, targetoptions={}) - invocable = _FakeInvocable() - requests = [] - preparations = [] - prepare_calls_and_storage = CoopSinglePhaseRewrite.prepare_calls_and_storage - - def prepare(rewrite, func_ir): - before = { - label: tuple(block.body) for label, block in func_ir.blocks.items() - } - prepared = prepare_calls_and_storage(rewrite, func_ir) - preparations.append((rewrite, prepared)) - if not prepared: - assert { - label: tuple(block.body) - for label, block in func_ir.blocks.items() - } == before - return prepared - - def require_exact_launch(requested_state): - launch_config = { - "block": (64, 1, 1), - "grid": (1, 1, 1), - "cluster": None, - } - requests.append(requested_state) - requested_state.metadata["targetoptions"]["__launch_config__"] = ( - launch_config - ) - return launch_config - - monkeypatch.setattr(rewrites, "require_launch_config", require_exact_launch) - monkeypatch.setattr( - CoopSinglePhaseRewrite, "prepare_calls_and_storage", prepare - ) - monkeypatch.setattr( - CoopSinglePhaseRewrite, - "_prepare_ltoir_bundle_for_matches", - lambda self, matches: None, - ) - monkeypatch.setattr( - CoopSinglePhaseRewrite, - "_materialize_invocable", - lambda self, match: (invocable, False), - ) - assert CoopWholeFunctionPlanner(state).run() - assert requests == [state] - assert [prepared for _, prepared in preparations] == [False, True] - assert preparations[0][0] is not preparations[1][0] - assert state.typingctx.refresh_count == 2 - - -@pytest.mark.parametrize( - ("targetoptions", "launch_config", "expected_detail"), - [ - ({}, {"block": (0, 1, 1)}, "invalid block=(0, 1, 1)"), - ({}, None, "no __launch_config__ metadata"), - ({}, {}, "contains no block shape"), - ( - {"launch_bounds": 128}, - None, - "launch_bounds=128 is only an upper bound", - ), - ], -) -def test_kernel_planner_reports_an_unresolved_launch_after_retry( - monkeypatch, - targetoptions, - launch_config, - expected_detail, -): - from cuda.coop.numba_mlir._compiler import _rewrite as rewrites - - def kernel(source, output): - return block_load(source, output, dtype=types.int32) - - state = _state(kernel, targetoptions=targetoptions) - requests = [] - - def require_unresolved_launch(requested_state): - requests.append(requested_state) - if launch_config is not None: - requested_state.metadata["targetoptions"]["__launch_config__"] = ( - launch_config - ) - return launch_config - - monkeypatch.setattr( - rewrites, "require_launch_config", require_unresolved_launch - ) - - with pytest.raises(CoopSinglePhaseRewriteError) as exc_info: - CoopWholeFunctionPlanner(state).run() - - message = str(exc_info.value) - assert "coop operation 'load'" in message - assert "could not infer an exact positive threads_per_block" in message - assert expected_detail in message - assert "compile-time constant launch shape" in message - assert requests == [state] - - -def test_device_planner_defers_without_requesting_a_launch(monkeypatch): - from cuda.coop.numba_mlir._compiler import _rewrite as rewrites - - def device_function(source, output): - return block_load(source, output, dtype=types.int32) - - state = _state(device_function, targetoptions={"device": True}) - before = { - label: tuple(block.body) - for label, block in state.func_ir.blocks.items() - } - - def unexpected_launch_request(_state): - pytest.fail("device-function planning requested kernel launch metadata") - - monkeypatch.setattr( - rewrites, "require_launch_config", unexpected_launch_request - ) - - assert not CoopWholeFunctionPlanner(state).run() - assert { - label: tuple(block.body) - for label, block in state.func_ir.blocks.items() - } == before diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_nvrtc.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_nvrtc.py deleted file mode 100644 index 3bf32f02b7f3..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_nvrtc.py +++ /dev/null @@ -1,733 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check NVRTC context, cache, and layout queries with fake compilers. - -Small header trees check that header order and content change the context -identity. Fake library records and NVRTC bindings check library load order, -version checks, cache-key arguments, and source dumps without depending on a -CUDA installation. A fake NVRTC program checks that storage-layout queries -share the provider compilation, separate cache entries, and destroy the -program after errors. The compile suite runs real compilation. -""" - -from __future__ import annotations - -import inspect -import os -from dataclasses import replace -from pathlib import Path -from types import SimpleNamespace - -import pytest - -from cuda.coop import _headers -from cuda.coop._headers._identity import include_dirs_identity -from cuda.coop._headers._toolkit import ToolkitCompilerLibraries -from cuda.coop.numba_mlir._compiler import _nvrtc - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - - -def _fake_nvrtc( - actual_version: tuple[int, int], - *, - result: int = 0, -) -> SimpleNamespace: - """Provide the version query needed for context validation.""" - return SimpleNamespace( - nvrtcResult=SimpleNamespace(NVRTC_SUCCESS=0), - nvrtcVersion=lambda: (result, *actual_version), - ) - - -def _libraries(**changes) -> ToolkitCompilerLibraries: - """Create a coherent fake toolchain with the requested differences.""" - libraries = ToolkitCompilerLibraries( - toolkit_root="/cuda/toolkit", - toolkit_version=(13, 3), - nvrtc_path="/cuda/toolkit/lib/libnvrtc.so.13", - nvrtc_builtins_path="/cuda/toolkit/lib/libnvrtc-builtins.so.13.3", - nvjitlink_path="/cuda/toolkit/lib/libnvJitLink.so.13", - nvrtc_version=(13, 3), - nvjitlink_version=(13, 4), - ) - return replace(libraries, **changes) - - -def _context(**changes) -> _nvrtc.CompileContext: - """Build a fixed compile context without resolving installed libraries.""" - context = _nvrtc.CompileContext( - toolkit_root="/cuda/toolkit", - toolkit_version=(13, 3), - nvrtc_path="/cuda/toolkit/lib/libnvrtc.so.13", - nvrtc_builtins_path="/cuda/toolkit/lib/libnvrtc-builtins.so.13.3", - nvjitlink_path="/cuda/toolkit/lib/libnvJitLink.so.13", - nvrtc_version=_nvrtc.version(13, 3), - nvjitlink_version=(13, 4), - include_dirs=("/cccl/include", "/cuda/toolkit/include"), - header_identity="headers-a", - ) - return replace(context, **changes) - - -def test_include_options_use_filesystem_encoding() -> None: - include_dir = "/tmp/cuda-coop-π/include" - - assert _nvrtc._include_options((include_dir,)) == [ - os.fsencode(f"--include-path={include_dir}") - ] - - -def test_required_headers_cover_executable_provider_includes() -> None: - assert { - "cub/block/block_load.cuh", - "cub/block/block_store.cuh", - } <= set(_nvrtc._REQUIRED_HEADERS) - - -def test_required_headers_reject_a_partial_configured_bundle( - tmp_path: Path, -) -> None: - bundle = tmp_path / "partial-bundle" - missing = "cuda/hierarchy" - present_headers = (set(_nvrtc._REQUIRED_HEADERS) - {missing}) | { - "cub/version.cuh" - } - for header in present_headers: - destination = bundle / header - destination.parent.mkdir(parents=True, exist_ok=True) - destination.touch() - - with pytest.raises(_headers.HeaderResolutionError, match=missing): - _headers.resolve_include_paths( - start=Path(__file__), - configured_roots=(bundle,), - required_headers=_nvrtc._REQUIRED_HEADERS, - ) - - -def test_resolve_context_captures_exact_toolchain_after_lazy_binding_load( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - cccl_include = tmp_path / "cccl" / "include" - cuda_include = tmp_path / "cuda" / "include" - cccl_include.mkdir(parents=True) - cuda_include.mkdir(parents=True) - (cccl_include / "primitive.cuh").write_bytes(b"cccl header") - (cuda_include / "cuda_runtime.h").write_bytes(b"cuda header") - include_dirs = (cccl_include, cuda_include) - libraries = _libraries() - events: list[object] = [] - real_identity = include_dirs_identity - - def resolve_paths(**kwargs): - events.append("headers") - assert kwargs["required_headers"] == _nvrtc._REQUIRED_HEADERS - return SimpleNamespace( - as_tuple=lambda: include_dirs, - cuda=(cuda_include,), - ) - - def preload(paths): - events.append(("preload", paths)) - return libraries - - def load_nvrtc(): - events.append("binding") - return _fake_nvrtc((13, 3)) - - def identify(paths): - events.append(("identity", paths)) - return real_identity(paths) - - monkeypatch.setattr(_nvrtc, "resolve_include_paths", resolve_paths) - monkeypatch.setattr(_nvrtc, "preload_toolkit_compiler_libraries", preload) - monkeypatch.setattr(_nvrtc, "_load_nvrtc", load_nvrtc) - monkeypatch.setattr(_nvrtc, "include_dirs_identity", identify) - - context = _nvrtc.resolve_compile_context() - expected_identity = real_identity(tuple(str(path) for path in include_dirs)) - - assert events == [ - "headers", - ("preload", (cuda_include,)), - "binding", - ("identity", tuple(str(path) for path in include_dirs)), - ] - assert expected_identity.recursive_walks == 2 - assert context == _nvrtc.CompileContext( - toolkit_root=libraries.toolkit_root, - toolkit_version=libraries.toolkit_version, - nvrtc_path=libraries.nvrtc_path, - nvrtc_builtins_path=libraries.nvrtc_builtins_path, - nvjitlink_path=libraries.nvjitlink_path, - nvrtc_version=_nvrtc.version(13, 3), - nvjitlink_version=libraries.nvjitlink_version, - include_dirs=tuple(str(path) for path in include_dirs), - header_identity=expected_identity.digest, - ) - - -def test_resolve_context_preserves_order_and_recursive_header_identity( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - first = tmp_path / "first" - second = tmp_path / "second" - cuda = tmp_path / "cuda" - for include_dir in (first, second, cuda): - include_dir.mkdir() - (include_dir / "header.cuh").write_bytes(b"same content") - ordered = [first, second, cuda] - monkeypatch.setattr( - _nvrtc, - "resolve_include_paths", - lambda **kwargs: SimpleNamespace( - as_tuple=lambda: tuple(ordered), - cuda=(cuda,), - ), - ) - monkeypatch.setattr( - _nvrtc, - "preload_toolkit_compiler_libraries", - lambda paths: _libraries(), - ) - monkeypatch.setattr( - _nvrtc, - "_load_nvrtc", - lambda: _fake_nvrtc((13, 3)), - ) - - forward = _nvrtc.resolve_compile_context() - ordered[:] = [second, first, cuda] - reverse = _nvrtc.resolve_compile_context() - (first / "header.cuh").write_bytes(b"mutated content") - mutated = _nvrtc.resolve_compile_context() - - assert forward.include_dirs == tuple( - str(path) for path in (first, second, cuda) - ) - assert reverse.include_dirs == tuple( - str(path) for path in (second, first, cuda) - ) - assert forward.header_identity != reverse.header_identity - assert reverse.header_identity != mutated.header_identity - - -def test_shared_context_scope_is_lazy_ordered_and_restored_after_errors( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> None: - first, second = tmp_path / "first", tmp_path / "second" - for path in (first, second): - path.mkdir() - (path / "header.cuh").write_text("original") - ordered = [first, second] - selections = [] - - def resolve_paths(**kwargs): - selections.append(tuple(ordered)) - return SimpleNamespace(as_tuple=lambda: tuple(ordered), cuda=(second,)) - - monkeypatch.setattr(_nvrtc, "resolve_include_paths", resolve_paths) - monkeypatch.setattr( - _nvrtc, "preload_toolkit_compiler_libraries", lambda paths: _libraries() - ) - monkeypatch.setattr(_nvrtc, "_load_nvrtc", lambda: _fake_nvrtc((13, 3))) - - with _nvrtc.compile_context_scope(): - pass - assert not selections - - with _nvrtc.compile_context_scope(): - original = _nvrtc.resolve_shared_compile_context() - assert _nvrtc.resolve_shared_compile_context() is original - # Re-select roots even when a context already exists in this attempt. - ordered.reverse() - reversed_context = _nvrtc.resolve_shared_compile_context() - assert reversed_context.header_identity != original.header_identity - ordered.reverse() - assert _nvrtc.resolve_shared_compile_context() is original - - with ( - pytest.raises(RuntimeError, match="failed attempt"), - _nvrtc.compile_context_scope(), - ): - nested = _nvrtc.resolve_shared_compile_context() - assert nested == original - assert nested is not original - raise RuntimeError("failed attempt") - assert _nvrtc.resolve_shared_compile_context() is original - - (first / "header.cuh").write_text("edited between attempts") - with _nvrtc.compile_context_scope(): - edited = _nvrtc.resolve_shared_compile_context() - assert edited.header_identity != original.header_identity - assert _nvrtc.resolve_shared_compile_context() is edited - - # Leaving the outermost scope restores ordinary fresh resolution. - unscoped = _nvrtc.resolve_shared_compile_context() - assert unscoped == edited - assert unscoped is not edited - assert _nvrtc.resolve_shared_compile_context() is not unscoped - - -@pytest.mark.parametrize( - ("field", "replacement"), - ( - ("toolkit_root", "/other/toolkit"), - ("toolkit_version", (13, 4)), - ("nvrtc_path", "/other/libnvrtc.so.13"), - ("nvrtc_builtins_path", "/other/libnvrtc-builtins.so.13.3"), - ("nvjitlink_path", "/other/libnvJitLink.so.13"), - ("nvrtc_version", _nvrtc.version(13, 4)), - ("nvjitlink_version", (13, 5)), - ("include_dirs", ("/cuda/toolkit/include", "/cccl/include")), - ("header_identity", "headers-b"), - ), -) -def test_compile_context_symbol_suffix_covers_every_identity_field( - field: str, - replacement: object, -) -> None: - context = _context() - changed = replace(context, **{field: replacement}) - - assert context.symbol_suffix == _context().symbol_suffix - assert context.symbol_suffix != changed.symbol_suffix - - -def test_compile_forwards_complete_context_into_cache_key( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[dict[str, object]] = [] - context = _context() - monkeypatch.setattr( - _nvrtc, - "resolve_compile_context", - lambda: pytest.fail( - "compile must not resolve a supplied context again" - ), - ) - monkeypatch.setattr(_nvrtc, "_dump_source", lambda *args: None) - monkeypatch.setattr( - _nvrtc, - "compile_impl", - lambda **kwargs: calls.append(kwargs) or b"ltoir", - ) - - runtime_version, result = _nvrtc.compile( - cpp="int provider;", - cc=90, - rdc=True, - code="lto", - context=context, - ) - - assert runtime_version == context.nvrtc_version - assert result == b"ltoir" - assert calls == [ - { - "cpp": "int provider;", - "cc": 90, - "rdc": True, - "code": "lto", - "toolkit_root": context.toolkit_root, - "toolkit_version": context.toolkit_version, - "nvrtc_path": context.nvrtc_path, - "nvrtc_builtins_path": context.nvrtc_builtins_path, - "nvjitlink_path": context.nvjitlink_path, - "nvrtc_version": context.nvrtc_version, - "nvjitlink_version": context.nvjitlink_version, - "include_dirs": context.include_dirs, - "header_identity": context.header_identity, - "compiler_options": _nvrtc._compiler_options( - cc=90, - rdc=True, - code="lto", - include_dirs=context.include_dirs, - ), - } - ] - - -@pytest.mark.parametrize( - ("targetoptions", "expected"), - [ - ({}, None), - ({"chip": "sm_90"}, _nvrtc.CompileTarget(cc=90)), - ({"chip": "sm_90a"}, _nvrtc.CompileTarget(cc=90, arch_suffix="a")), - ({"chip": "sm_100f"}, _nvrtc.CompileTarget(cc=100, arch_suffix="f")), - ({"chip": "sm_120"}, _nvrtc.CompileTarget(cc=120)), - ( - {"chip": "sm_120", "lineinfo": True}, - _nvrtc.CompileTarget(cc=120, lineinfo=True), - ), - ( - {"chip": "sm_120", "lineinfo": True, "debug": True}, - _nvrtc.CompileTarget(cc=120), - ), - ], -) -def test_compile_target_reads_the_resolved_kernel_target( - targetoptions, expected -) -> None: - assert _nvrtc.CompileTarget.from_targetoptions(targetoptions) == expected - - -@pytest.mark.parametrize("chip", ["90", "sm_9", "sm_90b", "compute_90", 90]) -def test_compile_target_rejects_an_unresolved_chip(chip) -> None: - with pytest.raises(RuntimeError, match="sm_XY"): - _nvrtc.CompileTarget.from_targetoptions({"chip": chip}) - - -def test_compile_target_selects_architecture_and_line_options() -> None: - target = _nvrtc.CompileTarget(cc=100, arch_suffix="f", lineinfo=True) - options = _nvrtc._compiler_options( - rdc=True, code="lto", include_dirs=(), **target.nvrtc_kwargs() - ) - identity = _nvrtc.compiler_identity( - context=_context(), rdc=True, code="lto", **target.nvrtc_kwargs() - ) - - assert b"--gpu-architecture=compute_100f" in options - assert b"-lineinfo" in options - assert identity.target == target - assert identity != _nvrtc.compiler_identity( - context=_context(), rdc=True, code="lto", cc=100, arch_suffix="f" - ) - - -def test_resolve_context_rejects_loaded_nvrtc_version_mismatch( - monkeypatch: pytest.MonkeyPatch, -) -> None: - include_dir = Path("/cuda/toolkit/include") - events: list[str] = [] - monkeypatch.setattr( - _nvrtc, - "resolve_include_paths", - lambda **kwargs: SimpleNamespace( - as_tuple=lambda: (include_dir,), - cuda=(include_dir,), - ), - ) - monkeypatch.setattr( - _nvrtc, - "preload_toolkit_compiler_libraries", - lambda paths: events.append("preload") or _libraries(), - ) - monkeypatch.setattr( - _nvrtc, - "_load_nvrtc", - lambda: events.append("binding") or _fake_nvrtc((12, 8)), - ) - monkeypatch.setattr( - _nvrtc, - "include_dirs_identity", - lambda paths: pytest.fail("mismatch must fail before header identity"), - ) - - with pytest.raises( - RuntimeError, - match=( - r"headers report Toolkit 13\.3, but loaded NVRTC .* reports 12\.8" - ), - ): - _nvrtc.resolve_compile_context() - - assert events == ["preload", "binding"] - - -def test_compile_impl_rejects_nvrtc_version_change_before_compilation( - monkeypatch: pytest.MonkeyPatch, -) -> None: - implementation = inspect.unwrap(_nvrtc.compile_impl) - monkeypatch.setattr( - _nvrtc, - "_load_nvrtc", - lambda: _fake_nvrtc((13, 2)), - ) - - with pytest.raises( - RuntimeError, - match=( - r"loaded NVRTC version changed after compile-context resolution: " - r"expected version\(major=13, minor=3\), got " - r"version\(major=13, minor=2\)" - ), - ): - implementation( - cpp="int provider;", - cc=90, - rdc=True, - code="lto", - toolkit_root="/cuda/toolkit", - toolkit_version=(13, 3), - nvrtc_path="/cuda/toolkit/lib/libnvrtc.so.13", - nvrtc_builtins_path="/cuda/toolkit/lib/libnvrtc-builtins.so.13.3", - nvjitlink_path="/cuda/toolkit/lib/libnvJitLink.so.13", - nvrtc_version=_nvrtc.version(13, 3), - nvjitlink_version=(13, 4), - include_dirs=("/cccl/include", "/cuda/toolkit/include"), - header_identity="headers-a", - compiler_options=_nvrtc._compiler_options( - cc=90, - rdc=True, - code="lto", - include_dirs=("/cccl/include", "/cuda/toolkit/include"), - ), - ) - - -def test_shared_source_dump_runs_before_provider_cache_lookup( - tmp_path, monkeypatch -): - source = 'extern "C" __device__ int probe() { return 1; }\n' - monkeypatch.setenv("CUDA_COOP_SOURCE_DUMP_DIR", str(tmp_path)) - - def cached(**kwargs): - dumped = tuple(tmp_path.glob("cuda_coop_numba_mlir_*.cu")) - assert len(dumped) == 1 - assert dumped[0].read_bytes() == source.encode("utf-8") - return b"cached" - - monkeypatch.setattr(_nvrtc, "compile_impl", cached) - _, result = _nvrtc.compile( - cpp=source, cc=90, rdc=True, code="lto", context=_context() - ) - assert result == b"cached" - - -def test_shared_source_dump_survives_compile_failure(tmp_path, monkeypatch): - source = "invalid CUDA source\n" - monkeypatch.setenv("CUDA_COOP_SOURCE_DUMP_DIR", str(tmp_path)) - - def fail(**kwargs): - raise RuntimeError("NVRTC compilation failed") - - monkeypatch.setattr(_nvrtc, "compile_impl", fail) - with pytest.raises(RuntimeError, match="NVRTC compilation failed"): - _nvrtc.compile( - cpp=source, cc=90, rdc=True, code="lto", context=_context() - ) - - dumped = tuple(tmp_path.glob("cuda_coop_numba_mlir_*.cu")) - assert len(dumped) == 1 - assert dumped[0].read_bytes() == source.encode("utf-8") - - -@pytest.mark.parametrize("shared_value", (None, "")) -def test_retired_source_dump_setting_is_ignored( - tmp_path, monkeypatch, shared_value -): - monkeypatch.setenv("CUDA_COOP_NUMBA_MLIR_NVRTC_DUMP_DIR", str(tmp_path)) - if shared_value is None: - monkeypatch.delenv("CUDA_COOP_SOURCE_DUMP_DIR", raising=False) - else: - monkeypatch.setenv("CUDA_COOP_SOURCE_DUMP_DIR", shared_value) - monkeypatch.setattr(_nvrtc, "compile_impl", lambda **kwargs: b"cached") - - _, result = _nvrtc.compile( - cpp="// source", cc=90, rdc=True, code="lto", context=_context() - ) - - assert result == b"cached" - assert not list(tmp_path.iterdir()) - - -def _layout_nvrtc(*, failure=None, lowered_name=None): - """Record the NVRTC program lifecycle and supply controlled layout names. - - Return the fake bindings with event, expression, and source logs. Names - encode fixed layouts in registration order; the fake does not evaluate - C++ types. Fail a chosen API phase or override a lowered name so tests can - check program cleanup on both compiler and decoder errors. - """ - - events = [] - expressions = [] - sources = [] - program = object() - layouts = ((24, 8), (64, 16)) - - def create(source, *_args): - events.append("create") - sources.append(source.decode()) - return 0, program - - def register(prog, expression): - assert prog is program - events.append("register") - expressions.append(expression.decode()) - return (1 if failure == "register" else 0,) - - def compile_program(prog, *_args): - assert prog is program - events.append("compile") - return (1 if failure == "compile" else 0,) - - def get_name(prog, expression): - assert prog is program - assert "compile" in events - events.append("name") - if failure == "name": - return 1, b"" - if lowered_name is not None: - return 0, lowered_name - expression = expression.decode() - symbol = expression.split("<", 1)[0].lstrip("&") - size, alignment = layouts[expressions.index(expression) % len(layouts)] - return 0, f"_Z{len(symbol)}{symbol}ILy{size}ELy{alignment}EE".encode() - - def get_image(prog, image): - assert prog is program - events.append("image") - image[:] = b"ltoir" - return (0,) - - def destroy(prog): - assert prog is program - events.append("destroy") - return (0,) - - nvrtc = _fake_nvrtc((13, 3)) - nvrtc.nvrtcCreateProgram = create - nvrtc.nvrtcAddNameExpression = register - nvrtc.nvrtcCompileProgram = compile_program - nvrtc.nvrtcGetLoweredName = get_name - nvrtc.nvrtcGetLTOIRSize = lambda prog: (0, 5) - nvrtc.nvrtcGetLTOIR = get_image - nvrtc.nvrtcDestroyProgram = destroy - nvrtc.nvrtcGetProgramLogSize = lambda prog: (0, 0) - nvrtc.nvrtcGetProgramLog = lambda prog, log: (0,) - return nvrtc, events, expressions, sources - - -@pytest.mark.parametrize("repeated_type", (False, True)) -def test_layout_queries_share_the_provider_compilation( - monkeypatch, repeated_type -): - nvrtc, events, expressions, sources = _layout_nvrtc() - monkeypatch.setattr(_nvrtc, "_load_nvrtc", lambda: nvrtc) - monkeypatch.setattr( - _nvrtc, "compile_impl", inspect.unwrap(_nvrtc.compile_impl) - ) - - runtime_version, result = _nvrtc.compile_with_layouts( - cpp="struct Small {}; struct Large {};", - layout_types=("Small", "Large", "Small") - if repeated_type - else ("Small", "Large"), - cc=90, - rdc=True, - context=_context(), - ) - - assert runtime_version == _context().nvrtc_version - expected = ( - ((24, 8), (64, 16), (24, 8)) if repeated_type else ((24, 8), (64, 16)) - ) - assert result == (b"ltoir", expected) - assert events == [ - "create", - "register", - "register", - "compile", - "name", - "name", - "image", - "destroy", - ] - assert len(sources) == 1 - assert "sizeof(Small)" in expressions[0] - assert "alignof(Small)" in expressions[0] - assert "sizeof(Large)" in expressions[1] - assert "alignof(Large)" in expressions[1] - - -@pytest.mark.parametrize("failure", ("register", "compile", "name")) -def test_layout_query_failures_destroy_the_program(monkeypatch, failure): - nvrtc, events, _, _ = _layout_nvrtc(failure=failure) - monkeypatch.setattr(_nvrtc, "_load_nvrtc", lambda: nvrtc) - monkeypatch.setattr( - _nvrtc, "compile_impl", inspect.unwrap(_nvrtc.compile_impl) - ) - - with pytest.raises(RuntimeError, match="NVRTC"): - _nvrtc.compile_with_layouts( - cpp="struct Storage {};", - layout_types=("Storage",), - cc=90, - rdc=True, - context=_context(), - ) - - assert events[-1] == "destroy" - assert events.count("destroy") == 1 - assert "image" not in events - - -def test_malformed_layout_name_destroys_the_program(monkeypatch): - nvrtc, events, _, _ = _layout_nvrtc(lowered_name=b"unexpected_lowered_name") - monkeypatch.setattr(_nvrtc, "_load_nvrtc", lambda: nvrtc) - monkeypatch.setattr( - _nvrtc, "compile_impl", inspect.unwrap(_nvrtc.compile_impl) - ) - - with pytest.raises(ValueError, match="(?i)(layout|lowered)"): - _nvrtc.compile_with_layouts( - cpp="struct Storage {};", - layout_types=("Storage",), - cc=90, - rdc=True, - context=_context(), - ) - - assert events[-1] == "destroy" - assert events.count("destroy") == 1 - - -@pytest.mark.parametrize("layout_types", ((), ("Storage",))) -def test_layout_queries_partition_the_memory_cache(monkeypatch, layout_types): - from cuda.coop.numba_mlir._compiler import _caching - - nvrtc, events, _, sources = _layout_nvrtc() - monkeypatch.setattr(_nvrtc, "_load_nvrtc", lambda: nvrtc) - monkeypatch.setattr(_caching, "_CACHE_USABLE", False) - _nvrtc.compile_impl.cache_clear() - arguments = { - "cpp": "struct Storage {}; struct Other {};", - "cc": 90, - "rdc": True, - "context": _context(), - } - try: - first = _nvrtc.compile_with_layouts( - layout_types=layout_types, **arguments - ) - repeated = _nvrtc.compile_with_layouts( - layout_types=layout_types, **arguments - ) - _nvrtc.compile_with_layouts(layout_types=("Other",), **arguments) - assert first == repeated - assert events.count("compile") == 2 - assert sources[0] != sources[1] - if not layout_types: - assert first[1] == (b"ltoir", ()) - assert events[:4] == ["create", "compile", "image", "destroy"] - finally: - _nvrtc.compile_impl.cache_clear() - - -@pytest.mark.parametrize("size,alignment", ((0, 8), (8, 0), (24, 3), (24, 16))) -def test_invalid_layout(size, alignment): - from cuda.coop.numba_mlir._compiler._layout import decode_layout_name - - with pytest.raises(ValueError, match="(?i)layout"): - decode_layout_name( - f"_Z5probeILy{size}ELy{alignment}EE", - symbol="probe", - expression="&probe<(sizeof(Storage)), (alignof(Storage))>", - ) diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_parameter_names.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_parameter_names.py deleted file mode 100644 index 8e6e66eb06de..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_parameter_names.py +++ /dev/null @@ -1,176 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from dataclasses import replace - -import pytest - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - - -def _materialize(specialization, *, adapter=None): - """Adapt a specialization without scratch or synchronization operands.""" - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir._compiler._operations import StorageABI - from cuda.coop.numba_mlir._lowering._core import NumbaMlirCoreAdapter - - return (adapter or NumbaMlirCoreAdapter()).materialize( - specialization, - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.NONE, - synchronization_scope=SynchronizationScope.NONE, - ) - - -def _source(algorithm): - """Generate wrapper text with a fixed identity, without compiling it.""" - from cuda.coop.numba_mlir._compiler._nvrtc import CompilerIdentity - - return algorithm._source_code( - compile_identity=CompilerIdentity( - cc=90, rdc=True, code="lto", compiler_options=() - ) - )[0] - - -@pytest.mark.parametrize("operation", ("load", "store")) -def test_data_movement_retains_names_through_dependent_specialization( - operation, -): - from cuda.coop._core import INT32, ArgumentBinding - from cuda.coop._core.block import ( - make_block_load_specialization, - make_block_store_specialization, - ) - - factory = ( - make_block_load_specialization - if operation == "load" - else make_block_store_specialization - ) - specialization = factory( - dtype=INT32, - block_dim=(32, 1, 1), - items_per_thread=2, - algorithm="direct", - valid_items=ArgumentBinding.omitted(), - include_pointer_offset=ArgumentBinding.runtime(), - ) - algorithm = _materialize(specialization) - method = algorithm.parameters[0] - assert [parameter.parameter_name for parameter in method] == ( - ["src", "dst", "offset"] - if operation == "load" - else ["dst", "src", "offset"] - ) - source = _source(algorithm) - abi_signature = next( - line for line in source.splitlines() if "__abi(void *__ret," in line - ) - assert "void *src" in abi_signature - assert "void *dst" in abi_signature - assert "::cuda::std::int64_t offset" in abi_signature - memory_name = "src" if operation == "load" else "dst" - assert f".{operation.title()}(({memory_name} + offset)," in source - - -def test_input_transform_retains_the_source_parameter_name(): - from cuda.coop._core import INT32, UINT8, Algorithm, Array - from cuda.coop.numba_mlir._lowering._core import ( - NumbaMlirArrayInputTransform, - NumbaMlirCoreAdapter, - ) - - specialization = Algorithm( - struct_name="Provider", - method_name="Run", - c_name="transformed_provider", - includes=(), - template_parameters=(), - parameters=((Array(INT32, 2, name="flags"),),), - template_arguments={}, - ) - algorithm = _materialize( - specialization, - adapter=NumbaMlirCoreAdapter( - input_transforms={ - "flags": NumbaMlirArrayInputTransform(UINT8, "{value} != 0") - } - ), - ) - - assert algorithm.parameters[0][0].parameter_name == "flags" - source = _source(algorithm) - assert "::cuda::std::uint8_t *flags" in source - assert "transformed_flags[0] = flags[0] != 0;" in source - assert ".Run(transformed_flags);" in source - assert "void *flags" in source - - -def test_parameter_renaming_preserves_provider_identity(): - from cuda.coop._core import INT32, Algorithm, Reference, Value - from cuda.coop.numba_mlir._types import algo_coalesce_key - - definition = Algorithm( - struct_name="Provider", - method_name="Run", - c_name="named_provider", - includes=(), - template_parameters=(), - parameters=( - (Value(INT32, name="value"), Reference(INT32, name="result")), - ), - template_arguments={}, - ) - renamed_definition = replace( - definition, - parameters=( - (Value(INT32, name="input"), Reference(INT32, name="output")), - ), - ) - original = _materialize(definition) - renamed = _materialize(renamed_definition) - - assert algo_coalesce_key(original) == algo_coalesce_key(renamed) - assert original.mangled_name( - original.parameters[0] - ) == renamed.mangled_name(renamed.parameters[0]) - original_source = _source(original) - renamed_source = _source(renamed) - assert "__abi(void *__ret, ::cuda::std::int32_t value," in original_source - assert "__abi(void *__ret, ::cuda::std::int32_t input," in renamed_source - assert original._private_symbol_key == renamed._private_symbol_key - assert original._private_symbol_digest == renamed._private_symbol_digest - - -def test_shared_value_abi_names_do_not_mutate_other_parameters(): - from numba_cuda_mlir import types - - from cuda.coop._core import INT32, Algorithm, Value - from cuda.coop.numba_mlir import _types - from cuda.coop.numba_mlir._lowering._core import NumbaMlirCoreAdapter - - shared_abi = _types.Value(types.int32) - definition = Algorithm( - struct_name="Provider", - method_name="Run", - c_name="shared_abi_provider", - includes=(), - template_parameters=(), - parameters=((Value(INT32, name="first"), Value(INT32, name="second")),), - template_arguments={}, - ) - algorithm = _materialize( - definition, - adapter=NumbaMlirCoreAdapter( - value_abis={"first": shared_abi, "second": shared_abi} - ), - ) - first, second = algorithm.parameters[0] - assert first is not second - assert first is not shared_abi - assert second is not shared_abi - assert first.parameter_name == "first" - assert second.parameter_name == "second" - assert getattr(shared_abi, "parameter_name", None) is None diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_public_foundation.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_public_foundation.py deleted file mode 100644 index 6bd0a8eb16d9..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_public_foundation.py +++ /dev/null @@ -1,325 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import ast -import importlib -import importlib.util -import inspect -import sys -from pathlib import Path -from typing import get_args - -import numpy as np -import pytest - -import cuda.coop as common_coop -import cuda.coop.numba_mlir as coop -from cuda.coop.numba_mlir import _temp_storage, _thread_data - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - -_COMMON_EXPORTS = [ - "__version__", - "register", - "Hierarchy", - "TempStorage", - "TempStorageLike", - "ThreadData", - "ThreadDataLike", - "ThreadGroup", - "ThreadHierarchy", - "this_block", - "this_cluster", - "this_grid", - "this_thread", - "this_warp", - "load", - "store", -] -_QUALIFIED_EXPORTS = [ - *( - name - for name in _COMMON_EXPORTS - if name not in {"__version__", "register"} - ), - "local", - "shared", -] -_EXCLUDED_BACKEND_MODULES = ( - "cuda.coop.numba_mlir._dataclass", - "cuda.coop.numba_mlir._enums", - "cuda.coop.numba_mlir._group._reduce", - "cuda.coop.numba_mlir._group._scan", - "cuda.coop.numba_mlir._stateful_function", - "cuda.coop.numba_mlir._compiler._group_reduce", - "cuda.coop.numba_mlir._compiler._group_scan", - "cuda.coop.numba_mlir._compiler._rewrite_reduce", - "cuda.coop.numba_mlir._compiler._rewrite_scan", - "cuda.coop.numba_mlir._lowering._reduce", - "cuda.coop.numba_mlir._lowering._scan", - "cuda.coop.numba_mlir._lowering._thread_group", - "cuda.coop.numba_mlir._lowering._warp", -) - - -def test_public_exports_are_only_the_load_store_foundation(): - assert sorted(common_coop.__all__) == sorted(_COMMON_EXPORTS) - assert dir(common_coop) == sorted(_COMMON_EXPORTS) - assert sorted(coop.__all__) == sorted(_QUALIFIED_EXPORTS) - assert dir(coop) == sorted(_QUALIFIED_EXPORTS) - - excluded_exports = { - "BlockLoadAlgorithm", - "BlockReduceAlgorithm", - "BlockScanAlgorithm", - "BlockStoreAlgorithm", - "StatefulFunction", - "exclusive_scan", - "gpu_dataclass", - "inclusive_scan", - "reduce", - "scan", - "sum", - "WarpLoadAlgorithm", - "WarpStoreAlgorithm", - } - assert excluded_exports.isdisjoint(common_coop.__all__) - assert excluded_exports.isdisjoint(coop.__all__) - - loaded = set(sys.modules) - assert "cuda.coop.numba_mlir._group._load_store" in loaded - assert "cuda.coop.numba_mlir._compiler._rewrite" in loaded - assert set(_EXCLUDED_BACKEND_MODULES).isdisjoint(loaded) - assert ( - importlib.import_module("cuda.coop.numba_mlir._lowering").__all__ == () - ) - - coop_root = Path(common_coop.__file__).resolve().parent - assert not (coop_root / "cutlass").exists() - - -def test_qualified_surface_is_common_plus_backend_extensions(): - assert set(coop.__all__) - set(common_coop.__all__) == {"local", "shared"} - assert set(common_coop.__all__) - set(coop.__all__) == { - "__version__", - "register", - } - - def call_shape(function): - return tuple( - (name, parameter.kind, parameter.default) - for name, parameter in inspect.signature( - function - ).parameters.items() - ) - - for operation in ("load", "store"): - assert call_shape(getattr(coop, operation)) == call_shape( - getattr(common_coop, operation) - ) - - assert call_shape(coop.TempStorage) == call_shape(common_coop.TempStorage) - for constructor in ( - "this_thread", - "this_warp", - "this_block", - "this_cluster", - "this_grid", - ): - assert call_shape(getattr(coop, constructor)) == call_shape( - getattr(common_coop, constructor) - ) - - common_thread_data = call_shape(common_coop.ThreadData) - qualified_thread_data = call_shape(coop.ThreadData) - assert qualified_thread_data == common_thread_data - assert qualified_thread_data[-1] == ( - "alignment", - inspect.Parameter.KEYWORD_ONLY, - None, - ) - assert call_shape(coop.TempStorage) == ( - ("size_in_bytes", inspect.Parameter.POSITIONAL_OR_KEYWORD, None), - ("alignment", inspect.Parameter.KEYWORD_ONLY, None), - ("auto_sync", inspect.Parameter.KEYWORD_ONLY, False), - ("sharing", inspect.Parameter.KEYWORD_ONLY, "shared"), - ) - assert call_shape(coop.ThreadGroup.group_by) == call_shape( - common_coop.ThreadGroup.group_by - ) - - assert coop.ThreadDataLike is common_coop.ThreadDataLike - assert coop.TempStorageLike is common_coop.TempStorageLike - assert coop.ThreadHierarchy is common_coop.ThreadHierarchy - assert coop.Hierarchy is common_coop.Hierarchy - - common_load_annotations = inspect.get_annotations( - common_coop.load, - eval_str=True, - ) - qualified_load_annotations = inspect.get_annotations( - coop.load, - eval_str=True, - ) - assert set(get_args(qualified_load_annotations["output"])) == { - common_load_annotations["output"], - np.ndarray, - } - assert ( - qualified_load_annotations["return"] - == common_load_annotations["return"] - ) - - coop_root = Path(common_coop.__file__).resolve().parent - - def stub_signatures(path): - module = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) - return [ - (node.name, ast.dump(node.args), ast.dump(node.returns)) - for node in module.body - if isinstance(node, ast.FunctionDef) - and node.name in {"load", "store"} - ] - - assert stub_signatures( - coop_root / "numba_mlir" / "_group" / "_load_store.pyi" - ) == (stub_signatures(coop_root / "_core" / "api" / "load_store.pyi")) - - -def test_group_descriptors_expose_only_canonical_extent_names(): - assert not hasattr(common_coop.ThreadHierarchy, "thread_count") - for group_type in (common_coop.ThreadGroup, coop.ThreadGroup): - assert not hasattr(group_type, "static_thread_count") - assert not hasattr(group_type, "thread_count") - assert hasattr(group_type, "static_size") - assert hasattr(group_type, "group_thread_count") - - -def test_stub_only_group_aliases_exist_at_runtime_without_becoming_exports(): - from cuda.coop._core.api import thread_group as common_groups - from cuda.coop.numba_mlir import _thread_group as qualified_groups - - for name in ("MemoryGroup", "ReductionGroup", "BlockGroup", "WarpGroup"): - assert getattr(common_groups, name) is common_groups.ThreadGroup - assert name not in common_groups.__all__ - for name in ("ReductionGroup", "BlockGroup", "WarpGroup"): - assert getattr(qualified_groups, name) is qualified_groups.ThreadGroup - assert name not in qualified_groups.__all__ - - -@pytest.mark.parametrize("module_name", _EXCLUDED_BACKEND_MODULES) -def test_excluded_backend_implementation_modules_remain_absent(module_name): - assert importlib.util.find_spec(module_name) is None - - coop_root = Path(common_coop.__file__).resolve().parent - relative_module = module_name.removeprefix("cuda.coop.") - module_path = coop_root.joinpath(*relative_module.split(".")) - assert not module_path.with_suffix(".py").exists() - assert not module_path.is_dir() - - -def test_python_operator_compilation_remains_absent(): - from cuda.coop.numba_mlir import _types - - assert not hasattr(_types, "_compile_device_ltoir") - assert tuple( - inspect.signature(_types.numba_type_to_wrapper).parameters - ) == ("numba_type",) - - -@pytest.mark.parametrize("operation", ("load", "store")) -def test_group_markers_use_exact_callable_identity(operation): - from cuda.coop.numba_mlir._compiler._operations import group_operation_name - - marker = getattr(coop, operation) - assert group_operation_name(marker) == operation - - def impostor(*args, **kwargs): - del args, kwargs - - impostor.__module__ = marker.__module__ - impostor.__name__ = marker.__name__ - impostor.__cuda_coop_backend_member__ = operation - assert group_operation_name(impostor) is None - - -@pytest.mark.parametrize("operation", ("load", "store")) -def test_lowering_factories_use_exact_callable_identity(operation): - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir import _lowering - from cuda.coop.numba_mlir._compiler._operations import ( - FactoryOperation, - StorageABI, - factory_operation, - ) - - factory = getattr(_lowering, operation) - assert factory_operation(factory) == FactoryOperation( - operation=operation, - namespace="block", - storage_abi=StorageABI.NONE, - execution_scope=SynchronizationScope.BLOCK, - synchronization_scope=SynchronizationScope.NONE, - ) - - def impostor(*args, **kwargs): - del args, kwargs - - impostor.__module__ = factory.__module__ - impostor.__name__ = factory.__name__ - assert factory_operation(impostor) is None - - -@pytest.mark.parametrize("operation", ("load", "store")) -def test_physical_warp_factories_use_exact_callable_identity(operation): - from cuda.coop._core import SynchronizationScope - from cuda.coop.numba_mlir._compiler._operations import ( - FactoryOperation, - StorageABI, - factory_operation, - ) - from cuda.coop.numba_mlir._lowering import _load_store - - for storage_bearing in (False, True): - factory = getattr( - _load_store, - ( - f"_warp_{operation}_with_storage" - if storage_bearing - else f"warp_{operation}" - ), - ) - assert factory_operation(factory) == FactoryOperation( - operation=operation, - namespace="warp", - storage_abi=( - StorageABI.LEADING_POINTER - if storage_bearing - else StorageABI.NONE - ), - execution_scope=SynchronizationScope.WARP, - synchronization_scope=( - SynchronizationScope.WARP - if storage_bearing - else SynchronizationScope.NONE - ), - ) - - def impostor(*args, **kwargs): - del args, kwargs - - impostor.__module__ = factory.__module__ - impostor.__name__ = factory.__name__ - assert factory_operation(impostor) is None - - -def test_public_runtime_helpers_have_semantic_module_owners(): - assert coop.local is importlib.import_module("numba_cuda_mlir.cuda").local - assert coop.shared is importlib.import_module("numba_cuda_mlir.cuda").shared - assert coop.ThreadData is _thread_data.ThreadData - assert coop.TempStorage is _temp_storage.TempStorage - assert coop.ThreadDataLike is common_coop.ThreadDataLike - assert coop.TempStorageLike is common_coop.TempStorageLike - assert coop.ThreadData.__module__ == "cuda.coop.numba_mlir._thread_data" - assert coop.TempStorage.__module__ == "cuda.coop.numba_mlir._temp_storage" diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_runtime_activation.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_runtime_activation.py deleted file mode 100644 index c3d43a8ccbc6..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_runtime_activation.py +++ /dev/null @@ -1,484 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check backend activation and import failures in fresh interpreters. - -Separate processes prevent one import order or registry state from masking -a later case. Embedded probes also exercise missing and unsupported runtimes. -""" - -import os -import subprocess -import sys -import textwrap -from pathlib import Path - -import pytest - -from cuda.coop.numba_mlir._compiler import _activation - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - -PACKAGE_ROOT = Path(__file__).parents[4] - - -def _run_import_probe(script: str) -> None: - """Run a fresh import against this package with default registration.""" - env = os.environ.copy() - env.pop("CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION", None) - env["PYTHONPATH"] = os.pathsep.join( - filter(None, (str(PACKAGE_ROOT), env.get("PYTHONPATH"))) - ) - result = subprocess.run( - [sys.executable, "-B", "-c", textwrap.dedent(script)], - check=False, - capture_output=True, - env=env, - text=True, - timeout=180, - ) - - assert result.returncode == 0, result.stdout + result.stderr - - -def test_numba_first_root_import_automatically_activates_backend(): - _run_import_probe( - """ - import sys - - import numba_cuda_mlir # noqa: F401 - import cuda.coop # noqa: F401 - - expected = { - "cuda.coop.numba_mlir", - "cuda.coop.numba_mlir._compiler._planner", - } - assert expected <= set(sys.modules), expected - set(sys.modules) - """ - ) - - -def test_numba_first_root_import_does_not_require_refresh_registries(): - _run_import_probe( - """ - import sys - - import numba_cuda_mlir - import numba_cuda_mlir.extending as extending - - del extending.refresh_registries - import cuda.coop # noqa: F401 - - assert "cuda.coop.numba_mlir" in sys.modules - assert numba_cuda_mlir.__version__.startswith("0.5.") - """ - ) - - -def test_root_first_qualified_import_explicitly_activates_backend(): - _run_import_probe( - """ - import sys - - before = set(sys.modules) - import cuda.coop # noqa: F401 - loaded = set(sys.modules) - before - - assert "cuda.coop.numba_mlir" not in sys.modules - assert not any( - name == "numba_cuda_mlir" or name.startswith("numba_cuda_mlir.") - for name in loaded - ), loaded - assert not any( - name == "cuda.bindings" or name.startswith("cuda.bindings.") - for name in loaded - ), loaded - - import cuda.coop.numba_mlir # noqa: F401 - - expected = { - "cuda.coop.numba_mlir", - "cuda.coop.numba_mlir._compiler._planner", - } - assert expected <= set(sys.modules), expected - set(sys.modules) - """ - ) - - -@pytest.mark.parametrize("backend", ["numba-cuda-mlir", "numba_cuda_mlir"]) -@pytest.mark.parametrize("disable_auto_registration", ["0", "1"]) -def test_root_first_register_explicitly_activates_backend_once( - backend, disable_auto_registration -): - _run_import_probe( - f""" - import os - import sys - - os.environ["CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION"] = ( - {disable_auto_registration!r} - ) - from cuda import coop - - assert "numba_cuda_mlir" not in sys.modules - assert "cuda.bindings" not in sys.modules - assert "cuda.coop.numba_mlir" not in sys.modules - assert coop.register({backend!r}) is None - - expected = {{ - "cuda.coop.numba_mlir", - "cuda.coop.numba_mlir._compiler._planner", - }} - assert expected <= set(sys.modules), expected - set(sys.modules) - - from numba_cuda_mlir._whole_function_planners import _planner_registry - from numba_cuda_mlir.numba_cuda.core.rewrites import rewrite_registry - from cuda.coop.numba_mlir._compiler._planner import ( - CoopWholeFunctionPlanner, - ) - - from cuda.coop.numba_mlir._compiler._activation import ( - _initialize_runtime_hooks, - ) - - assert coop.register("numba-cuda-mlir") is None - assert coop.register("numba_cuda_mlir") is None - _initialize_runtime_hooks() - _initialize_runtime_hooks() - assert [ - planner for planner in _planner_registry._planners - if planner.__module__.startswith("cuda.coop.") - ] == [CoopWholeFunctionPlanner] - assert not any( - rewrite.__module__.startswith("cuda.coop.") - for rewrites in rewrite_registry.rewrites.values() - for rewrite in rewrites - ) - """ - ) - - -def test_runtime_loading_retries_after_a_failed_qualified_import(monkeypatch): - runtime = object() - error = _activation.NumbaMlirBackendImportError( - "backend-runtime-missing", - "runtime unavailable", - ) - outcomes = iter(((None, error), (runtime, None))) - monkeypatch.setattr(_activation, "_cuda_module", None) - monkeypatch.setattr(_activation, "_load_runtime", lambda: next(outcomes)) - - with pytest.raises(_activation.NumbaMlirBackendImportError) as exc_info: - _activation._require_runtime() - - assert exc_info.value is error - assert _activation._require_runtime() is runtime - assert _activation._cuda_module is runtime - - -@pytest.mark.parametrize("error_type", ["RuntimeError", "AttributeError"]) -def test_failed_planner_import_can_retry_without_partial_registration( - error_type, -): - _run_import_probe( - """ - import importlib - import os - - os.environ["CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION"] = "1" - - # Initialize runtime-owned registrations before importing the backend. - from numba_cuda_mlir.compiler import ExternFunction # noqa: F401 - from numba_cuda_mlir.extending import typeof_impl - from numba_cuda_mlir._whole_function_planners import _planner_registry - from numba_cuda_mlir.numba_cuda.core.rewrites import rewrite_registry - - baseline = dict(typeof_impl.registry) - real_import_module = importlib.import_module - fail_once = True - injected_error = ERROR_TYPE("injected planner import failure") - - def fail_after_planner_import(name, package=None): - global fail_once - module = real_import_module(name, package) - if fail_once and name.endswith("._compiler._planner"): - fail_once = False - raise injected_error - return module - - importlib.import_module = fail_after_planner_import - try: - import cuda.coop.numba_mlir # noqa: F401 - except ERROR_TYPE as error: - assert error is injected_error - else: - raise AssertionError("injected activation failure did not occur") - finally: - importlib.import_module = real_import_module - - assert not any( - planner.__module__.startswith("cuda.coop.") - for planner in _planner_registry._planners - ) - assert not any( - rewrite.__module__.startswith("cuda.coop.") - for rewrites in rewrite_registry.rewrites.values() - for rewrite in rewrites - ) - assert dict(typeof_impl.registry) == baseline - - import cuda.coop.numba_mlir # noqa: F401 - from cuda.coop.numba_mlir._compiler._planner import ( - CoopWholeFunctionPlanner, - ) - - assert [ - planner for planner in _planner_registry._planners - if planner.__module__.startswith("cuda.coop.") - ] == [CoopWholeFunctionPlanner] - assert dict(typeof_impl.registry) == baseline - """.replace("ERROR_TYPE", error_type) - ) - - -@pytest.mark.parametrize( - "activate", - [ - "import cuda.coop.numba_mlir", - 'from cuda import coop; coop.register("numba-cuda-mlir")', - ], -) -@pytest.mark.parametrize( - "version", ["0.4.9", "0.5.0rc1", "0.6.0rc1", "0.6.0", "invalid", None] -) -def test_backend_registration_rejects_unsupported_runtime_version( - activate, version -): - _run_import_probe( - """ - import importlib.metadata - import os - import sys - from types import ModuleType - - os.environ["CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION"] = "1" - - runtime = ModuleType("numba_cuda_mlir") - runtime.__version__ = VERSION - sys.modules[runtime.__name__] = runtime - real_version = importlib.metadata.version - - def distribution_version(name): - if name == "numba-cuda-mlir": - raise importlib.metadata.PackageNotFoundError(name) - return real_version(name) - - importlib.metadata.version = distribution_version - try: - ACTIVATE_BACKEND - except ImportError as exc: - assert exc.backend == "numba-cuda-mlir" - assert exc.reason_code == "unsupported-runtime-version" - assert exc.details["detected_version"] == VERSION - assert exc.details["required_version"] == ">=0.5.0,<0.6" - message = str(exc) - assert "numba-cuda-mlir>=0.5.0,<0.6" in message - assert "cuda-coop[numba-cuda-mlir-cu12]" in message - assert "cuda-coop[numba-cuda-mlir-cu13]" in message - if VERSION is not None: - assert str(VERSION) in message - else: - raise AssertionError("unsupported runtime unexpectedly activated") - - assert "numba_cuda_mlir.cuda" not in sys.modules - assert "cuda.coop.numba_mlir._compiler._planner" not in sys.modules - """.replace("ACTIVATE_BACKEND", activate).replace( - "VERSION", repr(version) - ) - ) - - -@pytest.mark.parametrize("version", ["0.5.0+local", "0.5.1rc1", None]) -def test_backend_registration_accepts_supported_runtime_version(version): - _run_import_probe( - f""" - import os - import sys - - os.environ["CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION"] = "1" - - import numba_cuda_mlir - - # An absent module version falls back to installed package metadata. - version = {version!r} - if version is None: - del numba_cuda_mlir.__version__ - else: - numba_cuda_mlir.__version__ = version - - from cuda import coop - - coop.register("numba-cuda-mlir") - assert "cuda.coop.numba_mlir._compiler._planner" in sys.modules - """ - ) - - -def test_root_import_ignores_an_installed_but_unused_old_runtime(tmp_path): - runtime = tmp_path / "numba_cuda_mlir" - runtime.mkdir() - (runtime / "__init__.py").write_text('__version__ = "0.4.9"\n') - _run_import_probe( - f""" - import importlib.util - import sys - import warnings - - sys.path.insert(0, {str(tmp_path)!r}) - assert importlib.util.find_spec("numba_cuda_mlir") is not None - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - from cuda import coop - - assert callable(coop.ThreadData) - assert not caught, caught - assert "numba_cuda_mlir" not in sys.modules - assert "cuda.coop.numba_mlir" not in sys.modules - """ - ) - - -def test_root_import_warns_when_an_old_runtime_was_already_imported(): - _run_import_probe( - """ - import sys - import warnings - from types import ModuleType - - runtime = ModuleType("numba_cuda_mlir") - runtime.__version__ = "0.4.9" - sys.modules[runtime.__name__] = runtime - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - from cuda import coop - - from cuda.coop._core._auto_registration import ( - CudaCoopAutoRegistrationWarning, - ) - - assert callable(coop.ThreadData) - assert len(caught) == 1, caught - assert caught[0].category is CudaCoopAutoRegistrationWarning - assert "0.4.9" in str(caught[0].message) - assert "numba-cuda-mlir>=0.5.0,<0.6" in str(caught[0].message) - assert "cuda.coop.numba_mlir" not in sys.modules - assert "numba_cuda_mlir.cuda" not in sys.modules - """ - ) - - -def test_missing_packaging_is_reported_before_compiler_import(): - _run_import_probe( - """ - import importlib.abc - import os - import sys - from types import ModuleType - - os.environ["CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION"] = "1" - runtime = ModuleType("numba_cuda_mlir") - runtime.__version__ = "0.5.0" - sys.modules[runtime.__name__] = runtime - missing = ModuleNotFoundError( - "No module named 'packaging'", name="packaging" - ) - - class BlockPackaging(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname == "packaging": - raise missing - return None - - assert "packaging" not in sys.modules - sys.meta_path.insert(0, BlockPackaging()) - from cuda import coop - - try: - coop.register("numba-cuda-mlir") - except ImportError as exc: - assert exc.reason_code == "backend-dependency-missing" - assert exc.details["missing"] == "packaging" - assert exc.__cause__ is missing - message = str(exc) - assert "packaging" in message - assert "cuda-coop[numba-cuda-mlir-cu12]" in message - assert "cuda-coop[numba-cuda-mlir-cu13]" in message - else: - raise AssertionError( - "activation without packaging unexpectedly succeeded" - ) - - assert "numba_cuda_mlir.cuda" not in sys.modules - assert "cuda.coop.numba_mlir._compiler._planner" not in sys.modules - """ - ) - - -@pytest.mark.parametrize( - "activate", - [ - "import cuda.coop.numba_mlir", - 'from cuda import coop; coop.register("numba-cuda-mlir")', - ], -) -def test_backend_registration_reports_a_missing_public_runtime(activate): - script = textwrap.dedent( - """ - import importlib.abc - import os - import sys - - os.environ["CUDA_COOP_DISABLE_AUTO_DSL_REGISTRATION"] = "1" - - class BlockRuntime(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - del path, target - if fullname == "numba_cuda_mlir" or fullname.startswith( - "numba_cuda_mlir." - ): - raise ImportError("blocked runtime", name=fullname) - return None - - sys.meta_path.insert(0, BlockRuntime()) - - try: - ACTIVATE_BACKEND - except ImportError as exc: - assert exc.backend == "numba-cuda-mlir" - assert exc.reason_code == "backend-runtime-missing" - assert exc.details["missing"] == "numba_cuda_mlir" - assert isinstance(exc.__cause__, ImportError) - message = str(exc) - assert "numba-cuda-mlir>=0.5.0" in message - assert "cuda-coop[numba-cuda-mlir-cu12]" in message - assert "cuda-coop[numba-cuda-mlir-cu13]" in message - else: - raise AssertionError("qualified import unexpectedly succeeded") - """.replace("ACTIVATE_BACKEND", activate) - ) - env = os.environ.copy() - env["PYTHONPATH"] = os.pathsep.join( - filter(None, (str(PACKAGE_ROOT), env.get("PYTHONPATH"))) - ) - result = subprocess.run( - [sys.executable, "-S", "-B", "-c", script], - check=False, - capture_output=True, - env=env, - text=True, - timeout=180, - ) - - assert result.returncode == 0, result.stderr diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_storage_rewrites.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_storage_rewrites.py deleted file mode 100644 index f8b96d3d9460..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_storage_rewrites.py +++ /dev/null @@ -1,2423 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check payload and scratch rewrites by inspecting Numba's function IR. - -The kernels below are parsed and rewritten; their array accesses are not run. -Fake provider metadata isolates alias handling, allocation layout, and barrier -insertion from CUDA compilation. Tests use both ordinary frontend IR and -single-assignment form to check descriptors across branches and loops. -""" - -from enum import Enum -from types import SimpleNamespace - -import numpy as np -import pytest -from numba_cuda_mlir import cuda, types -from numba_cuda_mlir.numba_cuda.compiler import run_frontend -from numba_cuda_mlir.numbair_transforms import ir - -import cuda.coop as common_coop -import cuda.coop.numba_mlir as coop -from cuda.coop._core import SynchronizationScope -from cuda.coop.numba_mlir._compiler import _operations -from cuda.coop.numba_mlir._compiler._planner import CoopWholeFunctionPlanner -from cuda.coop.numba_mlir._compiler._rewrite import ( - CoopSinglePhaseRewrite, - CoopSinglePhaseRewriteError, -) -from cuda.coop.numba_mlir._compiler._rewrite_support import ( - _DEFAULT_STATIC_SHARED_MEMORY_BYTES, - _TempStorageCtorSpecification, - _TempStoragePlan, - _TempStorageRequirementSummary, - _TempStorageUseRequirement, -) - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - - -class _StringSharing(str, Enum): - SHARED = "shared" - - -class _TypingContext: - """Count refresh requests without a full compiler typing context.""" - - def __init__(self): - self.refresh_count = 0 - - def refresh(self): - self.refresh_count += 1 - - -@pytest.fixture(autouse=True) -def _restore_rewrite_registries(): - """Remove temporary test providers after each test.""" - - operation_prefix = "_test_storage_rewrite_family_" - try: - yield - finally: - for factory, metadata in tuple(_operations._FACTORY_OPERATIONS.items()): - if metadata.operation.startswith(operation_prefix): - del _operations._FACTORY_OPERATIONS[factory] - for operation in tuple(_operations._REWRITE_OPERATIONS): - if operation.startswith(operation_prefix): - del _operations._REWRITE_OPERATIONS[operation] - - -def _rewrite(function, *, arg_types=()): - """Rewrite payload constructors and fail if a block keeps matching. - - Supply only the compiler state these rewrites need. The iteration bound - turns an accidental rewrite cycle into a test failure instead of a hang. - """ - - func_ir = run_frontend(function) - typingctx = _TypingContext() - state = SimpleNamespace( - func_ir=func_ir, - args=arg_types, - typingctx=typingctx, - typemap={}, - calltypes={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - assert rewrite.prepare_calls_and_storage(func_ir) - rewrite.begin_rewrite() - for label in sorted(func_ir.blocks): - block = func_ir.blocks[label] - for _ in range(10): - if not rewrite.match( - func_ir, block, state.typemap, state.calltypes - ): - break - block = rewrite.apply() - func_ir.blocks[label] = block - else: - pytest.fail("payload rewrite did not reach a fixed point") - rewrite.finish_rewrite() - return func_ir, typingctx - - -def _call_targets(func_ir): - """Resolve emitted callees to Python functions for assertions.""" - - rewrite = object.__new__(CoopSinglePhaseRewrite) - rewrite._func_ir = func_ir - targets = [] - for block in func_ir.blocks.values(): - rewrite._block_defs = { - inst.target.name: inst.value - for inst in block.body - if isinstance(inst, ir.Assign) - } - for inst in block.body: - value = getattr(inst, "value", None) - if isinstance(value, ir.Expr) and value.op == "call": - targets.append(rewrite._resolve_python_value(value.func)) - return targets - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_qualified_thread_data_lowers_to_a_compiler_array(items_per_thread): - def kernel(items_per_thread): - data = coop.ThreadData(items_per_thread, types.int32, alignment=16) - return data[0] - - func_ir, typingctx = _rewrite( - kernel, arg_types=(types.IntegerLiteral(items_per_thread),) - ) - targets = _call_targets(func_ir) - - assert cuda.local.array in targets - assert typingctx.refresh_count == 1 - - -@pytest.mark.parametrize( - "module", (common_coop, coop), ids=("root", "qualified") -) -def test_thread_data_constructor_alias_can_feed_multiple_blocks(module): - def kernel(flag): - constructor = module.ThreadData - alias = constructor - if flag: - first = alias(1, types.int32) - return first[0] - second = alias(2, types.int32) - return second[1] - - func_ir, _ = _rewrite(kernel) - - assert _call_targets(func_ir).count(cuda.local.array) == 2 - constructor_bindings = { - inst.target.name: inst.value - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) - and inst.target.name in {"constructor", "alias"} - } - assert constructor_bindings.keys() == {"constructor", "alias"} - assert all( - isinstance(value, ir.Const) and value.value is None - for value in constructor_bindings.values() - ) - - -@pytest.mark.parametrize( - "module", (common_coop, coop), ids=("root", "qualified") -) -def test_thread_data_constructor_alias_can_feed_a_loop(module): - def kernel(): - constructor = module.ThreadData - result = 0 - for _ in range(2): - data = constructor(1, types.int32) - result += data[0] - return result - - func_ir, _ = _rewrite(kernel) - - assert _call_targets(func_ir).count(cuda.local.array) == 1 - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "module", (common_coop, coop), ids=("root", "qualified") -) -def test_thread_data_item_extent_is_a_compile_time_constant( - module, items_per_thread -): - def kernel(items_per_thread): - payload = module.ThreadData(items_per_thread, types.int32) - alias = payload - return alias.items_per_thread - - func_ir, _ = _rewrite( - kernel, arg_types=(types.IntegerLiteral(items_per_thread),) - ) - assignments = [ - stmt - for block in func_ir.blocks.values() - for stmt in block.body - if isinstance(stmt, ir.Assign) - ] - - assert not any( - isinstance(stmt.value, ir.Expr) - and stmt.value.op == "getattr" - and stmt.value.attr == "items_per_thread" - for stmt in assignments - ) - definitions = {stmt.target.name: stmt.value for stmt in assignments} - returned = next( - stmt - for block in func_ir.blocks.values() - for stmt in block.body - if isinstance(stmt, ir.Return) - ) - result = definitions[returned.value.name] - assert result.op == "cast" - extent = definitions[result.value.name] - assert isinstance(extent, ir.Const) - assert extent.value == items_per_thread - - -def test_thread_data_rejects_rebound_dtype_in_a_branch(): - def kernel(flag): - if flag: - dtype = np.int32 - first = coop.ThreadData(items_per_thread=2, dtype=dtype) - first[0] = 16777217 - dtype = np.float32 - second = coop.ThreadData(items_per_thread=2, dtype=dtype) - second[0] = 1.5 - return first[0] + second[0] - return 0 - - with pytest.raises(CoopSinglePhaseRewriteError, match="dtype must resolve"): - _rewrite(kernel) - - -def test_thread_data_rejects_rebound_alignment_in_a_loop(): - def kernel(count): - result = 0 - for _ in range(count): - alignment = 64 - first = coop.ThreadData( - items_per_thread=2, dtype=types.int32, alignment=alignment - ) - alignment = 16 - second = coop.ThreadData( - items_per_thread=2, dtype=types.int32, alignment=alignment - ) - result += first[0] + second[0] - return result - - with pytest.raises( - CoopSinglePhaseRewriteError, match="alignment must be a compile-time" - ): - _rewrite(kernel) - - -def test_rebound_shared_array_shape_cannot_hide_dynamic_allocation(): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value, count): - for _ in range(count): - size = 0 - dynamic = cuda.shared.array(size, types.int32) - size = 32 - static = cuda.shared.array(size, types.int32) - dynamic[0] = value - static[0] = value - value = provider(dynamic[0] + static[0]) - return value - - with pytest.raises(CoopSinglePhaseRewriteError, match="would alias"): - _rewrite_registered_provider(kernel) - - -def test_native_local_array_item_extent_is_not_rewritten(): - def kernel(): - payload = cuda.local.array(3, types.int32) - return payload.items_per_thread - - func_ir, _ = _rewrite(kernel) - assignments = [ - stmt - for block in func_ir.blocks.values() - for stmt in block.body - if isinstance(stmt, ir.Assign) - ] - - assert any( - isinstance(stmt.value, ir.Expr) - and stmt.value.op == "getattr" - and stmt.value.attr == "items_per_thread" - for stmt in assignments - ) - - -def test_mixed_thread_data_item_extent_is_not_rewritten(): - def kernel(flag): - thread_data = coop.ThreadData(items_per_thread=3, dtype=types.int32) - native = cuda.local.array(3, types.int32) - payload = thread_data if flag else native - return payload.items_per_thread - - func_ir, _ = _rewrite(kernel) - assignments = [ - stmt - for block in func_ir.blocks.values() - for stmt in block.body - if isinstance(stmt, ir.Assign) - ] - - assert any( - isinstance(stmt.value, ir.Expr) - and stmt.value.op == "getattr" - and stmt.value.attr == "items_per_thread" - for stmt in assignments - ) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "alignment", - [None, 1, 2, 4, 8, 16, np.int64(32)], -) -@pytest.mark.parametrize( - "module", (common_coop, coop), ids=("root", "qualified") -) -def test_thread_data_rewrite_uses_alignment( - alignment, module, items_per_thread -): - def kernel(items_per_thread): - data = module.ThreadData( - items_per_thread, - types.int32, - alignment=alignment, - ) - return data[0] - - func_ir, _ = _rewrite( - kernel, arg_types=(types.IntegerLiteral(items_per_thread),) - ) - - assert cuda.local.array in _call_targets(func_ir) - definitions = { - stmt.target.name: stmt.value - for block in func_ir.blocks.values() - for stmt in block.body - if isinstance(stmt, ir.Assign) - } - calls = [ - value - for value in definitions.values() - if isinstance(value, ir.Expr) and value.op == "call" - ] - assert len(calls) == 1 - alignment_refs = [ - value for name, value in calls[0].kws if name == "alignment" - ] - if alignment is None: - assert alignment_refs == [] - else: - assert len(alignment_refs) == 1 - assert definitions[alignment_refs[0].name].value == max(8, alignment) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -@pytest.mark.parametrize( - "module", (common_coop, coop), ids=("root", "qualified") -) -@pytest.mark.parametrize( - ("first_alignment", "second_alignment"), - [(16, 32), (32, 16), (None, 32), (32, None)], -) -def test_merged_thread_data_preserves_minimum_alignment( - module, first_alignment, second_alignment, items_per_thread -): - def kernel(flag, items_per_thread): - if flag: - data = module.ThreadData( - items_per_thread, types.int32, alignment=first_alignment - ) - else: - data = module.ThreadData( - items_per_thread, types.int32, alignment=second_alignment - ) - return data[0] - - func_ir, _ = _rewrite( - kernel, - arg_types=( - types.boolean, - types.IntegerLiteral(items_per_thread), - ), - ) - - assert _call_targets(func_ir).count(cuda.local.array) == 2 - definitions = { - stmt.target.name: stmt.value - for block in func_ir.blocks.values() - for stmt in block.body - if isinstance(stmt, ir.Assign) - } - calls = [ - stmt.value - for block in func_ir.blocks.values() - for stmt in block.body - if isinstance(stmt, ir.Assign) - and isinstance(stmt.value, ir.Expr) - and stmt.value.op == "call" - ] - calls = [ - call - for call, target in zip(calls, _call_targets(func_ir), strict=True) - if target is cuda.local.array - ] - assert len(calls) == 2 - for call, required_alignment in zip( - calls, (first_alignment, second_alignment), strict=True - ): - if required_alignment is None: - continue - alignment_refs = [ - value for name, value in call.kws if name == "alignment" - ] - assert len(alignment_refs) == 1 - assert definitions[alignment_refs[0].name].value >= required_alignment - - -@pytest.mark.parametrize( - "module", (common_coop, coop), ids=("root", "qualified") -) -@pytest.mark.parametrize( - ("alignment", "message"), - [ - (True, "alignment must be an integer or None"), - (1.5, "alignment must be an integer or None"), - (0, "alignment must be a positive integer"), - (-1, "alignment must be a positive integer"), - (3, "alignment must be a power of 2"), - ], -) -def test_thread_data_rewrite_rejects_invalid_alignment( - module, alignment, message -): - def kernel(): - return module.ThreadData( - items_per_thread=2, dtype=types.int32, alignment=alignment - ) - - with pytest.raises(CoopSinglePhaseRewriteError, match=message): - _rewrite(kernel) - - -@pytest.mark.parametrize( - "module", (common_coop, coop), ids=("root", "qualified") -) -def test_thread_data_rewrite_rejects_dynamic_alignment(module): - def kernel(alignment): - return module.ThreadData( - items_per_thread=2, dtype=types.int32, alignment=alignment - ) - - with pytest.raises( - CoopSinglePhaseRewriteError, match="alignment must be a compile-time" - ): - _rewrite(kernel) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_whole_function_planner_reuses_the_foundation_rewrite(items_per_thread): - def kernel(items_per_thread): - return coop.ThreadData(items_per_thread, types.int32) - - func_ir = run_frontend(kernel) - typingctx = _TypingContext() - state = SimpleNamespace( - func_ir=func_ir, - typingctx=typingctx, - typemap={}, - calltypes={}, - metadata={"targetoptions": {}}, - args=(types.IntegerLiteral(items_per_thread),), - ) - - assert CoopWholeFunctionPlanner(state).run() - assert cuda.local.array in _call_targets(func_ir) - assert typingctx.refresh_count == 1 - - -def _unused_temp_storage(): - _storage = coop.TempStorage(32) - return 0 - - -@pytest.mark.parametrize( - "kernel", - [ - pytest.param(_unused_temp_storage, id="no-consumer"), - pytest.param(lambda: coop.TempStorage(32), id="return"), - pytest.param(lambda: coop.TempStorage(32)[0], id="index"), - pytest.param(lambda: len(coop.TempStorage(32)), id="other-call"), - ], -) -def test_temp_storage_is_an_opaque_primitive_descriptor(kernel): - with pytest.raises( - CoopSinglePhaseRewriteError, - match="opaque compile-time descriptors", - ): - _rewrite(kernel) - - -def test_temp_storage_rewrite_rejects_starred_constructor_arguments(): - # Numba's frontend rejects **kwargs outright, so only *args can reach the - # rewrite. Starred options must not silently become TempStorage() - # and discard the caller's synchronization policy. - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - options = (None, None, False) - - def kernel(value): - storage = coop.TempStorage(*options) - return provider(value, temp_storage=storage) - - with pytest.raises(CoopSinglePhaseRewriteError, match="does not accept"): - _rewrite_registered_provider(kernel) - - -def test_temp_storage_rewrite_rejects_string_enum_sharing(): - def kernel(): - return coop.TempStorage(sharing=_StringSharing.SHARED) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="TempStorage sharing must be a string", - ): - _rewrite(kernel) - - -class _FakeInvocable: - """Supply provider storage metadata without compiling device code.""" - - files = ("storage-rewrite-test.ltoir",) - specialization = None - storage_abi = "leading_pointer" - execution_scope = "block" - synchronization_scope = "block" - - def __init__(self, *, size_in_bytes=64, alignment=16): - self.temp_storage_bytes = size_in_bytes - self.temp_storage_alignment = alignment - - def __call__(self, *args): - del args - - -def _register_leading_pointer_provider( - invocable, - *, - execution_scope=SynchronizationScope.BLOCK, - synchronization_scope=SynchronizationScope.BLOCK, -): - """Register a test provider whose generated call takes a scratch pointer. - - Its factory accepts no options and records each construction. Tests check - that invalid storage requests fail before the provider is constructed. - """ - - operation = f"_test_storage_rewrite_family_{id(invocable)}" - calls = [] - - def provider(*runtime_args, **factory_kwargs): - assert not runtime_args - assert not factory_kwargs - calls.append((runtime_args, factory_kwargs)) - return invocable - - provider.calls = calls - - _operations.register_factory( - provider, - operation=operation, - namespace="storage_test", - storage_abi=_operations.StorageABI.LEADING_POINTER, - execution_scope=execution_scope, - synchronization_scope=synchronization_scope, - ) - _operations.register_rewrite_operation( - operation, - _operations.RewriteOperationSpecification( - factory_namespaces=frozenset({"storage_test"}), - dtype_factory_kwargs=frozenset(), - runtime_arg_counts=frozenset({1}), - runtime_factory_kwargs=(), - runtime_factory_kw_prerequisites=(), - allowed_factory_kwargs=frozenset(), - required_factory_kwargs=frozenset(), - accepts_temp_storage=True, - scalar_binding_kwargs=frozenset(), - runtime_offset_kwarg=None, - infer_payload=lambda *_args: None, - ), - ) - return provider - - -def _frontend(function, *, ssa=False): - """Parse the kernel and optionally give each assignment a unique variable. - - Single-assignment form adds merge nodes at control-flow joins. Both forms - must retain descriptor origins despite their different variable names. - """ - - func_ir = run_frontend(function) - if ssa: - from numba_cuda_mlir.numba_cuda.core.ir_utils import build_definitions - from numba_cuda_mlir.numba_cuda.core.ssa import reconstruct_ssa - - func_ir = reconstruct_ssa(func_ir) - func_ir._definitions = build_definitions(func_ir.blocks) - return func_ir - - -def _rewrite_registered_provider(function, *, ssa=False, lifo=False): - """Rewrite test providers and optionally visit later blocks first. - - Revisit each changed block until it no longer matches. Reverse the visit - order to check that storage is emitted before all possible consumers, - even when a later block is rewritten first. - """ - - func_ir = _frontend(function, ssa=ssa) - typingctx = _TypingContext() - state = SimpleNamespace( - func_ir=func_ir, - args=(), - typingctx=typingctx, - typemap={}, - calltypes={}, - metadata={"targetoptions": {}}, - ) - rewrite = CoopSinglePhaseRewrite(state) - assert rewrite.prepare_calls_and_storage(func_ir) - rewrite.begin_rewrite() - items = list(func_ir.blocks.items()) - if not lifo: - items.reverse() - while items: - label, block = items.pop() - if rewrite.match(func_ir, block, state.typemap, state.calltypes): - new_block = rewrite.apply() - func_ir.blocks[label] = new_block - items.append((label, new_block)) - rewrite.finish_rewrite() - return func_ir, rewrite, state - - -def _rewrite_preflight(function): - """Prepare IR for tests that must fail before provider construction.""" - - func_ir = _frontend(function) - state = SimpleNamespace( - func_ir=func_ir, - args=(), - typingctx=_TypingContext(), - typemap={}, - calltypes={}, - metadata={"targetoptions": {}}, - ) - return func_ir, state, CoopSinglePhaseRewrite(state) - - -@pytest.mark.parametrize( - ("execution_scope", "synchronization_scope", "accepted"), - [ - pytest.param( - SynchronizationScope.BLOCK, - SynchronizationScope.BLOCK, - True, - id="block-block", - ), - pytest.param( - SynchronizationScope.BLOCK, - SynchronizationScope.NONE, - False, - id="block-none", - ), - pytest.param( - SynchronizationScope.WARP, - SynchronizationScope.WARP, - False, - id="warp-warp", - ), - pytest.param( - SynchronizationScope.NONE, - SynchronizationScope.NONE, - False, - id="none-none", - ), - pytest.param( - SynchronizationScope.GROUP, - SynchronizationScope.GROUP, - False, - id="group-group", - ), - ], -) -def test_storage_provider_without_plan_requires_block_scope( - execution_scope, - synchronization_scope, - accepted, -): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider( - invocable, - execution_scope=execution_scope, - synchronization_scope=synchronization_scope, - ) - - def kernel(value): - return provider(value) - - func_ir, _state, rewrite = _rewrite_preflight(kernel) - if accepted: - prepared = [] - rewrite._prepare_ltoir_bundle_for_matches = lambda matches: ( - prepared.append(tuple(matches)) - ) - - assert rewrite.prepare_calls_and_storage(func_ir) - assert len(prepared) == 1 - assert len(prepared[0]) == 1 - assert provider.calls == [((), {})] - return - - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: pytest.fail( - "invalid no-plan provider reached bundle preparation" - ) - rewrite._materialize_invocable = lambda _match: pytest.fail( - "invalid no-plan provider reached provider materialization" - ) - with pytest.raises( - CoopSinglePhaseRewriteError, - match="require block execution and block synchronization scopes", - ): - rewrite.prepare_calls_and_storage(func_ir) - assert provider.calls == [] - - -@pytest.mark.parametrize( - ("case", "with_descriptor", "message"), - [ - pytest.param( - "group", - False, - "execution scope 'group' has no storage emitter", - id="group", - ), - pytest.param( - "address-space", - False, - "shared-address-space TempStorage", - id="address-space", - ), - pytest.param( - "implementation-with-descriptor", - True, - "TempStorage ownership disagrees with its runtime arguments", - id="implementation-with-descriptor", - ), - pytest.param( - "caller-without-descriptor", - False, - "TempStorage ownership disagrees with its runtime arguments", - id="caller-without-descriptor", - ), - pytest.param( - "caller-none", - True, - "caller-owned TempStorage is supported only for " - "single-instance block", - id="caller-none", - ), - pytest.param( - "caller-group", - True, - "execution scope 'group' has no storage emitter", - id="caller-group", - ), - ], -) -def test_planned_storage_guardrails_fail_before_materialization( - case, - with_descriptor, - message, -): - from dataclasses import replace - - from cuda.coop._core import ( - GroupLoadStoreAlgorithm, - GroupTopologyRequirements, - LaunchFacts, - ParticipationRequirements, - StorageOwnership, - SynchronizationRequirements, - make_group_primitive_call, - resolve_thread_group, - this_block, - this_thread, - ) - from cuda.coop._core.group._execution_requirements import ( - _build_execution_requirements, - ) - from tests._group_planning import _load_store, _plan - - plan = _plan( - this_block(), - _load_store(algorithm=GroupLoadStoreAlgorithm.TRANSPOSE), - ) - execution_scope = SynchronizationScope.BLOCK - synchronization_scope = SynchronizationScope.BLOCK - if case in {"group", "caller-group"}: - execution_scope = SynchronizationScope.GROUP - synchronization_scope = SynchronizationScope.GROUP - launch = LaunchFacts(exact_block_dim=(64, 1, 1)) - resolved_group = resolve_thread_group( - this_block().group_by(2), launch - ).require_supported() - requirements = _build_execution_requirements( - resolved_group, - launch, - storage_ownership=StorageOwnership.IMPLEMENTATION, - cpp_type=plan.temp_storage.cpp_type, - ) - plan = replace( - plan, - call=make_group_primitive_call(resolved_group, plan.call.operation), - resolved_group=resolved_group, - topology=requirements.topology, - participation=requirements.participation, - synchronization=requirements.synchronization, - temp_storage=requirements.temp_storage, - ) - elif case == "address-space": - plan = replace( - plan, - temp_storage=replace(plan.temp_storage, address_space="local"), - ) - elif case == "caller-none": - execution_scope = SynchronizationScope.NONE - synchronization_scope = SynchronizationScope.NONE - launch = LaunchFacts(exact_block_dim=(64, 1, 1)) - resolved_thread = resolve_thread_group( - this_thread(), launch - ).require_supported() - operation = plan.call.operation - plan = replace( - plan, - call=make_group_primitive_call(resolved_thread, operation), - resolved_group=resolved_thread, - topology=GroupTopologyRequirements( - group_kind="thread", - logical_width=1, - instances=64, - instance_index="linear_thread_rank", - thread_rank="0", - execution_scope=SynchronizationScope.NONE, - ), - participation=ParticipationRequirements( - group_kind="thread", - exact_group_size=1, - exact_block_dim=(64, 1, 1), - complete_membership=True, - contiguous=True, - aligned=True, - converged_entry=True, - complete_parent_partition=True, - ), - synchronization=SynchronizationRequirements( - converged_entry=True, - storage_reuse_barrier=SynchronizationScope.NONE, - ), - temp_storage=replace( - plan.temp_storage, - instances=64, - instance_index="linear_thread_rank", - ), - ) - if case.startswith("caller-"): - plan = replace( - plan, - temp_storage=replace( - plan.temp_storage, - ownership=StorageOwnership.CALLER, - exact_layout_required=True, - sharing="shared", - ), - ) - - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider( - invocable, - execution_scope=execution_scope, - synchronization_scope=synchronization_scope, - ) - if with_descriptor: - - def kernel(value): - storage = coop.TempStorage() - return provider( - value, - temp_storage=storage, - __cuda_coop_group_lowering_plan__=plan, - ) - - else: - - def kernel(value): - return provider( - value, - __cuda_coop_group_lowering_plan__=plan, - ) - - func_ir, _state, rewrite = _rewrite_preflight(kernel) - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: pytest.fail( - "invalid planned provider reached bundle preparation" - ) - rewrite._materialize_invocable = lambda _match: pytest.fail( - "invalid planned provider reached provider materialization" - ) - - with pytest.raises(CoopSinglePhaseRewriteError, match=message): - rewrite.prepare_calls_and_storage(func_ir) - assert provider.calls == [] - - -@pytest.mark.parametrize( - "descriptor_auto_sync", [False, True], ids=["drift", "agree"] -) -def test_planned_caller_storage_contract_must_match_the_descriptor( - descriptor_auto_sync, -): - from cuda.coop._core import ( - GroupLoadStoreAlgorithm, - StorageOwnership, - this_block, - ) - from tests._group_planning import _load_store, _plan - - plan = _plan( - this_block(), - _load_store( - algorithm=GroupLoadStoreAlgorithm.TRANSPOSE, - storage_ownership=StorageOwnership.CALLER, - storage_sharing="shared", - storage_auto_sync=True, - ), - ) - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value): - storage = coop.TempStorage(auto_sync=descriptor_auto_sync) - return provider( - value, - temp_storage=storage, - __cuda_coop_group_lowering_plan__=plan, - ) - - func_ir, _state, rewrite = _rewrite_preflight(kernel) - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: None - rewrite._materialize_invocable = lambda _match: (invocable, False) - - if descriptor_auto_sync: - assert rewrite.prepare_calls_and_storage(func_ir) - return - # The planner parsed auto_sync=True into the plan while the rewrite sees - # auto_sync=False: neither parser may silently win. - with pytest.raises( - CoopSinglePhaseRewriteError, - match="disagrees between the group lowering plan", - ): - rewrite.prepare_calls_and_storage(func_ir) - - -def test_apply_refuses_a_plan_whose_auto_sync_disagrees_with_implicit_storage(): - from dataclasses import replace - - from cuda.coop._core import GroupLoadStoreAlgorithm, this_block - from tests._group_planning import _load_store, _plan - - plan = _plan( - this_block(), - _load_store(algorithm=GroupLoadStoreAlgorithm.TRANSPOSE), - ) - plan = replace( - plan, - temp_storage=replace(plan.temp_storage, auto_sync=False), - synchronization=replace( - plan.synchronization, - storage_reuse_barrier=SynchronizationScope.NONE, - ), - ) - invocable = _FakeInvocable() - invocable.synchronization_scope = "none" - provider = _register_leading_pointer_provider( - invocable, - synchronization_scope=SynchronizationScope.NONE, - ) - - def kernel(value): - return provider(value, __cuda_coop_group_lowering_plan__=plan) - - # Implementation-owned storage always carries the trailing barrier; a plan - # that claims otherwise must be rejected rather than drop the barrier. - with pytest.raises( - CoopSinglePhaseRewriteError, - match="disagrees between the group lowering plan and the descriptor", - ): - _rewrite_registered_provider(kernel) - - -def _resolved_calls(func_ir): - """Locate each emitted callee's block and assignment for layout checks.""" - - resolver = object.__new__(CoopSinglePhaseRewrite) - resolver._func_ir = func_ir - calls = [] - for label, block in func_ir.blocks.items(): - resolver._block_defs = { - inst.target.name: inst.value - for inst in block.body - if isinstance(inst, ir.Assign) - } - for inst in block.body: - value = getattr(inst, "value", None) - if isinstance(value, ir.Expr) and value.op == "call": - calls.append( - (label, inst, resolver._resolve_python_value(value.func)) - ) - return calls - - -def _rewrite_with_fake_invocable( - function, invocable, *, ssa=False, arg_types=() -): - """Exercise a real provider's rewrite using supplied compilation metadata. - - Skip device-code preparation while retaining argument, descriptor, and - storage analysis. Return the rewrite object so tests can inspect its plan. - """ - - func_ir = _frontend(function, ssa=ssa) - typingctx = _TypingContext() - state = SimpleNamespace( - func_ir=func_ir, - args=arg_types, - typingctx=typingctx, - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: None - rewrite._materialize_invocable = lambda _match: (invocable, False) - assert rewrite.prepare_calls_and_storage(func_ir) - rewrite.begin_rewrite() - for label in sorted(func_ir.blocks): - block = func_ir.blocks[label] - while rewrite.match(func_ir, block, state.typemap, state.calltypes): - block = rewrite.apply() - func_ir.blocks[label] = block - rewrite.finish_rewrite() - return func_ir, rewrite - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_temp_storage_aliases_from_one_constructor_are_accepted( - items_per_thread, -): - from cuda.coop.numba_mlir._lowering._load_store import load as provider_load - - def kernel(source, choose_first, items_per_thread): - storage = coop.TempStorage() - if choose_first: - selected = storage - else: - selected = storage - payload = coop.ThreadData(items_per_thread, types.int32) - return provider_load( - source, - payload, - dtype=types.int32, - threads_per_block=32, - items_per_thread=items_per_thread, - temp_storage=selected, - ) - - func_ir, _ = _rewrite_with_fake_invocable( - kernel, - _FakeInvocable(), - ssa=True, - arg_types=( - types.int32, - types.boolean, - types.IntegerLiteral(items_per_thread), - ), - ) - - assert any( - isinstance(inst, ir.Assign) - and isinstance(inst.value, ir.Expr) - and inst.value.op == "phi" - for block in func_ir.blocks.values() - for inst in block.body - ) - targets = _call_targets(func_ir) - assert cuda.shared.array not in targets - assert cuda.syncthreads not in targets - - -def test_temp_storage_phi_canonicalizes_equivalent_constructor_contracts(): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value, choose_first): - if choose_first: - selected = coop.TempStorage(auto_sync=True) - else: - selected = coop.TempStorage(auto_sync=True, sharing=" SHARED ") - return provider(value, temp_storage=selected) - - func_ir, rewrite, _ = _rewrite_registered_provider(kernel, ssa=True) - calls = _resolved_calls(func_ir) - - assert sum(target is cuda.shared.array for _, _, target in calls) == 1 - assert sum(target is invocable for _, _, target in calls) == 1 - assert rewrite._temp_storage_global_plan.total_size == 64 - assert len(set(rewrite._temp_storage_ctor_roots.values())) == 1 - - -def test_leading_pointer_storage_can_disable_external_reuse_sync(): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value): - storage = coop.TempStorage(auto_sync=False) - return provider(value, temp_storage=storage) - - func_ir, _, _ = _rewrite_registered_provider(kernel) - calls = _resolved_calls(func_ir) - targets = [target for _, _, target in calls] - invocable_calls = [ - inst.value for _, inst, target in calls if target is invocable - ] - - assert targets.count(cuda.shared.array) == 1 - assert len(invocable_calls) == 1 - assert len(invocable_calls[0].args) == 2 - assert cuda.syncthreads not in targets - assert cuda.syncwarp not in targets - - -@pytest.mark.parametrize( - "explicit_descriptor", [True, False], ids=["caller", "implicit"] -) -def test_temp_storage_backing_rejects_user_dynamic_shared_arrays( - explicit_descriptor, -): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - if explicit_descriptor: - - def kernel(value): - mine = cuda.shared.array(0, types.int32) - storage = coop.TempStorage() - mine[0] = value - return provider(value, temp_storage=storage) - - else: - - def kernel(value): - mine = cuda.shared.array(shape=0, dtype=types.int32) - mine[0] = value - return provider(value) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match=r"dynamic/runtime-sized cuda\.shared\.array.*would alias", - ): - _rewrite_registered_provider(kernel) - - -def test_temp_storage_backing_coexists_with_static_user_shared_arrays(): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value): - mine = cuda.shared.array(4, types.int32) - storage = coop.TempStorage() - mine[0] = value - return provider(value, temp_storage=storage) - - func_ir, _, _ = _rewrite_registered_provider(kernel) - targets = [target for _, _, target in _resolved_calls(func_ir)] - - assert targets.count(cuda.shared.array) == 2 - assert targets.count(invocable) == 1 - - -@pytest.mark.parametrize( - ("left", "right"), - [ - ((64, 16, True, "shared"), (96, 16, True, "shared")), - ((64, 16, True, "shared"), (64, 32, True, "shared")), - ((64, 16, True, "shared"), (64, 16, False, "shared")), - ((64, 16, None, "shared"), (64, 16, None, "exclusive")), - ], - ids=["capacity", "alignment", "auto-sync", "sharing"], -) -def test_temp_storage_phi_rejects_incompatible_contracts_before_compile( - left, right -): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - left_size, left_alignment, left_auto_sync, left_sharing = left - right_size, right_alignment, right_auto_sync, right_sharing = right - - def kernel(value, choose_first): - if choose_first: - selected = coop.TempStorage( - left_size, - alignment=left_alignment, - auto_sync=left_auto_sync, - sharing=left_sharing, - ) - else: - selected = coop.TempStorage( - right_size, - alignment=right_alignment, - auto_sync=right_auto_sync, - sharing=right_sharing, - ) - return provider(value, temp_storage=selected) - - func_ir = _frontend(kernel, ssa=True) - state = SimpleNamespace( - func_ir=func_ir, - args=(), - typingctx=_TypingContext(), - typemap={}, - calltypes={}, - metadata={"targetoptions": {}}, - ) - rewrite = CoopSinglePhaseRewrite(state) - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: pytest.fail( - "incompatible TempStorage reached provider preparation" - ) - rewrite._materialize_invocable = lambda _match: pytest.fail( - "incompatible TempStorage reached provider compilation" - ) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="TempStorage aliases have inconsistent contracts", - ): - rewrite.prepare_calls_and_storage(func_ir) - - -def test_group_planning_rejects_mixed_descriptor_phi(): - from cuda.coop.numba_mlir._compiler._group_planner import _GroupCallPlanner - from cuda.coop.numba_mlir._compiler._group_planner_support import ( - GroupRewriteError, - ) - - def consume(value): - del value - - def kernel(choose_storage): - if choose_storage: - selected = coop.TempStorage() - else: - selected = None - consume(selected) - - func_ir = _frontend(kernel, ssa=True) - phi_assign = next( - inst - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) - and isinstance(inst.value, ir.Expr) - and inst.value.op == "phi" - ) - planner = _GroupCallPlanner( - SimpleNamespace(func_ir=func_ir, args=(types.boolean,)), - {"block": (32, 1, 1), "grid": (1, 1, 1)}, - ) - - with pytest.raises(GroupRewriteError, match="on every path"): - planner.context.temp_storage(phi_assign.target) - - -@pytest.mark.parametrize("aliases", ["direct", "chain", "conditional"]) -def test_temp_storage_passed_to_a_non_inlined_device_function_is_rejected( - aliases, -): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - @cuda.jit(device=True, inline="never") - def helper(value, storage): - return provider(value, temp_storage=storage) - - def kernel(value, flag): - storage = coop.TempStorage() - alias = storage - chained = alias - if aliases == "direct": - selected = storage - elif aliases == "chain": - selected = chained - else: - selected = alias if flag else chained - first = helper(value, selected) - return helper(first, selected) - - func_ir, _state, rewrite = _rewrite_preflight(kernel) - with pytest.raises( - CoopSinglePhaseRewriteError, - match="device function that was not inlined", - ): - rewrite.prepare_calls_and_storage(func_ir) - - -@pytest.mark.parametrize( - "shape", - ["ternary", "initializer"], -) -def test_descriptor_joined_with_none_is_rejected_before_type_inference(shape): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - if shape == "ternary": - - def kernel(value, flag): - storage = coop.TempStorage() if flag else None - return provider(value, temp_storage=storage) - - else: - - def kernel(value, flag): - storage = None - if flag: - storage = coop.TempStorage() - return provider(value, temp_storage=storage) - - with pytest.raises(CoopSinglePhaseRewriteError, match="on every path"): - _rewrite_registered_provider(kernel) - - -def test_group_planning_rejects_descriptor_joined_with_none_without_ssa(): - from cuda.coop.numba_mlir._compiler._group_planner import _GroupCallPlanner - from cuda.coop.numba_mlir._compiler._group_planner_support import ( - GroupRewriteError, - ) - - def consume(value): - del value - - def kernel(choose_storage): - selected = None - if choose_storage: - selected = coop.TempStorage() - consume(selected) - - func_ir = _frontend(kernel) - consume_call = next( - inst - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) - and isinstance(inst.value, ir.Expr) - and inst.value.op == "call" - and getattr(func_ir.get_definition(inst.value.func), "value", None) - is consume - ) - planner = _GroupCallPlanner( - SimpleNamespace(func_ir=func_ir, args=(types.boolean,)), - {"block": (32, 1, 1), "grid": (1, 1, 1)}, - ) - - with pytest.raises(GroupRewriteError, match="on every path"): - planner.context.temp_storage(consume_call.value.args[0]) - - -@pytest.mark.parametrize( - "auto_sync", [False, None, True], ids=["manual", "compat-none", "auto"] -) -def test_rebound_constructor_name_collapses_only_with_auto_sync(auto_sync): - # Production IR is not in SSA form: a name rebound inside a branch keeps - # one name for two constructor sites. Collapsing them into one region is - # safe only while every use carries the trailing reuse barrier. - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value, flag): - storage = coop.TempStorage(auto_sync=auto_sync) - first = provider(value, temp_storage=storage) - if flag: - storage = coop.TempStorage(auto_sync=auto_sync) - return provider(first, temp_storage=storage) - - if auto_sync is not True: - with pytest.raises( - CoopSinglePhaseRewriteError, match="exactly one site" - ): - _rewrite_registered_provider(kernel) - return - - func_ir, rewrite, _ = _rewrite_registered_provider(kernel) - targets = [target for _, _, target in _resolved_calls(func_ir)] - - assert targets.count(cuda.shared.array) == 1 - assert targets.count(invocable) == 2 - assert targets.count(cuda.syncthreads) == 2 - assert rewrite._temp_storage_global_plan.total_size == 64 - - -def test_temp_storage_phi_rejects_merging_manual_sync_constructors(): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value, choose_first): - if choose_first: - selected = coop.TempStorage(auto_sync=False) - else: - selected = coop.TempStorage(auto_sync=False) - return provider(value, temp_storage=selected) - - with pytest.raises(CoopSinglePhaseRewriteError, match="exactly one site"): - _rewrite_registered_provider(kernel, ssa=True) - - -@pytest.mark.parametrize( - "auto_sync", [False, None, True], ids=["manual", "compat-none", "auto"] -) -def test_group_planning_collapses_constructor_sites_only_with_auto_sync( - auto_sync, -): - from cuda.coop.numba_mlir._compiler._group_planner import _GroupCallPlanner - from cuda.coop.numba_mlir._compiler._group_planner_support import ( - GroupRewriteError, - ) - - def consume(value): - del value - - def kernel(choose_storage): - if choose_storage: - selected = coop.TempStorage(auto_sync=auto_sync) - else: - selected = coop.TempStorage(auto_sync=auto_sync) - consume(selected) - - func_ir = _frontend(kernel) - consume_call = next( - inst - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) - and isinstance(inst.value, ir.Expr) - and inst.value.op == "call" - and getattr(func_ir.get_definition(inst.value.func), "value", None) - is consume - ) - planner = _GroupCallPlanner( - SimpleNamespace(func_ir=func_ir, args=(types.boolean,)), - {"block": (32, 1, 1), "grid": (1, 1, 1)}, - ) - - if auto_sync is not True: - with pytest.raises(GroupRewriteError, match="exactly one site"): - planner.context.temp_storage(consume_call.value.args[0]) - else: - assert planner.context.temp_storage(consume_call.value.args[0]) == ( - None, - None, - True, - "shared", - ) - - -def test_temp_storage_backing_dominates_nonentry_call_under_lifo_rewrite(): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value, take_provider): - if take_provider: - result = provider(value) - else: - result = value - return result - - func_ir, rewrite, _ = _rewrite_registered_provider(kernel, lifo=True) - calls = _resolved_calls(func_ir) - entry_label = min(func_ir.blocks) - backing_calls = [ - (label, inst) - for label, inst, target in calls - if target is cuda.shared.array - ] - invocable_calls = [ - (label, inst) for label, inst, target in calls if target is invocable - ] - - assert len(backing_calls) == 1 - assert len(invocable_calls) == 1 - backing_label, backing_assign = backing_calls[0] - call_label, call_assign = invocable_calls[0] - assert backing_label == entry_label - assert call_label != entry_label - storage_arg = call_assign.value.args[0] - storage_definition = next( - inst.value - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) and inst.target.name == storage_arg.name - ) - assert isinstance(storage_definition, ir.Expr) - assert storage_definition.op == "getitem" - assert storage_definition.value.name == backing_assign.target.name - assert rewrite._temp_storage_backing_emitted - - -def test_equivalent_temp_storage_phi_escape_is_rejected_before_compile(): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value, choose_first): - if choose_first: - selected = coop.TempStorage(auto_sync=True) - else: - selected = coop.TempStorage(auto_sync=True) - provider(value, temp_storage=selected) - return selected - - func_ir = _frontend(kernel, ssa=True) - state = SimpleNamespace( - func_ir=func_ir, - args=(), - typingctx=_TypingContext(), - typemap={}, - calltypes={}, - metadata={"targetoptions": {}}, - ) - rewrite = CoopSinglePhaseRewrite(state) - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: pytest.fail( - "escaping TempStorage reached provider preparation" - ) - rewrite._materialize_invocable = lambda _match: pytest.fail( - "escaping TempStorage reached provider compilation" - ) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="would escape to runtime", - ): - rewrite.prepare_calls_and_storage(func_ir) - - -def test_leading_pointer_provider_stages_one_dynamic_backing(monkeypatch): - from cuda.coop.numba_mlir._compiler import _rewrite_storage - - size_in_bytes = 64 * 1024 - invocable = _FakeInvocable(size_in_bytes=size_in_bytes) - provider = _register_leading_pointer_provider(invocable) - monkeypatch.setattr( - _rewrite_storage, - "_query_device_shared_memory_limits", - lambda: { - "max_default_shared_memory_per_block": 48 * 1024, - "max_optin_shared_memory_per_block": 96 * 1024, - }, - ) - - def kernel(value): - return provider(value) - - func_ir, rewrite, state = _rewrite_registered_provider(kernel, lifo=True) - backing_calls = [ - inst.value - for _, inst, target in _resolved_calls(func_ir) - if target is cuda.shared.array - ] - - assert len(backing_calls) == 1 - size_var = backing_calls[0].args[0] - size_definition = next( - inst.value - for block in func_ir.blocks.values() - for inst in block.body - if isinstance(inst, ir.Assign) and inst.target.name == size_var.name - ) - assert isinstance(size_definition, ir.Const) - assert size_definition.value == 0 - assert ( - rewrite._temp_storage_global_plan.dynamic_shared_bytes == size_in_bytes - ) - assert state.metadata["required_dynamic_shared_memory"] == size_in_bytes - - -def test_mixed_temp_storage_primitive_and_escape_fails_before_compile(): - from cuda.coop.numba_mlir._lowering._load_store import load as provider_load - - def kernel(source): - storage = coop.TempStorage() - payload = coop.ThreadData(items_per_thread=2, dtype=types.int32) - provider_load( - source, - payload, - dtype=types.int32, - threads_per_block=32, - items_per_thread=2, - temp_storage=storage, - ) - return storage - - func_ir = run_frontend(kernel) - state = SimpleNamespace( - func_ir=func_ir, - args=(), - typingctx=_TypingContext(), - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: pytest.fail( - "escaping TempStorage reached provider preparation" - ) - rewrite._materialize_invocable = lambda _match: pytest.fail( - "escaping TempStorage reached provider compilation" - ) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="would escape to runtime", - ): - rewrite.prepare_calls_and_storage(func_ir) - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_direct_provider_uses_no_implicit_storage_or_automatic_sync( - items_per_thread, -): - from cuda.coop.numba_mlir._lowering._load_store import load as provider_load - - def kernel(source, items_per_thread): - payload = coop.ThreadData(items_per_thread, types.int32) - return provider_load( - source, - payload, - dtype=types.int32, - threads_per_block=32, - items_per_thread=items_per_thread, - ) - - invocable = _FakeInvocable() - func_ir, rewrite = _rewrite_with_fake_invocable( - kernel, - invocable, - arg_types=( - types.int32, - types.IntegerLiteral(items_per_thread), - ), - ) - targets = _call_targets(func_ir) - - assert cuda.shared.array not in targets - assert cuda.syncthreads not in targets - assert rewrite._temp_storage_global_plan is None - resolver = object.__new__(CoopSinglePhaseRewrite) - resolver._func_ir = func_ir - invocable_calls = [] - for block in func_ir.blocks.values(): - resolver._block_defs = { - inst.target.name: inst.value - for inst in block.body - if isinstance(inst, ir.Assign) - } - invocable_calls.extend( - inst.value - for inst in block.body - if isinstance(inst, ir.Assign) - and isinstance(inst.value, ir.Expr) - and inst.value.op == "call" - and resolver._resolve_python_value(inst.value.func) is invocable - ) - assert len(invocable_calls) == 1 - assert len(invocable_calls[0].args) == 2 - - -@pytest.mark.parametrize("items_per_thread", [1, 4]) -def test_direct_provider_ignores_implicit_and_explicit_storage( - items_per_thread, -): - from cuda.coop.numba_mlir._lowering._load_store import load as provider_load - - def kernel(source_a, source_b, items_per_thread): - storage = coop.TempStorage() - payload_a = coop.ThreadData(items_per_thread, types.int32) - payload_b = coop.ThreadData(items_per_thread, types.int32) - provider_load( - source_a, - payload_a, - dtype=types.int32, - threads_per_block=32, - items_per_thread=items_per_thread, - temp_storage=storage, - ) - return provider_load( - source_b, - payload_b, - dtype=types.int32, - threads_per_block=32, - items_per_thread=items_per_thread, - ) - - invocable = _FakeInvocable() - func_ir, rewrite = _rewrite_with_fake_invocable( - kernel, - invocable, - arg_types=( - types.int32, - types.int32, - types.IntegerLiteral(items_per_thread), - ), - ) - targets = _call_targets(func_ir) - - assert cuda.shared.array not in targets - assert cuda.syncthreads not in targets - assert rewrite._temp_storage_global_plan is None - - -def test_getitem_temp_storage_syntax_is_not_an_accepted_descriptor_use(): - from cuda.coop.numba_mlir._lowering._load_store import load as provider_load - - def kernel(source): - storage = coop.TempStorage() - payload = coop.ThreadData(items_per_thread=2, dtype=types.int32) - return provider_load[storage]( - source, - payload, - dtype=types.int32, - threads_per_block=32, - items_per_thread=2, - ) - - func_ir = run_frontend(kernel) - state = SimpleNamespace( - func_ir=func_ir, - args=(), - typingctx=_TypingContext(), - typemap={}, - calltypes={}, - metadata={}, - ) - rewrite = CoopSinglePhaseRewrite(state) - rewrite._prepare_ltoir_bundle_for_matches = lambda _matches: pytest.fail( - "getitem TempStorage reached provider preparation" - ) - rewrite._materialize_invocable = lambda _match: pytest.fail( - "getitem TempStorage reached provider compilation" - ) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="may only be passed as temp_storage=", - ): - rewrite.prepare_calls_and_storage(func_ir) - - -def _planner_for_storage_policy(specification, use_specifications): - """Build one descriptor's use list from explicit size/alignment pairs. - - Stable call identities let tests inspect each assigned scratch slice - without first constructing and parsing a kernel. - """ - - rewrite = object.__new__(CoopSinglePhaseRewrite) - calls = [object() for _ in use_specifications] - rewrite._temp_storage_plans = {} - rewrite._temp_storage_ctor_specifications = {"storage": specification} - rewrite._func_temp_storage_requirements = { - "storage": _TempStorageRequirementSummary( - uses=[ - _TempStorageUseRequirement( - call_assign=call, - order=index, - size_in_bytes=size, - alignment=alignment, - ) - for index, (call, (size, alignment)) in enumerate( - zip(calls, use_specifications) - ) - ] - ) - } - return rewrite, calls - - -def test_shared_storage_reuses_one_aligned_slice_and_defaults_to_manual_sync(): - rewrite, calls = _planner_for_storage_policy( - _TempStorageCtorSpecification(None, None, None, "shared"), - [(24, 8), (64, 16)], - ) - - plan = rewrite._finalize_temp_storage_plan_for_var("storage") - - assert (plan.size_in_bytes, plan.alignment, plan.auto_sync) == ( - 64, - 16, - False, - ) - assert [plan.slices_by_call_id[id(call)].offset for call in calls] == [0, 0] - - -def test_shared_storage_can_delegate_synchronization_to_the_caller(): - rewrite, _ = _planner_for_storage_policy( - _TempStorageCtorSpecification(None, None, False, "shared"), - [(24, 8), (64, 16)], - ) - - plan = rewrite._finalize_temp_storage_plan_for_var("storage") - - assert not plan.auto_sync - - -@pytest.mark.parametrize( - "auto_sync", [None, False, True], ids=["compat-none", "manual", "auto"] -) -def test_exclusive_storage_assigns_distinct_aligned_slices_with_auto_sync( - auto_sync, -): - rewrite, calls = _planner_for_storage_policy( - _TempStorageCtorSpecification(None, None, auto_sync, "exclusive"), - [(24, 8), (16, 16)], - ) - - plan = rewrite._finalize_temp_storage_plan_for_var("storage") - - assert (plan.size_in_bytes, plan.alignment, plan.auto_sync) == ( - 48, - 16, - auto_sync is True, - ) - assert [plan.slices_by_call_id[id(call)].offset for call in calls] == [ - 0, - 32, - ] - - -def test_exclusive_storage_can_delegate_synchronization_to_the_caller(): - rewrite, _ = _planner_for_storage_policy( - _TempStorageCtorSpecification(None, None, False, "exclusive"), - [(24, 8)], - ) - - plan = rewrite._finalize_temp_storage_plan_for_var("storage") - - assert not plan.auto_sync - - -def test_exclusive_leading_pointer_storage_emits_a_barrier_per_call_site(): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value): - storage = coop.TempStorage(sharing="exclusive", auto_sync=True) - first = provider(value, temp_storage=storage) - second = provider(first, temp_storage=storage) - return second - - func_ir, _, _ = _rewrite_registered_provider(kernel) - targets = [target for _, _, target in _resolved_calls(func_ir)] - - assert targets.count(invocable) == 2 - assert targets.count(cuda.syncthreads) == 2 - - -@pytest.mark.parametrize("auto_sync", [None, False, True]) -def test_unused_storage_plan_requires_default_policy(auto_sync): - rewrite, _ = _planner_for_storage_policy( - _TempStorageCtorSpecification(128, None, auto_sync, "shared"), - [], - ) - - if auto_sync is True: - with pytest.raises( - CoopSinglePhaseRewriteError, - match="requires a cooperative primitive", - ): - rewrite._finalize_temp_storage_plan_for_var("storage") - else: - plan = rewrite._finalize_temp_storage_plan_for_var("storage") - assert plan.size_in_bytes == 128 - assert plan.auto_sync is False - - -def test_storage_capacity_is_validated_before_codegen(): - undersized, _ = _planner_for_storage_policy( - _TempStorageCtorSpecification(15, 8, None, "shared"), - [(16, 16)], - ) - - with pytest.raises( - CoopSinglePhaseRewriteError, match="smaller than required" - ): - undersized._finalize_temp_storage_plan_for_var("storage") - - -@pytest.mark.parametrize("alignment", [None, 1, 8, 16, 32]) -@pytest.mark.parametrize("sharing", ["shared", "exclusive"]) -def test_storage_alignment_satisfies_all_uses(alignment, sharing): - rewrite, calls = _planner_for_storage_policy( - _TempStorageCtorSpecification(None, alignment, None, sharing), - [(24, 8), (16, 16)], - ) - - plan = rewrite._finalize_temp_storage_plan_for_var("storage") - - assert plan.alignment == max(16, alignment or 1) - assert plan.size_in_bytes == (24 if sharing == "shared" else 48) - assert [plan.slices_by_call_id[id(call)].offset for call in calls] == ( - [0, 0] if sharing == "shared" else [0, 32] - ) - - -@pytest.mark.parametrize( - "module", [common_coop, coop], ids=["root", "qualified"] -) -def test_temp_storage_rewrite_normalizes_numpy_size(module): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - size_in_bytes = np.int64(4096) - - def kernel(value): - storage = module.TempStorage(size_in_bytes) - return provider(value, temp_storage=storage) - - func_ir, rewrite, _ = _rewrite_registered_provider(kernel) - - assert cuda.shared.array in _call_targets(func_ir) - assert rewrite._temp_storage_global_plan.total_size == 4096 - - -@pytest.mark.parametrize( - "module", [common_coop, coop], ids=["root", "qualified"] -) -@pytest.mark.parametrize("alignment", [None, 1, 2, 4, 8, 16, np.int64(32)]) -def test_temp_storage_rewrite_normalizes_minimum_alignment(module, alignment): - invocable = _FakeInvocable(alignment=16) - provider = _register_leading_pointer_provider(invocable) - - def kernel(value): - storage = module.TempStorage(alignment=alignment) - return provider(value, temp_storage=storage) - - func_ir, rewrite, _ = _rewrite_registered_provider(kernel) - - assert cuda.shared.array in _call_targets(func_ir) - assert rewrite._temp_storage_global_plan.max_alignment == max( - 16, alignment or 1 - ) - - -@pytest.mark.parametrize( - "module", [common_coop, coop], ids=["root", "qualified"] -) -@pytest.mark.parametrize( - ("alignment", "message"), - [ - (True, "alignment must be an integer or None"), - (1.5, "alignment must be an integer or None"), - (0, "alignment must be a positive integer"), - (-1, "alignment must be a positive integer"), - (3, "alignment must be a power of 2"), - ], -) -def test_temp_storage_rewrite_rejects_invalid_alignment( - module, alignment, message -): - def kernel(): - return module.TempStorage(alignment=alignment) - - with pytest.raises(CoopSinglePhaseRewriteError, match=message): - _rewrite(kernel) - - -@pytest.mark.parametrize( - "module", [common_coop, coop], ids=["root", "qualified"] -) -def test_temp_storage_rewrite_requires_keyword_options(module): - def kernel(): - storage = module.TempStorage(64, 16) - return storage - - with pytest.raises( - CoopSinglePhaseRewriteError, match="only size_in_bytes positionally" - ): - _rewrite(kernel) - - -def test_small_static_storage_does_not_require_a_device_query(monkeypatch): - from cuda.coop.numba_mlir._compiler import _rewrite_storage - - rewrite = object.__new__(CoopSinglePhaseRewrite) - monkeypatch.setattr( - _rewrite_storage, - "_query_device_shared_memory_limits", - lambda: pytest.fail("small static storage queried a device"), - ) - - assert rewrite._get_device_shared_memory_limits( - _DEFAULT_STATIC_SHARED_MEMORY_BYTES - ) == ( - _DEFAULT_STATIC_SHARED_MEMORY_BYTES, - _DEFAULT_STATIC_SHARED_MEMORY_BYTES, - ) - - -def test_large_storage_requires_an_exact_device_limit(monkeypatch): - from cuda.coop.numba_mlir._compiler import _rewrite_storage - - rewrite = object.__new__(CoopSinglePhaseRewrite) - monkeypatch.setattr( - _rewrite_storage, - "_query_device_shared_memory_limits", - lambda: (_ for _ in ()).throw(RuntimeError("no current device")), - ) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="require an exact current-device shared-memory query", - ): - rewrite._get_device_shared_memory_limits( - _DEFAULT_STATIC_SHARED_MEMORY_BYTES + 1 - ) - - -def _global_storage_planner( - monkeypatch, - *, - size, - default_limit, - optin_limit, - implicit_use_specifications=(), - alignment=16, -): - """Supply device limits and record dynamic-storage requests for one plan. - - Replace the descriptor's own plan with a fixed size and alignment, and - fix the device limits. The real function-wide planner combines that plan - with implicit uses and chooses static or dynamic shared memory. - """ - - from cuda.coop.numba_mlir._compiler import _rewrite_storage - - rewrite = object.__new__(CoopSinglePhaseRewrite) - rewrite._state = object() - rewrite._temp_storage_global_plan = None - rewrite._temp_storage_ctor_specifications = {"storage": object()} - rewrite._temp_storage_ctor_order = {"storage": 0} - rewrite._temp_storage_plans = {} - rewrite._func_temp_storage_requirements = { - "storage": _TempStorageRequirementSummary() - } - implicit_calls = [object() for _ in implicit_use_specifications] - rewrite._implicit_temp_storage_requirements = ( - _TempStorageRequirementSummary( - max_size_in_bytes=max( - (use_size for use_size, _ in implicit_use_specifications), - default=0, - ), - max_alignment=max( - (alignment for _, alignment in implicit_use_specifications), - default=1, - ), - uses=[ - _TempStorageUseRequirement( - call_assign=call, - order=index, - size_in_bytes=use_size, - alignment=alignment, - ) - for index, (call, (use_size, alignment)) in enumerate( - zip(implicit_calls, implicit_use_specifications) - ) - ], - ) - ) - rewrite._implicit_temp_storage_plan = None - rewrite._finalize_temp_storage_plan_for_var = lambda _key: _TempStoragePlan( - size_in_bytes=size, - alignment=alignment, - sharing="shared", - auto_sync=True, - slices_by_call_id={}, - ) - monkeypatch.setattr( - _rewrite_storage, - "_query_device_shared_memory_limits", - lambda: { - "max_default_shared_memory_per_block": default_limit, - "max_optin_shared_memory_per_block": optin_limit, - }, - ) - requested = [] - monkeypatch.setattr( - _rewrite_storage, - "set_required_dynamic_shared_memory", - lambda state, value: requested.append((state, value)), - ) - return rewrite, requested, implicit_calls - - -def test_storage_above_default_requests_exact_dynamic_shared_memory( - monkeypatch, -): - rewrite, requested, _ = _global_storage_planner( - monkeypatch, - size=64 * 1024, - default_limit=48 * 1024, - optin_limit=96 * 1024, - ) - - plan = rewrite._ensure_temp_storage_global_plan() - - assert plan.uses_dynamic_smem - assert plan.dynamic_shared_bytes == 64 * 1024 - assert requested == [(rewrite._state, 64 * 1024)] - - -def test_dynamic_backing_rejects_alignment_above_the_window_guarantee( - monkeypatch, -): - # The static path declares the shared array with the requested alignment; - # the dynamic window only guarantees 16 bytes, so a larger request must - # fail instead of becoming a false alignment assumption. - rewrite, requested, _ = _global_storage_planner( - monkeypatch, - size=64 * 1024, - default_limit=48 * 1024, - optin_limit=96 * 1024, - alignment=32, - ) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="guarantees only 16-byte alignment", - ): - rewrite._ensure_temp_storage_global_plan() - assert requested == [] - - -def test_static_backing_honors_alignment_above_the_dynamic_guarantee( - monkeypatch, -): - rewrite, requested, _ = _global_storage_planner( - monkeypatch, - size=32 * 1024, - default_limit=48 * 1024, - optin_limit=96 * 1024, - alignment=32, - ) - - plan = rewrite._ensure_temp_storage_global_plan() - - assert not plan.uses_dynamic_smem - assert plan.max_alignment == 32 - assert requested == [] - - -def test_storage_above_device_optin_limit_is_rejected(monkeypatch): - rewrite, requested, _ = _global_storage_planner( - monkeypatch, - size=100 * 1024, - default_limit=48 * 1024, - optin_limit=96 * 1024, - ) - - with pytest.raises(CoopSinglePhaseRewriteError, match="device max opt-in"): - rewrite._ensure_temp_storage_global_plan() - assert requested == [] - - -def test_implicit_and_explicit_storage_share_one_dynamic_backing(monkeypatch): - rewrite, requested, implicit_calls = _global_storage_planner( - monkeypatch, - size=40 * 1024, - implicit_use_specifications=((16 * 1024, 16), (8 * 1024, 8)), - default_limit=48 * 1024, - optin_limit=96 * 1024, - ) - - plan = rewrite._ensure_temp_storage_global_plan() - implicit = rewrite._implicit_temp_storage_plan - - assert plan.total_size == 56 * 1024 - assert plan.dynamic_shared_bytes == 56 * 1024 - assert requested == [(rewrite._state, 56 * 1024)] - assert implicit is not None - assert implicit.base_offset == 40 * 1024 - assert implicit.size_in_bytes == 16 * 1024 - assert set(implicit.slices_by_call_id) == { - id(call) for call in implicit_calls - } - - -def test_implicit_storage_counts_toward_the_optin_limit(monkeypatch): - rewrite, requested, _ = _global_storage_planner( - monkeypatch, - size=80 * 1024, - implicit_use_specifications=((32 * 1024, 16),), - default_limit=48 * 1024, - optin_limit=96 * 1024, - ) - - with pytest.raises(CoopSinglePhaseRewriteError, match="device max opt-in"): - rewrite._ensure_temp_storage_global_plan() - assert requested == [] - - -@pytest.mark.parametrize("ssa", [False, True], ids=["before-ssa", "ssa"]) -@pytest.mark.parametrize("rebound", [False, True], ids=["selected", "rebound"]) -def test_descriptor_none_alias_is_rejected_by_both_parsers(ssa, rebound): - from cuda.coop.numba_mlir._compiler._group_planner import _GroupCallPlanner - from cuda.coop.numba_mlir._compiler._group_planner_support import ( - GroupRewriteError, - ) - - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - if rebound: - - def kernel(value, flag): - empty = None - alias = empty - storage = coop.TempStorage() - if flag: - storage = alias - return provider(value, temp_storage=storage) - - else: - - def kernel(value, flag): - empty = None - alias = empty - storage = coop.TempStorage() - selected = storage if flag else alias - return provider(value, temp_storage=selected) - - with pytest.raises(CoopSinglePhaseRewriteError, match="on every path"): - _rewrite_registered_provider(kernel, ssa=ssa) - - func_ir = _frontend(kernel, ssa=ssa) - provider_call = next( - inst.value - for _, inst, target in _resolved_calls(func_ir) - if target is provider - ) - planner = _GroupCallPlanner( - SimpleNamespace(func_ir=func_ir, args=(types.int32, types.boolean)), - {"block": (32, 1, 1), "grid": (1, 1, 1)}, - ) - with pytest.raises(GroupRewriteError, match="on every path"): - planner.context.temp_storage(dict(provider_call.kws)["temp_storage"]) - - -@pytest.mark.parametrize("ssa", [False, True], ids=["before-ssa", "ssa"]) -def test_descriptor_loop_alias_retains_constructor_provenance(ssa): - from cuda.coop.numba_mlir._compiler._group_planner import _GroupCallPlanner - - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value, count): - storage = coop.TempStorage(auto_sync=True) - selected = storage - for _ in range(count): - alias = selected - selected = alias - value = provider(value, temp_storage=selected) - return value - - func_ir = _frontend(kernel, ssa=ssa) - provider_call = next( - inst.value - for _, inst, target in _resolved_calls(func_ir) - if target is provider - ) - planner = _GroupCallPlanner( - SimpleNamespace(func_ir=func_ir, args=(types.int32, types.int32)), - {"block": (32, 1, 1), "grid": (1, 1, 1)}, - ) - assert planner.context.temp_storage( - dict(provider_call.kws)["temp_storage"] - ) == ( - None, - None, - True, - "shared", - ) - func_ir, _, _ = _rewrite_registered_provider(kernel, ssa=ssa) - assert _call_targets(func_ir).count(cuda.syncthreads) == 1 - - -@pytest.mark.parametrize( - "explicit", [False, True], ids=["implicit", "explicit"] -) -@pytest.mark.parametrize("user_shape", ["static", "runtime", "zero"]) -def test_dynamic_coop_backing_rejects_all_user_shared_allocations( - monkeypatch, explicit, user_shape -): - from cuda.coop.numba_mlir._compiler import _rewrite_storage - - invocable = _FakeInvocable(size_in_bytes=64 * 1024) - provider = _register_leading_pointer_provider(invocable) - monkeypatch.setattr( - _rewrite_storage, - "_query_device_shared_memory_limits", - lambda: { - "max_default_shared_memory_per_block": 48 * 1024, - "max_optin_shared_memory_per_block": 96 * 1024, - }, - ) - - shape = (4, 4) if user_shape == "static" else 0 - if explicit: - if user_shape == "runtime": - - def kernel(value, size): - array = cuda.shared.array - alias = array - mine = alias(size, types.int32) - mine[0] = value - storage = coop.TempStorage() - return provider(value, temp_storage=storage) - - else: - - def kernel(value, size): - array = cuda.shared.array - alias = array - mine = alias(shape, types.int32) - mine[0] = value - storage = coop.TempStorage() - return provider(value, temp_storage=storage) - - elif user_shape == "runtime": - - def kernel(value, size): - array = cuda.shared.array - alias = array - mine = alias(size, types.int32) - mine[0] = value - return provider(value) - - else: - - def kernel(value, size): - array = cuda.shared.array - alias = array - mine = alias(shape, types.int32) - mine[0] = value - return provider(value) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="65536-byte dynamic shared-memory backing.*would alias", - ): - _rewrite_registered_provider(kernel) - - -def test_static_coop_backing_rejects_runtime_user_shared_allocation(): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value, size): - mine = cuda.shared.array(size, types.int32) - mine[0] = value - return provider(value) - - with pytest.raises( - CoopSinglePhaseRewriteError, - match="64-byte static.*dynamic/runtime-sized.*would alias", - ): - _rewrite_registered_provider(kernel) - - -def test_static_multidimensional_user_array_coexists_with_static_coop_backing(): - invocable = _FakeInvocable() - provider = _register_leading_pointer_provider(invocable) - - def kernel(value): - mine = cuda.shared.array((4, 4), types.int32) - mine[0, 0] = value - return provider(value) - - func_ir, _, _ = _rewrite_registered_provider(kernel) - assert _call_targets(func_ir).count(cuda.shared.array) == 2 diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_thread_data.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_thread_data.py deleted file mode 100644 index e17a07c24151..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_thread_data.py +++ /dev/null @@ -1,121 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from enum import Enum - -import numpy as np -import pytest - -import cuda.coop.numba_mlir as coop - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - - -class _StringSharing(str, Enum): - SHARED = "shared" - - -@pytest.mark.parametrize( - "call", - [ - lambda: coop.ThreadData(shape=4), - lambda: coop.ThreadData(4, address_space="local"), - lambda: coop.ThreadData(4, items_per_thread=4), - ], -) -def test_thread_data_rejects_unknown_or_duplicate_arguments(call): - with pytest.raises(TypeError): - call() - - -@pytest.mark.parametrize("alignment", [None, 1, 2, 4, 8, 16, np.int64(32)]) -def test_temp_storage_accepts_minimum_alignment(alignment): - storage = coop.TempStorage(alignment=alignment) - assert storage.alignment == alignment - assert storage.alignment is None or type(storage.alignment) is int - - -@pytest.mark.parametrize("size_in_bytes", [4096, np.int64(4096)]) -def test_temp_storage_accepts_integer_sizes(size_in_bytes): - storage = coop.TempStorage(size_in_bytes) - assert storage.size_in_bytes == 4096 - assert type(storage.size_in_bytes) is int - - -def test_temp_storage_uses_canonical_defaults_and_normalization(): - shared = coop.TempStorage(sharing=" SHARED ") - exclusive = coop.TempStorage(sharing=" Exclusive ") - - assert shared.sharing == "shared" - assert shared.auto_sync is False - assert exclusive.sharing == "exclusive" - assert exclusive.auto_sync is False - assert coop.TempStorage(auto_sync=None).auto_sync is False - assert ( - coop.TempStorage(sharing="exclusive", auto_sync=True).auto_sync is True - ) - assert ( - coop.TempStorage(sharing="exclusive", auto_sync=False).auto_sync - is False - ) - - -def test_temp_storage_rejects_string_enum_sharing(): - with pytest.raises(TypeError, match="sharing must be a string"): - coop.TempStorage(sharing=_StringSharing.SHARED) - - -@pytest.mark.parametrize( - ("kwargs", "error_type", "message"), - [ - ( - {"size_in_bytes": True}, - TypeError, - "TempStorage size_in_bytes must be an integer or None", - ), - ( - {"size_in_bytes": 0}, - ValueError, - "TempStorage size_in_bytes must be a positive integer", - ), - ( - {"size_in_bytes": 1.5}, - TypeError, - "TempStorage size_in_bytes must be an integer or None", - ), - ( - {"size_in_bytes": np.int64(0)}, - ValueError, - "TempStorage size_in_bytes must be a positive integer", - ), - ( - {"size_in_bytes": np.int64(-1)}, - ValueError, - "TempStorage size_in_bytes must be a positive integer", - ), - ( - {"alignment": False}, - TypeError, - "alignment must be an integer or None", - ), - ( - {"alignment": 3}, - ValueError, - "alignment must be a power of 2", - ), - ( - {"auto_sync": 1}, - TypeError, - "TempStorage auto_sync must be None/True/False", - ), - ( - {"sharing": 1}, - TypeError, - "TempStorage sharing must be a string", - ), - ], -) -def test_temp_storage_validation(kwargs, error_type, message): - with pytest.raises(error_type, match=message): - coop.TempStorage(**kwargs) diff --git a/python/cuda_coop/tests/backends/numba_mlir/unit/test_thread_groups.py b/python/cuda_coop/tests/backends/numba_mlir/unit/test_thread_groups.py deleted file mode 100644 index 54944efb05d2..000000000000 --- a/python/cuda_coop/tests/backends/numba_mlir/unit/test_thread_groups.py +++ /dev/null @@ -1,96 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import pytest - -import cuda.coop as common_coop -import cuda.coop.numba_mlir as coop -from cuda.coop._core import ThreadHierarchy as CoreThreadHierarchy - -pytestmark = [pytest.mark.backend_numba_mlir, pytest.mark.unit] - -_UNSUPPORTED_THREAD_GROUP_METHODS = frozenset( - { - "count", - "count_as", - "is_member", - "rank", - "rank_as", - "sync", - "sync_aligned", - } -) - - -def test_group_exports_use_the_shared_hierarchy_contract(): - assert coop.Hierarchy is coop.ThreadHierarchy - assert coop.ThreadHierarchy is CoreThreadHierarchy - - for name in ( - "Hierarchy", - "ThreadGroup", - "ThreadHierarchy", - "this_thread", - "this_warp", - "this_block", - "this_cluster", - "this_grid", - ): - assert name in coop.__all__ - - -def test_current_group_construction_preserves_backend_type(): - current = coop.this_block() - - assert type(current) is coop.ThreadGroup - assert current.is_current - assert repr(current).startswith("ThreadGroup(kind='block'") - assert current.block_dim is None - assert current.static_size is None - assert coop.this_cluster().static_size is None - assert coop.this_grid().static_size is None - assert coop.this_thread().static_size == 1 - assert coop.this_warp().static_size == 32 - - -def test_group_equality_hashing_and_group_by_use_numba_mlir_type(): - first = coop.this_block() - second = coop.this_block() - mapped = first.group_by(2) - - assert first == second - assert hash(first) == hash(second) - assert type(mapped) is coop.ThreadGroup - assert type(mapped.parent) is coop.ThreadGroup - assert mapped.kind == "warps_within_block" - assert mapped.static_size == 64 - assert mapped.groups_per_parent is None - assert mapped.complete_membership is None - - -def test_group_constructors_preserve_shared_validation(): - with pytest.raises(TypeError): - coop.ThreadHierarchy(block_dim=64) - with pytest.raises(TypeError): - coop.this_warp(16) - with pytest.raises(TypeError): - coop.this_warp(block_dim=64) - with pytest.raises(ValueError, match="requires the count to divide"): - coop.this_warp().group_by(12) - with pytest.raises(NotImplementedError, match="nested"): - coop.this_warp().group_by(8).group_by(2) - - -@pytest.mark.parametrize( - "api", (common_coop, coop), ids=("common", "qualified") -) -def test_thread_group_runtime_surface_is_descriptor_only(api): - group = api.this_block() - - assert callable(group.group_by) - assert _UNSUPPORTED_THREAD_GROUP_METHODS.isdisjoint(dir(group)) - assert all( - not hasattr(group, method) - for method in _UNSUPPORTED_THREAD_GROUP_METHODS - ) diff --git a/python/cuda_coop/tests/packaging/test_typing.py b/python/cuda_coop/tests/packaging/test_typing.py deleted file mode 100644 index cfbfaebaf3d8..000000000000 --- a/python/cuda_coop/tests/packaging/test_typing.py +++ /dev/null @@ -1,129 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check public stubs through isolated valid and invalid consumer programs. - -Copy only the imported package's ``.pyi`` stubs and ``py.typed`` marker to a -temporary root. Then implementation modules and an unrelated installed wheel -cannot satisfy missing public declarations. -""" - -from __future__ import annotations - -import importlib.util -import os -import re -import shutil -import subprocess -import sys -from pathlib import Path - -import pytest - -from cuda import coop - -_CONSUMER_ROOT = Path(__file__).with_name("typing") -_VALID_CONSUMERS = ("common_consumer.py", "numba_consumer.py") - - -def _mypy_args(cache_dir: Path) -> list[str]: - return [ - sys.executable, - "-m", - "mypy", - "--strict", - "--python-version", - f"{sys.version_info.major}.{sys.version_info.minor}", - "--disallow-any-unimported", - "--show-error-codes", - "--cache-dir", - str(cache_dir), - ] - - -def _run_mypy( - mypy_args: list[str], - consumers: list[Path], - *, - stub_root: Path, -) -> subprocess.CompletedProcess[str]: - """Type-check copied consumers against only the selected stub root.""" - environment = os.environ.copy() - environment["MYPYPATH"] = str(stub_root) - environment.pop("PYTHONPATH", None) - return subprocess.run( - [*mypy_args, *(str(consumer) for consumer in consumers)], - cwd=stub_root.parent, - env=environment, - capture_output=True, - text=True, - check=False, - ) - - -def _expected_diagnostics(consumer: Path) -> set[tuple[int, str]]: - """Read expected error codes at their current consumer source lines.""" - return { - (line_number, error_code) - for line_number, line in enumerate( - consumer.read_text(encoding="utf-8").splitlines(), - start=1, - ) - if (match := re.search(r"# expected-error: (\[[a-z-]+\])", line)) - for error_code in (match.group(1),) - } - - -def test_public_stubs_pass_strict_consumer_type_checks(tmp_path: Path) -> None: - if importlib.util.find_spec("mypy") is None: - pytest.skip("mypy is not installed") - - package_root = Path(coop.__file__).parent - stub_root = tmp_path / "stubs" / "cuda" / "coop" - for source in package_root.rglob("*.pyi"): - destination = stub_root / source.relative_to(package_root) - destination.parent.mkdir(parents=True, exist_ok=True) - shutil.copyfile(source, destination) - shutil.copyfile(package_root / "py.typed", stub_root / "py.typed") - # Package markers keep an installed wheel from winning namespace lookup. - (stub_root.parent / "__init__.pyi").touch() - (stub_root / "_core" / "__init__.pyi").touch() - - consumer_root = tmp_path / "consumers" - consumer_root.mkdir() - valid_consumers = [] - for name in _VALID_CONSUMERS: - destination = consumer_root / name - shutil.copyfile(_CONSUMER_ROOT / name, destination) - valid_consumers.append(destination) - - mypy_args = _mypy_args(tmp_path / "mypy-cache") - result = _run_mypy( - mypy_args, - valid_consumers, - stub_root=tmp_path / "stubs", - ) - assert result.returncode == 0, result.stdout + result.stderr - - invalid_consumer = consumer_root / "invalid_consumer.py" - shutil.copyfile(_CONSUMER_ROOT / "invalid_consumer.py", invalid_consumer) - invalid_result = _run_mypy( - mypy_args, - [invalid_consumer], - stub_root=tmp_path / "stubs", - ) - invalid_output = invalid_result.stdout + invalid_result.stderr - assert invalid_result.returncode == 1, invalid_output - - diagnostics = { - (int(line), code) - for line, code in re.findall( - r"invalid_consumer\.py:(\d+): error: .* (\[[a-z-]+\])$", - invalid_output, - flags=re.MULTILINE, - ) - } - assert diagnostics == _expected_diagnostics(invalid_consumer), ( - invalid_output - ) diff --git a/python/cuda_coop/tests/packaging/typing/common_consumer.py b/python/cuda_coop/tests/packaging/typing/common_consumer.py deleted file mode 100644 index 093ae8a28419..000000000000 --- a/python/cuda_coop/tests/packaging/typing/common_consumer.py +++ /dev/null @@ -1,76 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from __future__ import annotations - -from typing import Literal - -import numpy as np -from typing_extensions import assert_type - -from cuda import coop - - -def check_common_surface(source: object, destination: object) -> None: - """Exercise public declarations without importing package internals.""" - - block = coop.this_block() - warp = coop.this_warp() - logical_warp = warp.group_by(8) - values = coop.ThreadData(items_per_thread=2, dtype=np.int16, alignment=16) - storage = coop.TempStorage(64, alignment=16, sharing="shared") - - assert_type(block, coop.ThreadGroup[Literal["block"]]) - assert_type(warp, coop.ThreadGroup[Literal["warp"]]) - assert_type( - logical_warp, - coop.ThreadGroup[Literal["threads_within_warp"]], - ) - assert_type(values, coop.ThreadDataLike[np.int16]) - assert_type(storage, coop.TempStorageLike) - assert_type( - coop.load( - block, - source, - values, - algorithm="direct", - valid_items=15, - oob_default=np.int16(0), - offset=1, - temp_storage=storage, - ), - None, - ) - assert_type( - coop.load(block, source, values, algorithm="warp_transpose_timesliced"), - None, - ) - assert_type( - coop.store( - block, - destination, - values, - algorithm="direct", - valid_items=15, - offset=1, - temp_storage=storage, - ), - None, - ) - assert_type( - coop.load(warp, source, values, algorithm="transpose"), - None, - ) - assert_type( - coop.store(warp, destination, values, algorithm="vectorize"), - None, - ) - assert_type( - coop.load(logical_warp, source, values, algorithm="transpose"), - None, - ) - assert_type( - coop.store(logical_warp, destination, values, algorithm="striped"), - None, - ) diff --git a/python/cuda_coop/tests/packaging/typing/invalid_consumer.py b/python/cuda_coop/tests/packaging/typing/invalid_consumer.py deleted file mode 100644 index b8b0b9691b2f..000000000000 --- a/python/cuda_coop/tests/packaging/typing/invalid_consumer.py +++ /dev/null @@ -1,134 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from __future__ import annotations - -import numpy as np - -import cuda.coop as common -import cuda.coop.numba_mlir as coop - -common.register("numba") # expected-error: [arg-type] - -coop.TempStorage(64, 16) # expected-error: [call-arg] -common.TempStorage(64, 16) # expected-error: [call-arg] -values = coop.ThreadData(items_per_thread=2, dtype=np.int32) -common_values = common.ThreadData(items_per_thread=2, dtype=np.int32) -common_block = common.this_block() -common_block.rank() # expected-error: [attr-defined] -common_block.count() # expected-error: [attr-defined] -common_block.rank_as(np.uint32) # expected-error: [attr-defined] -common_block.count_as(np.uint32) # expected-error: [attr-defined] -common_block.sync() # expected-error: [attr-defined] -common_block.sync_aligned() # expected-error: [attr-defined] -common_block.is_member() # expected-error: [attr-defined] -qualified_block = coop.this_block() -qualified_block.rank() # expected-error: [attr-defined] -qualified_block.count() # expected-error: [attr-defined] -qualified_block.rank_as(np.uint32) # expected-error: [attr-defined] -qualified_block.count_as(np.uint32) # expected-error: [attr-defined] -qualified_block.sync() # expected-error: [attr-defined] -qualified_block.sync_aligned() # expected-error: [attr-defined] -qualified_block.is_member() # expected-error: [attr-defined] -common.load( # expected-error: [call-overload] - common.this_block(), - object(), - common_values, - algorithm="stripd", -) -common.load( # expected-error: [call-overload] - common.this_warp(), - object(), - common_values, - algorithm="warp_transpose", -) -common.load( - common.this_warp(), # expected-error: [arg-type] - object(), - common_values, - temp_storage=common.TempStorage(), -) -common.load( - common.this_warp().group_by(8), # expected-error: [arg-type] - object(), - common_values, - temp_storage=common.TempStorage(), -) -common.store( # expected-error: [call-overload] - common.this_warp(), - object(), - common_values, - algorithm="warp_transpose", -) -coop.load( # expected-error: [call-overload] - coop.this_warp(), - object(), - values, - algorithm="warp_transpose", -) -coop.load( - coop.this_warp(), # expected-error: [arg-type] - object(), - values, - temp_storage=coop.TempStorage(), -) -coop.load( - coop.this_warp().group_by(8), # expected-error: [arg-type] - object(), - values, - temp_storage=coop.TempStorage(), -) -coop.load( # expected-error: [call-overload] - coop.this_warp(), - object(), - values, - algorithm=0, -) -coop.store( # expected-error: [call-overload] - coop.this_warp(), - object(), - values, - algorithm=True, -) -coop.load( # expected-error: [call-overload] - coop.this_block(), - object(), - values, - algorithm=0, -) -coop.store( # expected-error: [call-overload] - coop.this_block(), - object(), - values, - algorithm=True, -) -coop.load( # expected-error: [call-overload] - coop.this_block(), - object(), - values, - algorithm="stripd", -) -coop.load( # expected-error: [call-overload] - coop.this_block(), - object(), - values, - valid_items=1.5, -) -coop.load( # expected-error: [call-overload] - coop.this_block(), - object(), - values, - oob_default=0, -) -coop.store( # expected-error: [call-overload] - coop.this_block(), - object(), - values, - offset="1", -) -# Test rejected attributes. -coop.BlockLoadAlgorithm # expected-error: [attr-defined] # noqa: B018 -coop.BlockStoreAlgorithm # expected-error: [attr-defined] # noqa: B018 -coop.WarpLoadAlgorithm # expected-error: [attr-defined] # noqa: B018 -coop.WarpStoreAlgorithm # expected-error: [attr-defined] # noqa: B018 diff --git a/python/cuda_coop/tests/packaging/typing/numba_consumer.py b/python/cuda_coop/tests/packaging/typing/numba_consumer.py deleted file mode 100644 index 18e7d0f8de7c..000000000000 --- a/python/cuda_coop/tests/packaging/typing/numba_consumer.py +++ /dev/null @@ -1,119 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from __future__ import annotations - -from typing import Literal - -import numpy as np -from typing_extensions import assert_type - -import cuda.coop.numba_mlir as coop - - -def check_numba_surface(source: object, destination: object) -> None: - """Exercise Numba declarations through their public package.""" - - block = coop.this_block() - warp = coop.this_warp() - logical_warp = warp.group_by(8) - byte_values = coop.ThreadData(items_per_thread=1, dtype=np.int8) - values = coop.ThreadData(items_per_thread=2, dtype=np.uint16, alignment=16) - storage = coop.TempStorage(alignment=16, sharing="shared") - - assert_type(block, coop.ThreadGroup[Literal["block"]]) - assert_type(warp, coop.ThreadGroup[Literal["warp"]]) - assert_type( - logical_warp, - coop.ThreadGroup[Literal["threads_within_warp"]], - ) - assert_type(byte_values, coop.ThreadDataLike[np.int8]) - assert_type(values, coop.ThreadDataLike[np.uint16]) - assert_type(storage, coop.TempStorage) - common_storage: coop.TempStorageLike = storage - assert_type(common_storage, coop.TempStorageLike) - assert_type( - coop.load( - block, - source, - values, - algorithm="direct", - temp_storage=storage, - ), - None, - ) - assert_type( - coop.load(block, source, byte_values), - None, - ) - assert_type( - coop.load( - block, - source, - byte_values, - valid_items=1, - oob_default=0, - ), - None, - ) - assert_type( - coop.load(block, source, values, algorithm="vectorize"), - None, - ) - assert_type( - coop.load( - block, - source, - values, - algorithm="warp_transpose_timesliced", - ), - None, - ) - assert_type( - coop.store( - block, - destination, - values, - algorithm="direct", - temp_storage=storage, - ), - None, - ) - assert_type(coop.store(block, destination, byte_values), None) - assert_type( - coop.load( - warp, - source, - values, - algorithm="transpose", - ), - None, - ) - assert_type( - coop.store( - warp, - destination, - values, - algorithm="striped", - ), - None, - ) - assert_type( - coop.load( - logical_warp, - source, - values, - algorithm="transpose", - ), - None, - ) - assert_type( - coop.store( - logical_warp, - destination, - values, - algorithm="striped", - ), - None, - ) diff --git a/python/cuda_coop/tests/test_bindings.py b/python/cuda_coop/tests/test_bindings.py deleted file mode 100644 index 2a44bee6e7d7..000000000000 --- a/python/cuda_coop/tests/test_bindings.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import numpy as np -import pytest - -from cuda.coop._core import ( - INT32, - CxxFunction, - RuntimeValue, - Value, - binding, - i32_parameter, -) - - -def test_i32_binding_selects_omitted_constant_or_runtime_parameter(): - assert i32_parameter(binding(None), name="value") is None - assert i32_parameter( - binding(None), name="value", omitted_value=3 - ) == CxxFunction("3", dtype=INT32, name="value") - assert i32_parameter(binding(np.int32(5)), name="value") == CxxFunction( - "5", dtype=INT32, name="value" - ) - assert i32_parameter( - binding(RuntimeValue("payload")), name="value" - ) == Value(INT32, name="value") - - -@pytest.mark.parametrize("value", [-(1 << 31), (1 << 31) - 1]) -def test_i32_binding_accepts_signed_boundaries(value): - assert i32_parameter(binding(value), name="value") == CxxFunction( - str(value), dtype=INT32, name="value" - ) - - -@pytest.mark.parametrize( - ("value", "error", "message"), - [ - (True, TypeError, "must be an integer"), - (1.5, TypeError, "must be an integer"), - (1 << 31, ValueError, "fit a signed 32-bit integer"), - ], -) -def test_i32_binding_rejects_unrepresentable_constants(value, error, message): - with pytest.raises(error, match=message): - i32_parameter(binding(value), name="value") - - -@pytest.mark.parametrize(("left", "right"), [(True, 1), (0.0, -0.0)]) -def test_static_bindings_do_not_alias_in_caches(left, right): - assert len({binding(left), binding(right)}) == 2 diff --git a/python/cuda_coop/tests/test_block_load_store.py b/python/cuda_coop/tests/test_block_load_store.py deleted file mode 100644 index 4b409d43e0c6..000000000000 --- a/python/cuda_coop/tests/test_block_load_store.py +++ /dev/null @@ -1,202 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check Block Load and Store descriptors before code generation. - -The parameter lists describe the CUB call and its runtime arguments. These -tests inspect those lists and validate static controls without a compiler. -""" - -import numpy as np -import pytest - -from cuda.coop._core import ArgumentBinding, ArgumentKind -from cuda.coop._core.block import ( - make_block_load_specialization, - make_block_load_store_semantics, - make_block_store_specialization, -) - - -def test_block_load_partial_default_and_pointer_offset_overloads(): - specialization = make_block_load_specialization( - dtype="f32", - block_dim=(32, 1, 1), - items_per_thread=2, - algorithm="transpose", - valid_items=True, - oob_default=True, - include_full_tile=True, - include_pointer_offset=True, - ) - - # Each optional overload must preserve CUB's argument order. - assert [ - [parameter.name for parameter in method[3:]] - for method in specialization.parameters - ] == [ - [], - ["num_valid_items", "oob_default"], - ["num_valid_items", "oob_default", "offset"], - ["offset"], - ] - - -def test_block_load_static_controls_stay_out_of_the_runtime_abi(): - specialization = make_block_load_specialization( - dtype="f32", - block_dim=(32, 1, 1), - items_per_thread=2, - algorithm="direct", - valid_items=ArgumentBinding.static(17), - oob_default=ArgumentBinding.static(0), - include_pointer_offset=ArgumentBinding.static(4), - ) - - assert len(specialization.parameters) == 1 - controls = specialization.parameters[0][-3:] - assert all(p.argument_kind is ArgumentKind.STATIC for p in controls) - assert controls[-1].static_value == 4 - - -@pytest.mark.parametrize( - "make_specialization", - [make_block_load_specialization, make_block_store_specialization], -) -def test_numpy_integer_controls_share_the_same_specialization( - make_specialization, -): - def build(integer): - return make_specialization( - dtype="i32", - block_dim=(32, 1, 1), - items_per_thread=integer(2), - algorithm="direct", - valid_items=ArgumentBinding.static(integer(5)), - include_pointer_offset=ArgumentBinding.static(integer(7)), - ) - - assert build(int).semantic_key == build(np.int64).semantic_key - - -@pytest.mark.parametrize("items_per_thread", [0, -1, True, "two"]) -def test_block_load_store_rejects_invalid_item_count(items_per_thread): - with pytest.raises( - ValueError, match="items_per_thread must be a positive integer" - ): - make_block_load_store_semantics( - kind="load", - dtype="i32", - items_per_thread=items_per_thread, - algorithm="direct", - ) - - -@pytest.mark.parametrize( - ("overrides", "message"), - [ - ({"algorithm": "unknown"}, "unsupported BlockLoad algorithm"), - ( - {"kind": "store", "valid_items": True, "oob_default": True}, - "only valid for BlockLoad", - ), - ({"oob_default": True}, "requires a valid_items"), - ({"include_full_tile": True}, "include_full_tile"), - ], -) -def test_block_load_store_rejects_invalid_options(overrides, message): - options = { - "kind": "load", - "dtype": "i32", - "items_per_thread": 1, - "algorithm": "direct", - } - options.update(overrides) - with pytest.raises(ValueError, match=message): - make_block_load_store_semantics(**options) - - -@pytest.mark.parametrize( - "make_specialization", - [make_block_load_specialization, make_block_store_specialization], -) -def test_block_valid_items_bounds_use_all_three_dimensions(make_specialization): - def build(value): - return make_specialization( - dtype="i32", - block_dim=(8, 2, 2), - items_per_thread=2, - algorithm="direct", - valid_items=ArgumentBinding.static(value), - ) - - for value in (0, 64): - build(value) - for value in (-1, 65): - with pytest.raises(ValueError, match="block tile size"): - build(value) - - -@pytest.mark.parametrize( - "make_specialization", - [make_block_load_specialization, make_block_store_specialization], -) -@pytest.mark.parametrize( - "algorithm", ["warp_transpose", "warp_transpose_timesliced"] -) -def test_direct_block_warp_transpose_requires_complete_physical_warps( - make_specialization, algorithm -): - with pytest.raises(ValueError, match="multiple of 32"): - make_specialization( - dtype="i32", - block_dim=(16, 3, 1), - items_per_thread=2, - algorithm=algorithm, - ) - - -@pytest.mark.parametrize( - ("value", "literal"), - [ - (True, "true"), - (7, "7"), - (1.5, "1.5"), - (np.int64(-(1 << 63)), "(-9223372036854775807LL - 1LL)"), - (np.uint64((1 << 64) - 1), "18446744073709551615ULL"), - ], -) -def test_block_load_renders_static_oob_default_as_cpp_scalar(value, literal): - specialization = make_block_load_specialization( - dtype="f32", - block_dim=(32, 1, 1), - items_per_thread=2, - algorithm="direct", - valid_items=ArgumentBinding.static(1), - oob_default=ArgumentBinding.static(value), - ) - - assert specialization.parameters[0][-1].cpp == literal - - -@pytest.mark.parametrize( - ("value", "message"), - [ - (float("inf"), "must be finite"), - (float("-inf"), "must be finite"), - (float("nan"), "must be finite"), - (-(1 << 63) - 1, "fit a 64-bit integer"), - (1 << 64, "fit a 64-bit integer"), - ], -) -def test_block_load_rejects_unrepresentable_static_oob_default(value, message): - with pytest.raises(ValueError, match=message): - make_block_load_specialization( - dtype="f32", - block_dim=(32, 1, 1), - items_per_thread=2, - algorithm="direct", - valid_items=ArgumentBinding.static(1), - oob_default=ArgumentBinding.static(value), - ) diff --git a/python/cuda_coop/tests/test_group_load_store.py b/python/cuda_coop/tests/test_group_load_store.py deleted file mode 100644 index 48a03a4e1a27..000000000000 --- a/python/cuda_coop/tests/test_group_load_store.py +++ /dev/null @@ -1,295 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check the compiler contracts produced by shared Load and Store planning. - -Plans select a CUB implementation and describe its arguments, participation, -scratch storage, and reuse barriers. These tests inspect plans on the host; -they do not compile or execute the selected device code. -""" - -import numpy as np -import pytest - -from cuda.coop._core import ( - INT64, - ArgumentBinding, - ArgumentKind, - Array, - GroupLoadStoreAlgorithm, - GroupLoweringTarget, - Pointer, - PreconditionEnforcement, - StorageOwnership, - SynchronizationScope, - UnsupportedReasonCode, - this_block, - this_warp, -) -from tests._group_planning import _load_store, _plan - - -@pytest.mark.parametrize("kind", ["load", "store"]) -def test_block_plan_selects_cub_with_exact_launch_dimensions(kind): - plan = _plan(this_block(), _load_store(kind), (8, 4, 1)) - - assert plan.target is GroupLoweringTarget.CUB_BLOCK - assert plan.provenance.cpp_class == f"cub::Block{kind.title()}" - arguments = plan.implementation.template_arguments - assert tuple(arguments[f"BLOCK_DIM_{axis}"] for axis in "XYZ") == (8, 4, 1) - - -@pytest.mark.parametrize("kind", ["load", "store"]) -@pytest.mark.parametrize("width", [1, 2, 4, 8, 16, 32]) -def test_warp_plan_selects_cub_for_the_group_width(kind, width): - group = this_warp() if width == 32 else this_warp().group_by(width) - plan = _plan(group, _load_store(kind), 64) - - assert plan.target is GroupLoweringTarget.CUB_WARP - assert plan.provenance.cpp_class == f"cub::Warp{kind.title()}" - assert ( - plan.implementation.template_arguments["LOGICAL_WARP_THREADS"] == width - ) - - -@pytest.mark.parametrize("kind", ["load", "store"]) -@pytest.mark.parametrize("group", [this_block(), this_warp()]) -def test_load_store_abi_offsets_the_global_pointer(kind, group): - plan = _plan(group, _load_store(kind, offset=ArgumentBinding.runtime())) - _, pointer, items, offset = plan.implementation.parameters[0] - - assert isinstance(pointer, Pointer) and isinstance(items, Array) - assert (pointer.is_output, items.is_output) == ( - kind == "store", - kind == "load", - ) - assert offset.pointer_arg_index == 0 - assert offset.dtype == INT64 - - -@pytest.mark.parametrize("algorithm", list(GroupLoadStoreAlgorithm)) -def test_block_scratch_and_reuse_barrier_follow_the_algorithm(algorithm): - plan = _plan(this_block(), _load_store(algorithm=algorithm)) - uses_scratch = algorithm in { - GroupLoadStoreAlgorithm.TRANSPOSE, - GroupLoadStoreAlgorithm.WARP_TRANSPOSE, - GroupLoadStoreAlgorithm.WARP_TRANSPOSE_TIMESLICED, - } - - assert plan.implementation.template_arguments["ALGORITHM"] == ( - f"::cub::BLOCK_LOAD_{algorithm.value.upper()}" - ) - assert plan.temp_storage.ownership is ( - StorageOwnership.IMPLEMENTATION - if uses_scratch - else StorageOwnership.NONE - ) - assert plan.synchronization.storage_reuse_barrier is ( - SynchronizationScope.BLOCK - if uses_scratch - else SynchronizationScope.NONE - ) - - -@pytest.mark.parametrize("width", [8, 32]) -@pytest.mark.parametrize( - "algorithm", ["direct", "striped", "vectorize", "transpose"] -) -def test_warp_scratch_is_per_group(width, algorithm): - group = this_warp() if width == 32 else this_warp().group_by(width) - plan = _plan(group, _load_store(algorithm=algorithm), 64) - uses_scratch = algorithm == "transpose" - - assert plan.implementation.template_arguments["ALGORITHM"] == ( - f"::cub::WARP_LOAD_{algorithm.upper()}" - ) - assert plan.temp_storage.instances == ( - 64 // width if uses_scratch else None - ) - assert plan.synchronization.storage_reuse_barrier is ( - SynchronizationScope.WARP if uses_scratch else SynchronizationScope.NONE - ) - - -@pytest.mark.parametrize("width", [8, 32]) -def test_warp_effective_offset_leaves_room_for_every_group_tile(width): - group = this_warp() if width == 32 else this_warp().group_by(width) - maximum_offset = (1 << 63) - 1 - (64 // width - 1) * width * 2 - plan = _plan( - group, _load_store(offset=ArgumentBinding.static(maximum_offset)), 64 - ) - - # The adapter adds the group tile origin to this preserved user offset. - assert plan.call.operation.offset.value == maximum_offset - assert plan.implementation.metadata["effective_offset_stride"] == width * 2 - assert ( - plan.implementation.parameters[0][-1].argument_kind - is ArgumentKind.RUNTIME - ) - assert ( - plan.participation.argument_preconditions[0].maximum == maximum_offset - ) - with pytest.raises(ValueError, match="tile origin must fit"): - _plan( - group, - _load_store(offset=ArgumentBinding.static(maximum_offset + 1)), - 64, - ) - - -@pytest.mark.parametrize( - ("group", "tile_size"), - [(this_block(), 128), (this_warp(), 64), (this_warp().group_by(8), 16)], -) -def test_valid_items_bounds_use_the_group_tile(group, tile_size): - for value in (0, tile_size): - _plan(group, _load_store(valid_items=ArgumentBinding.static(value)), 64) - for value in (-1, tile_size + 1): - with pytest.raises(ValueError, match="group tile size"): - _plan( - group, - _load_store(valid_items=ArgumentBinding.static(value)), - 64, - ) - - runtime = _plan( - group, _load_store(valid_items=ArgumentBinding.runtime()), 64 - ) - condition = runtime.participation.argument_preconditions[0] - assert (condition.minimum, condition.maximum) == (0, tile_size) - assert condition.enforcement is PreconditionEnforcement.CALLER - - -def test_block_caller_scratch_auto_sync_changes_the_barrier_and_artifact(): - def plan(auto_sync): - return _plan( - this_block(), - _load_store( - algorithm="transpose", - storage_ownership=StorageOwnership.CALLER, - storage_sharing="shared", - storage_auto_sync=auto_sync, - ), - ) - - synchronized = plan(True) - unsynchronized = plan(False) - - assert synchronized.temp_storage.exact_layout_required - assert ( - synchronized.synchronization.storage_reuse_barrier - is SynchronizationScope.BLOCK - ) - assert ( - unsynchronized.synchronization.storage_reuse_barrier - is SynchronizationScope.NONE - ) - assert synchronized.artifact_key != unsynchronized.artifact_key - - -@pytest.mark.parametrize("kind", ["load", "store"]) -@pytest.mark.parametrize("group", [this_warp(), this_warp().group_by(8)]) -@pytest.mark.parametrize("algorithm", ["direct", "transpose"]) -def test_warp_groups_reject_caller_owned_scratch(kind, group, algorithm): - plan = _plan( - group, - _load_store( - kind, - algorithm=algorithm, - storage_ownership=StorageOwnership.CALLER, - storage_sharing="shared", - ), - ) - - assert plan.target is GroupLoweringTarget.UNSUPPORTED - assert plan.unsupported.code is UnsupportedReasonCode.OPERATION_VARIANT - assert "does not accept explicit temp_storage" in plan.unsupported.message - - -def test_unused_scratch_does_not_change_the_artifact(): - implicit = _plan(this_block(), _load_store()) - explicit = _plan( - this_block(), - _load_store( - storage_ownership=StorageOwnership.CALLER, - storage_sharing="shared", - storage_size_in_bytes=256, - storage_alignment=16, - ), - ) - - assert explicit.temp_storage.ownership is StorageOwnership.NONE - assert implicit.artifact_key == explicit.artifact_key - - -@pytest.mark.parametrize( - "algorithm", ["warp_transpose", "warp_transpose_timesliced"] -) -@pytest.mark.parametrize("group", [this_warp(), this_warp().group_by(8)]) -def test_warp_groups_reject_block_only_algorithms(group, algorithm): - plan = _plan(group, _load_store(algorithm=algorithm), 64) - - assert plan.unsupported.code is UnsupportedReasonCode.OPERATION_VARIANT - - -@pytest.mark.parametrize( - ("group", "block_size", "reason"), - [ - ( - this_warp().group_by(3, exhaustive=False), - 64, - UnsupportedReasonCode.GROUP_KIND, - ), - (this_warp(), 48, UnsupportedReasonCode.PARTIAL_PHYSICAL_WARP), - ( - this_warp().group_by(8), - 48, - UnsupportedReasonCode.PARTIAL_PHYSICAL_WARP, - ), - ], -) -def test_unsupported_warp_membership(group, block_size, reason): - plan = _plan(group, _load_store(), block_size) - - assert plan.unsupported.code is reason - - -def test_different_algorithms_and_warp_widths_do_not_share_artifacts(): - plans = [ - *( - _plan(this_block(), _load_store(algorithm=algorithm)) - for algorithm in GroupLoadStoreAlgorithm - ), - *( - _plan(this_warp().group_by(width), _load_store()) - for width in (8, 16, 32) - ), - ] - - assert len({plan.artifact_key for plan in plans}) == len(plans) - - -def test_numpy_static_controls_share_the_same_artifact(): - plain = _load_store( - valid_items=ArgumentBinding.static(5), offset=ArgumentBinding.static(7) - ) - numpy = _load_store( - valid_items=ArgumentBinding.static(np.int32(5)), - offset=ArgumentBinding.static(np.int64(7)), - ) - - assert ( - _plan(this_block(), plain).artifact_key - == _plan(this_block(), numpy).artifact_key - ) - - -def test_offsets_reject_negative_constants_and_record_runtime_preconditions(): - with pytest.raises(ValueError, match="static offset must be nonnegative"): - _load_store(offset=ArgumentBinding.static(-1)) - - runtime = _plan(this_block(), _load_store(offset=ArgumentBinding.runtime())) - condition = runtime.participation.argument_preconditions[0] - assert condition.minimum == 0 - assert condition.enforcement is PreconditionEnforcement.CALLER diff --git a/python/cuda_coop/tests/test_group_resolution.py b/python/cuda_coop/tests/test_group_resolution.py deleted file mode 100644 index e393d30e61ca..000000000000 --- a/python/cuda_coop/tests/test_group_resolution.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import pytest - -from cuda.coop._core import ( - LaunchFacts, - UnsupportedReasonCode, - merge_launch_facts, - this_block, - this_cluster, - this_grid, - this_thread, -) -from tests._group_planning import _load_store, _plan - - -def test_maximum_block_dimensions_do_not_substitute_for_exact_dimensions(): - maximum = LaunchFacts(max_block_dim=(16, 8, 2)) - unsupported = _plan(this_block(), _load_store(), maximum) - assert ( - unsupported.unsupported.code - is UnsupportedReasonCode.MISSING_EXACT_BLOCK_DIM - ) - - merged = merge_launch_facts(maximum, LaunchFacts(exact_block_dim=(8, 4, 2))) - supported = _plan(this_block(), _load_store(), merged) - arguments = supported.implementation.template_arguments - assert tuple(arguments[f"BLOCK_DIM_{axis}"] for axis in "XYZ") == (8, 4, 2) - - -@pytest.mark.parametrize("group", [this_thread(), this_cluster(), this_grid()]) -def test_non_load_store_targets_are_typed_unsupported_before_resolution(group): - plan = _plan(group, _load_store(), LaunchFacts(exact_block_dim=48)) - - assert plan.unsupported.code is UnsupportedReasonCode.GROUP_KIND - assert plan.artifact_key is None diff --git a/python/cuda_coop/tests/test_header_bundle.py b/python/cuda_coop/tests/test_header_bundle.py deleted file mode 100644 index 96270c09ef81..000000000000 --- a/python/cuda_coop/tests/test_header_bundle.py +++ /dev/null @@ -1,309 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check how the wheel build records the source of its bundled headers. - -Temporary source trees exercise the real package CMake rules with small -header fixtures. The tests configure the package without building a wheel -or modifying the checkout that runs the tests. -""" - -from __future__ import annotations - -import json -import os -import shutil -import subprocess -from pathlib import Path - -import pytest - -_PACKAGE_ROOT = Path(__file__).parents[1] - - -def _require_packaging_tools() -> None: - required_tools = ("cmake", "git") - missing_tools = [ - tool for tool in required_tools if shutil.which(tool) is None - ] - if missing_tools: - pytest.skip( - "required packaging tools are unavailable: " - + ", ".join(missing_tools) - ) - - -def _isolated_git_env() -> dict[str, str]: - """Keep external Git settings and repository paths out of the fixture.""" - env = os.environ.copy() - for name in ( - "GIT_CEILING_DIRECTORIES", - "GIT_DIR", - "GIT_INDEX_FILE", - "GIT_WORK_TREE", - ): - env.pop(name, None) - env.update( - { - "GIT_CONFIG_GLOBAL": os.devnull, - "GIT_CONFIG_NOSYSTEM": "1", - "GIT_CONFIG_SYSTEM": os.devnull, - } - ) - return env - - -def _prepare_minimal_cccl_source(source_root: Path) -> Path: - """Create the source layout needed to configure the real package rules. - - Empty CCCL install rules avoid staging the full header tree. Individual - tests add only the header paths needed for their provenance check. - """ - package_root = source_root / "python" / "cuda_coop" - install_rules = source_root / "cmake" / "CCCLInstallRules.cmake" - package_root.mkdir(parents=True) - install_rules.parent.mkdir(parents=True) - shutil.copyfile( - _PACKAGE_ROOT / "CMakeLists.txt", package_root / "CMakeLists.txt" - ) - install_rules.touch() - return package_root - - -def _initialize_git_repository(source_root: Path, env: dict[str, str]) -> str: - """Commit the fixture's initial state and return its source revision. - - An empty template directory avoids inherited hooks. Explicit identity and - signing settings let the temporary repository work without user setup. - """ - subprocess.run( - [ - "git", - "-C", - source_root, - "-c", - "init.templateDir=", - "init", - "--quiet", - ], - check=True, - env=env, - ) - subprocess.run( - ["git", "-C", source_root, "add", "--all"], - check=True, - env=env, - ) - subprocess.run( - [ - "git", - "-C", - source_root, - "-c", - "commit.gpgsign=false", - "-c", - "user.name=cuda-coop packaging test", - "-c", - "user.email=cuda-coop-packaging-test@example.invalid", - "commit", - "--quiet", - "--message=initial source", - ], - check=True, - env=env, - ) - return subprocess.run( - ["git", "-C", source_root, "rev-parse", "HEAD"], - check=True, - env=env, - capture_output=True, - text=True, - ).stdout.strip() - - -def _configure_cuda_coop( - package_root: Path, - build_root: Path, - env: dict[str, str], - *definitions: str, -) -> subprocess.CompletedProcess[str]: - return subprocess.run( - [ - "cmake", - "-S", - package_root, - "-B", - build_root, - *definitions, - ], - check=False, - env=env, - capture_output=True, - text=True, - ) - - -def _bundle_provenance(build_root: Path) -> dict[str, str]: - return json.loads( - (build_root / "cccl-bundle-provenance.json").read_text(encoding="utf-8") - ) - - -def test_gitless_archive_does_not_inherit_enclosing_repository_revision( - tmp_path: Path, -) -> None: - _require_packaging_tools() - - outer_repository = tmp_path / "outer" - source_root = outer_repository / "cccl-archive" - package_root = _prepare_minimal_cccl_source(source_root) - build_root = tmp_path / "build" - - git_env = _isolated_git_env() - - subprocess.run( - ["git", "-c", "init.templateDir=", "init", "--quiet", outer_repository], - check=True, - env=git_env, - ) - subprocess.run( - [ - "git", - "-C", - outer_repository, - "-c", - "commit.gpgsign=false", - "-c", - "user.name=cuda-coop packaging test", - "-c", - "user.email=cuda-coop-packaging-test@example.invalid", - "commit", - "--quiet", - "--allow-empty", - "--message=outer repository", - ], - check=True, - env=git_env, - ) - - result = _configure_cuda_coop(package_root, build_root, git_env) - result.check_returncode() - - assert _bundle_provenance(build_root) == {"cccl_source_commit": "unknown"} - - -@pytest.mark.parametrize( - "change,relative_path", - [ - ("modified", Path("cub/cub/test.cuh")), - ("untracked", Path("libcudacxx/include/cuda/functional")), - ("ignored", Path("libcudacxx/include/cuda/hierarchy")), - ("deleted", Path("cmake/install/cub.cmake")), - ("modified", Path("python/cuda_coop/CMakeLists.txt")), - ], -) -def test_changed_header_bundle_fails_closed( - tmp_path: Path, change: str, relative_path: Path -) -> None: - _require_packaging_tools() - - source_root = tmp_path / "cccl" - package_root = _prepare_minimal_cccl_source(source_root) - header = source_root / relative_path - header.parent.mkdir(parents=True, exist_ok=True) - if change in {"modified", "deleted"} and not header.exists(): - header.write_text("// original\n", encoding="utf-8") - elif change == "ignored": - (source_root / ".gitignore").write_text( - f"/{relative_path.as_posix()}\n", encoding="utf-8" - ) - - git_env = _isolated_git_env() - _initialize_git_repository(source_root, git_env) - if change == "modified": - with header.open("a", encoding="utf-8") as stream: - stream.write("\n# modified\n") - elif change == "deleted": - header.unlink() - else: - header.write_text("// local\n", encoding="utf-8") - - result = _configure_cuda_coop( - package_root, - tmp_path / "build", - git_env, - f"-DCUDA_COOP_CCCL_SOURCE_REVISION={'2' * 40}", - ) - - assert result.returncode != 0 - diagnostic = result.stdout + result.stderr - assert "CCCL header bundle inputs contain local changes" in diagnostic - assert "CUDA_COOP_ALLOW_DIRTY_HEADER_BUNDLE=ON" in diagnostic - - -def test_allow_dirty_header_bundle_forces_unknown_provenance( - tmp_path: Path, -) -> None: - _require_packaging_tools() - - source_root = tmp_path / "cccl" - package_root = _prepare_minimal_cccl_source(source_root) - header = source_root / "cub" / "cub" / "test.cuh" - header.parent.mkdir(parents=True, exist_ok=True) - header.write_text("// original\n", encoding="utf-8") - git_env = _isolated_git_env() - _initialize_git_repository(source_root, git_env) - header.write_text("// modified\n", encoding="utf-8") - - claimed_revision = "1" * 40 - build_root = tmp_path / "build" - result = _configure_cuda_coop( - package_root, - build_root, - git_env, - "-DCUDA_COOP_ALLOW_DIRTY_HEADER_BUNDLE=ON", - f"-DCUDA_COOP_CCCL_SOURCE_REVISION={claimed_revision}", - ) - result.check_returncode() - - assert _bundle_provenance(build_root) == {"cccl_source_commit": "unknown"} - - -def test_unrelated_dirty_file_preserves_head_revision(tmp_path: Path) -> None: - _require_packaging_tools() - - source_root = tmp_path / "cccl" - package_root = _prepare_minimal_cccl_source(source_root) - git_env = _isolated_git_env() - revision = _initialize_git_repository(source_root, git_env) - unrelated = source_root / "docs" / "notes.md" - unrelated.parent.mkdir() - unrelated.write_text("local notes\n", encoding="utf-8") - - build_root = tmp_path / "build" - result = _configure_cuda_coop(package_root, build_root, git_env) - result.check_returncode() - - assert _bundle_provenance(build_root) == {"cccl_source_commit": revision} - - -def test_gitless_archive_accepts_explicit_source_revision( - tmp_path: Path, -) -> None: - _require_packaging_tools() - - source_root = tmp_path / "cccl-archive" - package_root = _prepare_minimal_cccl_source(source_root) - git_env = _isolated_git_env() - revision = "v1.2.3+archive" - build_root = tmp_path / "build" - result = _configure_cuda_coop( - package_root, - build_root, - git_env, - f"-DCUDA_COOP_CCCL_SOURCE_REVISION={revision}", - ) - result.check_returncode() - - assert _bundle_provenance(build_root) == {"cccl_source_commit": revision} diff --git a/python/cuda_coop/tests/test_header_identity.py b/python/cuda_coop/tests/test_header_identity.py deleted file mode 100644 index cfa791466f1f..000000000000 --- a/python/cuda_coop/tests/test_header_identity.py +++ /dev/null @@ -1,118 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check that cache identities track the selected headers and their contents. - -Header order, file contents, and source paths can affect generated code. -The fixtures isolate those changes without invoking a compiler. -""" - -from __future__ import annotations - -import os -from pathlib import Path - -import pytest - -from cuda.coop._headers import _identity -from cuda.coop._headers._identity import ( - HeaderIdentityError, - include_dirs_identity, -) - - -def _include_root(path: Path, content: bytes) -> Path: - path.mkdir() - (path / "primitive.cuh").write_bytes(content) - return path - - -def test_identity_preserves_include_order_and_provenance( - tmp_path: Path, -) -> None: - first = _include_root(tmp_path / "first", b"same header") - second = _include_root(tmp_path / "second", b"same header") - - forward = include_dirs_identity((str(first), str(second))) - reverse = include_dirs_identity((str(second), str(first))) - first_only = include_dirs_identity((str(first),)) - second_only = include_dirs_identity((str(second),)) - - assert tuple(root.path for root in forward.roots) == ( - str(first.resolve()), - str(second.resolve()), - ) - assert forward.digest != reverse.digest - assert first_only.roots[0].digest == second_only.roots[0].digest - assert first_only.digest != second_only.digest - - -@pytest.mark.parametrize("mutation", ["content", "addition", "deletion"]) -def test_identity_changes_with_header_tree( - tmp_path: Path, mutation: str -) -> None: - include = _include_root(tmp_path / "include", b"before") - header = include / "primitive.cuh" - original_stat = header.stat() - before = include_dirs_identity((str(include),)) - - if mutation == "content": - header.write_bytes(b"after!") - # Equal size and timestamps must not hide changed compiler inputs. - os.utime( - header, ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns) - ) - elif mutation == "addition": - (include / "added.cuh").write_bytes(b"added") - else: - header.unlink() - - assert include_dirs_identity((str(include),)).digest != before.digest - - -def test_recursive_identity_frames_file_contents_unambiguously( - tmp_path: Path, -) -> None: - combined = tmp_path / "combined" - combined.mkdir() - # One file contains bytes that resemble the next file's hash record. - # Length framing must distinguish this from two separate files. - (combined / "a").write_bytes(b"A\0" + b"1:b\0file\0B") - - split = tmp_path / "split" - split.mkdir() - (split / "a").write_bytes(b"A") - (split / "b").write_bytes(b"B") - - combined_identity = include_dirs_identity((str(combined),)) - split_identity = include_dirs_identity((str(split),)) - - assert combined_identity.roots[0].digest != split_identity.roots[0].digest - - -def test_identity_rejects_missing_include_root(tmp_path: Path) -> None: - missing = tmp_path / "missing" - - with pytest.raises(HeaderIdentityError, match="not a directory"): - include_dirs_identity((str(missing),)) - - -def test_missing_git_executable_falls_back_to_recursive_identity( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - repository = tmp_path / "repository" - (repository / ".git").mkdir(parents=True) - include = repository / "include" - include.mkdir() - (include / "header.cuh").write_bytes(b"header") - - def missing_git(*args: object, **kwargs: object) -> None: - raise FileNotFoundError("git") - - monkeypatch.setattr(_identity.subprocess, "run", missing_git) - - identity = include_dirs_identity((str(include),)) - - assert identity.roots[0].method == "recursive-content" diff --git a/python/cuda_coop/tests/test_headers.py b/python/cuda_coop/tests/test_headers.py deleted file mode 100644 index 0f58e7eca435..000000000000 --- a/python/cuda_coop/tests/test_headers.py +++ /dev/null @@ -1,180 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check header selection for source checkouts and installed packages. - -CCCL headers must come from one source. CUDA Toolkit headers are a separate -requirement when compiling; importing the shared API must remain possible -without loading CUDA bindings. -""" - -from __future__ import annotations - -import subprocess -import sys -from pathlib import Path - -import pytest - -import cuda.coop._headers as headers -from cuda.coop._headers import CoopIncludePaths, resolve_include_paths - -_PACKAGE_ROOT = Path(__file__).parents[1] - - -def _write_source_checkout(checkout: Path) -> None: - for path in ( - checkout / "thrust", - checkout / "cub" / "cub", - checkout / "libcudacxx" / "include", - ): - path.mkdir(parents=True) - (checkout / "cub" / "cub" / "version.cuh").write_text( - "// source probe\n", - encoding="utf-8", - ) - - -def test_environment_inside_checkout_uses_installed_headers( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - checkout = tmp_path / "cccl" - _write_source_checkout(checkout) - installed_module = ( - checkout - / ".venv" - / "lib" - / "python3.14" - / "site-packages" - / "cuda" - / "coop" - / "_headers" - / "__init__.py" - ) - installed_module.parent.mkdir(parents=True) - installed_module.touch() - installed_bundle = tmp_path / "installed-bundle" - installed_bundle.mkdir() - expected = CoopIncludePaths( - cccl=(installed_bundle,), - cuda=(), - origin="installed test bundle", - ) - monkeypatch.setattr(headers, "_installed_include_paths", lambda: expected) - - paths = resolve_include_paths(start=installed_module) - - assert paths == expected - - -def test_source_resolution_uses_one_coherent_header_set() -> None: - paths = resolve_include_paths( - start=Path(__file__), - required_headers=( - "cub/block/block_load.cuh", - "thrust/detail/raw_pointer_cast.h", - "cuda/hierarchy", - "cuda/std/cstdint", - ), - ) - checkout = _PACKAGE_ROOT.parents[1] - assert paths.cccl == ( - checkout / "thrust", - checkout / "cub", - checkout / "libcudacxx" / "include", - ) - - -def test_incomplete_source_checkout_reports_missing_headers( - tmp_path: Path, -) -> None: - checkout = tmp_path / "cccl" - probe = checkout / "cub" / "cub" / "version.cuh" - probe.parent.mkdir(parents=True) - probe.touch() - source_module = checkout / "python" / "cuda_coop" / "cuda" / "coop" - source_module.mkdir(parents=True) - - with pytest.raises( - headers.HeaderResolutionError, match="missing source include roots" - ): - resolve_include_paths(start=source_module) - - -def test_import_cuda_coop_does_not_import_cuda_bindings() -> None: - script = f""" -import sys -sys.path.insert(0, {str(_PACKAGE_ROOT)!r}) -import cuda.coop -unexpected = sorted( - name for name in sys.modules - if name == "cuda.bindings" or name.startswith("cuda.bindings.") -) -if unexpected: - raise RuntimeError(f"cuda.coop eagerly imported {{unexpected}}") -""" - - result = subprocess.run( - [sys.executable, "-I", "-c", script], - check=False, - capture_output=True, - text=True, - ) - - assert result.returncode == 0, result.stderr - - -@pytest.mark.parametrize( - "platform,configured", [("linux", False), ("win32", True)] -) -def test_cuda_header_fallback_respects_platform( - tmp_path, monkeypatch, platform, configured -): - from types import SimpleNamespace - - import cuda.pathfinder - - monkeypatch.setattr(headers, "sys", SimpleNamespace(platform=platform)) - monkeypatch.setattr( - cuda.pathfinder, "find_nvidia_header_directory", lambda _: None - ) - for name in ("CUDA_PATH", "CUDA_HOME", "CUDA_ROOT"): - monkeypatch.delenv(name, raising=False) - include = tmp_path / "CUDA Toolkit" / "include" - if configured: - include.mkdir(parents=True) - (include / "cuda_runtime.h").touch() - monkeypatch.setenv("CUDA_PATH", str(include.parent)) - calls = [] - select = headers._select_cuda_include_path - - def record(paths): - paths = tuple(paths) - calls.append(paths) - # Avoid relying on this host's fallback toolkit installation. - return select(path for path in paths if path == include) - - monkeypatch.setattr(headers, "_select_cuda_include_path", record) - assert headers._cuda_include_paths() == ((include,) if configured else ()) - fallback = (Path("/usr/local/cuda/include"),) if platform == "linux" else () - assert calls[-1] == ((include,) if configured else ()) + fallback - - -def test_required_header_diagnostic_never_falls_back_to_toolkit() -> None: - with pytest.raises( - headers.HeaderResolutionError, match="does not fall back" - ): - resolve_include_paths( - start=Path(__file__), - required_headers=("cub/block/not_a_primitive.cuh",), - ) - - -def test_cuda_headers_are_required_only_when_compiling() -> None: - paths = CoopIncludePaths( - cccl=(Path("cccl/include"),), cuda=(), origin="test" - ) - with pytest.raises(headers.HeaderResolutionError, match="cuda_runtime.h"): - paths.as_tuple() diff --git a/python/cuda_coop/tests/test_registration.py b/python/cuda_coop/tests/test_registration.py deleted file mode 100644 index 8c29b7f39911..000000000000 --- a/python/cuda_coop/tests/test_registration.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import pytest - -from cuda import coop -from cuda.coop import _registration - - -def test_register_rejects_unsupported_backend_without_importing(monkeypatch): - def unexpected_import(name): - pytest.fail(f"unsupported backend attempted to import {name}") - - monkeypatch.setattr( - _registration.importlib, "import_module", unexpected_import - ) - - with pytest.raises(ValueError, match="Unsupported cuda.coop backend"): - coop.register("cutlass") - - -@pytest.mark.parametrize("backend", ["numba-cuda-mlir", "numba_cuda_mlir"]) -def test_register_reports_an_unavailable_adapter(monkeypatch, backend): - def missing_adapter(name): - raise ModuleNotFoundError("missing adapter", name=name) - - monkeypatch.setattr( - _registration.importlib, "import_module", missing_adapter - ) - - with pytest.raises( - ImportError, match="does not include the Numba-CUDA-MLIR adapter" - ): - coop.register(backend) - - -def test_register_preserves_missing_backend_dependency(monkeypatch): - error = ModuleNotFoundError("missing dependency", name="backend_dependency") - - def broken_adapter(name): - raise error - - monkeypatch.setattr( - _registration.importlib, "import_module", broken_adapter - ) - - with pytest.raises(ModuleNotFoundError) as exc_info: - coop.register("numba-cuda-mlir") - - assert exc_info.value is error diff --git a/python/cuda_coop/tests/test_source_dump.py b/python/cuda_coop/tests/test_source_dump.py deleted file mode 100644 index af46871b53f7..000000000000 --- a/python/cuda_coop/tests/test_source_dump.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from concurrent.futures import ThreadPoolExecutor - -from cuda.coop._core._source_dump import dump_source - - -def test_source_dump_distinguishes_source_target_and_backend( - tmp_path, monkeypatch -): - monkeypatch.setenv("CUDA_COOP_SOURCE_DUMP_DIR", str(tmp_path)) - source = "// CUDA source with a Unicode comment: π\n" - first = dump_source(source, backend="numba_mlir", identity=(90, "lto")) - repeated = dump_source(source, backend="numba_mlir", identity=(90, "lto")) - changed_source = dump_source( - "// different\n", backend="numba_mlir", identity=(90, "lto") - ) - changed_target = dump_source( - source, backend="numba_mlir", identity=(120, "lto") - ) - changed_backend = dump_source( - source, backend="cutlass", identity=(90, "lto") - ) - - assert first == repeated - assert len({first, changed_source, changed_target, changed_backend}) == 4 - assert first.read_bytes() == source.encode("utf-8") - - -def test_source_dump_disabled_by_default(tmp_path, monkeypatch): - monkeypatch.delenv("CUDA_COOP_SOURCE_DUMP_DIR", raising=False) - monkeypatch.chdir(tmp_path) - assert dump_source("// source", backend="numba_mlir") is None - assert not list(tmp_path.iterdir()) - - -def test_concurrent_source_dumps_leave_complete_files(tmp_path, monkeypatch): - monkeypatch.setenv("CUDA_COOP_SOURCE_DUMP_DIR", str(tmp_path)) - source = "// shared translation unit\n" * 1000 - - def write(_): - return dump_source(source, backend="numba_mlir") - - with ThreadPoolExecutor(max_workers=4) as executor: - paths = set(executor.map(write, range(12))) - - assert len(paths) == 1 - assert set(tmp_path.iterdir()) == paths - assert paths.pop().read_bytes() == source.encode("utf-8") diff --git a/python/cuda_coop/tests/test_thread_group.py b/python/cuda_coop/tests/test_thread_group.py deleted file mode 100644 index 5bf59838788a..000000000000 --- a/python/cuda_coop/tests/test_thread_group.py +++ /dev/null @@ -1,127 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -from typing import Literal, get_type_hints - -import pytest - -import cuda.coop._core.thread_group as _thread_group -from cuda.coop._core import ( - ThreadGroup, - ThreadHierarchy, - make_thread_group, - this_block, - this_warp, -) - - -def test_thread_group_annotations_resolve_at_runtime(): - def kernel(group: ThreadGroup[Literal["block"]]): - pass - - assert get_type_hints(kernel)["group"] == ThreadGroup[Literal["block"]] - - -def test_resolved_hierarchy_counts_threads_at_each_level(): - hierarchy = ThreadHierarchy._resolved( - block_dim=(8, 4), grid_dim=3, cluster_dim=2 - ) - - assert ThreadGroup(kind="block", hierarchy=hierarchy).static_size == 32 - assert ThreadGroup(kind="cluster", hierarchy=hierarchy).static_size == 64 - assert ThreadGroup(kind="grid", hierarchy=hierarchy).static_size == 192 - - -def test_current_group_resolution_preserves_backend_type_and_cache_identity(): - class BackendGroup(ThreadGroup): - pass - - current = make_thread_group("block", group_type=BackendGroup) - assert current.static_size is None - - hierarchy = ThreadHierarchy._resolved(block_dim=(8, 4)) - resolved = current.with_hierarchy(hierarchy, source="inferred_launch") - assert type(resolved) is BackendGroup - assert resolved.static_size == 32 - assert current not in {resolved} - assert BackendGroup(kind="block", hierarchy=hierarchy) in {resolved} - - -def test_physical_group_rendering_uses_current_or_resolved_hierarchy(): - current = this_block() - assert "implicit_hierarchy()" in _thread_group.render_group_decl(current) - - hierarchy = ThreadHierarchy._resolved(block_dim=(8, 4)) - resolved = current.with_hierarchy(hierarchy) - assert "block_dims<8, 4>()" in "\n".join( - _thread_group.render_hierarchy_decl(hierarchy) - ) - assert "this_block group{hierarchy}" in _thread_group.render_group_decl( - resolved - ) - - -@pytest.mark.parametrize( - ("kind", "count", "exhaustive", "synchronizer"), - [ - ("warp", 8, True, "lane_synchronizer{}"), - ("block", 2, False, "barrier_synchronizer{group_barriers}"), - ], -) -def test_mapped_group_rendering_uses_membership_and_synchronizer( - kind, count, exhaustive, synchronizer -): - group = ThreadGroup( - kind=kind, hierarchy=ThreadHierarchy._resolved(block_dim=160) - ).group_by(count, exhaustive=exhaustive) - source = "\n".join(_thread_group.render_group_decl_lines(group)) - - assert f"group_by<{count}, {str(exhaustive).lower()}>" in source - assert synchronizer in source - if kind == "block": - assert "barrier<::cuda::thread_scope_block>[2]" in source - - -def test_partial_warp_mapping_excludes_remainder_lanes(): - group = this_warp().group_by(12, exhaustive=False) - - assert group.static_size == 12 - assert group.groups_per_parent == 2 - assert group.remainder_count == 8 - - -def test_block_mapping_waits_for_launch_dimensions(): - current = this_block().group_by(3, exhaustive=False) - assert current.static_size == 96 - assert current.groups_per_parent is None - - resolved = current.with_hierarchy(ThreadHierarchy._resolved(block_dim=320)) - assert resolved.groups_per_parent == 3 - assert resolved.remainder_count == 1 - - -@pytest.mark.parametrize( - ("block_threads", "count", "exhaustive", "message"), - [ - (48, 1, False, "complete warps"), - (64, 3, False, "cannot exceed the parent warp count"), - (128, 3, True, "requires the count to divide"), - ], -) -def test_group_by_revalidates_deferred_constraints( - block_threads, count, exhaustive, message -): - current = this_block().group_by(count, exhaustive=exhaustive) - - with pytest.raises(ValueError, match=message): - current.with_hierarchy( - ThreadHierarchy._resolved(block_dim=block_threads) - ) - - -def test_group_by_requires_exhaustive_non_nested_membership(): - with pytest.raises(ValueError, match="requires the count to divide"): - this_warp().group_by(12) - with pytest.raises(NotImplementedError, match="nested"): - this_warp().group_by(8).group_by(2) diff --git a/python/cuda_coop/tests/test_toolkit.py b/python/cuda_coop/tests/test_toolkit.py deleted file mode 100644 index bc14a7441822..000000000000 --- a/python/cuda_coop/tests/test_toolkit.py +++ /dev/null @@ -1,522 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -"""Check CUDA Toolkit library selection with temporary files and fake loaders. - -The selected libraries must come from the header installation and have -compatible versions. The fakes record load order and emulate version queries -so the tests can cover Linux and Windows without loading CUDA. -""" - -from __future__ import annotations - -import os -from pathlib import Path -from types import SimpleNamespace - -import cuda.pathfinder -import pytest - -from cuda.coop._headers import _toolkit - - -class _VersionFunction: - """Emulate a C version query that writes through two integer pointers. - - Assert the ctypes signature because NVRTC and nvJitLink use different - integer types for their version outputs. - """ - - def __init__(self, version, integer_type): - self.version = version - self.pointer_type = _toolkit.ctypes.POINTER(integer_type) - self.argtypes = None - self.restype = None - - def __call__(self, major, minor): - assert self.argtypes == (self.pointer_type, self.pointer_type) - assert self.restype is _toolkit.ctypes.c_int - _toolkit.ctypes.cast(major, self.pointer_type)[0] = self.version[0] - _toolkit.ctypes.cast(minor, self.pointer_type)[0] = self.version[1] - return 0 - - -@pytest.fixture(autouse=True) -def _isolated_process_toolkit_state(): - """Reset fake library handles and toolkit selection around each test.""" - with _toolkit._PRELOAD_LOCK: - _toolkit._EXACT_LIBRARY_HANDLES.clear() - _toolkit._PROCESS_TOOLKIT_SELECTION = None - yield - with _toolkit._PRELOAD_LOCK: - _toolkit._EXACT_LIBRARY_HANDLES.clear() - _toolkit._PROCESS_TOOLKIT_SELECTION = None - - -def _write_cuda_header(include_dir: Path, encoded_version: int) -> None: - include_dir.mkdir(parents=True) - (include_dir / "cuda_runtime_api.h").write_text( - f"#define CUDART_VERSION {encoded_version}\n", - encoding="utf-8", - ) - - -def _library_path(directory: Path, kind: str, major: int) -> Path: - if _toolkit.os.name == "nt": - name = ( - f"nvrtc64_{major}0_0.dll" - if kind == "nvrtc" - else f"nvJitLink_{major}0_0.dll" - ) - else: - name = f"lib{kind}.so.{major}" - return directory / name - - -def _builtins_path(directory: Path, major: int, minor: int) -> Path: - name = ( - f"nvrtc-builtins64_{major}{minor}.dll" - if _toolkit.os.name == "nt" - else f"libnvrtc-builtins.so.{major}.{minor}" - ) - return directory / name - - -def _write_complete_toolkit( - root: Path, - encoded_version: int, - *, - library_dir_name: str = "lib", -) -> dict[str, Path]: - """Create a version header and empty libraries under one toolkit root.""" - major = encoded_version // 1000 - minor = (encoded_version % 1000) // 10 - include_dir = root / "include" - lib_dir = root / library_dir_name - _write_cuda_header(include_dir, encoded_version) - lib_dir.mkdir(parents=True) - paths = { - "root": root, - "include": include_dir, - "nvrtc": _library_path(lib_dir, "nvrtc", major), - "builtins": _builtins_path(lib_dir, major, minor), - "nvjitlink": _library_path(lib_dir, "nvJitLink", major), - } - for name in ("nvrtc", "builtins", "nvjitlink"): - paths[name].touch() - return paths - - -def _write_split_wheel_toolkit( - nvidia_root: Path, - encoded_version: int, - *, - library_dir_name: str = "lib", -) -> dict[str, Path]: - """Create sibling runtime, NVRTC, and nvJitLink wheel directories.""" - major = encoded_version // 1000 - minor = (encoded_version % 1000) // 10 - include_dir = nvidia_root / "cuda_runtime" / "include" - nvrtc_dir = nvidia_root / "cuda_nvrtc" / library_dir_name - nvjitlink_dir = nvidia_root / "nvjitlink" / library_dir_name - _write_cuda_header(include_dir, encoded_version) - nvrtc_dir.mkdir(parents=True) - nvjitlink_dir.mkdir(parents=True) - paths = { - "root": nvidia_root, - "include": include_dir, - "nvrtc": _library_path(nvrtc_dir, "nvrtc", major), - "builtins": _builtins_path(nvrtc_dir, major, minor), - "nvjitlink": _library_path(nvjitlink_dir, "nvJitLink", major), - } - for name in ("nvrtc", "builtins", "nvjitlink"): - paths[name].touch() - return paths - - -def _patch_compiler_loaders( - monkeypatch: pytest.MonkeyPatch, - paths: dict[str, Path], - *, - nvrtc_version: tuple[int, int], - nvjitlink_version: tuple[int, int], - actual_nvrtc: Path | None = None, - actual_nvjitlink: Path | None = None, - failures: dict[Path, str] | None = None, -) -> tuple[list[Path], list[str]]: - """Replace compiler loaders and return their call logs for assertions. - - The ctypes fake uses the requested versions and optional load failures. - The pathfinder fake can report a different loaded path to test a mismatch - with the selected toolkit. Neither fake opens a native library. - """ - attempts: list[Path] = [] - pathfinder_calls: list[str] = [] - failures = failures or {} - - def load_library(path, *, mode): - assert mode == getattr(_toolkit.ctypes, "RTLD_GLOBAL", 0) - candidate = Path(path) - attempts.append(candidate) - if candidate in failures: - raise OSError(failures[candidate]) - if candidate == paths["nvrtc"]: - return SimpleNamespace( - nvrtcVersion=_VersionFunction( - nvrtc_version, _toolkit.ctypes.c_int - ) - ) - if candidate.name == paths["nvjitlink"].name: - return SimpleNamespace( - nvJitLinkVersion=_VersionFunction( - nvjitlink_version, _toolkit.ctypes.c_uint - ) - ) - return object() - - def load_with_pathfinder(kind: str): - pathfinder_calls.append(kind) - return SimpleNamespace( - abs_path=str( - (actual_nvrtc or paths["nvrtc"]) - if kind == "nvrtc" - else (actual_nvjitlink or paths["nvjitlink"]) - ) - ) - - monkeypatch.setattr(_toolkit.ctypes, "CDLL", load_library) - monkeypatch.setattr( - cuda.pathfinder, - "load_nvidia_dynamic_lib", - load_with_pathfinder, - ) - return attempts, pathfinder_calls - - -@pytest.mark.parametrize( - "os_name,library_dir_name", - [ - ("posix", "lib"), - ("nt", "bin"), - ("nt", "bin/x64"), - ("nt", "bin/x86_64"), - ], -) -def test_preload_loads_one_same_root_monolithic_set_and_reuses_exact_handles( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - os_name: str, - library_dir_name: str, -) -> None: - monkeypatch.setattr( - _toolkit, "os", SimpleNamespace(name=os_name, path=os.path) - ) - paths = _write_complete_toolkit( - tmp_path / "toolkit", 13020, library_dir_name=library_dir_name - ) - attempts, _ = _patch_compiler_loaders( - monkeypatch, - paths, - nvrtc_version=(13, 2), - nvjitlink_version=(13, 4), - ) - - first = _toolkit.preload_toolkit_compiler_libraries((paths["include"],)) - second = _toolkit.preload_toolkit_compiler_libraries((paths["include"],)) - - assert attempts == [paths["builtins"], paths["nvrtc"], paths["nvjitlink"]] - assert first == second - assert first.toolkit_version == (13, 2) - assert first.nvjitlink_version == (13, 4) - - -@pytest.mark.parametrize( - "os_name,library_dir_name", [("posix", "lib"), ("nt", "bin")] -) -def test_preload_loads_split_wheel_set_from_one_nvidia_anchor( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - os_name: str, - library_dir_name: str, -) -> None: - monkeypatch.setattr( - _toolkit, "os", SimpleNamespace(name=os_name, path=os.path) - ) - paths = _write_split_wheel_toolkit( - tmp_path / "site-packages" / "nvidia", - 12090, - library_dir_name=library_dir_name, - ) - attempts, _ = _patch_compiler_loaders( - monkeypatch, - paths, - nvrtc_version=(12, 9), - nvjitlink_version=(12, 9), - ) - - libraries = _toolkit.preload_toolkit_compiler_libraries((paths["include"],)) - - assert attempts == [paths["builtins"], paths["nvrtc"], paths["nvjitlink"]] - assert libraries.toolkit_version == (12, 9) - - -def test_preload_rejects_split_wheel_components_from_later_nvidia_anchor( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - first = _write_split_wheel_toolkit( - tmp_path / "first" / "site-packages" / "nvidia", - 12090, - ) - second = _write_split_wheel_toolkit( - tmp_path / "second" / "site-packages" / "nvidia", - 12090, - ) - first["nvjitlink"].unlink() - attempts: list[Path] = [] - monkeypatch.setattr( - _toolkit.ctypes, - "CDLL", - lambda path, *, mode: attempts.append(Path(path)), - ) - monkeypatch.setattr( - cuda.pathfinder, - "load_nvidia_dynamic_lib", - lambda kind: pytest.fail(f"unexpected fallback load for {kind}"), - ) - - with pytest.raises( - RuntimeError, - match=( - "require NVRTC, nvrtc-builtins, and nvJitLink " - "from one CUDA Toolkit root" - ), - ): - _toolkit.preload_toolkit_compiler_libraries( - (first["include"], second["include"]) - ) - - assert attempts == [] - - -def test_preload_rejects_nvrtc_without_adjacent_builtins( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - root = tmp_path / "toolkit" - include_dir = root / "include" - lib_dir = root / "lib" - _write_cuda_header(include_dir, 13020) - lib_dir.mkdir() - _library_path(lib_dir, "nvrtc", 13).touch() - _library_path(lib_dir, "nvJitLink", 13).touch() - monkeypatch.setattr( - _toolkit.ctypes, - "CDLL", - lambda path, *, mode: pytest.fail(f"unexpected load of {path}"), - ) - - with pytest.raises(RuntimeError, match="without adjacent builtins"): - _toolkit.preload_toolkit_compiler_libraries((include_dir,)) - - -def test_preload_stops_if_nvrtc_fails_after_builtins_becomes_global( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - paths = _write_complete_toolkit(tmp_path / "toolkit", 13020) - attempts, pathfinder_calls = _patch_compiler_loaders( - monkeypatch, - paths, - nvrtc_version=(13, 2), - nvjitlink_version=(13, 2), - failures={paths["nvrtc"]: "NVRTC load failed"}, - ) - - with pytest.raises( - RuntimeError, - match="failed loading same-root NVRTC after its builtins library", - ): - _toolkit.preload_toolkit_compiler_libraries((paths["include"],)) - - assert attempts == [paths["builtins"], paths["nvrtc"]] - assert pathfinder_calls == [] - assert _toolkit._PROCESS_TOOLKIT_SELECTION == ( - str(paths["root"].resolve()), - (13, 2), - ) - - -def test_preload_tries_nvjitlink_candidates_only_within_selected_root( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - paths = _write_complete_toolkit(tmp_path / "toolkit", 13020) - lib64 = paths["root"] / "lib64" - lib64.mkdir() - second_nvjitlink = _library_path(lib64, "nvJitLink", 13) - second_nvjitlink.touch() - attempts, _ = _patch_compiler_loaders( - monkeypatch, - paths, - nvrtc_version=(13, 2), - nvjitlink_version=(13, 3), - actual_nvjitlink=second_nvjitlink, - failures={paths["nvjitlink"]: "first candidate unavailable"}, - ) - - libraries = _toolkit.preload_toolkit_compiler_libraries((paths["include"],)) - - assert attempts == [ - paths["builtins"], - paths["nvrtc"], - paths["nvjitlink"], - second_nvjitlink, - ] - assert libraries.nvjitlink_path == str(second_nvjitlink) - - -@pytest.mark.parametrize("kind", ["nvrtc", "nvJitLink"]) -def test_preload_rejects_pathfinder_library_from_another_root( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - kind: str, -) -> None: - paths = _write_complete_toolkit(tmp_path / "toolkit", 13020) - version = (13, 2) - mismatched = ( - tmp_path - / "other" - / (paths["nvrtc"].name if kind == "nvrtc" else paths["nvjitlink"].name) - ) - _patch_compiler_loaders( - monkeypatch, - paths, - nvrtc_version=version, - nvjitlink_version=version, - actual_nvrtc=mismatched if kind == "nvrtc" else None, - actual_nvjitlink=mismatched if kind == "nvJitLink" else None, - ) - - with pytest.raises( - RuntimeError, - match=rf"process uses .* for {kind}", - ): - _toolkit.preload_toolkit_compiler_libraries((paths["include"],)) - - -def test_preload_rejects_nvrtc_version_mismatched_with_headers( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - paths = _write_complete_toolkit(tmp_path / "toolkit", 13020) - _patch_compiler_loaders( - monkeypatch, - paths, - nvrtc_version=(13, 1), - nvjitlink_version=(13, 2), - ) - - with pytest.raises( - RuntimeError, - match=( - r"headers report Toolkit 13\.2, " - r"but loaded NVRTC .* reports 13\.1" - ), - ): - _toolkit.preload_toolkit_compiler_libraries((paths["include"],)) - - -@pytest.mark.parametrize("actual_version", [(12, 8), (13, 1)]) -def test_preload_rejects_incompatible_same_root_nvjitlink_version( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - actual_version: tuple[int, int], -) -> None: - paths = _write_complete_toolkit(tmp_path / "toolkit", 13020) - _patch_compiler_loaders( - monkeypatch, - paths, - nvrtc_version=(13, 2), - nvjitlink_version=actual_version, - ) - - with pytest.raises( - RuntimeError, - match=( - r"headers report Toolkit 13\.2, but loaded nvJitLink .* reports " - rf"{actual_version[0]}\.{actual_version[1]}" - ), - ): - _toolkit.preload_toolkit_compiler_libraries((paths["include"],)) - - -def test_process_global_toolkit_selection_rejects_a_later_root( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - first = _write_complete_toolkit(tmp_path / "first", 13020) - second = _write_complete_toolkit(tmp_path / "second", 13020) - _patch_compiler_loaders( - monkeypatch, - first, - nvrtc_version=(13, 2), - nvjitlink_version=(13, 2), - ) - _toolkit.preload_toolkit_compiler_libraries((first["include"],)) - - with pytest.raises(RuntimeError, match="resolved headers select different"): - _toolkit.preload_toolkit_compiler_libraries((second["include"],)) - - -def test_toolkit_version_rejects_disagreeing_header_roots( - tmp_path: Path, -) -> None: - first = tmp_path / "first" / "include" - second = tmp_path / "second" / "include" - _write_cuda_header(first, 13020) - _write_cuda_header(second, 12080) - - with pytest.raises(RuntimeError, match="disagree on Toolkit version"): - _toolkit.preload_toolkit_compiler_libraries((first, second)) - - -def test_toolkit_version_rejects_unparseable_selected_header( - tmp_path: Path, -) -> None: - include_dir = tmp_path / "toolkit" / "include" - include_dir.mkdir(parents=True) - (include_dir / "cuda_runtime_api.h").write_text( - "// no CUDART_VERSION\n", - encoding="utf-8", - ) - - with pytest.raises(RuntimeError, match="failed parsing CUDART_VERSION"): - _toolkit.preload_toolkit_compiler_libraries((include_dir,)) - - -@pytest.mark.skipif(os.name == "nt", reason="requires POSIX symbolic links") -def test_preload_rejects_library_symlink_escaping_toolkit_root( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - paths = _write_complete_toolkit(tmp_path / "toolkit", 13020) - version = (13, 2) - outside = tmp_path / "outside" / paths["nvjitlink"].name - outside.parent.mkdir() - outside.touch() - paths["nvjitlink"].unlink() - paths["nvjitlink"].symlink_to(outside) - attempts, pathfinder_calls = _patch_compiler_loaders( - monkeypatch, - paths, - nvrtc_version=version, - nvjitlink_version=version, - ) - - with pytest.raises(RuntimeError, match="library escapes"): - _toolkit.preload_toolkit_compiler_libraries((paths["include"],)) - - assert attempts == [paths["builtins"], paths["nvrtc"]] - assert pathfinder_calls == [] diff --git a/python/cuda_coop/tests/test_types.py b/python/cuda_coop/tests/test_types.py deleted file mode 100644 index adbd30519b5d..000000000000 --- a/python/cuda_coop/tests/test_types.py +++ /dev/null @@ -1,127 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import struct -from dataclasses import replace -from functools import partial - -import pytest - -from cuda.coop._core import ( - Algorithm, - ArgumentKind, - Dependency, - TemplateParameter, - Value, - semantic_token, -) - - -def _algorithm(*, template_arguments, metadata=None): - return Algorithm( - template_arguments=template_arguments, - metadata={} if metadata is None else metadata, - struct_name="BlockExample", - method_name="Run", - c_name="block_example", - includes=(), - template_parameters=(TemplateParameter("T"),), - parameters=((Value(Dependency("T"), name="value"),),), - ) - - -def test_algorithm_identity_distinguishes_specializations(): - base = _algorithm( - template_arguments={"T": "int"}, metadata={"mode": "base"} - ) - variants = ( - replace(base, template_arguments={"T": "float"}), - replace(base, metadata={"mode": "alternate"}), - replace(base, method_name="Other"), - ) - - cache = {base: "compiled"} - equivalent = _algorithm( - template_arguments={"T": "int"}, metadata={"mode": "base"} - ) - assert cache[equivalent] == "compiled" - for variant in variants: - assert variant not in cache - - -@pytest.mark.parametrize( - "callback", - [ - partial(lambda value: value, 1), - staticmethod(lambda value: value), - classmethod(lambda cls, value: value), - ], -) -def test_wrapped_callbacks_raise(callback): - with pytest.raises(TypeError, match="callable values are not supported"): - _algorithm(template_arguments={"T": "int", "callback": callback}) - - -@pytest.mark.parametrize("value", (False, True)) -def test_algorithm_identity_distinguishes_boolean_and_integer_settings(value): - boolean = _algorithm( - template_arguments={"T": "int", "settings": {"flag": (value,)}} - ) - integer = _algorithm( - template_arguments={"T": "int", "settings": {"flag": (int(value),)}} - ) - - assert len({boolean: "boolean", integer: "integer"}) == 2 - - -def test_algorithm_template_arguments_freeze_nested_semantic_containers(): - nested = {"values": [1], "modes": {"direct"}} - specialization = _algorithm( - template_arguments={"T": "int", "settings": nested} - ) - equivalent = _algorithm(template_arguments={"T": "int", "settings": nested}) - cache = {specialization: "compiled"} - - nested["values"].append(2) - nested["modes"].add("striped") - - assert cache[equivalent] == "compiled" - assert specialization.template_arguments["settings"] == { - "values": (1,), - "modes": frozenset({"direct"}), - } - - -def test_algorithm_template_arguments_reject_container_cycles(): - cyclic = [] - cyclic.append(cyclic) - - with pytest.raises(ValueError, match="container cycles"): - _algorithm(template_arguments={"T": "int", "settings": cyclic}) - - -def test_algorithm_requires_all_template_arguments(): - with pytest.raises( - ValueError, match="Template argument\\(s\\) not provided: T" - ): - _algorithm(template_arguments={}) - - -@pytest.mark.parametrize( - ("left", "right"), - [ - (ArgumentKind.STATIC, "static"), - (False, 0), - (True, 1), - (0.0, -0.0), - (1.0, 1), - (float("nan"), -float("nan")), - ( - struct.unpack(">d", bytes.fromhex("7ff8000000000001"))[0], - struct.unpack(">d", bytes.fromhex("7ff8000000000002"))[0], - ), - ], -) -def test_semantic_token_preserves_scalar_identity(left, right): - assert semantic_token(left) != semantic_token(right) diff --git a/python/cuda_coop/tests/test_warp_load_store.py b/python/cuda_coop/tests/test_warp_load_store.py deleted file mode 100644 index 99942215d419..000000000000 --- a/python/cuda_coop/tests/test_warp_load_store.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -import numpy as np -import pytest - -from cuda.coop._core import ArgumentBinding -from cuda.coop._core.warp import ( - make_warp_load_specialization, - make_warp_load_store_semantics, - make_warp_store_specialization, -) - - -@pytest.mark.parametrize( - "threads_in_warp", [True, 0, 3, 64, 8.0, np.int64(3), np.bool_(True)] -) -def test_warp_load_store_rejects_unsupported_widths(threads_in_warp): - with pytest.raises(ValueError, match="threads_in_warp in"): - make_warp_load_specialization( - dtype="i32", - items_per_thread=2, - threads_in_warp=threads_in_warp, - algorithm="direct", - ) - - -@pytest.mark.parametrize( - "make_specialization", - [make_warp_load_specialization, make_warp_store_specialization], -) -def test_numpy_integer_controls_share_the_same_specialization( - make_specialization, -): - def build(integer): - return make_specialization( - dtype="i32", - items_per_thread=integer(2), - threads_in_warp=integer(8), - algorithm="direct", - ) - - assert build(int).semantic_key == build(np.int64).semantic_key - - -@pytest.mark.parametrize( - "make_specialization", - [make_warp_load_specialization, make_warp_store_specialization], -) -@pytest.mark.parametrize("valid_items", [-1, 17]) -def test_warp_load_store_rejects_static_valid_items_outside_tile( - make_specialization, valid_items -): - with pytest.raises(ValueError, match="warp tile size"): - make_specialization( - dtype="i32", - items_per_thread=2, - threads_in_warp=8, - algorithm="direct", - valid_items=ArgumentBinding.static(valid_items), - ) - - -@pytest.mark.parametrize( - ("overrides", "message"), - [ - ({"algorithm": "warp_transpose"}, "unsupported WarpLoad algorithm"), - ( - {"items_per_thread": True}, - "items_per_thread must be a positive integer", - ), - ( - {"items_per_thread": 2.0}, - "items_per_thread must be a positive integer", - ), - ( - {"kind": "store", "valid_items": True, "oob_default": True}, - "only valid for WarpLoad", - ), - ({"oob_default": True}, "requires a valid_items"), - ], -) -def test_warp_load_store_rejects_invalid_options(overrides, message): - options = { - "kind": "load", - "dtype": "i32", - "items_per_thread": 1, - "algorithm": "direct", - } - options.update(overrides) - with pytest.raises(ValueError, match=message): - make_warp_load_store_semantics(**options) diff --git a/python/cuda_coop/tests/typing/pyright-unused.json b/python/cuda_coop/tests/typing/pyright-unused.json deleted file mode 100644 index 1402fda975d2..000000000000 --- a/python/cuda_coop/tests/typing/pyright-unused.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "include": [ - "../../cuda/coop" - ], - "extraPaths": [ - "../.." - ], - "typeCheckingMode": "off", - "pythonVersion": "3.10", - "reportUnusedFunction": "error", - "reportMissingImports": "none", - "reportMissingModuleSource": "none" -} From 1ff52757b0177bf9891dc7ac4ba33d0f1a8ff065 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 27 Oct 2026 12:54:25 -0700 Subject: [PATCH 2/2] Keep nested C++ license changes out of CI builds Match the remaining root and C++ license paths explicitly after removing the cooperative package. Preserve Python package license coverage. Signed-off-by: Trent Nelson --- ci/project_files_and_dependencies.yaml | 7 ++++++- ci/test/inspect_changes/ignored_only.dirty_files | 6 ++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/ci/project_files_and_dependencies.yaml b/ci/project_files_and_dependencies.yaml index d2b7a763e596..ae29192a15ac 100644 --- a/ci/project_files_and_dependencies.yaml +++ b/ci/project_files_and_dependencies.yaml @@ -305,4 +305,9 @@ ignore_regexes: - 'ci/rapids/' - '^ci/util/pre-commit/' - 'docs/' - - 'LICENSE' + - '^LICENSE$' + - '^cub/LICENSE\.TXT$' + - '^cudax/LICENSE\.TXT$' + - '^libcudacxx/LICENSE\.TXT$' + - '^libcudacxx/test/LICENSE\.TXT$' + - '^thrust/LICENSE$' diff --git a/ci/test/inspect_changes/ignored_only.dirty_files b/ci/test/inspect_changes/ignored_only.dirty_files index 7d2fa473a692..daf5748dddb0 100644 --- a/ci/test/inspect_changes/ignored_only.dirty_files +++ b/ci/test/inspect_changes/ignored_only.dirty_files @@ -1 +1,7 @@ docs/index.rst +LICENSE +cub/LICENSE.TXT +cudax/LICENSE.TXT +libcudacxx/LICENSE.TXT +libcudacxx/test/LICENSE.TXT +thrust/LICENSE