diff --git a/cuda_core/cuda/core/_cpp/resource_handles.hpp b/cuda_core/cuda/core/_cpp/resource_handles.hpp index 60a9f3c53a9..f2415cd23ac 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.hpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.hpp @@ -754,6 +754,10 @@ inline CUlibrary as_cu(const LibraryHandle& h) noexcept { return h ? *h : nullptr; } +inline CUmodule as_cu(const CUmodule& h) noexcept { + return h; +} + inline CUkernel as_cu(const KernelHandle& h) noexcept { return h ? *h : nullptr; } @@ -838,6 +842,10 @@ inline std::intptr_t as_intptr(const LibraryHandle& h) noexcept { return reinterpret_cast(as_cu(h)); } +inline std::intptr_t as_intptr(const CUmodule& h) noexcept { + return reinterpret_cast(as_cu(h)); +} + inline std::intptr_t as_intptr(const KernelHandle& h) noexcept { return reinterpret_cast(as_cu(h)); } @@ -968,6 +976,10 @@ inline PyObject* as_py(const LibraryHandle& h) noexcept { return detail::make_py("cuda.bindings.driver", "CUlibrary", as_intptr(h)); } +inline PyObject* as_py(const CUmodule& h) noexcept { + return detail::make_py("cuda.bindings.driver", "CUmodule", as_intptr(h)); +} + inline PyObject* as_py(const KernelHandle& h) noexcept { return detail::make_py("cuda.bindings.driver", "CUkernel", as_intptr(h)); } diff --git a/cuda_core/cuda/core/_module.pyi b/cuda_core/cuda/core/_module.pyi index f51b4cb2817..9ff758bb6c7 100644 --- a/cuda_core/cuda/core/_module.pyi +++ b/cuda_core/cuda/core/_module.pyi @@ -458,7 +458,7 @@ class ObjectCode: """ - def get_module(self) -> object: + def get_module(self) -> driver.CUmodule: """Return a context-dependent :obj:`~driver.CUmodule` for legacy interop. Bridges the native :obj:`~driver.CUlibrary` (see :attr:`handle`) to a diff --git a/cuda_core/cuda/core/_module.pyx b/cuda_core/cuda/core/_module.pyx index 704f6d2f856..95e149065bf 100644 --- a/cuda_core/cuda/core/_module.pyx +++ b/cuda_core/cuda/core/_module.pyx @@ -6,7 +6,6 @@ from __future__ import annotations cimport cython from libc.stddef cimport size_t -from libc.stdint cimport intptr_t from libcpp.mutex cimport py_safe_call_once from collections import namedtuple @@ -811,7 +810,7 @@ cdef class ObjectCode: HANDLE_RETURN(get_last_error()) return Kernel._from_handle(h_kernel) - def get_module(self) -> object: + def get_module(self) -> driver.CUmodule: """Return a context-dependent :obj:`~driver.CUmodule` for legacy interop. Bridges the native :obj:`~driver.CUlibrary` (see :attr:`handle`) to a @@ -828,7 +827,7 @@ cdef class ObjectCode: cdef cydriver.CUmodule mod with nogil: HANDLE_RETURN(cydriver.cuLibraryGetModule(&mod, as_cu(self._h_library))) - return driver.CUmodule(mod) + return as_py(mod) @property def code(self) -> CodeTypeT: diff --git a/cuda_core/cuda/core/_resource_handles.pxd b/cuda_core/cuda/core/_resource_handles.pxd index b0ae65d1666..168aec3d72e 100644 --- a/cuda_core/cuda/core/_resource_handles.pxd +++ b/cuda_core/cuda/core/_resource_handles.pxd @@ -85,6 +85,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": cydriver.CUmemoryPool as_cu(MemoryPoolHandle h) noexcept nogil cydriver.CUdeviceptr as_cu(DevicePtrHandle h) noexcept nogil cydriver.CUlibrary as_cu(LibraryHandle h) noexcept nogil + cydriver.CUmodule as_cu(cydriver.CUmodule h) noexcept nogil cydriver.CUkernel as_cu(KernelHandle h) noexcept nogil cydriver.CUgraph as_cu(GraphHandle h) noexcept nogil cydriver.CUgraphExec as_cu(GraphExecHandle h) noexcept nogil @@ -107,6 +108,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": intptr_t as_intptr(MemoryPoolHandle h) noexcept nogil intptr_t as_intptr(DevicePtrHandle h) noexcept nogil intptr_t as_intptr(LibraryHandle h) noexcept nogil + intptr_t as_intptr(const cydriver.CUmodule& h) noexcept nogil intptr_t as_intptr(KernelHandle h) noexcept nogil intptr_t as_intptr(GraphHandle h) noexcept nogil intptr_t as_intptr(GraphExecHandle h) noexcept nogil @@ -130,6 +132,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": object as_py(MemoryPoolHandle h) object as_py(DevicePtrHandle h) object as_py(LibraryHandle h) + object as_py(const cydriver.CUmodule& h) object as_py(KernelHandle h) object as_py(GraphHandle h) object as_py(GraphExecHandle h)