nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/bindings/nvshmem.pyi

@@ -0,0 +1,179 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import _cython_3_2_4
+import enum
+import numpy.dtypes
+from _typeshed import Incomplete
+from typing import Any, Callable, ClassVar
+
+STATUS_FULL_MPG: int
+STATUS_INVALID: int
+STATUS_IS_BOOTSTRAPPED: int
+STATUS_IS_INITIALIZED: int
+STATUS_LIMITED_MPG: int
+STATUS_NOT_INITIALIZED: int
+TEAMS_MIN: int
+TEAM_GPU_LEADERS: int
+TEAM_GPU_LEADERS_INDEX: int
+TEAM_INDEX_MAX: int
+TEAM_INVALID: int
+TEAM_NODE: int
+TEAM_NODE_INDEX: int
+TEAM_SAME_GPU: int
+TEAM_SAME_GPU_INDEX: int
+TEAM_SAME_MYPE_NODE: int
+TEAM_SAME_MYPE_NODE_INDEX: int
+TEAM_SHARED: int
+TEAM_SHARED_INDEX: int
+TEAM_WORLD: int
+TEAM_WORLD_INDEX: int
+__pyx_capi__: dict
+__test__: dict
+align: _cython_3_2_4.cython_function_or_method
+barrier_all_on_stream: _cython_3_2_4.cython_function_or_method
+calloc: _cython_3_2_4.cython_function_or_method
+check_status: _cython_3_2_4.cython_function_or_method
+free: _cython_3_2_4.cython_function_or_method
+get_uniqueid: _cython_3_2_4.cython_function_or_method
+hostlib_finalize: _cython_3_2_4.cython_function_or_method
+hostlib_init_attr: _cython_3_2_4.cython_function_or_method
+init_attr_dtype: numpy.dtypes.VoidDType
+init_status: _cython_3_2_4.cython_function_or_method
+int_p: _cython_3_2_4.cython_function_or_method
+malloc: _cython_3_2_4.cython_function_or_method
+my_pe: _cython_3_2_4.cython_function_or_method
+n_pes: _cython_3_2_4.cython_function_or_method
+ptr: _cython_3_2_4.cython_function_or_method
+set_attr_uniqueid_args: _cython_3_2_4.cython_function_or_method
+sync_all_on_stream: _cython_3_2_4.cython_function_or_method
+team_my_pe: _cython_3_2_4.cython_function_or_method
+uniqueid_dtype: numpy.dtypes.VoidDType
+
+class Flags(enum.IntEnum):
+    """See `flags`."""
+    __new__: ClassVar[Callable] = ...
+    INIT_MAX: ClassVar[Flags] = ...
+    INIT_THREAD_PES: ClassVar[Flags] = ...
+    INIT_WITH_MPI_COMM: ClassVar[Flags] = ...
+    INIT_WITH_SHMEM: ClassVar[Flags] = ...
+    INIT_WITH_UNIQUEID: ClassVar[Flags] = ...
+    _generate_next_value_: ClassVar[Callable] = ...
+    _member_map_: ClassVar[dict] = ...
+    _member_names_: ClassVar[list] = ...
+    _member_type_: ClassVar[type[int]] = ...
+    _value2member_map_: ClassVar[dict] = ...
+
+class InitAttr:
+    """InitAttr()"""
+    mpi_comm: Incomplete
+    ptr: Incomplete
+    version: Incomplete
+    def __init__(self) -> Any:
+        """Initialize self.  See help(type(self)) for accurate signature."""
+    def __reduce__(self):
+        """InitAttr.__reduce_cython__(self)"""
+    def __reduce_cython__(self) -> Any:
+        """InitAttr.__reduce_cython__(self)"""
+    def __setstate_cython__(self, __pyx_state) -> Any:
+        """InitAttr.__setstate_cython__(self, __pyx_state)"""
+
+class NvshmemError(Exception):
+    def __init__(self, status) -> Any:
+        """NvshmemError.__init__(self, status)"""
+    def __reduce__(self) -> Any:
+        """NvshmemError.__reduce__(self)"""
+
+class Status(enum.IntEnum):
+    """See `nvshmemx_status`."""
+    __new__: ClassVar[Callable] = ...
+    ERROR_COLLECTIVE_LAUNCH_FAILED: ClassVar[Status] = ...
+    ERROR_GPU_NOT_SELECTED: ClassVar[Status] = ...
+    ERROR_INTERNAL: ClassVar[Status] = ...
+    ERROR_INVALID_VALUE: ClassVar[Status] = ...
+    ERROR_NOT_SUPPORTED: ClassVar[Status] = ...
+    ERROR_OUT_OF_MEMORY: ClassVar[Status] = ...
+    ERROR_SENTINEL: ClassVar[Status] = ...
+    ERROR_SYMMETRY: ClassVar[Status] = ...
+    SUCCESS: ClassVar[Status] = ...
+    _generate_next_value_: ClassVar[Callable] = ...
+    _member_map_: ClassVar[dict] = ...
+    _member_names_: ClassVar[list] = ...
+    _member_type_: ClassVar[type[int]] = ...
+    _value2member_map_: ClassVar[dict] = ...
+
+class UniqueId(uniqueid):
+    @classmethod
+    def __init__(cls, *args, **kwargs) -> None:
+        """Create and return a new object.  See help(type) for accurate signature."""
+    def __reduce__(self):
+        """UniqueId.__reduce_cython__(self)"""
+    def __reduce_cython__(self) -> Any:
+        """UniqueId.__reduce_cython__(self)"""
+    def __setstate_cython__(self, __pyx_state) -> Any:
+        """UniqueId.__setstate_cython__(self, __pyx_state)"""
+
+class uniqueid:
+    """uniqueid(size=1)
+
+    Empty-initialize an array of `nvshmemx_uniqueid_v1`.
+
+    The resulting object is of length `size` and of dtype `uniqueid_dtype`.
+    If default-constructed, the instance represents a single struct.
+
+    Args:
+        size (int): number of structs, default=1.
+
+
+    .. seealso:: `nvshmemx_uniqueid_v1`"""
+    ptr: Incomplete
+    version: Incomplete
+    def __init__(self, size=...) -> Any:
+        """Initialize self.  See help(type(self)) for accurate signature."""
+    @staticmethod
+    def from_data(data) -> Any:
+        """uniqueid.from_data(data)
+
+        Create an uniqueid instance wrapping the given NumPy array.
+
+        Args:
+            data (_numpy.ndarray): a 1D array of dtype `uniqueid_dtype` holding the data."""
+    @staticmethod
+    def from_ptr(intptr_tptr, size_tsize=..., boolreadonly=...) -> Any:
+        """uniqueid.from_ptr(intptr_t ptr, size_t size=1, bool readonly=False)
+
+        Create an uniqueid instance wrapping the given pointer.
+
+        Args:
+            ptr (intptr_t): pointer address as Python :class:`int` to the data.
+            size (int): number of structs, default=1.
+            readonly (bool): whether the data is read-only (to the user). default is `False`."""
+    def __delitem__(self, other) -> None:
+        """Delete self[key]."""
+    def __eq__(self, other: object) -> bool:
+        """Return self==value."""
+    def __ge__(self, other: object) -> bool:
+        """Return self>=value."""
+    def __getitem__(self, index):
+        """Return self[key]."""
+    def __gt__(self, other: object) -> bool:
+        """Return self>value."""
+    def __int__(self) -> int:
+        """int(self)"""
+    def __le__(self, other: object) -> bool:
+        """Return self<=value."""
+    def __len__(self) -> int:
+        """Return len(self)."""
+    def __lt__(self, other: object) -> bool:
+        """Return self<value."""
+    def __ne__(self, other: object) -> bool:
+        """Return self!=value."""
+    def __reduce__(self):
+        """uniqueid.__reduce_cython__(self)"""
+    def __reduce_cython__(self) -> Any:
+        """uniqueid.__reduce_cython__(self)"""
+    def __setitem__(self, index, object) -> None:
+        """Set self[key] to value."""
+    def __setstate_cython__(self, __pyx_state) -> Any:
+        """uniqueid.__setstate_cython__(self, __pyx_state)"""

