Skip to content

Commit 600176e

Browse files
committed
add an api to return module in ObjectCode to support calling legacy driver APIs
1 parent e1f5d97 commit 600176e

3 files changed

Lines changed: 41 additions & 0 deletions

File tree

‎cuda_core/cuda/core/_module.pyi‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -458,6 +458,16 @@ class ObjectCode:
458458
459459
"""
460460

461+
def get_module(self) -> object:
462+
"""Return the :obj:`~driver.CUmodule` associated with this object code.
463+
464+
Returns
465+
-------
466+
:obj:`~driver.CUmodule`
467+
Module handle for the current CUDA context, suitable for legacy
468+
driver APIs that accept ``CUmodule``.
469+
"""
470+
461471
@property
462472
def code(self) -> CodeTypeT:
463473
"""Return the underlying code object."""

‎cuda_core/cuda/core/_module.pyx‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ from __future__ import annotations
66

77
cimport cython
88
from libc.stddef cimport size_t
9+
from libc.stdint cimport intptr_t
910

1011
from collections import namedtuple
1112
from os import fsencode, fspath, PathLike
@@ -796,6 +797,21 @@ cdef class ObjectCode:
796797
HANDLE_RETURN(get_last_error())
797798
return Kernel._from_handle(h_kernel)
798799

800+
def get_module(self) -> object:
801+
"""Return the :obj:`~driver.CUmodule` associated with this object code.
802+
803+
Returns
804+
-------
805+
:obj:`~driver.CUmodule`
806+
Module handle for the current CUDA context, suitable for legacy
807+
driver APIs that accept ``CUmodule``.
808+
"""
809+
self._lazy_load_module()
810+
cdef cydriver.CUmodule mod
811+
with nogil:
812+
HANDLE_RETURN(cydriver.cuLibraryGetModule(&mod, as_cu(self._h_library)))
813+
return driver.CUmodule(<intptr_t>mod)
814+
799815
@property
800816
def code(self) -> CodeTypeT:
801817
"""Return the underlying code object."""

‎cuda_core/tests/test_module.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,21 @@ def test_get_kernel(init_cuda):
226226
assert object_code.get_kernel(b"ABC").handle is not None
227227

228228

229+
def test_object_code_get_module_for_legacy_integration(init_cuda):
230+
src = """
231+
extern "C" __global__ void ABC() { }
232+
extern "C" __global__ void DEF() { }
233+
"""
234+
object_code = Program(src, "c++").compile("cubin")
235+
236+
# Bridge: CUlibrary (new) → CUmodule (legacy)
237+
module = object_code.get_module()
238+
239+
# Legacy module-only API consumes it directly
240+
count = handle_return(driver.cuModuleGetFunctionCount(module))
241+
assert count == 2
242+
243+
229244
@pytest.mark.parametrize(
230245
"attr, expected_type",
231246
[

0 commit comments

Comments
 (0)