nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/bindings/nvpl/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# type: ignore
+
+from . import fft
+
+__all__ = [
+    "fft",
+]

nvmath/bindings/nvpl/_internal/fft.pxd

@@ -0,0 +1,36 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# This code was automatically generated with version 0.3.0. Do not modify it directly.
+
+from ..cyfft cimport *
+
+
+###############################################################################
+# Wrapper functions
+###############################################################################
+
+cdef int _nvpl_fft_get_version() except?-42 nogil
+cdef fftw_plan _fftw_plan_many_dft(int rank, const int* n, int batch, fftw_complex* in_, const int* inembed, int istride, int idist, fftw_complex* out, const int* onembed, int ostride, int odist, int sign, unsigned flags) except?NULL nogil
+cdef fftw_plan _fftw_plan_many_dft_r2c(int rank, const int* n, int batch, double* in_, const int* inembed, int istride, int idist, fftw_complex* out, const int* onembed, int ostride, int odist, unsigned flags) except?NULL nogil
+cdef fftw_plan _fftw_plan_many_dft_c2r(int rank, const int* n, int batch, fftw_complex* in_, const int* inembed, int istride, int idist, double* out, const int* onembed, int ostride, int odist, unsigned flags) except?NULL nogil
+cdef void _fftw_execute_dft(const fftw_plan plan, fftw_complex* idata, fftw_complex* odata) except* nogil
+cdef void _fftw_execute_dft_r2c(const fftw_plan plan, double* idata, fftw_complex* odata) except* nogil
+cdef void _fftw_execute_dft_c2r(const fftw_plan plan, fftw_complex* idata, double* odata) except* nogil
+cdef fftwf_plan _fftwf_plan_many_dft(int rank, const int* n, int batch, fftwf_complex* in_, const int* inembed, int istride, int idist, fftwf_complex* out, const int* onembed, int ostride, int odist, int sign, unsigned flags) except?NULL nogil
+cdef fftwf_plan _fftwf_plan_many_dft_r2c(int rank, const int* n, int batch, float* in_, const int* inembed, int istride, int idist, fftwf_complex* out, const int* onembed, int ostride, int odist, unsigned flags) except?NULL nogil
+cdef fftwf_plan _fftwf_plan_many_dft_c2r(int rank, const int* n, int batch, fftwf_complex* in_, const int* inembed, int istride, int idist, float* out, const int* onembed, int ostride, int odist, unsigned flags) except?NULL nogil
+cdef void _fftwf_execute_dft(const fftwf_plan plan, fftwf_complex* idata, fftwf_complex* odata) except* nogil
+cdef void _fftwf_execute_dft_r2c(const fftwf_plan plan, float* idata, fftwf_complex* odata) except* nogil
+cdef void _fftwf_execute_dft_c2r(const fftwf_plan plan, fftwf_complex* idata, float* odata) except* nogil
+cdef int _fftw_init_threads() except?-42 nogil
+cdef int _fftwf_init_threads() except?-42 nogil
+cdef void _fftw_plan_with_nthreads(int nthreads) except* nogil
+cdef void _fftwf_plan_with_nthreads(int nthreads) except* nogil
+cdef int _fftw_planner_nthreads() except?-42 nogil
+cdef int _fftwf_planner_nthreads() except?-42 nogil
+cdef void _fftw_cleanup_threads() except* nogil
+cdef void _fftwf_cleanup_threads() except* nogil
+cdef void _fftw_destroy_plan(fftw_plan plan) except* nogil
+cdef void _fftwf_destroy_plan(fftwf_plan plan) except* nogil

nvmath/bindings/nvpl/cyfft.pxd

