expose ObjectCode to public + fix file loading

Author: leofang

cuda_core/cuda/core/experimental/__init__.py

@@ -7,6 +7,7 @@
 from cuda.core.experimental._event import EventOptions
 from cuda.core.experimental._launcher import LaunchConfig, launch
 from cuda.core.experimental._linker import Linker, LinkerOptions
+from cuda.core.experimental._module import ObjectCode
 from cuda.core.experimental._program import Program, ProgramOptions
 from cuda.core.experimental._stream import Stream, StreamOptions
 from cuda.core.experimental._system import System

cuda_core/cuda/core/experimental/_module.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
-from typing import Union
+from typing import Optional, Union
 from warnings import warn
 
 from cuda.core.experimental._utils import driver, get_binding_version, handle_return, precondition
@@ -236,12 +236,16 @@ class ObjectCode:
     module : Union[bytes, str]
         Either a bytes object containing the cubin to load, or
         a file path string pointing to the cubin to load.
+    symbol_mapping : Optional[dict]
+        A dictionary specifying how the unmangled symbol names (as keys)
+        should be mapped to the mangled names before trying to retrieve
+        them (default to no mappings).
     """
 
     __slots__ = ("_handle", "_backend_version", "_code_type", "_module", "_loader", "_sym_map")
     _supported_code_type = ("cubin", "ptx", "ltoir", "fatbin")
 
-    def __init__(self, module: Union[bytes, str]):
+    def __init__(self, module: Union[bytes, str], *, symbol_mapping: Optional[dict]=None):
         _lazy_init()
 
         # handle is assigned during _lazy_load
@@ -250,9 +254,9 @@ def __init__(self, module: Union[bytes, str]):
         self._loader = _backend[self._backend_version]
         self._code_type = "cubin"
         self._module = module
-        self._sym_map = {}
+        self._sym_map = {} if symbol_mapping is None else symbol_mapping
 
-    def _init(module, code_type, *, symbol_mapping=None):
+    def _init(module, code_type, *, symbol_mapping: Optional[dict]=None):
         self = ObjectCode.__new__(ObjectCode)
         if code_type not in self._supported_code_type:
             raise ValueError
@@ -278,9 +282,9 @@ def _lazy_load_module(self, *args, **kwargs):
         module = self._module
         if isinstance(module, str):
             if self._backend_version == "new":
-                self._handle = handle_return(self._loader["file"](module, [], [], 0, [], [], 0))
+                self._handle = handle_return(self._loader["file"](module.encode(), [], [], 0, [], [], 0))
             else:  # "old" backend
-                self._handle = handle_return(self._loader["file"](module))
+                self._handle = handle_return(self._loader["file"](module.encode()))
         else:
             assert isinstance(module, bytes)
             if self._backend_version == "new":

cuda_core/docs/source/release/0.2.0-notes.rst

@@ -12,15 +12,19 @@ Highlights
 - Add :class:`~DeviceProperties` to provide pythonic access to device properties.
 - Add kernel attributes to :class:`~Kernel`
 
-Limitations
------------
-
-- <TODO>
-
 Breaking Changes
 ----------------
 
 - Change ``__cuda_stream__`` from attribute to method
 - The :meth:`~Program.compile` method no longer accepts the `options` argument. Instead, you can optionally pass an instance of :class:`~ProgramOptions` to the constructor of :obj:`~Program`.
-- The internal constructor of :class:`~ObjectCode` no longer accepts the jit_options argument. Options are provided to upstream :class:`~ProgramOptions` or :class:`~LinkerOptions` instead.
-- :meth: `~Device.properties` now provides an instance of :class:`~DeviceProperties` instead of a dictionary. 
+- :meth: `~Device.properties` now provides an instance of :class:`~DeviceProperties` instead of a dictionary.
+
+New features
+------------
+
+- Expose :class:`ObjectCode` as a public API, which allows loading cubins from memory or disk. For loading other kinds of code types, please continue using :class:`Program`.
+
+Limitations
+-----------
+
+- <TODO>

cuda_core/tests/test_module.py

@@ -10,7 +10,7 @@
 import pytest
 from conftest import can_load_generated_ptx
 
-from cuda.core.experimental import Program, ProgramOptions, system
+from cuda.core.experimental import ObjectCode, Program, ProgramOptions, system
 
 
 @pytest.fixture(scope="function")
@@ -37,7 +37,7 @@ def get_saxpy_kernel(init_cuda):
     )
 
     # run in single precision
-    return mod.get_kernel("saxpy<float>")
+    return mod.get_kernel("saxpy<float>"), mod
 
 
 @pytest.mark.xfail(not can_load_generated_ptx(), reason="PTX version too new")
@@ -72,7 +72,7 @@ def test_get_kernel(init_cuda):
     ],
 )
 def test_read_only_kernel_attributes(get_saxpy_kernel, attr, expected_type):
-    kernel = get_saxpy_kernel
+    kernel, _ = get_saxpy_kernel
     method = getattr(kernel.attributes, attr)
     # get the value without providing a device ordinal
     value = method()
@@ -82,3 +82,23 @@ def test_read_only_kernel_attributes(get_saxpy_kernel, attr, expected_type):
     for device in system.devices:
         value = method(device.device_id)
     assert isinstance(value, expected_type), f"Expected {attr} to be of type {expected_type}, but got {type(value)}"
+
+
+def test_object_code_load_cubin(get_saxpy_kernel):
+    _, mod = get_saxpy_kernel
+    cubin = mod._module
+    sym_map = mod._sym_map
+    assert isinstance(cubin, bytes)
+    mod = ObjectCode(cubin, symbol_mapping=sym_map)
+    ker = mod.get_kernel("saxpy<double>")
+
+
+def test_object_code_load_cubin_from_file(get_saxpy_kernel, tmp_path):
+    _, mod = get_saxpy_kernel
+    cubin = mod._module
+    sym_map = mod._sym_map
+    assert isinstance(cubin, bytes)
+    cubin_file = tmp_path / "test.cubin"
+    cubin_file.write_bytes(cubin)
+    mod = ObjectCode(str(cubin_file), symbol_mapping=sym_map)
+    ker = mod.get_kernel("saxpy<double>")