support getting native handle in C++/Cython

Author: leofang

.github/workflows/test-wheel-linux.yml

@@ -223,6 +223,19 @@ jobs:
           pytest -rxXs tests/
           popd
 
+          # It is a bit convoluted to run the Cython tests against CTK wheels,
+          # so let's just skip them.
+          if [[ "${{ inputs.local-ctk }}" == 1 ]]; then
+            if [[ "${{ inputs.host-platform }}" == linux* ]]; then
+              bash tests/cython/build_tests.sh
+            elif [[ "${{ inputs.host-platform }}" == win* ]]; then
+              # TODO: enable this once win-64 runners are up
+              exit 1
+            fi
+            pytest -rxXs tests/cython
+            popd
+          fi
+
       - name: Ensure cuda-python installable
         run: |
           if [[ "${{ inputs.local-ctk }}" == 1 ]]; then

cuda_core/README.md

@@ -41,3 +41,14 @@ for more details, including how to sign your commits.
 To run these tests:
 * `python -m pytest tests/` against editable installations
 * `pytest tests/` against installed packages
+
+### Cython Unit Tests
+
+Cython tests are located in `tests/cython` and need to be built. These builds have the same CUDA Toolkit header requirements as [those of cuda.bindings](https://nvidia.github.io/cuda-python/cuda-bindings/latest/install.html#requirements) where the major.minor version must match `cuda.bindings`. To build them:
+
+1. Setup environment variable `CUDA_HOME` with the path to the CUDA Toolkit installation.
+2. Run `build_tests` script located in `test/cython` appropriate to your platform. This will both cythonize the tests and build them.
+
+To run these tests:
+* `python -m pytest tests/cython/` against editable installations
+* `pytest tests/cython/` against installed packages

cuda_core/cuda/core/experimental/_event.py

@@ -130,6 +130,6 @@ def is_done(self) -> bool:
             raise CUDAError(f"unexpected error: {result}")
 
     @property
-    def handle(self) -> int:
+    def handle(self) -> "CUevent":
         """Return the underlying cudaEvent_t pointer address as Python int."""
-        return int(self._mnff.handle)
+        return self._mnff.handle

cuda_core/cuda/core/experimental/_linker.py

@@ -323,6 +323,10 @@ def _exception_manager(self):
         raise e
 
 
+nvJitLinkHandleT = int
+LinkerHandleT = Union[nvJitLinkHandleT, "CUlinkState"]
+
+
 class Linker:
     """Represent a linking machinery to link one or multiple object codes into
     :obj:`~cuda.core.experimental._module.ObjectCode` with the specified options.
@@ -483,7 +487,7 @@ def _input_type_from_code_type(self, code_type: str):
         return input_type
 
     @property
-    def handle(self):
+    def handle(self) -> LinkerHandleT:
         """Return the underlying handle object."""
         return self._mnff.handle
 

cuda_core/cuda/core/experimental/_program.py

@@ -8,7 +8,7 @@
 from warnings import warn
 
 from cuda.core.experimental._device import Device
-from cuda.core.experimental._linker import Linker, LinkerOptions
+from cuda.core.experimental._linker import Linker, LinkerHandleT, LinkerOptions
 from cuda.core.experimental._module import ObjectCode
 from cuda.core.experimental._utils import (
     _handle_boolean_option,
@@ -498,6 +498,6 @@ def backend(self) -> str:
         return self._backend
 
     @property
-    def handle(self):
+    def handle(self) -> Union["nvrtcProgram", LinkerHandleT]:
         """Return the underlying handle object."""
         return self._mnff.handle

cuda_core/cuda/core/experimental/_stream.py

@@ -147,9 +147,9 @@ def __cuda_stream__(self) -> Tuple[int, int]:
         return (0, self.handle)
 
     @property
-    def handle(self) -> int:
+    def handle(self) -> "CUstream":
         """Return the underlying cudaStream_t pointer address as Python int."""
-        return int(self._mnff.handle)
+        return self._mnff.handle
 
     @property
     def is_nonblocking(self) -> bool:

cuda_core/cuda/core/experimental/include/utility.hpp

@@ -0,0 +1,17 @@
+#include <type_traits>
+
+// In cuda.bindings 12.8, the private member name was renamed from "_ptr" to "_pvt_ptr".
+// We want to have the C++ layer supporting all past 12.x versions, so some tricks are needed.
+// Since there's no std::has_member<T, member_name> so we use SFINAE to create the same effect.
+
+template <typename T,
+          std::enable_if_t<std::is_pointer_v<decltype(std::remove_pointer_t<T>::_pvt_ptr)>, bool> = true>
+inline auto& get_cuda_native_handle(const T& obj) {
+    return *(obj->_pvt_ptr);
+}
+
+template <typename T,
+          std::enable_if_t<std::is_pointer_v<decltype(std::remove_pointer_t<T>::_ptr)>, bool> = true>
+inline auto& get_cuda_native_handle(const T& obj) {
+    return *(obj->_ptr);
+}

cuda_core/tests/cython/build_tests.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+SCRIPTPATH=$(dirname $(realpath "$0"))
+CPLUS_INCLUDE_PATH=$SCRIPTPATH/../../cuda/core/experimental/include:$CUDA_HOME/include:$CPLUS_INCLUDE_PATH cythonize -3 -i $(dirname "$0")/test_*.pyx

cuda_core/tests/cython/test_cython.py

@@ -0,0 +1,42 @@
+# Copyright 2021-2024 NVIDIA Corporation.  All rights reserved.
+#
+# Please refer to the NVIDIA end user license agreement (EULA) associated
+# with this source code for terms and conditions that govern your use of
+# this software. Any use, reproduction, disclosure, or distribution of
+# this software and related documentation outside the terms of the EULA
+# is strictly prohibited.
+import functools
+import importlib
+import sys
+
+
+def py_func(func):
+    """
+    Wraps func in a plain Python function.
+    """
+
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        return func(*args, **kwargs)
+
+    return wrapped
+
+
+cython_test_modules = [
+    "test_get_cuda_native_handle",
+]
+
+
+for mod in cython_test_modules:
+    try:
+        # For each callable in `mod` with name `test_*`,
+        # wrap the callable in a plain Python function
+        # and set the result as an attribute of this module.
+        mod = importlib.import_module(mod)
+        for name in dir(mod):
+            item = getattr(mod, name)
+            if callable(item) and name.startswith("test_"):
+                item = py_func(item)
+                setattr(sys.modules[__name__], name, item)
+    except ImportError:
+        raise

cuda_core/tests/cython/test_get_cuda_native_handle.pyx

@@ -0,0 +1,32 @@
+# distutils: language = c++
+# distutils: extra_compile_args = -std=c++17
+
+from libc.stdint cimport intptr_t
+
+from cuda.bindings.driver cimport (CUstream as pyCUstream,
+                                   CUevent as pyCUevent)
+from cuda.bindings.cydriver cimport CUstream, CUevent
+
+from cuda.core.experimental import Device
+
+
+cdef extern from "utility.hpp":
+    void* get_cuda_native_handle[T](T)
+
+
+def test_get_cuda_native_handle():
+    dev = Device(0)
+    dev.set_current()
+    
+    s = dev.create_stream()
+    cdef pyCUstream s_py = s.handle
+    cdef CUstream s_c = <CUstream>get_cuda_native_handle(s_py)
+    assert <intptr_t>(s_c) == <intptr_t>(int(s_py))
+
+    e = s.record()
+    cdef pyCUevent e_py = e.handle
+    cdef CUevent e_c = <CUevent>get_cuda_native_handle(e_py)
+    assert <intptr_t>(e_c) == <intptr_t>(int(e_py))
+
+    e.close()
+    s.close()

cuda_core/tests/pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+norecursedirs = cython