nvmath/device/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .common import *  # noqa: F403
+from .common_cuda import *  # noqa: F403
+from .cublasdx import *  # noqa: F403
+from .cublasdx_backend import *  # noqa: F403
+from .cufftdx import *  # noqa: F403
+from .cusolverdx import *  # noqa: F403
+from .types import *  # noqa: F403
+from .vector_types_numba import *  # noqa: F403
+
+# isort: split
+
+# register models in numba; must occur after imports above
+from . import (
+    cublasdx_numba,  # noqa: F401
+    cufftdx_numba,  # noqa: F401
+    cusolverdx_numba,  # noqa: F401
+)

nvmath/device/_deprecated.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import sys
+from warnings import warn
+
+
+def _deprecated(
+    msg,
+    category: type[Warning] | None = DeprecationWarning,
+    stacklevel: int = 1,
+):
+    """Dropin replacement for @warnings.deprecated(...)"""
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            warn(msg, category=category, stacklevel=stacklevel + 1)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+deprecated = _deprecated
+
+if (sys.version_info.major, sys.version_info.minor) >= (3, 13):
+    from warnings import deprecated as _warnings_deprecated  # type: ignore
+
+    deprecated = _warnings_deprecated

nvmath/device/common.py

@@ -0,0 +1,313 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import tempfile
+from abc import abstractmethod
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+
+from .common_cuda import MAX_SUPPORTED_CC, MIN_SUPPORTED_CC, CodeType, ComputeCapability, get_current_device_cc
+
+__all__ = [
+    "make_tensor",
+    "make_fragment_like",
+    "OpaqueTensor",
+    "Layout",
+    "axpby",
+    "copy",
+    "copy_fragment",
+    "clear",
+    "copy_wait",
+]
+
+SHARED_DEVICE_DOCSTRINGS = {
+    "compiler": "A string to specify the compiler for the device code, currently supports ``None`` (default) and ``'numba'``",
+    #
+    "precision": """\
+The computation precision specified as a numpy float dtype, currently supports ``numpy.float16``, ``numpy.float32`` and
+``numpy.float64``.""".replace("\n", " "),
+    #
+    "code_type": "The target GPU code and compute-capability.",
+    "sm": "Target mathdx compute-capability.",
+    #
+    "execution": "A string specifying the execution method, can be ``'Block'`` or ``'Thread'``.",
+}
+
+
+# TODO: maybe pre-compile regular expression
+def make_binary_tempfile(content, suffix: str) -> tempfile._TemporaryFileWrapper:
+    """Write `content` to a temporary file with the given `suffix`.
+
+    A closed file object returned; it is the user's responsibility to delete the file when
+    finished.
+
+    .. seealso:: :py:func:`delete_binary_tempfiles`
+
+    """
+    with tempfile.NamedTemporaryFile(mode="w+b", suffix=suffix, delete=False) as tmp:
+        tmp.write(content)
+        tmp.flush()
+    return tmp
+
+
+def delete_binary_tempfiles(filenames: list[str]):
+    for name in filenames:
+        if os.path.isfile(name):
+            os.remove(name)
+
+
+def check_in(name, value, coll, format="{name} must be in {coll_str} ; got {name} = {value}"):
+    if value not in coll:
+        coll_str = ", ".join(f'"{t}"' for t in coll)
+        msg = format.format(name=name, value=value, coll_str=coll_str)
+        raise ValueError(msg)
+
+
+def check_not_in(name, value, coll, format="{name} must not be any of those value {coll_str} ; got {name} = {value}"):
+    if value in coll:
+        coll_str = ", ".join(f'"{t}"' for t in coll)
+        msg = format.format(name=name, value=value, coll_str=coll_str)
+        raise ValueError(msg)
+
+
+def check_contains(set, key):
+    check_in("", key, set, "{arg} must be in {set}")
+
+
+def parse_sm(sm: Any) -> ComputeCapability:
+    if sm is None:
+        sm = get_current_device_cc()
+    else:
+        if not isinstance(sm, ComputeCapability) and not isinstance(sm, int):
+            raise ValueError(f"sm should be a ComputeCapability or an int; got sm = {sm}")
+        if isinstance(sm, int):
+            sm = ComputeCapability(sm // 10, sm % 10)
+
+    return sm
+
+
+def check_sm(sm, library_name: str, var_name: str = "sm"):
+    if not isinstance(sm, ComputeCapability):
+        raise ValueError(f"{var_name} should be an instance of ComputeCapability ; got {var_name} = {sm}")
+    if sm < MIN_SUPPORTED_CC:
+        raise RuntimeError(f"Minimal compute capability {MIN_SUPPORTED_CC} is required by {library_name}, got {sm}")
+    if sm > MAX_SUPPORTED_CC:
+        raise RuntimeError(f"The maximum compute capability currently supported by device APIs is {MAX_SUPPORTED_CC}, got {sm}")
+    if sm.minor < 0:
+        raise ValueError(f"{var_name}.minor must be >= 0 ; got {var_name}.minor = {sm.minor}")
+
+
+def parse_code_type(code_type: Any) -> CodeType:
+    if not isinstance(code_type, Sequence) or len(code_type) != 2:
+        raise ValueError(f"code_type should be an instance of CodeType or a 2-tuple ; got code_type = {code_type}")
+
+    return CodeType(code_type[0], ComputeCapability(*code_type[1]))
+
+
+def check_code_type(code_type, library_name):
+    if not isinstance(code_type, CodeType):
+        raise ValueError(f"code_type should be an instance of CodeType ; got code_type = {code_type}")
+    check_sm(code_type.cc, library_name, "code_type.cc")
+    check_in("code_type.kind", code_type.kind, ["lto"])
+
+
+def pad_or_truncate(list, target_len):
+    return list[:target_len] + [0] * (target_len - len(list))
+
+
+class Layout:
+    """
+    Layout for the :py:class:`nvmath.device.OpaqueTensor`.
+
+    .. note:: Do not create directly, use appropriate method from
+        :py:func:`nvmath.device.Matmul`. Refer to
+        :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+        for guidance on which method to use.
+    """
+
+    def __init__(self):
+        raise RuntimeError("Layout should not be called directly.")
+
+    @property
+    @abstractmethod
+    def size(self) -> int:
+        """
+        Number of valid elements in a tensor. This is simply a product of all
+        shape dimensions.
+
+        Refer to the cuBLASDx documentation for more details on how to use this attribute:
+        :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def cosize(self) -> int:
+        """
+        Returns a distance from last element of a tensor to its first element.
+        It describes how many elements does the argument layout span.
+
+        Refer to the cuBLASDx documentation for more details on how to use this attribute:
+        :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def alignment(self) -> int:
+        """
+        Returns the required alignment (in bytes) for the tensor data buffer.
+
+        Refer to the cuBLASDx documentation for more details on how to use this attribute:
+        :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+        """
+        pass
+
+
+class OpaqueTensor:
+    """
+    Abstraction over the cuBLASDx tensor type (an alias of the CuTe tensor type).
+    The CuTe tensor layout is powerful and supports layouts not provided by NumPy,
+    so this a bridge to add this functionality to Python.
+
+    .. note:: Do not create directly, use :py:func:`nvmath.device.make_tensor`.
+
+    Refer to the cuBLASDx documentation for more details on how to use this class:
+    :cublasdx_doc:`api/other_tensors.html#tensors`
+    """
+
+    buffer: np.ndarray
+    layout: Layout
+    leading_dimension: int | None
+
+    def __init__(self, *args):
+        raise RuntimeError("OpaqueTensor should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def make_tensor(array: np.ndarray, layout: Layout) -> OpaqueTensor:
+    """
+    make_tensor is a helper function for creating
+    :py:class:`nvmath.device.OpaqueTensor` objects.
+
+    Args:
+        array (np.ndarray): The input array to be wrapped as an OpaqueTensor.
+        layout (Layout): The layout of the tensor, which describes how the data is
+            organized in memory.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#create-tensor-other-label`
+    """
+    raise RuntimeError("make_tensor should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def make_fragment_like(tensor: OpaqueTensor, dtype) -> OpaqueTensor:
+    """
+    make_fragment_like is a helper function for creating register fragments with
+    the same layout as input tensor, but different dtype.
+
+    Args:
+        tensor (OpaqueTensor): The input tensor to be used as a template for the fragment.
+        dtype: The data type of the fragment to be created.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+    """
+    raise RuntimeError("make_fragment_like should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def axpby(alpha: float, x_tensor: OpaqueTensor, beta: float, y_tensor: OpaqueTensor) -> None:
+    """
+    AXPBY operation: y = alpha * x + beta * y
+
+    Args:
+        alpha (float): Scalar multiplier for x_tensor.
+        x_tensor (OpaqueTensor): Input tensor x.
+        beta (float): Scalar multiplier for y_tensor.
+        y_tensor (OpaqueTensor): Input/output tensor y, which will be updated
+            with the result.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+    """
+    raise RuntimeError("axpby should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def copy(src: OpaqueTensor, dst: OpaqueTensor, alignment=None):
+    """
+    Copies data from the source tensor to the destination tensor.
+
+    Args:
+        src (OpaqueTensor): The source tensor to copy from.
+        dst (OpaqueTensor): The destination tensor to copy to.
+        alignment (int, optional): The alignment (in bytes) for the copy operation.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#cooperative-global-shared-copying`
+    """
+    raise RuntimeError("copy should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def copy_fragment(src: OpaqueTensor, dst: OpaqueTensor, alignment=None):
+    """
+    A bidirectional copying method to copy data between register fragments and
+    global memory tensors.
+
+    Args:
+        src (OpaqueTensor): The source tensor to copy from.
+        dst (OpaqueTensor): The destination tensor to copy to.
+        alignment (int, optional): The alignment (in bytes) for the copy operation.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#copying-registers-tensors`
+    """
+    raise RuntimeError("copy_fragment should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def clear(arr: OpaqueTensor):
+    """
+    Clears the contents of the given tensor by setting all elements to zero.
+
+    Args:
+        arr (OpaqueTensor): The tensor to be cleared.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#imported-tensor-utilities`
+    """
+    raise RuntimeError("clear should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def copy_wait():
+    """
+    Creates synchronization point. It has to be called after :py:func:`nvmath.device.copy`
+    to ensure that the copy operation has completed before any subsequent
+    operations are executed.
+
+    Refer to the cuBLASDx documentation for more details on how to use this function:
+    :cublasdx_doc:`api/other_tensors.html#cooperative-global-shared-copying`
+    """
+    raise RuntimeError("copy_wait should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def check_positive_integer_sequence(arg, arg_name, min_len, max_len):
+    if not isinstance(arg, Sequence):
+        raise ValueError(f'Parameter "{arg_name}" must be a sequence. Got: {type(arg).__name__}.')
+
+    for i, x in enumerate(arg):
+        if not isinstance(x, (int, np.integer)):
+            raise ValueError(
+                f'Parameter "{arg_name}" must be a sequence of positive integers. '
+                + f"Element {arg_name}[{i}] is {type(x).__name__}."
+            )
+        if x <= 0:
+            raise ValueError(f'Parameter "{arg_name}" values must be positive integers. Got {arg_name}[{i}] = {x}.')
+
+    if min_len == max_len:
+        if len(arg) != min_len:
+            raise ValueError(f'Parameter "{arg_name}" must be a sequence of length {min_len}. Got length: {len(arg)}.')
+    else:
+        if len(arg) > max_len or len(arg) < min_len:
+            raise ValueError(f'Parameter "{arg_name}" must define {min_len} to {max_len} values. Got {len(arg)} value(s).')

nvmath/device/common_backend.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import weakref
+from collections.abc import Callable
+
+import numpy as np
+
+from nvmath.bindings import mathdx
+from nvmath.device.common_cuda import ComputeCapability, ISAVersion
+
+from .types import complex32, complex64, complex128, half2, half4
+
+MATHDX_TYPES_TO_NP = {
+    mathdx.CommondxValueType.R_16F: np.float16,
+    mathdx.CommondxValueType.R_16F2: half2,
+    mathdx.CommondxValueType.R_32F: np.float32,
+    mathdx.CommondxValueType.R_64F: np.float64,
+    mathdx.CommondxValueType.C_16F: complex32,
+    mathdx.CommondxValueType.C_16F2: half4,
+    mathdx.CommondxValueType.C_32F: complex64,
+    mathdx.CommondxValueType.C_64F: complex128,
+    mathdx.CommondxValueType.R_8I: np.int8,
+    mathdx.CommondxValueType.R_16I: np.int16,
+    mathdx.CommondxValueType.R_32I: np.int32,
+    mathdx.CommondxValueType.R_64I: np.int64,
+    mathdx.CommondxValueType.R_8UI: np.uint8,
+    mathdx.CommondxValueType.R_16UI: np.uint16,
+    mathdx.CommondxValueType.R_32UI: np.uint32,
+    mathdx.CommondxValueType.R_64UI: np.uint64,
+}
+
+
+NP_TYPES_TO_MATHDX_TYPES = {np_type: mathdx_type for mathdx_type, np_type in MATHDX_TYPES_TO_NP.items()}
+
+NP_TYPES_TO_MATHDX_PRECISION = {
+    np.float16: mathdx.CommondxPrecision.F16,
+    np.float32: mathdx.CommondxPrecision.F32,
+    np.float64: mathdx.CommondxPrecision.F64,
+    np.int8: mathdx.CommondxPrecision.I8,
+    np.int16: mathdx.CommondxPrecision.I16,
+    np.int32: mathdx.CommondxPrecision.I32,
+    np.int64: mathdx.CommondxPrecision.I64,
+    np.uint8: mathdx.CommondxPrecision.UI8,
+    np.uint16: mathdx.CommondxPrecision.UI16,
+    np.uint32: mathdx.CommondxPrecision.UI32,
+    np.uint64: mathdx.CommondxPrecision.UI64,
+}
+
+EXECUTION_STR_TO_MATHDX = {
+    "Block": mathdx.CommondxExecution.BLOCK,
+    "Thread": mathdx.CommondxExecution.THREAD,
+}
+
+NVARG_GEN_OPT_LTO = "-gen-opt-lto"
+
+
+class DescriptorWrapper:
+    """
+    A smart pointer to a descriptor.
+
+    Lives as long as we need the descriptor (due to lru cache usage). Once the
+    descriptor is no longer needed/cached, it will trigger proper resource
+    deallocation.
+    """
+
+    def __init__(self, descriptor, destructor):
+        self.descriptor = descriptor
+        self._destructor = destructor
+        if destructor is not None:
+            weakref.finalize(self, self._destructor, self.descriptor)
+
+
+def get_lto(code_descriptor: int) -> bytes:
+    """Get lto binary from the mathdx common code descriptor."""
+    lto_size = mathdx.commondx_get_code_ltoir_size(code_descriptor)
+
+    lto_fn = bytearray(lto_size)
+    mathdx.commondx_get_code_ltoir(code_descriptor, lto_size, lto_fn)
+    return bytes(lto_fn)
+
+
+def get_isa_version(code_descriptor: int) -> ISAVersion:
+    """Parse isa version from the mathdx common code descriptor."""
+    isa = mathdx.commondx_get_code_option_int64(code_descriptor, mathdx.CommondxOption.CODE_ISA)
+    return ISAVersion.from_integer(isa)
+
+
+def build_get_int_traits(get_trait_int64s: Callable[[int, int, int, int], int]):
+    """Generate function that returns mathdx tuple of int traits."""
+
+    def get_int_traits(handle: int, trait_type: int, size: int) -> tuple:
+        int_buffer = np.zeros(size, np.int64)
+        get_trait_int64s(handle, trait_type, size, int_buffer.ctypes.data)
+        return tuple(map(int, int_buffer))
+
+    return get_int_traits
+
+
+def build_get_str_trait(
+    get_trait_str_size: Callable[[int, int], int],
+    get_trait_str: Callable[[int, int, int, bytearray], bytearray],
+):
+    """Generate function that returns mathdx string trait."""
+
+    def get_str_trait(handle: int, trait_type: int) -> str:
+        symbol_size = get_trait_str_size(handle, trait_type)
+
+        symbol = bytearray(symbol_size)
+
+        get_trait_str(handle, trait_type, symbol_size, symbol)
+
+        # terminate trailing 0 (indicator of c-string end)
+        symbol_str = symbol[:-1].decode()
+
+        return symbol_str
+
+    return get_str_trait
+
+
+def get_mathdx_sm(version: ComputeCapability) -> list:
+    sm = [version.integer]
+    if version.arch == "a":
+        sm.append(mathdx.CommondxArchModifier.ARCH_SPECIFIC)
+    return sm
+
+
+def set_code_target_sm(code: int, version: ComputeCapability):
+    sm = get_mathdx_sm(version)
+    mathdx.commondx_set_code_option_int64s(code, mathdx.CommondxOption.TARGET_SM, len(sm), sm)