[WIP] Add experimental support for compiling CUDA kernels with Clang

xobjects/context_cupy.py

@@ -4,6 +4,11 @@
 # ########################################### #
 
 import logging
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
 from typing import Dict, List, Tuple
 
 import numpy as np
@@ -358,7 +363,7 @@ def __invert__(self):
 typedef unsigned short     uint16_t; //only_for_context cuda
 typedef unsigned char      uint8_t;  //only_for_context cuda
 
-#if defined(__CUDACC__) || defined(__HIPCC_RTC__)
+#if defined(__CUDACC_RTC__) || defined(__HIPCC_RTC__)
 typedef signed long long   int64_t;
 typedef unsigned long long uint64_t;
 #endif
@@ -402,6 +407,7 @@ def __init__(
         default_block_size=256,
         default_shared_mem_size_bytes=0,
         device=None,
+        compiler='nvrtc',
     ):
         if device is not None:
             cupy.cuda.Device(device).use()
@@ -410,6 +416,7 @@ def __init__(
 
         self.default_block_size = default_block_size
         self.default_shared_mem_size_bytes = default_shared_mem_size_bytes
+        self.compiler = compiler
 
     def _make_buffer(self, capacity):
         return BufferCupy(capacity=capacity, context=self)
@@ -466,15 +473,26 @@ def build_kernels(
             *extra_compile_args,
             *include_flags,
             "-DXO_CONTEXT_CUDA",
-            # Skip heavy optimizations (e.g. involving cloning),
-            # which for us don't translate to a lot of runtime gains,
-            # but consume a lot of compile time and memory:
-            "--Ofast-compile=min",
         )
 
-        module = cupy.RawModule(
-            code=specialized_source, options=extra_compile_args
-        )
+        if self.compiler == 'clang':
+            module = self._build_module_with_clang(
+                specialized_source, extra_compile_args
+            )
+        else:
+            # NVRTC (default): add NVRTC-specific flags
+            nvrtc_args = (
+                *extra_compile_args,
+                # Skip heavy optimizations (e.g. involving cloning),
+                # which for us don't translate to a lot of runtime gains
+                "--Ofast-compile=min",
+            )
+            module = cupy.RawModule(
+                code=specialized_source, options=nvrtc_args
+            )
+
+        with open('out_source.cu', 'w+') as f:
+            f.write(specialized_source)
 
         out_kernels = {}
         for pyname, kernel in kernel_descriptions.items():
@@ -493,6 +511,74 @@ def build_kernels(
 
         return out_kernels
 
+    def _find_clang(self):
+        # Allow explicit override via a dedicated env var
+        override = os.environ.get('XSUITE_CLANG')
+        if override:
+            return override
+        # Look alongside the current Python executable (conda env bin dir)
+        bin_dir = os.path.dirname(sys.executable)
+        for name in ['clang++', 'clang++-22', 'clang++-21', 'clang++-20',
+                     'clang++-19', 'clang++-18']:
+            candidate = os.path.join(bin_dir, name)
+            if os.path.isfile(candidate):
+                return candidate
+        # Fall back to PATH
+        for name in ['clang++', 'clang++-22', 'clang++-21', 'clang++-20']:
+            found = shutil.which(name)
+            if found:
+                return found
+        raise RuntimeError(
+            "clang++ not found. Set XSUITE_CLANG to the clang++ binary path."
+        )
+
+    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
+        cuda_include = os.path.join(cupy.cuda.get_cuda_path(), "include")
+
+        # Keep -I and -D flags; skip NVRTC-specific flags (--Ofast-compile etc.)
+        clang_args = [
+            a for a in extra_compile_args
+            if a.startswith(('-I', '-D', '-O', '-std='))
+        ]
+
+        src_fd, src_path = tempfile.mkstemp(suffix='.cu')
+        ptx_fd, ptx_path = tempfile.mkstemp(suffix='.ptx')
+        os.close(src_fd)
+        os.close(ptx_fd)
+        try:
+            with open(src_path, 'w') as f:
+                f.write(source)
+
+            cmd = [
+                clang,
+                '-x', 'cuda',
+                f'--cuda-gpu-arch=sm_{cc}',
+                '--cuda-device-only',
+                '-S',
+                f'-I{cuda_include}',
+                '-O3',
+            ] + clang_args + ['-o', ptx_path, src_path]
+
+            log.debug("clang command: %s", ' '.join(cmd))
+            result = subprocess.run(cmd, capture_output=True, text=True)
+            if result.returncode != 0:
+                raise RuntimeError(
+                    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()
+        finally:
+            if os.path.exists(src_path):
+                os.unlink(src_path)
+            if os.path.exists(ptx_path):
+                os.unlink(ptx_path)
+
+        return module
+
     def __str__(self):
         return f"{type(self).__name__}:{cupy.cuda.get_device_id()}"
 

xobjects/headers/atomicadd.h

@@ -121,8 +121,24 @@ DEF_ATOMIC_ADD(double  , f64)
         typedef unsigned short     uint16_t;
         typedef unsigned int       uint32_t;
     #else
-        // Alternatively, NVCC path is fine with host headers
+        // NVCC and Clang have system headers available
         #include <stdint.h>
+        #if defined(__clang__)
+        // On LP64 Linux, stdint.h maps uint64_t -> unsigned long, but CUDA
+        // atomics only have overloads for unsigned long long.  Provide
+        // transparent forwarding overloads so that atomicCAS/atomicAdd
+        // called with uint64_t* compile and behave correctly.
+        __device__ static inline unsigned long
+        atomicCAS(volatile unsigned long* a, unsigned long c, unsigned long v) {
+            return (unsigned long)atomicCAS((unsigned long long*)(void*)a,
+                                           (unsigned long long)c, (unsigned long long)v);
+        }
+        __device__ static inline unsigned long
+        atomicAdd(volatile unsigned long* a, unsigned long v) {
+            return (unsigned long)atomicAdd((unsigned long long*)(void*)a,
+                                           (unsigned long long)v);
+        }
+        #endif // __clang__
     #endif // __CUDACC_RTC__
     #define GPUVOLATILE    GPUGLMEM
     #define __XT_CAS_U32(ptr, exp, val) atomicCAS((GPUVOLATILE uint32_t*)(ptr), (uint32_t)(exp), (uint32_t)(val))