Add more ObjectCode constructors (#652)

* add ctors

* updates

* add CU_JIT_INPUT_LIBRARY

* lib -> library everywhere

* fix

* Add missing `@staticmethod`

---------

Co-authored-by: Leo Fang <leo80042@gmail.com>
Co-authored-by: Ralf W. Grosse-Kunstleve <rwgkio@gmail.com>
Co-authored-by: Ralf W. Grosse-Kunstleve <rgrossekunst@nvidia.com>

Author: 3 people

cuda_core/cuda/core/experimental/_linker.py

@@ -77,13 +77,15 @@ def _lazy_init():
             "fatbin": _nvjitlink.InputType.FATBIN,
             "ltoir": _nvjitlink.InputType.LTOIR,
             "object": _nvjitlink.InputType.OBJECT,
+            "library": _nvjitlink.InputType.LIBRARY,
         }
     else:
         _driver_input_types = {
             "ptx": _driver.CUjitInputType.CU_JIT_INPUT_PTX,
             "cubin": _driver.CUjitInputType.CU_JIT_INPUT_CUBIN,
             "fatbin": _driver.CUjitInputType.CU_JIT_INPUT_FATBINARY,
             "object": _driver.CUjitInputType.CU_JIT_INPUT_OBJECT,
+            "library": _driver.CUjitInputType.CU_JIT_INPUT_LIBRARY,
         }
     _inited = True
 

cuda_core/cuda/core/experimental/_module.py

@@ -276,7 +276,7 @@ class ObjectCode:
     """
 
     __slots__ = ("_handle", "_backend_version", "_code_type", "_module", "_loader", "_sym_map")
-    _supported_code_type = ("cubin", "ptx", "ltoir", "fatbin")
+    _supported_code_type = ("cubin", "ptx", "ltoir", "fatbin", "object", "library")
 
     def __new__(self, *args, **kwargs):
         raise RuntimeError(
@@ -342,6 +342,70 @@ def from_ptx(module: Union[bytes, str], *, symbol_mapping: Optional[dict] = None
         """
         return ObjectCode._init(module, "ptx", symbol_mapping=symbol_mapping)
 
+    @staticmethod
+    def from_ltoir(module: Union[bytes, str], *, symbol_mapping: Optional[dict] = None) -> "ObjectCode":
+        """Create an :class:`ObjectCode` instance from an existing LTOIR.
+
+        Parameters
+        ----------
+        module : Union[bytes, str]
+            Either a bytes object containing the in-memory ltoir code to load, or
+            a file path string pointing to the on-disk ltoir file 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).
+        """
+        return ObjectCode._init(module, "ltoir", symbol_mapping=symbol_mapping)
+
+    @staticmethod
+    def from_fatbin(module: Union[bytes, str], *, symbol_mapping: Optional[dict] = None) -> "ObjectCode":
+        """Create an :class:`ObjectCode` instance from an existing fatbin.
+
+        Parameters
+        ----------
+        module : Union[bytes, str]
+            Either a bytes object containing the in-memory fatbin to load, or
+            a file path string pointing to the on-disk fatbin 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).
+        """
+        return ObjectCode._init(module, "fatbin", symbol_mapping=symbol_mapping)
+
+    @staticmethod
+    def from_object(module: Union[bytes, str], *, symbol_mapping: Optional[dict] = None) -> "ObjectCode":
+        """Create an :class:`ObjectCode` instance from an existing object code.
+
+        Parameters
+        ----------
+        module : Union[bytes, str]
+            Either a bytes object containing the in-memory object code to load, or
+            a file path string pointing to the on-disk object code 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).
+        """
+        return ObjectCode._init(module, "object", symbol_mapping=symbol_mapping)
+
+    @staticmethod
+    def from_library(module: Union[bytes, str], *, symbol_mapping: Optional[dict] = None) -> "ObjectCode":
+        """Create an :class:`ObjectCode` instance from an existing library.
+
+        Parameters
+        ----------
+        module : Union[bytes, str]
+            Either a bytes object containing the in-memory library to load, or
+            a file path string pointing to the on-disk library 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).
+        """
+        return ObjectCode._init(module, "library", symbol_mapping=symbol_mapping)
+
     # TODO: do we want to unload in a finalizer? Probably not..
 
     def _lazy_load_module(self, *args, **kwargs):