diff --git a/cuda_core/cuda/core/_module.pyi b/cuda_core/cuda/core/_module.pyi index 5125b99131..6c0cc85b50 100644 --- a/cuda_core/cuda/core/_module.pyi +++ b/cuda_core/cuda/core/_module.pyi @@ -458,6 +458,16 @@ class ObjectCode: """ + def get_module(self) -> object: + """Return the :obj:`~driver.CUmodule` associated with this object code. + + 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.""" diff --git a/cuda_core/cuda/core/_module.pyx b/cuda_core/cuda/core/_module.pyx index 91c8ad4389..990a2c8445 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,21 @@ cdef class ObjectCode: HANDLE_RETURN(get_last_error()) return Kernel._from_handle(h_kernel) + def get_module(self) -> object: + """Return the :obj:`~driver.CUmodule` associated with this object code. + + 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.""" 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", [