@@ -0,0 +1,87 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# This code was automatically generated with version 0.3.0. Do not modify it directly.
+# This layer exposes the C header to Cython as-is.
+
+###############################################################################
+# Types (structs, enums, ...)
+###############################################################################
+
+# enums
+
+
+
+# types
+cdef extern from *:
+    """
+    // Transform direction
+    #define FFTW_FORWARD -1
+    #define FFTW_INVERSE  1
+    #define FFTW_BACKWARD 1
+
+    // Planner flags
+    #define FFTW_ESTIMATE           0x01
+    #define FFTW_MEASURE            0x02
+    #define FFTW_PATIENT            0x03
+    #define FFTW_EXHAUSTIVE         0x04
+    #define FFTW_WISDOM_ONLY        0x05
+
+    // Algorithm restriction flags
+    #define FFTW_DESTROY_INPUT      0x08
+    #define FFTW_PRESERVE_INPUT     0x0C
+    #define FFTW_UNALIGNED          0x10
+
+    typedef double fftw_complex[2] __attribute__ ((aligned (16)));
+    typedef float fftwf_complex[2] __attribute__ ((aligned (8)));
+    """
+
+    cdef const int FFTW_FORWARD
+    cdef const int FFTW_INVERSE
+
+    cdef const int FFTW_ESTIMATE
+    cdef const int FFTW_MEASURE
+    cdef const int FFTW_PATIENT
+    cdef const int FFTW_EXHAUSTIVE
+    cdef const int FFTW_WISDOM_ONLY
+
+    cdef const int FFTW_DESTROY_INPUT
+    cdef const int FFTW_PRESERVE_INPUT
+    cdef const int FFTW_UNALIGNED
+
+    ctypedef double fftw_complex[2]
+    ctypedef float fftwf_complex[2]
+
+
+ctypedef void* fftw_plan 'fftw_plan'
+ctypedef void* fftwf_plan 'fftwf_plan'
+
+
+###############################################################################
+# Functions
+###############################################################################
+
+cdef int nvpl_fft_get_version() except?-42 nogil
+cdef fftw_plan fftw_plan_many_dft(int rank, const int* n, int batch, fftw_complex* in_, const int* inembed, int istride, int idist, fftw_complex* out, const int* onembed, int ostride, int odist, int sign, unsigned flags) except?NULL nogil
+cdef fftw_plan fftw_plan_many_dft_r2c(int rank, const int* n, int batch, double* in_, const int* inembed, int istride, int idist, fftw_complex* out, const int* onembed, int ostride, int odist, unsigned flags) except?NULL nogil
+cdef fftw_plan fftw_plan_many_dft_c2r(int rank, const int* n, int batch, fftw_complex* in_, const int* inembed, int istride, int idist, double* out, const int* onembed, int ostride, int odist, unsigned flags) except?NULL nogil
+cdef void fftw_execute_dft(const fftw_plan plan, fftw_complex* idata, fftw_complex* odata) except* nogil
+cdef void fftw_execute_dft_r2c(const fftw_plan plan, double* idata, fftw_complex* odata) except* nogil
+cdef void fftw_execute_dft_c2r(const fftw_plan plan, fftw_complex* idata, double* odata) except* nogil
+cdef fftwf_plan fftwf_plan_many_dft(int rank, const int* n, int batch, fftwf_complex* in_, const int* inembed, int istride, int idist, fftwf_complex* out, const int* onembed, int ostride, int odist, int sign, unsigned flags) except?NULL nogil
+cdef fftwf_plan fftwf_plan_many_dft_r2c(int rank, const int* n, int batch, float* in_, const int* inembed, int istride, int idist, fftwf_complex* out, const int* onembed, int ostride, int odist, unsigned flags) except?NULL nogil
+cdef fftwf_plan fftwf_plan_many_dft_c2r(int rank, const int* n, int batch, fftwf_complex* in_, const int* inembed, int istride, int idist, float* out, const int* onembed, int ostride, int odist, unsigned flags) except?NULL nogil
+cdef void fftwf_execute_dft(const fftwf_plan plan, fftwf_complex* idata, fftwf_complex* odata) except* nogil
+cdef void fftwf_execute_dft_r2c(const fftwf_plan plan, float* idata, fftwf_complex* odata) except* nogil
+cdef void fftwf_execute_dft_c2r(const fftwf_plan plan, fftwf_complex* idata, float* odata) except* nogil
+cdef int fftw_init_threads() except?-42 nogil
+cdef int fftwf_init_threads() except?-42 nogil
+cdef void fftw_plan_with_nthreads(int nthreads) except* nogil
+cdef void fftwf_plan_with_nthreads(int nthreads) except* nogil
+cdef int fftw_planner_nthreads() except?-42 nogil
+cdef int fftwf_planner_nthreads() except?-42 nogil
+cdef void fftw_cleanup_threads() except* nogil
+cdef void fftwf_cleanup_threads() except* nogil
+cdef void fftw_destroy_plan(fftw_plan plan) except* nogil
+cdef void fftwf_destroy_plan(fftwf_plan plan) except* nogil

