nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/linalg/_internal/matmul_desc_ifc.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface class to encapsulate low-level calls to get and set matmul descriptor attributes.
+"""
+
+__all__ = ["MatmulDescInterface"]
+
+import ctypes
+import logging
+
+import numpy as np
+
+from nvmath.bindings import cublasLt as cublaslt
+
+logger = logging.getLogger()
+
+DescEnum = cublaslt.MatmulDescAttribute
+
+
+def scalar_attributes():
+    return [e.name for e in DescEnum]
+
+
+DESC_ENUM_SCALAR_ATTR = scalar_attributes()
+
+
+def _get_attribute_ctype(name):
+    return np.ctypeslib.as_ctypes_type(cublaslt.get_matmul_desc_attribute_dtype(DescEnum[name]))
+
+
+DESC_ENUM_SCALAR_ATTR_INFO = {name: (DescEnum[name].value, _get_attribute_ctype(name)) for name in DESC_ENUM_SCALAR_ATTR}
+
+
+class MatmulDescInterface:
+    def __init__(self, matmul_desc):
+        """ """
+        self.matmul_desc = matmul_desc
+
+    def __getattr__(self, name):
+        _name = name.upper()
+        logging.debug("Getting Matmul Description attribute %s.", _name)
+        info = DESC_ENUM_SCALAR_ATTR_INFO.get(_name)
+        if info is None:
+            return super().__getattr__(name)
+        enum_value, ctype = info
+        name = _name
+        attribute_buffer = ctype()
+        size_written = ctypes.c_uint64()
+        cublaslt.matmul_desc_get_attribute(
+            self.matmul_desc,
+            enum_value,
+            ctypes.addressof(attribute_buffer),
+            ctypes.sizeof(attribute_buffer),
+            ctypes.addressof(size_written),
+        )
+        return attribute_buffer.value
+
+    def __setattr__(self, name, value):
+        _name = name.upper()
+        logging.debug("Setting Matmul Description attribute %s to %s.", _name, value)
+        info = DESC_ENUM_SCALAR_ATTR_INFO.get(_name)
+        if info is None:
+            return super().__setattr__(name, value)
+        enum_value, ctype = info
+        name = _name
+        ctypes_value = ctype(value)
+        cublaslt.matmul_desc_set_attribute(
+            self.matmul_desc, enum_value, ctypes.addressof(ctypes_value), ctypes.sizeof(ctypes_value)
+        )

nvmath/linalg/_internal/matmul_pref_ifc.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface class to encapsulate low-level calls to get and set matmul plan preference
+attributes.
+"""
+
+__all__ = ["MatmulPreferenceInterface"]
+
+import logging
+
+import numpy as np
+
+from nvmath.bindings import cublasLt as cublaslt
+
+
+logger = logging.getLogger()
+
+PreferenceEnum = cublaslt.MatmulPreferenceAttribute
+
+
+def scalar_attributes():
+    return [e.name for e in PreferenceEnum]
+
+
+PREF_ENUM_SCALAR_ATTR = scalar_attributes()
+
+
+class MatmulPreferenceInterface:
+    def __init__(self, matmul_pref):
+        """ """
+        self.matmul_pref = matmul_pref
+
+    def __getattr__(self, name):
+        _name = name.upper()
+        logging.debug("Getting Matmul Preference attribute %s.", _name)
+        if _name not in PREF_ENUM_SCALAR_ATTR:
+            return super().__getattr__(name)
+        name = _name
+        get_dtype = cublaslt.get_matmul_preference_attribute_dtype
+        attribute_buffer = np.zeros((1,), dtype=get_dtype(PreferenceEnum[name]))
+        size_written = np.zeros((1,), dtype=np.uint64)
+        cublaslt.matmul_preference_get_attribute(
+            self.matmul_pref,
+            PreferenceEnum[name].value,
+            attribute_buffer.ctypes.data,
+            attribute_buffer.itemsize,
+            size_written.ctypes.data,
+        )
+        return attribute_buffer[0]
+
+    def __setattr__(self, name, value):
+        _name = name.upper()
+        logging.debug("Setting Matmul Preference attribute %s to %s.", _name, value)
+        if _name not in PREF_ENUM_SCALAR_ATTR:
+            return super().__setattr__(name, value)
+        name = _name
+        get_dtype = cublaslt.get_matmul_preference_attribute_dtype
+        attribute_buffer = np.zeros((1,), dtype=get_dtype(PreferenceEnum[name]))
+        attribute_buffer[0] = value
+        cublaslt.matmul_preference_set_attribute(
+            self.matmul_pref, PreferenceEnum[name].value, attribute_buffer.ctypes.data, attribute_buffer.itemsize
+        )

