diff --git a/tests/test_settings.py b/tests/test_settings.py new file mode 100644 index 0000000..233094b --- /dev/null +++ b/tests/test_settings.py @@ -0,0 +1,27 @@ +# copyright ################################# # +# This file is part of the Xobjects Package. # +# Copyright (c) CERN, 2026. # +# ########################################### # + +import xobjects as xo +import pytest + + +def test_cuda_compute_capability(): + assert xo.ContextCupy().cuda_compute_capability is None + + with xo.settings.override(cuda_compute_capability=90): + assert xo.ContextCupy().cuda_compute_capability == 90 + assert ( + xo.ContextCupy(cuda_compute_capability=80).cuda_compute_capability + == 80 + ) + + with pytest.raises(ValueError): + xo.ContextCupy(cuda_compute_capability=0) + + with pytest.raises(ValueError): + xo.ContextCupy(cuda_compute_capability=-1) + + with pytest.raises(ValueError): + xo.ContextCupy(cuda_compute_capability="90") diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py index b687b6b..545c549 100644 --- a/xobjects/context_cupy.py +++ b/xobjects/context_cupy.py @@ -46,6 +46,8 @@ import cupyx.scipy.stats from cupyx.scipy import fftpack as cufftp from cupy_backends.cuda.libs import nvrtc + from cupy._core import core as cupy_core + from cupy.cuda import function as cupy_function _enabled = True except ImportError: @@ -402,10 +404,19 @@ class ContextCupy(XContext): environment variable ``XSUITE_CUDA_FAST_COMPILE=0``, to disable it. Args: - default_block_size (int): CUDA thread size that is used by default + default_block_size (int): + CUDA thread size that is used by default for kernel execution in case a block size is not specified directly in the kernel object. The default value is 256. - device (int): Identifier of the device to be used by the context. + device (int): + Identifier of the device to be used by the context. + cuda_compute_capability (int, optional): + CUDA virtual architecture used for kernel compilation. If omitted, + the architecture is selected automatically from the active GPU. + When specified, kernels are compiled to PTX for the requested + compute capability and JIT-compiled by the CUDA driver for the + actual device. For example, ``cuda_compute_capability=90`` targets + ``compute_90``. Returns: ContextCupy: context object. @@ -425,6 +436,7 @@ def __init__( default_shared_mem_size_bytes=0, device=None, backend: Literal[None, "nvrtc", "clang"] = None, + cuda_compute_capability=None, ): if device is not None: cupy.cuda.Device(device).use() @@ -444,6 +456,20 @@ def __init__( self.backend = backend + if cuda_compute_capability is None: + cuda_compute_capability = settings.cuda_compute_capability + + if cuda_compute_capability is not None and ( + not isinstance(cuda_compute_capability, int) + or cuda_compute_capability <= 0 + ): + raise ValueError( + "cuda_compute_capability must be a positive integer " + "such as 90 or 120." + ) + + self.cuda_compute_capability = cuda_compute_capability + def _make_buffer(self, capacity): return BufferCupy(capacity=capacity, context=self) @@ -524,9 +550,17 @@ def build_kernels( "Compilation time and memory usage might be high: if this is a problem, please update CUDA nvrtc." ) - module = cupy.RawModule( - code=specialized_source, options=nvrtc_args - ) + if self.cuda_compute_capability is None: + module = cupy.RawModule( + code=specialized_source, + options=nvrtc_args, + ) + else: + module = self._build_module_with_nvrtc_ptx( + specialized_source, + nvrtc_args, + ) + else: # Clang++ backend module = self._build_module_with_clang( @@ -567,9 +601,47 @@ def _find_clang(self): "executable." ) + def _build_module_with_nvrtc_ptx(self, source, extra_compile_args=()): + # Match the include paths/options that CuPy normally adds for RawModule. + options = cupy_core.assemble_cupy_compiler_options( + tuple(extra_compile_args) + ) + + options += ( + "-ftz=true", + f"-arch=compute_{self.cuda_compute_capability}", + "--device-as-default-execution-space", + ) + + program = nvrtc.createProgram( + source, + "xobjects.cu", + (), + (), + ) + + try: + nvrtc.compileProgram(program, options) + ptx = nvrtc.getPTX(program) + except nvrtc.NVRTCError as err: + log = nvrtc.getProgramLog(program) + raise RuntimeError( + f"NVRTC compilation failed for compute_{self.cuda_compute_capability}:\n" + f"{log}" + ) from err + finally: + nvrtc.destroyProgram(program) + + module = cupy_function.Module() + module.load(ptx) + + return module + def _build_module_with_clang(self, source, extra_compile_args=()): clang = self._find_clang() - cc = cupy.cuda.Device(cupy.cuda.get_device_id()).compute_capability + cc = self.cuda_compute_capability + if cc is None: + cc = cupy.cuda.Device(cupy.cuda.get_device_id()).compute_capability cuda_include = os.path.join(cupy.cuda.get_cuda_path(), "include") # Keep -I and -D flags; skip NVRTC-specific flags (--Ofast-compile etc.) @@ -608,9 +680,14 @@ def _build_module_with_clang(self, source, extra_compile_args=()): f"clang CUDA compilation failed:\n{result.stderr}" ) - module = cupy.RawModule(path=ptx_path) - # Force the driver to load the PTX now, before we delete the file - module.compile() + # Load the generated PTX directly. Avoid RawModule(path=...) + # because its memoisation key contains the temporary filename. + with open(ptx_path, "rb") as fid: + ptx = fid.read() + + module = cupy_function.Module() + module.load(ptx) + finally: if os.path.exists(src_path): os.unlink(src_path) diff --git a/xobjects/settings.py b/xobjects/settings.py index 21dac33..74715ec 100644 --- a/xobjects/settings.py +++ b/xobjects/settings.py @@ -51,6 +51,11 @@ class Settings: Values: ``None`` (default), executable path Selects the clang++ executable used by the CUDA clang backend. + ``cuda_compute_capability`` — ``XSUITE_CUDA_COMPUTE_CAPABILITY`` + Values: ``None`` (default), integer compute capability such as ``90`` + Selects the CUDA virtual architecture used for PTX compilation. + If unset, the compute capability is inferred from the active GPU. + Notes ----- The same object is exposed as ``xobjects.settings`` and @@ -249,6 +254,13 @@ def _parse_optional_string(value): return value if value else None +def _parse_optional_int(value): + value = value.strip() + if not value: + return None + return int(value) + + settings = Settings() settings._register( "print_mode", @@ -326,3 +338,10 @@ def _parse_optional_string(value): environment_parser=_parse_optional_string, value_type=(str, type(None)), ) +settings._register( + "cuda_compute_capability", + default=None, + environment_variable="XSUITE_CUDA_COMPUTE_CAPABILITY", + environment_parser=_parse_optional_int, + value_type=(int, type(None)), +)