nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/linalg/_internal/matrix_layout_ifc.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2024-2026, 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"]
+
+from nvmath._internal.attribute_ifc_factory import make_cublas_attribute_interface
+from nvmath.bindings import cublasLt as cublaslt
+
+# Create a class, MatrixLayoutInterface, such that each enum member in
+# MatrixLayoutAttribute is exposed as a lowercase property (getter + setter).
+# For example, if the class instance is stored as ``layout_ifc``,
+# then the enum member ``ROWS`` becomes the property
+# ``layout_ifc.rows``.
+MatrixLayoutInterface = make_cublas_attribute_interface(
+    class_module=__name__,
+    class_name="MatrixLayoutInterface",
+    attribute_enum=cublaslt.MatrixLayoutAttribute,
+    get_attribute_dtype_fn=cublaslt.get_matrix_layout_attribute_dtype,
+    get_attribute_fn=cublaslt.matrix_layout_get_attribute,
+    set_attribute_fn=cublaslt.matrix_layout_set_attribute,
+)

nvmath/linalg/_internal/typemaps.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2024-2026, 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 = {}
+    # FP4 scale types (https://docs.nvidia.com/cuda/cublas/index.html#id105)
+    scale_type_map["float4_e2m1fn_x2", "float4_e2m1fn_x2"] = dt.CUDA_R_32F
+    # FP8 scale types
+    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_32F_EMULATED_16BFX9] = f32
+    scale_type_map[ct.COMPUTE_64F] = f64
+    scale_type_map[ct.COMPUTE_64F_PEDANTIC] = f64
+    scale_type_map[ct.COMPUTE_64F_EMULATED_FIXEDPOINT] = 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 compute type.
+    """
+
+    ct = cublas.ComputeType
+
+    compute_type_map = {}
+    # FP4 compute types (https://docs.nvidia.com/cuda/cublas/index.html#id105)
+    compute_type_map["float4_e2m1fn_x2", "float4_e2m1fn_x2"] = ct.COMPUTE_32F
+    # FP8 compute types
+    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 = [
+    "float4_e2m1fn_x2",  # FP4 (supported on SM 10.0+)
+    "float8_e4m3fn",
+    "float8_e5m2",
+    "bfloat16",
+    "float16",
+    "float32",
+    "float64",
+    "complex32",
+    "complex64",
+    "complex128",
+]

nvmath/linalg/_internal/utils.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2024-2026, 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 atexit
+import typing
+
+from nvmath.bindings import cublas
+from nvmath.bindings import cublasLt as cublaslt
+from nvmath.internal import utils
+
+HANDLES: dict[str, dict[int, int]] = {
+    "cublas": {},
+    "cublaslt": {},
+}
+
+
+def create_handle(device_id: int, binding="cublaslt") -> int:
+    """
+    Currently for internal use only.
+    """
+    with utils.device_ctx(device_id):
+        match binding:
+            case "cublas":
+                handle = cublas.create()
+            case "cublaslt" | _:
+                handle = cublaslt.create()
+    return handle
+
+
+def destroy_handle(handle: int, binding="cublaslt"):
+    """
+    Currently for internal use only.
+    """
+    match binding:
+        case "cublas":
+            cublas.destroy(handle)
+        case "cublaslt" | _:
+            cublaslt.destroy(handle)
+
+
+@atexit.register
+def _cleanup_handles():
+    """
+    Cleanup all cached handles on program exit.
+    """
+    for binding, device_handles in HANDLES.items():
+        for handle in device_handles.values():
+            destroy_handle(handle, binding=binding)
+        device_handles.clear()
+
+
+def get_handle(device_id: int, binding="cublaslt") -> int:
+    """
+    Retrieve the cuBLAS[lt] library handle for the specified device. If one doesn't exist,
+    create, cache, and return the handle.
+
+    According to the docs for cublasLtHandle_t, any valid cublasHandle_t can be used in
+    place of cublasLtHandle_t with a simple cast, so we use the same handle for both APIs.
+
+    Handles are cached and automatically cleaned up on program exit via the atexit handler.
+    We expect to have exactly one handle per device / thread.
+    """
+    if device_id in HANDLES[binding]:
+        handle = HANDLES[binding][device_id]
+    else:
+        handle = create_handle(device_id, binding=binding)
+        HANDLES[binding][device_id] = handle
+    return handle
+
+
+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], min_stride: int = 1):
+    """
+    Calculate the strides for the provided shape and axis order.
+    """
+    assert len(axis_order) == len(shape), f"axis_order length ({len(axis_order)}) must equal shape length ({len(shape)})"
+    assert len(set(axis_order)) == len(axis_order), f"axis_order must not contain duplicates: {axis_order}"
+    assert set(axis_order) == set(range(len(shape))), f"axis_order must be permutation of range({len(shape)}): {axis_order}"
+
+    strides: list[None | int] = [None] * len(shape)
+
+    stride = min_stride
+    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-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from . import helpers as helpers  # noqa: F403
+from ._algorithmmod import *  # noqa: F403
+from ._configuration import *  # noqa: F403
+from .matmulmod import *  # noqa: F403

nvmath/linalg/advanced/_algorithmmod.py

@@ -0,0 +1,149 @@
+# Copyright (c) 2024-2026, 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.bindings.cublasLt import MatmulHeuristicResult  # type: ignore
+from nvmath.linalg._internal.algo_cap_ifc import AlgoCapInterface
+from nvmath.linalg._internal.algo_config_ifc import AlgoConfigInterface
+from nvmath.linalg._internal.enum_to_tuples import (
+    CLUSTER_SHAPE_TO_ENUM,
+    ENUM_TO_CLUSTER_SHAPE,
+    ENUM_TO_MATMUL_STAGE,
+    ENUM_TO_MATMUL_TILE,
+    MATMUL_STAGE_TO_ENUM,
+    MATMUL_TILE_TO_ENUM,
+)
+from nvmath.linalg.advanced._configuration import AlgorithmCapabilities
+
+
+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: MatmulHeuristicResult) -> None:
+        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) -> int:
+        "The ID of the algorithm (integer)."
+        return self.config_ifc.id
+
+    @property
+    def tile(self) -> tuple[int | str, ...] | str:
+        """A tuple representing the tile (see MatmulAlgoConfigAttribute.TILE_ID).
+        The value provided must be one of the `tile_ids` in the algorithm capabilities."""
+        return ENUM_TO_MATMUL_TILE[self.config_ifc.tile_id]
+
+    @tile.setter
+    def tile(self, tile: tuple[int | str, ...] | str) -> None:
+        self.config_ifc.tile_id = MATMUL_TILE_TO_ENUM[tile]
+
+    @property
+    def stages(self) -> tuple[int | str, ...] | str:
+        """A tuple representing the stages (see MatmulAlgoConfigAttribute.STAGES_ID).
+        The value provided must be one of the `stages_ids` in the algorithm capabilities."""
+        return ENUM_TO_MATMUL_STAGE[self.config_ifc.stages_id]
+
+    @stages.setter
+    def stages(self, stages: tuple[int | str, ...] | str) -> None:
+        self.config_ifc.stages_id = MATMUL_STAGE_TO_ENUM[stages]
+
+    @property
+    def split_k(self) -> int:
+        """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: int) -> None:
+        self.config_ifc.splitk_num = number
+
+    @property
+    def reduction_scheme(self) -> int:
+        """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: int) -> None:
+        self.config_ifc.reduction_scheme = scheme_id
+
+    @property
+    def cta_swizzling(self) -> bool:
+        """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) -> None:
+        self.config_ifc.cta_swizzling = flag
+
+    @property
+    def custom_option(self) -> int:
+        """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) -> None:
+        self.config_ifc.custom_option = value
+
+    @property
+    def inner_shape(self) -> int:
+        """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: int) -> None:
+        self.config_ifc.inner_shape_id = shape
+
+    @property
+    def cluster_shape(self) -> tuple[int | str, ...] | str:
+        """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 ENUM_TO_CLUSTER_SHAPE[self.config_ifc.cluster_shape_id]
+
+    @cluster_shape.setter
+    def cluster_shape(self, shape: tuple[int | str, ...] | str) -> None:
+        self.config_ifc.cluster_shape_id = CLUSTER_SHAPE_TO_ENUM[shape]