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

.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: |

.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

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

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 = 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",

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