Cuda compute

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

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)

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)),
+)