cuda.core: directory-aware C++ build rule and drop the dead top-level copy from the merged wheel

.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

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)

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/<stem>/, or the single legacy file cuda/core/_cpp/<stem>.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",

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" }

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")