diff --git a/cuda_core/cuda/core/_module.pyi b/cuda_core/cuda/core/_module.pyi index 5125b99131..f51b4cb281 100644 --- a/cuda_core/cuda/core/_module.pyi +++ b/cuda_core/cuda/core/_module.pyi @@ -458,6 +458,20 @@ class ObjectCode: """ + def get_module(self) -> object: + """Return a context-dependent :obj:`~driver.CUmodule` for legacy interop. + + Bridges the native :obj:`~driver.CUlibrary` (see :attr:`handle`) to a + ``CUmodule`` via ``cuLibraryGetModule``, for use with legacy driver APIs + that only accept ``CUmodule``. + + Returns + ------- + :obj:`~driver.CUmodule` + Module handle for the current CUDA context, suitable for legacy + driver APIs that accept ``CUmodule``. + """ + @property def code(self) -> CodeTypeT: """Return the underlying code object.""" @@ -476,7 +490,10 @@ class ObjectCode: @property def handle(self) -> object: - """Return the underlying handle object. + """Return the native, context-independent :obj:`~driver.CUlibrary` handle. + + Used by ``cuda.core`` and newer driver library APIs. For legacy APIs + that only accept a ``CUmodule``, use :meth:`get_module` instead. .. caution:: diff --git a/cuda_core/cuda/core/_module.pyx b/cuda_core/cuda/core/_module.pyx index 91c8ad4389..5734e31414 100644 --- a/cuda_core/cuda/core/_module.pyx +++ b/cuda_core/cuda/core/_module.pyx @@ -6,6 +6,7 @@ from __future__ import annotations cimport cython from libc.stddef cimport size_t +from libc.stdint cimport intptr_t from collections import namedtuple from os import fsencode, fspath, PathLike @@ -796,6 +797,25 @@ cdef class ObjectCode: HANDLE_RETURN(get_last_error()) return Kernel._from_handle(h_kernel) + def get_module(self) -> object: + """Return a context-dependent :obj:`~driver.CUmodule` for legacy interop. + + Bridges the native :obj:`~driver.CUlibrary` (see :attr:`handle`) to a + ``CUmodule`` via ``cuLibraryGetModule``, for use with legacy driver APIs + that only accept ``CUmodule``. + + Returns + ------- + :obj:`~driver.CUmodule` + Module handle for the current CUDA context, suitable for legacy + driver APIs that accept ``CUmodule``. + """ + self._lazy_load_module() + cdef cydriver.CUmodule mod + with nogil: + HANDLE_RETURN(cydriver.cuLibraryGetModule(&mod, as_cu(self._h_library))) + return driver.CUmodule(mod) + @property def code(self) -> CodeTypeT: """Return the underlying code object.""" @@ -818,7 +838,10 @@ cdef class ObjectCode: @property def handle(self) -> object: - """Return the underlying handle object. + """Return the native, context-independent :obj:`~driver.CUlibrary` handle. + + Used by ``cuda.core`` and newer driver library APIs. For legacy APIs + that only accept a ``CUmodule``, use :meth:`get_module` instead. .. caution:: diff --git a/cuda_core/tests/test_module.py b/cuda_core/tests/test_module.py index a060096661..b5c9b4bee7 100644 --- a/cuda_core/tests/test_module.py +++ b/cuda_core/tests/test_module.py @@ -226,6 +226,21 @@ def test_get_kernel(init_cuda): assert object_code.get_kernel(b"ABC").handle is not None +def test_object_code_get_module_for_legacy_integration(init_cuda): + src = """ + extern "C" __global__ void ABC() { } + extern "C" __global__ void DEF() { } + """ + object_code = Program(src, "c++").compile("cubin") + + # Bridge: CUlibrary (new) → CUmodule (legacy) + module = object_code.get_module() + + # Legacy module-only API consumes it directly + count = handle_return(driver.cuModuleGetFunctionCount(module)) + assert count == 2 + + @pytest.mark.parametrize( "attr, expected_type", [