From 7e7777a7ca18505de09f1ff7fe436bf4736ea21e Mon Sep 17 00:00:00 2001 From: "Frederik F. Van der Veken" Date: Thu, 24 Sep 2026 16:06:13 +0200 Subject: [PATCH 1/5] Add option to select CUDA compute capability --- xobjects/context_cupy.py | 102 ++++++++++++++++++++++++++++++++++++--- 1 file changed, 96 insertions(+), 6 deletions(-) diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py index b687b6b..8ad9777 100644 --- a/xobjects/context_cupy.py +++ b/xobjects/context_cupy.py @@ -46,6 +46,7 @@ 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 _enabled = True except ImportError: @@ -402,10 +403,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 or str, 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 +435,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() @@ -443,6 +454,25 @@ def __init__( ) self.backend = backend + if cuda_compute_capability is not None: + cuda_compute_capability = str(cuda_compute_capability) + + if cuda_compute_capability.startswith("compute_"): + cuda_compute_capability = cuda_compute_capability.removeprefix( + "compute_" + ) + elif cuda_compute_capability.startswith("sm_"): + cuda_compute_capability = cuda_compute_capability.removeprefix( + "sm_" + ) + + if not cuda_compute_capability.isdigit(): + raise ValueError( + f"Invalid CUDA architecture {cuda_compute_capability!r}. " + "Expected e.g. 90, '90', 'compute_90', or 'sm_90'." + ) + + self.cuda_arch = cuda_compute_capability def _make_buffer(self, capacity): return BufferCupy(capacity=capacity, context=self) @@ -524,9 +554,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_arch is None: + module = cupy.RawModule( + code=specialized_source, + options=nvrtc_args, + ) + else: + module = self._build_module_with_nvrtc( + specialized_source, + nvrtc_args, + ) + else: # Clang++ backend module = self._build_module_with_clang( @@ -567,9 +605,61 @@ def _find_clang(self): "executable." ) + def _build_module_with_nvrtc(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_arch}", + "--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_arch}:\n" + f"{log}" + ) from err + finally: + nvrtc.destroyProgram(program) + + # RawModule(path=...) gives us the same interface currently used by + # KernelCupy, while the CUDA driver JITs the PTX for the actual GPU. + fd, ptx_path = tempfile.mkstemp(suffix=".ptx") + os.close(fd) + + try: + with open(ptx_path, "wb") as fid: + fid.write(ptx) + + module = cupy.RawModule(path=ptx_path) + + # Force loading/JIT before deleting the temporary PTX. + module.compile() + finally: + if os.path.exists(ptx_path): + os.unlink(ptx_path) + + 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_arch + 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.) From 701a7e43ee14d265aefae0cc2b979c50974c5653 Mon Sep 17 00:00:00 2001 From: "Frederik F. Van der Veken" Date: Thu, 24 Sep 2026 16:24:58 +0200 Subject: [PATCH 2/5] Make cuda_compute_capability a setting (with flag XSUITE_CUDA_COMPUTE_CAPABILITY) --- xobjects/context_cupy.py | 25 ++++++++++--------------- xobjects/settings.py | 19 +++++++++++++++++++ 2 files changed, 29 insertions(+), 15 deletions(-) diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py index 8ad9777..224fafc 100644 --- a/xobjects/context_cupy.py +++ b/xobjects/context_cupy.py @@ -454,23 +454,18 @@ def __init__( ) self.backend = backend - if cuda_compute_capability is not None: - cuda_compute_capability = str(cuda_compute_capability) - if cuda_compute_capability.startswith("compute_"): - cuda_compute_capability = cuda_compute_capability.removeprefix( - "compute_" - ) - elif cuda_compute_capability.startswith("sm_"): - cuda_compute_capability = cuda_compute_capability.removeprefix( - "sm_" - ) + if cuda_compute_capability is None: + cuda_compute_capability = settings.cuda_compute_capability - if not cuda_compute_capability.isdigit(): - raise ValueError( - f"Invalid CUDA architecture {cuda_compute_capability!r}. " - "Expected e.g. 90, '90', 'compute_90', or 'sm_90'." - ) + 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_arch = cuda_compute_capability diff --git a/xobjects/settings.py b/xobjects/settings.py index 21dac33..1c21605 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)), +) From b49ca34cfc26c529a5ac215310fd6593c3116a31 Mon Sep 17 00:00:00 2001 From: "Frederik F. Van der Veken" Date: Thu, 24 Sep 2026 16:31:26 +0200 Subject: [PATCH 3/5] Internal rename and added test --- tests/test_settings.py | 20 ++++++++++++++++++++ xobjects/context_cupy.py | 10 +++++----- 2 files changed, 25 insertions(+), 5 deletions(-) create mode 100644 tests/test_settings.py diff --git a/tests/test_settings.py b/tests/test_settings.py new file mode 100644 index 0000000..c7d3bf6 --- /dev/null +++ b/tests/test_settings.py @@ -0,0 +1,20 @@ +# copyright ################################# # +# This file is part of the Xobjects Package. # +# Copyright (c) CERN, 2026. # +# ########################################### # + +import xobjects as xo +import numpy as np +import pytest + +from xobjects.test_helpers import for_all_test_contexts + + +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 + ) diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py index 224fafc..560359b 100644 --- a/xobjects/context_cupy.py +++ b/xobjects/context_cupy.py @@ -467,7 +467,7 @@ def __init__( "such as 90 or 120." ) - self.cuda_arch = cuda_compute_capability + self.cuda_compute_capability = cuda_compute_capability def _make_buffer(self, capacity): return BufferCupy(capacity=capacity, context=self) @@ -549,7 +549,7 @@ def build_kernels( "Compilation time and memory usage might be high: if this is a problem, please update CUDA nvrtc." ) - if self.cuda_arch is None: + if self.cuda_compute_capability is None: module = cupy.RawModule( code=specialized_source, options=nvrtc_args, @@ -608,7 +608,7 @@ def _build_module_with_nvrtc(self, source, extra_compile_args=()): options += ( "-ftz=true", - f"-arch=compute_{self.cuda_arch}", + f"-arch=compute_{self.cuda_compute_capability}", "--device-as-default-execution-space", ) @@ -625,7 +625,7 @@ def _build_module_with_nvrtc(self, source, extra_compile_args=()): except nvrtc.NVRTCError as err: log = nvrtc.getProgramLog(program) raise RuntimeError( - f"NVRTC compilation failed for compute_{self.cuda_arch}:\n" + f"NVRTC compilation failed for compute_{self.cuda_compute_capability}:\n" f"{log}" ) from err finally: @@ -652,7 +652,7 @@ def _build_module_with_nvrtc(self, source, extra_compile_args=()): def _build_module_with_clang(self, source, extra_compile_args=()): clang = self._find_clang() - cc = self.cuda_arch + 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") From d948e992e612537706142c90a26ea2b3f97beaef Mon Sep 17 00:00:00 2001 From: "Frederik F. Van der Veken" Date: Thu, 24 Sep 2026 20:32:36 +0200 Subject: [PATCH 4/5] Some improvements --- tests/test_settings.py | 13 ++++++++++--- xobjects/context_cupy.py | 6 +++--- xobjects/settings.py | 6 +++--- 3 files changed, 16 insertions(+), 9 deletions(-) diff --git a/tests/test_settings.py b/tests/test_settings.py index c7d3bf6..233094b 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -4,17 +4,24 @@ # ########################################### # import xobjects as xo -import numpy as np import pytest -from xobjects.test_helpers import for_all_test_contexts - 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 560359b..bd8d428 100644 --- a/xobjects/context_cupy.py +++ b/xobjects/context_cupy.py @@ -409,7 +409,7 @@ class ContextCupy(XContext): directly in the kernel object. The default value is 256. device (int): Identifier of the device to be used by the context. - cuda_compute_capability (int or str, optional): + 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 @@ -555,7 +555,7 @@ def build_kernels( options=nvrtc_args, ) else: - module = self._build_module_with_nvrtc( + module = self._build_module_with_nvrtc_ptx( specialized_source, nvrtc_args, ) @@ -600,7 +600,7 @@ def _find_clang(self): "executable." ) - def _build_module_with_nvrtc(self, source, extra_compile_args=()): + 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) diff --git a/xobjects/settings.py b/xobjects/settings.py index 1c21605..74715ec 100644 --- a/xobjects/settings.py +++ b/xobjects/settings.py @@ -52,9 +52,9 @@ class Settings: 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. + 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 ----- From d70c1e8f8c2cf0bafb14bb4f90644e3dba40a432 Mon Sep 17 00:00:00 2001 From: "Frederik F. Van der Veken" Date: Fri, 25 Sep 2026 10:29:59 +0200 Subject: [PATCH 5/5] Avoid recompiling temporary modules everytime --- xobjects/context_cupy.py | 30 +++++++++++------------------- 1 file changed, 11 insertions(+), 19 deletions(-) diff --git a/xobjects/context_cupy.py b/xobjects/context_cupy.py index bd8d428..545c549 100644 --- a/xobjects/context_cupy.py +++ b/xobjects/context_cupy.py @@ -47,6 +47,7 @@ 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: @@ -631,22 +632,8 @@ def _build_module_with_nvrtc_ptx(self, source, extra_compile_args=()): finally: nvrtc.destroyProgram(program) - # RawModule(path=...) gives us the same interface currently used by - # KernelCupy, while the CUDA driver JITs the PTX for the actual GPU. - fd, ptx_path = tempfile.mkstemp(suffix=".ptx") - os.close(fd) - - try: - with open(ptx_path, "wb") as fid: - fid.write(ptx) - - module = cupy.RawModule(path=ptx_path) - - # Force loading/JIT before deleting the temporary PTX. - module.compile() - finally: - if os.path.exists(ptx_path): - os.unlink(ptx_path) + module = cupy_function.Module() + module.load(ptx) return module @@ -693,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)