nvmath/bindings/nvpl/fft.pxd

@@ -0,0 +1,100 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# This code was automatically generated with version 0.3.0. Do not modify it directly.
+
+from libc.stdint cimport intptr_t
+
+from .cyfft cimport *
+
+
+###############################################################################
+# Types
+###############################################################################
+
+
+
+
+cdef class Plan:
+    cdef intptr_t forward_handle
+    cdef intptr_t inverse_handle
+    cdef int precision
+    cdef int kind
+
+
+###############################################################################
+# Enum
+###############################################################################
+
+
+# The `UNSPECIFIED` value is not valid sign in NVPL's header,
+# but is used here to allow for planning C2C that can handle
+# both directions
+cpdef enum Sign:
+    UNSPECIFIED = 0
+    FORWARD = FFTW_FORWARD
+    INVERSE = FFTW_INVERSE
+
+
+cpdef enum PlannerFlags:
+    ESTIMATE = FFTW_ESTIMATE
+    MEASURE = FFTW_MEASURE
+    PATIENT = FFTW_PATIENT
+    EXHAUSTIVE = FFTW_EXHAUSTIVE
+    WISDOM_ONLY = FFTW_WISDOM_ONLY
+
+
+cpdef enum Precision:
+    FLOAT = 32
+    DOUBLE = 64
+
+
+cpdef enum Kind:
+    C2C = 1
+    C2R = 2
+    R2C = 3
+
+
+
+
+
+###############################################################################
+# Convenience wrappers/adapters
+###############################################################################
+
+cpdef int plan_with_nthreads(int precision, int num_threads) except -1
+cpdef int planner_nthreads(int precision) except? -1
+cpdef int init_threads(int precision) except -1
+cpdef int cleanup_threads(int precision) except -1
+cpdef Plan plan_many(int precision, int kind, int sign, int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, unsigned flags)
+cpdef int destroy(Plan plan) except -1
+cpdef int execute(Plan plan, intptr_t idata, intptr_t odata, int sign) except -1
+
+###############################################################################
+# Functions
+###############################################################################
+
+cpdef int get_version() except? -1
+cpdef intptr_t plan_many_c2c_double(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, int sign, unsigned flags) except? 0
+cpdef intptr_t plan_many_r2c_double(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, unsigned flags) except? 0
+cpdef intptr_t plan_many_c2r_double(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, unsigned flags) except? 0
+cpdef void execute_c2c_double(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef void execute_r2c_double(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef void execute_c2r_double(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef intptr_t plan_many_c2c_float(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, int sign, unsigned flags) except? 0
+cpdef intptr_t plan_many_r2c_float(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, unsigned flags) except? 0
+cpdef intptr_t plan_many_c2r_float(int rank, n, int batch, intptr_t in_, inembed, int istride, int idist, intptr_t out, onembed, int ostride, int odist, unsigned flags) except? 0
+cpdef void execute_c2c_float(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef void execute_r2c_float(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef void execute_c2r_float(intptr_t plan, intptr_t idata, intptr_t odata) except*
+cpdef int init_threads_double() except 0
+cpdef int init_threads_float() except 0
+cpdef void plan_with_nthreads_double(int nthreads) except*
+cpdef void plan_with_nthreads_float(int nthreads) except*
+cpdef int planner_nthreads_double() except? 0
+cpdef int planner_nthreads_float() except? 0
+cpdef void cleanup_threads_double() except*
+cpdef void cleanup_threads_float() except*
+cpdef void destroy_plan_double(intptr_t plan) except*
+cpdef void destroy_plan_float(intptr_t plan) except*

nvmath/bindings/nvshmem.pxd

@@ -0,0 +1,54 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# This code was automatically generated with version 3.1.7. Do not modify it directly.
+
+from libc.stdint cimport intptr_t
+
+from .cynvshmem cimport *
+
+
+###############################################################################
+# Types
+###############################################################################
+
+ctypedef nvshmem_team_config_v1 team_config
+ctypedef nvshmemx_uniqueid_args_v1 uniqueid_args
+ctypedef nvshmemx_init_args_v1 init_args
+ctypedef nvshmemx_init_attr_v1 init_attr
+
+ctypedef cudaStream_t Stream
+
+
+###############################################################################
+# Enum
+###############################################################################
+
+ctypedef _anon_enum0 __anon_enum0
+ctypedef _anon_enum1 __anon_enum1
+ctypedef _anon_enum2 __anon_enum2
+ctypedef nvshmemx_status _Status
+ctypedef flags _Flags
+
+
+###############################################################################
+# Functions
+###############################################################################
+
+cpdef int init_status() except? -1
+cpdef int my_pe() except? -1
+cpdef int n_pes() except? -1
+cpdef intptr_t malloc(size_t size) except? 0
+cpdef intptr_t calloc(size_t count, size_t size) except? 0
+cpdef intptr_t align(size_t alignment, size_t size) except? 0
+cpdef void free(intptr_t ptr) except*
+cpdef intptr_t ptr(intptr_t dest, int pe) except? 0
+cpdef void int_p(intptr_t dest, int value, int pe) except*
+cpdef int team_my_pe(int32_t team) except? -1
+cpdef void barrier_all_on_stream(intptr_t stream) except*
+cpdef void sync_all_on_stream(intptr_t stream) except*
+cpdef hostlib_init_attr(unsigned int flags, intptr_t attr)
+cpdef void hostlib_finalize() except*
+cpdef set_attr_uniqueid_args(int myrank, int nranks, intptr_t uniqueid, intptr_t attr)
+cpdef get_uniqueid(intptr_t uniqueid)

nvmath/device/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .patch import patch_codegen
+from nvmath._utils import force_loading_libmathdx
+
+patch_codegen()
+force_loading_libmathdx("12")
+
+from .common_cuda import *  # noqa: E402, F403
+from .cufftdx import *  # noqa: E402, F403
+from .cublasdx import *  # noqa: E402, F403
+from .cublasdx_backend import *  # noqa: E402, F403
+from .vector_types_numba import *  # noqa: E402, F403
+from . import nvrtc  # noqa: E402, F403, F401
+from .common import make_tensor  # noqa: E402, F401
+
+# register models in numba
+from . import cublasdx_numba  # noqa: E402, F401
+
+del patch_codegen
+del force_loading_libmathdx

nvmath/device/_deprecated.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2024-2025, 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/caching.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import hashlib
+import json
+import logging
+import pickle
+
+_ENABLE_CACHE = "NVMATH_ENABLE_CACHE" in os.environ
+if _ENABLE_CACHE:
+    logging.warning("nvmath: NVMATH_ENABLE_CACHE is set in the environment, cache is enabled")
+    _CACHE_LOCATION = os.path.join(os.path.expanduser("~"), ".cache", "nvmath-device")
+    if not os.path.exists(_CACHE_LOCATION):
+        logging.debug(f"pymathdx: creating directory {_CACHE_LOCATION}")
+        os.makedirs(_CACHE_LOCATION)
+
+
+# We use
+# json.dumps to serialize args/kwargs to a string
+# hashlib to compute the hash
+def json_hash(*args, **kwargs):
+    hasher = hashlib.sha1()
+    hasher.update(json.dumps([args, kwargs]).encode("utf-8"))
+    return hasher.hexdigest()
+
+
+def disk_cache(func):
+    def cacher(*args, **kwargs):
+        if _ENABLE_CACHE:
+            # compute hash(args, kwargs)
+            h = json_hash(*args, **kwargs)
+            # if file exist...
+            if os.path.isfile(os.path.join(_CACHE_LOCATION, h)):
+                # open it
+                with open(os.path.join(_CACHE_LOCATION, h), "rb") as f:
+                    out = pickle.load(f)
+                # return cache
+                return out
+            else:
+                # compute output
+                out = func(*args, **kwargs)
+                # store to file
+                with open(os.path.join(_CACHE_LOCATION, h), "wb") as f:
+                    pickle.dump(out, f)
+                return out
+        else:
+            return func(*args, **kwargs)
+
+    return cacher

nvmath/device/common.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import abstractmethod
+import tempfile
+
+import numpy as np
+
+from .common_cuda import CodeType
+
+
+__all__ = [
+    "make_tensor",
+    "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.",
+    #
+    "execution": "A string specifying the execution method, can be ``'Block'`` or ``'Thread'``.",
+}
+
+
+# TODO: maybe pre-compile regular expression
+def make_binary_tempfile(content, suffix):
+    # TODO: may need to set it False for Windows? (refer to Python API doc)
+    tmp = tempfile.NamedTemporaryFile(mode="w+b", suffix=suffix, delete=True)  # noqa: SIM115
+    tmp.write(content)
+    tmp.flush()
+    return tmp
+
+
+def check_in(name, arg, set):
+    if arg not in set:
+        raise ValueError(f"{name} must be in {set} ; got {name} = {arg}")
+
+
+def check_not_in(name, arg, set):
+    if arg in set:
+        raise ValueError(f"{name} must not be any of those value {set} ; got {name} = {arg}")
+
+
+def check_contains(set, key):
+    if key not in set:
+        raise ValueError(f"{key} must be in {set}")
+
+
+def check_code_type(code_type):
+    if isinstance(code_type, CodeType):
+        if code_type.cc.major < 7:
+            raise ValueError(f"code_type.cc.major must be >= 7 ; got code_type.cc.major = {code_type.cc.major}")
+        if code_type.cc.minor < 0:
+            raise ValueError(f"code_type.cc.minor must be >= 0 ; got code_type.cc.minor = {code_type.cc.minor}")
+        check_in("code_type.kind", code_type.kind, ["lto"])
+    else:
+        raise ValueError(f"code_type should be an instance of CodeType ; got code_type = {code_type}")
+
+
+def pad_or_truncate(list, target_len):
+    return list[:target_len] + [0] * (target_len - len(list))
+
+
+class Layout:
+    """Layout for the OpaqueTensor"""
+
+    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.
+        """
+        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.
+        """
+        pass
+
+
+class OpaqueTensor:
+    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:
+    raise RuntimeError("make_tensor 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:
+    raise RuntimeError("axpby should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def copy(src: OpaqueTensor, dst: OpaqueTensor):
+    raise RuntimeError("copy should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def copy_fragment(src: OpaqueTensor, dst: OpaqueTensor):
+    raise RuntimeError("copy_fragment should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def clear(arr: OpaqueTensor):
+    raise RuntimeError("copy_c should not be called directly outside of a numba.cuda.jit(...) kernel.")
+
+
+def copy_wait():
+    raise RuntimeError("copy_wait should not be called directly outside of a numba.cuda.jit(...) kernel.")

nvmath/device/common_backend.py

@@ -0,0 +1,125 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+import numpy as np
+
+from nvmath.device.common_cuda import ISAVersion
+from .types import np_float16x2, np_float16x4
+
+from nvmath.bindings import mathdx
+
+MATHDX_TYPES_TO_NP = {
+    mathdx.CommondxValueType.R_16F: np.float16,
+    mathdx.CommondxValueType.R_16F2: np_float16x2,
+    mathdx.CommondxValueType.R_32F: np.float32,
+    mathdx.CommondxValueType.R_64F: np.float64,
+    mathdx.CommondxValueType.C_16F: np_float16x2,
+    mathdx.CommondxValueType.C_16F2: np_float16x4,
+    mathdx.CommondxValueType.C_32F: np.complex64,
+    mathdx.CommondxValueType.C_64F: np.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
+
+    def __del__(self):
+        if not self.descriptor:
+            return
+
+        self._destructor(self.descriptor)
+
+        # Safety clean up
+        self.descriptor = None
+        self._destructor = None
+
+
+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