Skip to content

Commit 3aefb73

Browse files
sebergseberg
committed
BLD: Insert CCCL also via FLAGS environment to work around conda nvcc
Conda's `cuda-nvcc` inserts it's own CCCL in a way that setuptools ends up placing them before our includes. However, using the system CCCL means we find a version that isn't (currently) compatible. That means that in such a setup, the old "TMA doesn't link" bug resurfaces. This is a somewhat weird work-around, but it was the best idea, I could think of...
1 parent 3f58a7b commit 3aefb73

2 files changed

Lines changed: 40 additions & 11 deletions

File tree

‎cuda_core/build_hooks.py‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,11 @@
1212
import importlib.util
1313
import os
1414
import re
15+
import shlex
1516
import sys
1617
import tempfile
1718
import zipfile
19+
from distutils import sysconfig
1820
from pathlib import Path
1921

2022
import Cython as _Cython
@@ -340,7 +342,8 @@ def _build_define_macros(cuda_major: str, bindings_cuda_version: int | None = No
340342

341343
def _build_config_key(cuda_major, toolchain, debug, coverage):
342344
"""Return a stable string key for the build configuration."""
343-
return f"cu{cuda_major}-{toolchain}-{'debug' if debug else 'opt'}{'-cov' if coverage else ''}"
345+
# Rebuild artifacts made before the pinned CCCL took precedence over environment flags.
346+
return f"cu{cuda_major}-{toolchain}-{'debug' if debug else 'opt'}{'-cov' if coverage else ''}-cccl-first"
344347

345348

346349
def _check_build_config(toolchain, debug, coverage):
@@ -445,6 +448,20 @@ def _get_cccl_include_dirs() -> list[str]:
445448
return [str(libcudacxx_include)]
446449

447450

451+
def _configure_cccl_include_flags(include_dirs):
452+
if sys.platform == "win32":
453+
return
454+
455+
# setuptools puts environment flags before Extension.include_dirs;
456+
# conda's cuda-nvcc activation adds a competing CCCL -I there.
457+
cccl_flags = shlex.join([f"-I{path}" for path in include_dirs])
458+
default_flags = sysconfig.get_config_var("CFLAGS") or ""
459+
for name in ("CFLAGS", "CXXFLAGS", "CPPFLAGS"):
460+
flags = os.environ.get(name, "" if name == "CPPFLAGS" else default_flags)
461+
if flags != cccl_flags and not flags.startswith(cccl_flags + " "):
462+
os.environ[name] = f"{cccl_flags} {flags}"
463+
464+
448465
def _build_cuda_core(debug=False):
449466
# Customizing the build hooks is needed because we must defer cythonization until cuda-bindings,
450467
# now a required build-time dependency that's dynamically installed via the other hook below,
@@ -491,7 +508,9 @@ def module_names():
491508
yield mod
492509

493510
# The pinned CCCL must precede the toolkit's bundled headers.
494-
all_include_dirs = [*_get_cccl_include_dirs(), os.path.join(cuda_path, "include")]
511+
cccl_include_dirs = _get_cccl_include_dirs()
512+
all_include_dirs = [*cccl_include_dirs, os.path.join(cuda_path, "include")]
513+
_configure_cccl_include_flags(cccl_include_dirs)
495514

496515
# Resolve the C/C++ toolchain (CUDA_PYTHON_TOOLCHAIN). The default (gnu on
497516
# Linux, msvc on Windows) reproduces the previous build behavior and does

‎cuda_core/tests/test_build_hooks.py‎

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,17 @@ def _load_module(name, path, *, register=False):
6363

6464
@pytest.fixture(autouse=True)
6565
def _isolate_toolchain_env():
66-
names = ("CUDA_PYTHON_TOOLCHAIN", "CC", "CXX", "LDSHARED", "LDCXXSHARED", "CUDA_PYTHON_CYTHON_CACHE_DIR")
66+
names = (
67+
"CUDA_PYTHON_TOOLCHAIN",
68+
"CC",
69+
"CXX",
70+
"CFLAGS",
71+
"CXXFLAGS",
72+
"CPPFLAGS",
73+
"LDSHARED",
74+
"LDCXXSHARED",
75+
"CUDA_PYTHON_CYTHON_CACHE_DIR",
76+
)
6777
original = {name: os.environ[name] for name in names if name in os.environ}
6878
for name in names:
6979
os.environ.pop(name, None)
@@ -236,38 +246,38 @@ def test_missing_stamp_forces_rebuild(self, stamp):
236246
# No stamp means the last build's config is unknown, so rebuild.
237247
cuda_major, key = build_hooks._check_build_config("gnu", False, False)
238248
assert cuda_major == "13"
239-
assert key == "cu13-gnu-opt"
249+
assert key == "cu13-gnu-opt-cccl-first"
240250
assert build_hooks.force_build_ext is True
241251

242252
@pytest.mark.agent_authored(model="glm-5.2")
243253
def test_same_config_does_not_force(self, stamp):
244-
_write_stamp(stamp, "cu13-gnu-opt")
254+
_write_stamp(stamp, "cu13-gnu-opt-cccl-first")
245255
cuda_major, key = build_hooks._check_build_config("gnu", False, False)
246256
assert cuda_major == "13"
247257
assert build_hooks.force_build_ext is False
248258

249259
@pytest.mark.agent_authored(model="glm-5.2")
250260
def test_changed_cuda_major_forces_rebuild(self, stamp):
251-
_write_stamp(stamp, "cu12-gnu-opt")
261+
_write_stamp(stamp, "cu12-gnu-opt-cccl-first")
252262
cuda_major, _ = build_hooks._check_build_config("gnu", False, False)
253263
assert cuda_major == "13"
254264
assert build_hooks.force_build_ext is True
255265

256266
@pytest.mark.agent_authored(model="glm-5.2")
257267
def test_changed_toolchain_forces_rebuild(self, stamp):
258-
_write_stamp(stamp, "cu13-gnu-opt")
268+
_write_stamp(stamp, "cu13-gnu-opt-cccl-first")
259269
build_hooks._check_build_config("llvm", False, False)
260270
assert build_hooks.force_build_ext is True
261271

262272
@pytest.mark.agent_authored(model="glm-5.2")
263273
def test_changed_debug_forces_rebuild(self, stamp):
264-
_write_stamp(stamp, "cu13-gnu-opt")
274+
_write_stamp(stamp, "cu13-gnu-opt-cccl-first")
265275
build_hooks._check_build_config("gnu", True, False)
266276
assert build_hooks.force_build_ext is True
267277

268278
@pytest.mark.agent_authored(model="glm-5.2")
269279
def test_changed_coverage_forces_rebuild(self, stamp):
270-
_write_stamp(stamp, "cu13-gnu-opt")
280+
_write_stamp(stamp, "cu13-gnu-opt-cccl-first")
271281
build_hooks._check_build_config("gnu", False, True)
272282
assert build_hooks.force_build_ext is True
273283

@@ -428,8 +438,8 @@ def test_majors_use_different_dirs(self, monkeypatch):
428438

429439
assert dir_12 != dir_13
430440
toolchain = "msvc" if sys.platform == "win32" else "gnu"
431-
assert dir_12.name == f"cu12-{toolchain}-opt"
432-
assert dir_13.name == f"cu13-{toolchain}-opt"
441+
assert dir_12.name == f"cu12-{toolchain}-opt-cccl-first"
442+
assert dir_13.name == f"cu13-{toolchain}-opt-cccl-first"
433443

434444
@pytest.mark.agent_authored(model="glm-5.2")
435445
def test_dir_is_anchored_not_relative_to_cwd(self, monkeypatch):

0 commit comments

Comments
 (0)