Release/0.6.7

.github/workflows/black.yml

@@ -7,4 +7,4 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
-      - uses: psf/black@26.1.0
+      - uses: pre-commit/action@v3.0.1

.pre-commit-config.yaml

@@ -5,8 +5,8 @@ repos:
     -   id: check-yaml
     -   id: end-of-file-fixer
     -   id: trailing-whitespace
-  - repo: https://github.com/psf/black
-    rev: 23.10.1
+  - repo: https://github.com/psf/black-pre-commit-mirror
+    rev: 26.3.1
     hooks:
       - id: black
-        language_version: python3 # Should be a command that runs python3.6+
+        language_version: python3

xobjects/context.py

@@ -354,25 +354,75 @@ def build_kernels(
     ) -> Dict[Tuple[str, tuple], KernelType]:
         pass
 
-    def get_installed_c_source_paths(self) -> List[str]:
-        """Returns a list of include paths registered in dependent packages.
+    def get_installed_c_source_and_library_paths(
+        self,
+    ) -> tuple[set[Path], set[str], set[Path]]:
+        """Returns a list of C paths registered in dependent packages.
 
-        In a package that depends on xobjects, you can register C source paths
-        using the entry point `xobjects.c_sources`. A path to the directory
-        containing the specified module will be added to the include path when
-        building kernels. For example, the following will allow to write
-        ``#include <xtrack/path/to/some/header.h>`` in kernel sources:
+        In a package that depends on xobjects, you can register C source and
+        library paths using the entry point `xobjects.build_info`. These paths
+        will be added to the C include path and the library path when building
+        kernels. For example, the following will allow to write
+        ``#include <xcoll/path/to/some/header.h>`` in kernel sources, and
+        allow to use functions from the library ``xcoll/lib/libFlukaIO.a``:
 
          .. code-block:: toml
             [project.entry-points.xobjects]