nvmath/linalg/_internal/matrix_layout_ifc.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface class to encapsulate low-level calls to get and set matrix layout attributes.
+"""
+
+__all__ = ["MatrixLayoutInterface"]
+
+
+import numpy as np
+
+from nvmath.bindings import cublasLt as cublaslt
+
+
+LayoutEnum = cublaslt.MatrixLayoutAttribute
+
+
+def scalar_attributes():
+    return [e.name for e in LayoutEnum]
+
+
+LAYOUT_ENUM_SCALAR_ATTR = scalar_attributes()
+
+
+class MatrixLayoutInterface:
+    def __init__(self, matrix_layout):
+        """ """
+        self.matrix_layout = matrix_layout
+
+    def __getattr__(self, name):
+        _name = name.upper()
+        if _name not in LAYOUT_ENUM_SCALAR_ATTR:
+            return super().__getattr__(name)
+        name = _name
+        get_dtype = cublaslt.get_matrix_layout_attribute_dtype
+        attribute_buffer = np.zeros((1,), dtype=get_dtype(LayoutEnum[name]))
+        size_written = np.zeros((1,), dtype=np.uint64)
+        cublaslt.matrix_layout_get_attribute(
+            self.matrix_layout,
+            LayoutEnum[name].value,
+            attribute_buffer.ctypes.data,
+            attribute_buffer.itemsize,
+            size_written.ctypes.data,
+        )
+        return attribute_buffer[0]
+
+    def __setattr__(self, name, value):
+        _name = name.upper()
+        if _name not in LAYOUT_ENUM_SCALAR_ATTR:
+            return super().__setattr__(name, value)
+        name = _name
+        get_dtype = cublaslt.get_matrix_layout_attribute_dtype
+        attribute_buffer = np.zeros((1,), dtype=get_dtype(LayoutEnum[name]))
+        attribute_buffer[0] = value
+        cublaslt.matrix_layout_set_attribute(
+            self.matrix_layout, LayoutEnum[name].value, attribute_buffer.ctypes.data, attribute_buffer.itemsize
+        )

nvmath/linalg/_internal/typemaps.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Functions to link type names with CUBLAS compute types.
+"""
+
+__all__ = [
+    "NAMES_TO_DEFAULT_SCALE_TYPE",
+    "NAMES_TO_DEFAULT_COMPUTE_TYPE",
+    "SUPPORTED_TYPES",
+    "COMPUTE_TYPE_TO_DEFAULT_SCALE_TYPE",
+    "SCALE_TYPE_TO_DEFAULT_COMPUTE_TYPE",
+]
+
+from nvmath.bindings import cublas  # type: ignore
+from nvmath.internal.typemaps import cudaDataType
+
+
+def create_default_scale_type_map():
+    """
+    Map the data type name to the corresponding CUDA data type that's appropriate for
+    default scale.
+    """
+
+    dt = cudaDataType
+
+    scale_type_map = {}
+    scale_type_map["float8_e4m3fn", "float8_e5m2"] = dt.CUDA_R_32F
+    scale_type_map["float8_e5m2", "float8_e4m3fn"] = dt.CUDA_R_32F
+    scale_type_map["float8_e4m3fn", "float8_e4m3fn"] = dt.CUDA_R_32F
+    scale_type_map["bfloat16", "bfloat16"] = dt.CUDA_R_32F
+    scale_type_map["float16", "float16"] = dt.CUDA_R_32F
+    scale_type_map["float32", "float32"] = dt.CUDA_R_32F
+    scale_type_map["float64", "float64"] = dt.CUDA_R_64F
+    scale_type_map["complex32", "complex32"] = dt.CUDA_C_32F
+    scale_type_map["complex64", "complex64"] = dt.CUDA_C_32F
+    scale_type_map["complex128", "complex128"] = dt.CUDA_C_64F
+
+    return scale_type_map
+
+
+def create_compute_type_to_scale_type_map(is_complex):
+    """
+    Map the compute type to the corresponding CUDA data type that's appropriate for
+    default scale.
+    """
+
+    dt = cudaDataType
+    ct = cublas.ComputeType
+
+    scale_type_map = {}
+    scale_type_map[ct.COMPUTE_16F] = dt.CUDA_R_16F
+    scale_type_map[ct.COMPUTE_16F_PEDANTIC] = dt.CUDA_R_16F
+
+    f32 = dt.CUDA_C_32F if is_complex else dt.CUDA_R_32F
+    f64 = dt.CUDA_C_64F if is_complex else dt.CUDA_R_64F
+
+    scale_type_map[ct.COMPUTE_32F] = f32
+    scale_type_map[ct.COMPUTE_32F_PEDANTIC] = f32
+    scale_type_map[ct.COMPUTE_32F_FAST_16F] = f32
+    scale_type_map[ct.COMPUTE_32F_FAST_16BF] = f32
+    scale_type_map[ct.COMPUTE_32F_FAST_TF32] = f32
+    scale_type_map[ct.COMPUTE_64F] = f64
+    scale_type_map[ct.COMPUTE_64F_PEDANTIC] = f64
+
+    return scale_type_map
+
+
+def create_scale_type_to_compute_type_map():
+    """
+    Map the scale type to the corresponding compute type that's an appropriate default.
+    """
+
+    dt = cudaDataType
+    ct = cublas.ComputeType
+
+    compute_type_map = {}
+    compute_type_map[dt.CUDA_R_16F] = ct.COMPUTE_16F
+    compute_type_map[dt.CUDA_R_16BF] = ct.COMPUTE_32F
+    compute_type_map[dt.CUDA_R_32F] = ct.COMPUTE_32F
+    compute_type_map[dt.CUDA_C_32F] = ct.COMPUTE_32F
+    compute_type_map[dt.CUDA_R_64F] = ct.COMPUTE_64F
+    compute_type_map[dt.CUDA_C_64F] = ct.COMPUTE_64F
+    return compute_type_map
+
+
+def create_compute_type_map():
+    """
+    Map the data type name to the corresponding CUDA data type that's appropriate for
+    default scale.
+    """
+
+    ct = cublas.ComputeType
+
+    compute_type_map = {}
+    compute_type_map["float8_e4m3fn", "float8_e5m2"] = ct.COMPUTE_32F
+    compute_type_map["float8_e5m2", "float8_e4m3fn"] = ct.COMPUTE_32F
+    compute_type_map["float8_e4m3fn", "float8_e4m3fn"] = ct.COMPUTE_32F
+    compute_type_map["bfloat16", "bfloat16"] = ct.COMPUTE_32F
+    compute_type_map["float16", "float16"] = ct.COMPUTE_32F
+    compute_type_map["float32", "float32"] = ct.COMPUTE_32F
+    compute_type_map["float64", "float64"] = ct.COMPUTE_64F
+    compute_type_map["complex32", "complex32"] = ct.COMPUTE_32F
+    compute_type_map["complex64", "complex64"] = ct.COMPUTE_32F
+    compute_type_map["complex128", "complex128"] = ct.COMPUTE_64F
+
+    return compute_type_map
+
+
+CUBLAS_COMPUTE_TYPE_TO_NAME = {
+    cublas.ComputeType.COMPUTE_32F: ("float32", "complex64"),
+    cublas.ComputeType.COMPUTE_64F: ("float64", "complex128"),
+}
+
+NAMES_TO_DEFAULT_SCALE_TYPE = create_default_scale_type_map()
+NAMES_TO_DEFAULT_COMPUTE_TYPE = create_compute_type_map()
+COMPUTE_TYPE_TO_DEFAULT_SCALE_TYPE = {
+    "real": create_compute_type_to_scale_type_map(is_complex=False),
+    "complex": create_compute_type_to_scale_type_map(is_complex=True),
+}
+SCALE_TYPE_TO_DEFAULT_COMPUTE_TYPE = create_scale_type_to_compute_type_map()
+
+SUPPORTED_TYPES = [
+    "float8_e4m3fn",
+    "float8_e5m2",
+    "bfloat16",
+    "float16",
+    "float32",
+    "float64",
+    "complex32",
+    "complex64",
+    "complex128",
+]

