diff --git a/.gitignore b/.gitignore index b8eaa5ca208..6938baaaff4 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..6e7f51e73b6 100644 --- a/ci/tools/merge_cuda_core_wheels.py +++ b/ci/tools/merge_cuda_core_wheels.py @@ -133,9 +133,11 @@ def merge_wheels(wheels: list[Path], output_dir: Path, show_wheel_contents: bool # Copy version-specific directories from each wheel into versioned subdirectories base_dir = Path("cuda") / "core" + versioned_dirs = set() for i, wheel_dir in enumerate(extracted_wheels): cuda_version = wheels[i].name.split(".cu")[1].split(".")[0] versioned_dir = base_wheel / base_dir / f"cu{cuda_version}" + versioned_dirs.add(versioned_dir.name) # Copy entire directory tree from source wheel to versioned directory print(f" Copying {wheel_dir / base_dir} to {versioned_dir}", file=sys.stderr) @@ -145,25 +147,17 @@ 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) - 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",) + # 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 the versioned tree. Anything else + # left at top level is a dead copy that nothing imports. + items_to_keep = {"__init__.py", "_version.py", *versioned_dirs} 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) @@ -172,6 +166,11 @@ def merge_wheels(wheels: list[Path], output_dir: Path, show_wheel_contents: bool os.remove(f_abspath) removed_count += 1 print(f"Removed {removed_count} items from cuda/core/ directory", file=sys.stderr) + remaining = {entry.name for entry in os.scandir(base_wheel / base_dir)} + if remaining != items_to_keep: + raise RuntimeError( + f"unexpected top level under cuda/core/: {sorted(remaining)} (expected {sorted(items_to_keep)})" + ) # Repack the merged wheel output_dir.mkdir(parents=True, exist_ok=True) diff --git a/cuda_core/build_hooks.py b/cuda_core/build_hooks.py index fc112b5a74e..f105c92005a 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 = Path("cuda", "core", "_cpp", mod_name.lstrip("_")) + if cpp_stem.is_dir(): + cpp_sources = sorted(str(path) for path in cpp_stem.rglob("*.cpp")) + if not cpp_sources: + raise RuntimeError(f"{cpp_stem}/ exists but contains no .cpp files") + sources.extend(cpp_sources) + elif cpp_stem.with_suffix(".cpp").is_file(): + sources.append(str(cpp_stem.with_suffix(".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")