From 600176e9de036d4c8685fd53f7e5082e21cb731c Mon Sep 17 00:00:00 2001
From: Jinfeng <jinfengl@nvidia.com>
Date: Fri, 10 Jul 2026 23:31:00 +0000
Subject: [PATCH] add an api to return module in ObjectCode to support calling
 legacy driver APIs

---
 cuda_core/cuda/core/_module.pyi | 10 ++++++++++
 cuda_core/cuda/core/_module.pyx | 16 ++++++++++++++++
 cuda_core/tests/test_module.py  | 15 +++++++++++++++
 3 files changed, 41 insertions(+)

diff --git a/cuda_core/cuda/core/_module.pyi b/cuda_core/cuda/core/_module.pyi
index 5125b99131a..6c0cc85b50a 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 91c8ad43895..990a2c8445c 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(<intptr_t>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 a0600966619..b5c9b4bee7a 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",
     [