nvmath/linalg/_internal/utils.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Utilities.
+"""
+
+__all__ = [
+    "axis_order_in_memory",
+    "calculate_strides",
+    "check_batch_tileable",
+    "create_handle",
+    "destroy_handle",
+    "get_handle",
+    "pointer_aligned_to",
+]
+
+import typing
+
+from nvmath.bindings import cublasLt as cublaslt
+from nvmath.internal import utils
+
+HANDLES: dict[int, int] = {}
+
+
+def create_handle(device_id: int) -> int:
+    """
+    Currently for internal use only.
+    """
+    with utils.device_ctx(device_id):
+        handle = cublaslt.create()
+
+    return handle
+
+
+def destroy_handle(handle: int):
+    """
+    Currently for internal use only.
+    """
+    cublaslt.destroy(handle)
+
+
+def get_handle(device_id: int) -> int:
+    """
+    Retrieve the BLAS library handle for the specified device. If one doesn't exist, create,
+    cache, and return the handle.
+    """
+    return HANDLES.setdefault(device_id, create_handle(device_id))
+
+
+def pointer_aligned_to(address):
+    """
+    Return the number of bytes the address is aligned to.
+    """
+    return address & ~(address - 1)
+
+
+def axis_order_in_memory(strides):
+    """
+    Compute the order in which the axes appear in memory.
+    """
+    if len(strides) == 0:
+        return ()
+
+    _, axis_order = zip(*sorted(zip(strides, range(len(strides)), strict=True)), strict=True)
+
+    return axis_order
+
+
+def calculate_strides(shape: typing.Sequence[int], axis_order: typing.Sequence[int]):
+    """
+    Calculate the strides for the provided shape and axis order.
+    """
+    strides: list[None | int] = [None] * len(shape)
+
+    stride = 1
+    for axis in axis_order:
+        strides[axis] = stride
+        stride *= shape[axis]
+
+    return strides
+
+
+def _contiguous_layout(sorted_shape, sorted_strides):
+    return all(sorted_shape[s - 1] * sorted_strides[s - 1] == sorted_strides[s] for s in range(1, len(sorted_strides)))
+
+
+def check_batch_tileable(batch_shape, batch_strides):
+    """
+    Check if the matrix layout is tileable across the specified batch layout.
+    """
+    sorted_batch_strides, sorted_batch_shape = zip(
+        *sorted((batch_strides[a], batch_shape[a]) for a in range(len(batch_shape))), strict=True
+    )
+    return _contiguous_layout(sorted_batch_shape, sorted_batch_strides)

nvmath/linalg/advanced/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from ._algorithmmod import *  # noqa: F403
+from ._configuration import *  # noqa: F403
+from .matmulmod import *  # noqa: F403
+from . import helpers as helpers  # noqa: F403

nvmath/linalg/advanced/_algorithmmod.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+An interface class to query algorithm capabilities and configure it.
+"""
+
+__all__ = ["Algorithm"]
+
+import dataclasses
+
+
+from nvmath.linalg.advanced._configuration import AlgorithmCapabilities
+from nvmath.linalg._internal.algo_cap_ifc import AlgoCapInterface
+from nvmath.linalg._internal.algo_config_ifc import AlgoConfigInterface
+from nvmath.bindings.cublasLt import MatmulHeuristicResult  # type: ignore
+
+
+class Algorithm:
+    """
+    An interface class to query algorithm capabilities and configure the algorithm.
+
+    Note that objects of this type should not be constructed directly by the user.
+    """
+
+    def __init__(self, algorithm):
+        assert isinstance(algorithm, MatmulHeuristicResult), "Internal error."
+        self.algorithm = algorithm
+        self.cap_ifc = AlgoCapInterface(algorithm)
+        self.config_ifc = AlgoConfigInterface(algorithm)
+
+    @property
+    def capabilities(self):
+        """
+        Return the capabilities of this algorithm as a
+        :class:`nvmath.linalg.advanced.AlgorithmCapabilities` dataclass.
+        """
+        names = [field.name for field in dataclasses.fields(AlgorithmCapabilities)]
+        _capabilities = {}
+        for name in names:
+            # Not all capabilities in the enum are supported in all CTK versions.
+            try:
+                _capabilities[name] = getattr(self.cap_ifc, name)
+            except:
+                pass
+        return AlgorithmCapabilities(**_capabilities)
+
+    @property
+    def algorithm_id(self):
+        "The ID of the algorithm (integer)."
+        return self.config_ifc.id
+
+    @property
+    def tile(self):
+        """A tuple representing the tile (see MatmulAlgoConfigAttribute.TILE_ID).
+        The value provided must be one of the `tile_ids` in the algorithm capabilities."""
+        return self.config_ifc.tile_id
+
+    @tile.setter
+    def tile(self, tile):
+        self.config_ifc.tile_id = tile
+
+    @property
+    def stages(self):
+        """A tuple representing the stages (see MatmulAlgoConfigAttribute.STAGES_ID).
+        The value provided must be one of the `stages_ids` in the algorithm capabilities."""
+        return self.config_ifc.stages_id
+
+    @stages.setter
+    def stages(self, stages):
+        self.config_ifc.stages_id = stages
+
+    @property
+    def split_k(self):
+        """The number of split-k steps (see MatmulAlgoConfigAttribute.SPLITK_NUM).
+
+        This can be set only if `splitk_support` is 1 in the algorithm capabilities."""
+        return self.config_ifc.splitk_num
+
+    @split_k.setter
+    def split_k(self, number):
+        self.config_ifc.splitk_num = number
+
+    @property
+    def reduction_scheme(self):
+        """The reduction scheme used (see MatmulAlgoConfigAttribute.REDUCTION_SCHEME).
+
+        The value provided must be consistent with the `reduction_scheme_mask` in the
+        algorithm capabilities."""
+        return self.config_ifc.reduction_scheme
+
+    @reduction_scheme.setter
+    def reduction_scheme(self, scheme_id):
+        self.config_ifc.reduction_scheme = scheme_id
+
+    @property
+    def cta_swizzling(self):
+        """A flag indicating CTA swizzling (see MatmulAlgoConfigAttribute.CTA_SWIZZLING).
+
+        This can be set only if `cta_swizzling` is 1 in the algorithm capabilities."""
+        return self.config_ifc.cta_swizzling
+
+    @cta_swizzling.setter
+    def cta_swizzling(self, flag: bool):
+        self.config_ifc.cta_swizzling = flag
+
+    @property
+    def custom_option(self):
+        """A value indicating the custom option (see
+        MatmulAlgoConfigAttribute.CUSTOM_OPTION).
+
+        The value provided must be less than `custom_option_max` in the algorithm
+        capabilities."""
+        return self.config_ifc.custom_option
+
+    @custom_option.setter
+    def custom_option(self, value: int):
+        self.config_ifc.custom_option = value
+
+    @property
+    def inner_shape(self):
+        """A value indicating the inner shape (see
+        MatmulAlgoConfigAttribute.INNER_SHAPE_ID)."""
+        return self.config_ifc.inner_shape_id
+
+    @inner_shape.setter
+    def inner_shape(self, shape):
+        self.config_ifc.inner_shape_id = shape
+
+    @property
+    def cluster_shape(self):
+        """A tuple representing the cluster shape (see
+        MatmulAlgoConfigAttribute.CLUSTER_SHAPE_ID).
+
+        The value provided must be one of the `cluster_shape_ids` in the algorithm
+        capabilities."""
+        return self.config_ifc.cluster_shape_id
+
+    @cluster_shape.setter
+    def cluster_shape(self, shape):
+        self.config_ifc.cluster_shape_id = shape