-            include = "xtrack"
+            build_info = "xcoll._xobjects:get_build_info"
+
+        and in the file ``xcoll/_xobjects.py``:
+
+         .. code-block:: python
+            from ..general import _pkg_root
+            def get_build_info():
+                return {
+                    "include_dirs": [_pkg_root.parent],
+                    "libraries": ["FlukaIO"],
+                    "library_dirs": [_pkg_root / "lib"],
+                }
         """
-        sources = []
+        sources = set()
+        libs = set()
+        lib_paths = set()
+
+        # Old entry point for backward compatibility
         for ep in entry_points(group="xobjects", name="include"):
             module = ep.load()
             path = Path(module.__file__).parents[1]
-            sources.append(str(path))
-        return sources
+            sources.add(path)
+
+        # New entry point
+        for ep in entry_points(group="xobjects", name="build_info"):
+            get_build_info = ep.load()
+            info = get_build_info()
+            include_dirs = info.get("include_dirs", [])
+            if not hasattr(include_dirs, "__iter__") or isinstance(
+                include_dirs, str
+            ):
+                include_dirs = [include_dirs]
+            include_dirs = [
+                Path(dd).expanduser().resolve() for dd in include_dirs
+            ]
+            sources.update(include_dirs)
+
+            library_dirs = info.get("library_dirs", [])
+            if not hasattr(library_dirs, "__iter__") or isinstance(
+                library_dirs, str
+            ):
+                library_dirs = [library_dirs]
+            library_dirs = [
+                Path(dd).expanduser().resolve() for dd in library_dirs
+            ]
+            lib_paths.update(library_dirs)
+
+            libraries = info.get("libraries", [])
+            if not hasattr(libraries, "__iter__") or isinstance(
+                libraries, str
+            ):
+                libraries = [libraries]
+            libs.update(libraries)
+
+        return sources, libs, lib_paths
 
     @abstractmethod
     def nparray_to_context_array(self, arr, copy=False):

xobjects/context_cpu.py

@@ -499,10 +499,11 @@ def compile_kernel(
             xtr_compile_args.append("-DXO_CONTEXT_CPU_SERIAL")
             xtr_link_args.append("-DXO_CONTEXT_CPU_SERIAL")
 
-        extra_include_paths = self.get_installed_c_source_paths()
-        include_flags = [f"-I{path}" for path in extra_include_paths]
-        xtr_compile_args.extend(include_flags)
-        xtr_link_args.extend(include_flags)
+        (
+            extra_include_paths,
+            extra_libraries,
+            extra_library_paths,
+        ) = self.get_installed_c_source_and_library_paths()
 
         if os.name == "nt":  # windows
             # TODO: to be handled properly
@@ -516,6 +517,9 @@ def compile_kernel(
         ffi_interface.set_source(
             module_name,
             specialized_source,
+            include_dirs=[path.as_posix() for path in extra_include_paths],
+            libraries=list(extra_libraries),
+            library_dirs=[path.as_posix() for path in extra_library_paths],
             extra_compile_args=xtr_compile_args,
             extra_link_args=xtr_link_args,
         )

xobjects/context_cupy.py

@@ -8,6 +8,7 @@
 import shutil
 import subprocess
 import tempfile
+import warnings
 from typing import Dict, List, Tuple, Literal
 
 import numpy as np
@@ -28,6 +29,14 @@
 
 log = logging.getLogger(__name__)
 
+no_fast_compile = False
+"""Disable NVRTC fast compile tuning when building CUDA kernels.
+
+When set to ``True``, ``ContextCupy`` does not pass ``--Ofast-compile=min``
+to NVRTC. The ``XO_CUDA_NO_FAST_COMPILE`` environment variable provides the
+same behavior.
+"""
+
 try:
     import cupy
     import cupyx.scipy
@@ -385,6 +394,13 @@ class ContextCupy(XContext):
     Creates a Cupy Context object, that allows performing the computations
     on nVidia GPUs.
 
+    The module-level flag ``xobjects.context_cupy.no_fast_compile`` controls
+    whether NVRTC fast compile tuning is disabled. By default it is ``False``,
+    so CUDA kernels built with NVRTC >= 12.9 use ``--Ofast-compile=min`` to
+    reduce compilation time and memory usage, at the cost of some runtime
+    performance. Set it to ``True`` to disable this option. The environment
+    variable ``XO_CUDA_NO_FAST_COMPILE`` also disables it.
+
     Args:
         default_block_size (int):  CUDA thread size that is used by default
             for kernel execution in case a block size is not specified
@@ -477,8 +493,15 @@ def build_kernels(
             with open(save_source_as, "w") as fid:
                 fid.write(specialized_source)
 
-        extra_include_paths = self.get_installed_c_source_paths()
-        include_flags = [f"-I{path}" for path in extra_include_paths]
+        (
+            # TODO: how to deal with CUDA libraries?
+            extra_include_paths,
+            _,
+            _,
+        ) = self.get_installed_c_source_and_library_paths()
+        include_flags = [
+            f"-I{path.as_posix()}" for path in extra_include_paths
+        ]
         extra_compile_args = (
             *extra_compile_args,
             *include_flags,
@@ -488,11 +511,20 @@ def build_kernels(
         if self.backend == "nvrtc":
             # NVRTC (default): add NVRTC-specific flags
             nvrtc_args = (*extra_compile_args,)
+            fast_compile = not (
+                no_fast_compile or os.environ.get("XO_CUDA_NO_FAST_COMPILE")
+            )
             if nvrtc and nvrtc.getVersion() >= (12, 9):
                 # If supported, skip prohibitively heavy optimisations (e.g.
                 # involving cloning). This it at the expense of <20%
                 # runtime performance, but gain of a lot of compile time and memory.
-                nvrtc_args += ("--Ofast-compile=min",)
+                if fast_compile:
+                    nvrtc_args += ("--Ofast-compile=min",)
+            elif fast_compile:
+                warnings.warn(
+                    "Detected nvrtc version < 12.9, which does not support compile-time optimisation tuning. "
+                    "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

xobjects/context_pyopencl.py

@@ -228,8 +228,15 @@ def build_kernels(
             with open(save_source_as, "w") as fid:
                 fid.write(specialized_source)
 
-        extra_include_paths = self.get_installed_c_source_paths()
-        include_flags = [f"-I{path}" for path in extra_include_paths]
+        (
+            # TODO: how to deal with OpenCL libraries?
+            extra_include_paths,
+            _,
+            _,
+        ) = self.get_installed_c_source_and_library_paths()
+        include_flags = [
+            f"-I{path.as_posix()}" for path in extra_include_paths
+        ]
 
         extra_compile_args = (
             *extra_compile_args,

xobjects/headers/atomicadd.h

@@ -102,7 +102,7 @@ DEF_ATOMIC_ADD(double  , f64)
 #if defined(XO_CONTEXT_CUDA)
     // CUDA compiler may not have <stdint.h>, so define the types if needed.
     #if defined(__CUDACC_RTC__) || defined(__HIPCC_RTC__)
-        // NVRTC and HIPRTC (CuPy RawModule default) can’t see <stdint.h> 
+        // NVRTC and HIPRTC (CuPy RawModule default) can’t see <stdint.h>
         // We detect via __CUDACC_RTC__ (Nvidia) or __HIPCC_RTC__ (ROCm)
         typedef signed char        int8_t;
         typedef short              int16_t;