diff --git a/.github/workflows/build-wheel.yml b/.github/workflows/build-wheel.yml index 14cebd633d0..01091a7dd79 100644 --- a/.github/workflows/build-wheel.yml +++ b/.github/workflows/build-wheel.yml @@ -696,6 +696,27 @@ jobs: "${{ env.CUDA_CORE_ARTIFACTS_DIR }}"/cu"${BUILD_PREV_CUDA_MAJOR}"/cuda_core*.whl \ --output-dir "${{ env.CUDA_CORE_ARTIFACTS_DIR }}" + - name: Check merged cuda.core wheel layout + if: ${{ env.BUILD_CORE == 'true' }} + run: | + # cuda/core/__init__.py rewrites __path__ to the active cuXX/ tree, so + # only the entry points and the two versioned trees belong at the top + # level of cuda/core/; anything else is a dead copy. + python - "${{ env.CUDA_CORE_ARTIFACTS_DIR }}"/cuda_core*.whl <<'EOF' + import os + import sys + import zipfile + + prefix = "cuda/core/" + members = zipfile.ZipFile(sys.argv[1]).namelist() + names = {n[len(prefix) :].split("/")[0] for n in members if n.startswith(prefix)} + names.discard("") + expected = {"__init__.py", "_version.py"} + expected |= {f"cu{os.environ[v]}" for v in ("BUILD_CUDA_MAJOR", "BUILD_PREV_CUDA_MAJOR")} + assert names == expected, f"unexpected top-level entries under cuda/core/: {sorted(names ^ expected)}" + print("merged wheel top level:", sorted(names)) + EOF + - name: Check cuda.core wheel if: ${{ env.BUILD_CORE == 'true' }} run: | diff --git a/.gitignore b/.gitignore index 91c3b74a105..338c72735d4 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,7 @@ __pycache__/ !*_impl.cpp !cuda_bindings/cuda/bindings/_lib/param_packer.cpp !cuda_bindings/cuda/bindings/_bindings/loader.cpp +!cuda_core/cuda/core/_cpp/**/*.cpp cache_driver cache_runtime cache_nvrtc diff --git a/ci/tools/merge_cuda_core_wheels.py b/ci/tools/merge_cuda_core_wheels.py index 23a8a21289f..47af0751186 100644 --- a/ci/tools/merge_cuda_core_wheels.py +++ b/ci/tools/merge_cuda_core_wheels.py @@ -145,25 +145,22 @@ def merge_wheels(wheels: list[Path], output_dir: Path, show_wheel_contents: bool os.truncate(versioned_dir / "__init__.py", 0) print("\n=== Removing files from cuda/core/ directory ===", file=sys.stderr) + # Only what cuda/core/__init__.py uses before it rewrites __path__ to the + # versioned subpackage stays at top level: it imports _version, then + # redirects every later import into cu12/ or cu13/. Anything else left at + # top level is a dead copy that nothing imports. items_to_keep = ( "__init__.py", "_version.py", - "_include", - "_cpp", # Headers for Cython development "cu12", "cu13", ) - # _resource_handles is shared (not CUDA-version-specific) and must stay - # at top level. It's imported early in __init__.py before versioned code. - items_to_keep_prefix = ("_resource_handles",) all_items = os.scandir(base_wheel / base_dir) removed_count = 0 for f in all_items: f_abspath = f.path if f.name in items_to_keep: continue - if any(f.name.startswith(prefix) for prefix in items_to_keep_prefix): - continue if f.is_dir(): print(f" Removing directory: {f.name}", file=sys.stderr) shutil.rmtree(f_abspath) diff --git a/cuda_core/build_hooks.py b/cuda_core/build_hooks.py index fc112b5a74e..01f6c756bc8 100644 --- a/cuda_core/build_hooks.py +++ b/cuda_core/build_hooks.py @@ -183,6 +183,22 @@ def _relativize_extension_sources(extensions) -> None: ] +def _extension_sources(mod_name): + """The module's .pyx plus its C++, if any: every .cpp under + cuda/core/_cpp//, or the single legacy file cuda/core/_cpp/.cpp. + Example: _tensor_map.pyx compiles _cpp/tensor_map.cpp.""" + sources = [f"cuda/core/{mod_name}.pyx"] + cpp_stem = os.path.join("cuda", "core", "_cpp", mod_name.lstrip("_")) + if os.path.isdir(cpp_stem): + cpp_sources = sorted(glob.glob(os.path.join(cpp_stem, "**", "*.cpp"), recursive=True)) + if not cpp_sources: + raise RuntimeError(f"{cpp_stem}/ exists but contains no .cpp files") + sources.extend(cpp_sources) + elif os.path.isfile(cpp_stem + ".cpp"): + sources.append(cpp_stem + ".cpp") + return sources + + def _build_cuda_core(debug=False): # Customizing the build hooks is needed because we must defer cythonization until cuda-bindings, # now a required build-time dependency that's dynamically installed via the other hook below, @@ -227,18 +243,6 @@ def module_names(): continue yield mod - def get_sources(mod_name): - """Get source files for a module, including any .cpp files.""" - sources = [f"cuda/core/{mod_name}.pyx"] - - # Add module-specific .cpp file from _cpp/ directory if it exists - # Example: _resource_handles.pyx finds _cpp/resource_handles.cpp. - cpp_file = f"cuda/core/_cpp/{mod_name.lstrip('_')}.cpp" - if os.path.exists(cpp_file): - sources.append(cpp_file) - - return sources - all_include_dirs = [os.path.join(cuda_path, "include")] extra_compile_args = [] extra_link_args = [] @@ -264,7 +268,7 @@ def get_sources(mod_name): ext_modules = tuple( Extension( f"cuda.core.{mod.replace(os.path.sep, '.')}", - sources=get_sources(mod), + sources=_extension_sources(mod), include_dirs=[ "cuda/core/_include", "cuda/core/_cpp", diff --git a/cuda_core/pyproject.toml b/cuda_core/pyproject.toml index 3eac3a4c0af..3df8536a32b 100644 --- a/cuda_core/pyproject.toml +++ b/cuda_core/pyproject.toml @@ -108,7 +108,7 @@ include-package-data = false [tool.setuptools.package-data] "*" = ["*.pxd", "*.pyi", "py.typed"] "cuda.core._include" = ["*.h", "*.hpp"] -"cuda.core._cpp" = ["*.h", "*.hpp"] +"cuda.core._cpp" = ["**/*.h", "**/*.hpp"] [tool.setuptools.dynamic] readme = { file = ["DESCRIPTION.rst"], content-type = "text/x-rst" } diff --git a/cuda_core/tests/test_build_hooks.py b/cuda_core/tests/test_build_hooks.py index 2f1b3211781..b59e059548a 100644 --- a/cuda_core/tests/test_build_hooks.py +++ b/cuda_core/tests/test_build_hooks.py @@ -325,3 +325,42 @@ def test_flag_set_forces_rebuild(self, monkeypatch): def test_flag_clear_leaves_default(self, monkeypatch): assert not self._finalized_build_ext(False, monkeypatch).force + + +class TestExtensionSources: + """_extension_sources: a directory of .cpp files, a single legacy .cpp, or nothing.""" + + @pytest.fixture + def tree(self, tmp_path, monkeypatch): + core = tmp_path / "cuda" / "core" + cpp = core / "_cpp" + (cpp / "a" / "nested").mkdir(parents=True) + (cpp / "d").mkdir() + for name in ("_a.pyx", "_b.pyx", "_c.pyx", "_d.pyx"): + (core / name).write_text("") + for name in ("a/x.cpp", "a/y.cpp", "a/nested/z.cpp", "a/notes.md", "b.cpp"): + (cpp / name).write_text("") + monkeypatch.chdir(tmp_path) + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_directory_of_sources(self, tree): + a = os.path.join("cuda", "core", "_cpp", "a") + assert build_hooks._extension_sources("_a") == [ + "cuda/core/_a.pyx", + os.path.join(a, "nested", "z.cpp"), + os.path.join(a, "x.cpp"), + os.path.join(a, "y.cpp"), + ] + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_legacy_single_file_and_no_cpp(self, tree): + assert build_hooks._extension_sources("_b") == [ + "cuda/core/_b.pyx", + os.path.join("cuda", "core", "_cpp", "b.cpp"), + ] + assert build_hooks._extension_sources("_c") == ["cuda/core/_c.pyx"] + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_empty_directory_is_an_error(self, tree): + with pytest.raises(RuntimeError, match="no .cpp files"): + build_hooks._extension_sources("_d")