nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/linalg/__init__.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from nvmath.bindings.cublas import ComputeType  # type: ignore
+
+from . import advanced
+from .generic import (
+    DiagonalMatrixQualifier,
+    DiagType,
+    ExecutionCPU,
+    ExecutionCUDA,
+    FillMode,
+    GeneralMatrixQualifier,
+    HermitianMatrixQualifier,
+    InvalidMatmulState,
+    Matmul,
+    MatmulOptions,
+    MatrixQualifier,
+    SideMode,
+    SymmetricMatrixQualifier,
+    TriangularMatrixQualifier,
+    matmul,
+    matrix_qualifiers_dtype,
+)
+from .generic.solvermod import (
+    DirectSolver,
+    DirectSolverOptions,
+    InvalidDirectSolverState,
+    direct_solver,
+)
+
+__all__ = [
+    "advanced",
+    "ComputeType",
+    "DiagonalMatrixQualifier",
+    "ExecutionCPU",
+    "ExecutionCUDA",
+    "GeneralMatrixQualifier",
+    "HermitianMatrixQualifier",
+    "InvalidMatmulState",
+    "InvalidDirectSolverState",
+    "DirectSolver",
+    "DirectSolverOptions",
+    "matmul",
+    "Matmul",
+    "MatmulOptions",
+    "MatrixQualifier",
+    "matrix_qualifiers_dtype",
+    "direct_solver",
+    "SymmetricMatrixQualifier",
+    "TriangularMatrixQualifier",
+    "SideMode",
+    "FillMode",
+    "DiagType",
+]

nvmath/linalg/_internal/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0

nvmath/linalg/_internal/algo_cap_ifc.py

@@ -0,0 +1,82 @@
+# 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 algorithm capabilities information.
+"""
+
+__all__ = ["AlgoCapInterface"]
+
+import numpy as np
+
+from nvmath._internal.attribute_ifc_factory import make_cublas_attribute_interface
+from nvmath.bindings import cublasLt as cublaslt
+
+from .enum_to_tuples import ENUM_TO_MATMUL_STAGE, ENUM_TO_MATMUL_TILE
+
+CapEnum = cublaslt.MatmulAlgoCapAttribute
+
+# Create a base class, _AlgoCapBase, such that each enum member in
+# MatmulAlgoCapAttribute (except TILE_IDS and STAGES_IDS) is exposed
+# as a read-only lowercase property. For example, if the class instance
+# is stored as ``cap_ifc``, then the enum member ``ATOMIC_SYNC``
+# becomes the property ``cap_ifc.atomic_sync``.
+#
+# The C binding expects a raw pointer to the algo sub-buffer, not the
+# handle dict. _get_attribute bridges that gap by extracting
+# `handle["algo"].ctypes.data` before forwarding the call.
+_raw_get = cublaslt.matmul_algo_cap_get_attribute
+
+
+def _get_attribute(handle, enum_value, buf_addr, buf_size, size_written_addr):
+    _raw_get(handle["algo"].ctypes.data, enum_value, buf_addr, buf_size, size_written_addr)
+
+
+# TILE_IDS and STAGES_IDS are excluded because they are variable-length
+# arrays that require a two-phase size-then-read query. They are implemented
+# as extra properties on the AlgoCapInterface subclass below.
+_AlgoCapBase = make_cublas_attribute_interface(
+    class_module=__name__,
+    class_name="_AlgoCapBase",
+    attribute_enum=CapEnum,
+    get_attribute_dtype_fn=cublaslt.get_matmul_algo_cap_attribute_dtype,
+    get_attribute_fn=_get_attribute,
+    exclude={"TILE_IDS", "STAGES_IDS"},
+)
+
+
+class AlgoCapInterface(_AlgoCapBase):  # type: ignore[valid-type, misc]
+    """The factory base class is extended with two extra properties:
+    ``tile_ids`` and ``stages_ids`` are added here because they return
+    variable-length arrays via a two-phase size-then-read query."""
+
+    def _get_array_attribute(self, name, enum_to_value_map):
+        dtype = cublaslt.get_matmul_algo_cap_attribute_dtype(CapEnum[name])
+        size_written = np.zeros((1,), dtype=np.uint64)
+
+        _raw_get(self._handle["algo"].ctypes.data, CapEnum[name].value, 0, 0, size_written.ctypes.data)
+
+        size = size_written[0]
+        if size == 0:
+            return ()
+
+        num = int(size // dtype().itemsize)
+        attribute_buffer = np.zeros((num,), dtype=dtype)
+        _raw_get(
+            self._handle["algo"].ctypes.data,
+            CapEnum[name].value,
+            attribute_buffer.ctypes.data,
+            size,
+            size_written.ctypes.data,
+        )
+
+        return tuple(enum_to_value_map[e] for e in attribute_buffer)
+
+    @property
+    def tile_ids(self):
+        return self._get_array_attribute("TILE_IDS", ENUM_TO_MATMUL_TILE)
+
+    @property
+    def stages_ids(self):
+        return self._get_array_attribute("STAGES_IDS", ENUM_TO_MATMUL_STAGE)

nvmath/linalg/_internal/algo_config_ifc.py

@@ -0,0 +1,43 @@
+# 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 algorithm configuration
+attributes.
+"""
+
+__all__ = ["AlgoConfigInterface"]
+
+from nvmath._internal.attribute_ifc_factory import make_cublas_attribute_interface
+from nvmath.bindings import cublasLt as cublaslt
+
+# Create a class, AlgoConfigInterface, such that each enum member in
+# MatmulAlgoConfigAttribute is exposed as a lowercase property (getter + setter).
+# For example, if the class instance is stored as ``config_ifc``,
+# then the enum member ``CLUSTER_SHAPE_ID`` becomes the property
+# ``config_ifc.cluster_shape_id``.
+#
+# The C bindings expect a raw pointer to the algo sub-buffer, not the
+# handle dict. _get_attribute and _set_attribute bridge that gap by
+# extracting `handle["algo"].ctypes.data` before forwarding the call.
+_raw_get = cublaslt.matmul_algo_config_get_attribute
+_raw_set = cublaslt.matmul_algo_config_set_attribute
+
+
+def _get_attribute(handle, enum_value, buf_addr, buf_size, size_written_addr):
+    _raw_get(handle["algo"].ctypes.data, enum_value, buf_addr, buf_size, size_written_addr)
+
+
+def _set_attribute(handle, enum_value, buf_addr, buf_size):
+    _raw_set(handle["algo"].ctypes.data, enum_value, buf_addr, buf_size)
+
+
+AlgoConfigInterface = make_cublas_attribute_interface(
+    class_module=__name__,
+    class_name="AlgoConfigInterface",
+    attribute_enum=cublaslt.MatmulAlgoConfigAttribute,
+    get_attribute_dtype_fn=cublaslt.get_matmul_algo_config_attribute_dtype,
+    get_attribute_fn=_get_attribute,
+    set_attribute_fn=_set_attribute,
+)

nvmath/linalg/_internal/batch.py

@@ -0,0 +1,234 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations  # allows typehint of class methods to return the self class
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from nvmath._internal.layout import is_overlapping_layout
+from nvmath.linalg._internal.utils import axis_order_in_memory, calculate_strides, check_batch_tileable
+
+
+@dataclass(slots=True, frozen=True)
+class BatchTraits:
+    """Represents the traits of a batched data.
+
+    A BatchTraits is valid if the non-batch dimensions are contiguous and dense (the data
+    span the entire extent of the tensor and there are no gaps between adjacent elements).
+    The batch dimensions may be optionally overlapping.
+
+    Attributes:
+        shape: The dimensions of the batch. An empty tuple `()` indicates no batching.
+        strides: The memory strides for each dimension of the batch.
+        overlap_allowed: Whether the batch dimensions are allowed to overlap in memory.
+
+    """
+
+    shape: Sequence[int]
+    strides: Sequence[int]
+    overlap_allowed: bool = False
+
+    @classmethod
+    def from_full_shape_and_strides(
+        cls,
+        shape: Sequence[int],
+        strides: Sequence[int],
+        num_trailing_dims: int,
+        overlap_allowed: bool = False,
+    ) -> BatchTraits:
+        """Create BatchTraits from the shape and strides of a tensor.
+
+        Args:
+            shape: The full shape of the tensor including both batch and non-batch
+                dimensions.
+            strides: The memory strides for each dimension of the full tensor.
+            num_trailing_dims: The number of trailing (non-batch) dimensions to exclude from
+                the batch. For example, for a batched matrix with shape (B1, B2, M, N),
+                num_trailing_dims=2 would extract (B1, B2) as batch dimensions.
+            overlap_allowed: Whether the batch dimensions are allowed to overlap in memory.
+                If False, raises ValueError if overlapping layout is detected. Defaults to
+                False.
+
+        Returns:
+            BatchTraits with the extracted batch shape and strides.
+
+        Raises:
+            ValueError: If the batch layout is not tileable (cannot be efficiently
+                processed).
+            ValueError: If overlap_allowed is False and the batch layout has overlapping
+                memory.
+
+        Example:
+            >>> # Extract batch info from a tensor of shape (3, 4, 5, 6)
+            >>> # treating last 2 dimensions as non-batch
+            >>> batch = BatchTraits.from_full_shape_and_strides(
+            ...     shape=(3, 4, 5, 6),
+            ...     strides=(120, 30, 6, 1),  # Row-major strides
+            ...     num_trailing_dims=2,  # Last 2 dims are non-batch
+            ... )
+            >>> batch.shape
+            (3, 4)
+            >>> batch.strides
+            (120, 30)
+        """
+        assert len(shape) == len(strides), (shape, strides)
+        num_dims = len(shape)
+        leading_shape = shape[: num_dims - num_trailing_dims]
+        leading_strides = strides[: num_dims - num_trailing_dims]
+        # Check for valid batching
+        # The samples in the batch must be tileable.
+        if leading_shape and not check_batch_tileable(leading_shape, leading_strides):
+            message = (
+                f"The batch layout for shape = {leading_shape} and strides = "
+                f"{leading_strides} is currently not supported because it is not tileable."
+            )
+            raise ValueError(message)
+
+        if not overlap_allowed and is_overlapping_layout(leading_shape, leading_strides):
+            raise ValueError(
+                "Only non-overlapping tensors are valid for batching. "
+                f"Shape {leading_shape} with strides {leading_strides} is not valid for batching."
+            )
+        # It's OK if we prune off the information about the trailing dimensions because the
+        # strides should already contain information about them.
+        return BatchTraits(
+            shape=leading_shape,
+            strides=leading_strides,
+            overlap_allowed=overlap_allowed,
+        )
+
+    @classmethod
+    def from_batch_shape_and_size(
+        cls,
+        shape: Sequence[int],
+        axis_order: Sequence[int],
+        batch_stride: int,
+    ) -> BatchTraits:
+        """Create BatchTraits from a batch shape and the size of one batch sample.
+
+        This method constructs BatchTraits by calculating the appropriate strides for a new
+        tensor with the given batch `shape`. The strides are computed based on the number of
+        elements of one batch sample and the desired layout order specified via
+        `axis_order`.
+
+        Args:
+            shape: The shape of the batch dimensions
+            axis_order: The order of batch axes in memory
+            batch_stride: Number of elements of one batch sample (e.g., M*N
+                for an MxN matrix)
+
+        Returns:
+            BatchTraits with the specified batch shape and computed strides.
+
+        Example:
+            >>> # Creating batch info for tensor of shape (3, 4, 5, 5)
+            >>> batch = BatchTraits.from_batch_shape_and_size(
+            ...     shape=(3, 4),  # 2D batch
+            ...     axis_order=(1, 0),  # ROW order,inner dimension (axis 1) comes first
+            ...     batch_stride=25,  # Each 5x5 matrix = 25 elements
+            ... )
+            >>> batch.shape
+            (3, 4)
+            >>> batch.strides
+            (100, 25)  # Jump 100 for outer dim, 25 for inner dim
+        """
+        strides = calculate_strides(
+            shape,
+            axis_order,
+            batch_stride,
+        )
+        return BatchTraits.from_full_shape_and_strides(
+            shape=shape,
+            strides=strides,
+            num_trailing_dims=0,
+            overlap_allowed=False,
+        )
+
+    @property
+    def count(self) -> int:
+        """The total number of elements in the batch. Returns 0 if non-batched."""
+        return 0 if not self.shape else math.prod(self.shape)
+
+    @property
+    def stride(self) -> int:
+        """The stride between elements the batch. Returns 0 if non-batched."""
+        return 0 if not self.shape else min(self.strides)
+
+    @property
+    def axis_order(self) -> tuple[int, ...]:
+        """The axis order of the batch. Returns an empty tuple if non-batched."""
+        return axis_order_in_memory(self.strides)
+
+    def __eq__(self, other) -> bool:
+        """Return whether two BatchTraits are equal.
+
+        Two BatchTraits are equal only if both the shape and strides of the batches are
+        equal.
+        """
+        if not isinstance(other, BatchTraits):
+            raise TypeError("Unsupported operand type(s) for ==.")
+        return self.shape == other.shape and self.strides == other.strides
+
+    def combine(
+        self,
+        other: BatchTraits | tuple[Sequence[int], Sequence[int]],
+        *,
+        allow_mixed_axis_order: bool = False,
+    ) -> tuple[Sequence[int], Sequence[int]]:
+        """Return the combined shape of two BatchTraits.
+
+        Currently two BatchTraits are only combinable if they have the same shape and axis
+        order or if one is an empty batch.
+
+        Parameters
+        ----------
+        other: The other BatchTraits or (shape, axis_order) tuple to combine with.
+        allow_mixed_axis_order: If True, skip the axis_order compatibility check and pick
+            the first non-empty axis_order. Useful for explicit batching where each
+            operand's pointer array is computed independently.
+
+        Returns
+        -------
+        shape: The combined shape of the two BatchTraits.
+        axis_order: The combined axis order of the two BatchTraits.
+
+        Raises
+        ------
+        ValueError: if the dimensions and strides are incompatible.
+        """
+        if isinstance(other, BatchTraits):
+            shape, axis_order = other.shape, other.axis_order
+        else:
+            shape, axis_order = other  # type: ignore[assignment]
+        if not self.shape or self.shape == shape:
+            new_shape = shape
+        elif not shape:
+            new_shape = self.shape
+        else:
+            msg = (
+                f"Batch dimensions {self.shape} and {shape} are not compatible. "
+                "Dimensions must be the same OR one input must be non-batched."
+            )
+            raise ValueError(msg)
+        if not self.axis_order or self.axis_order == axis_order:
+            new_axis_order = axis_order
+        elif not axis_order:
+            new_axis_order = self.axis_order
+        elif allow_mixed_axis_order:
+            new_axis_order = self.axis_order  # pick self's order; D will be allocated fresh anyway
+        else:
+            msg = (
+                f"Batch order {self.axis_order} and {axis_order} are not compatible. "
+                "Batch orders must be the same OR one input must be non-batched."
+            )
+            raise ValueError(msg)
+        return new_shape, new_axis_order
+
+    def __mul__(self, other: BatchTraits | tuple[Sequence[int], Sequence[int]]) -> tuple[Sequence[int], Sequence[int]]:
+        return self.combine(other)
+
+    def __str__(self) -> str:
+        return f"shape {self.shape} and strides {self.strides}"

nvmath/linalg/_internal/enum_to_tuples.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Utilities to convert tuple like enumerators to tuples, along with forward and inverse maps.
+"""
+
+__all__ = [
+    "CLUSTER_SHAPES",
+    "CLUSTER_SHAPE_TO_ENUM",
+    "ENUM_TO_CLUSTER_SHAPE",
+    "MATMUL_STAGES",
+    "MATMUL_STAGE_TO_ENUM",
+    "ENUM_TO_MATMUL_STAGE",
+    "MATMUL_TILES",
+    "MATMUL_TILE_TO_ENUM",
+    "ENUM_TO_MATMUL_TILE",
+]
+
+import re
+
+from nvmath.bindings import cublasLt as cublaslt
+
+
+def integer_or_string(value):
+    try:
+        value = int(value)
+    except ValueError:
+        ...
+    return value
+
+
+def create_valid_tuples_from_enum(enum, prefix, *, expr=r"(?:(\d+)x(\d+|\w+)(?:x(\d+))?|(AUTO|UNDEFINED))"):
+    """
+    Create a sequence of tuples representing the allowed combinations for the given
+    enumeration.
+    """
+
+    combinations = []
+    enumerator_to_value = {}
+    value_to_enumerator = {}
+    expr = prefix + expr
+    for e in enum:
+        m = re.match(expr, e.name)
+        if m:
+            # print(m.groups())
+            if m.group(4):
+                v = m.group(4)
+            else:
+                groups = m.groups()[: m.groups().index(None)]
+                v = tuple(map(integer_or_string, (g for g in groups)))
+            combinations.append(v)
+            value_to_enumerator[v] = e
+            enumerator_to_value[e] = v
+
+    return tuple(combinations), value_to_enumerator, enumerator_to_value
+
+
+CLUSTER_SHAPES, CLUSTER_SHAPE_TO_ENUM, ENUM_TO_CLUSTER_SHAPE = create_valid_tuples_from_enum(cublaslt.ClusterShape, "SHAPE_")
+
+MATMUL_STAGES, MATMUL_STAGE_TO_ENUM, ENUM_TO_MATMUL_STAGE = create_valid_tuples_from_enum(cublaslt.MatmulStages, "STAGES_")
+
+MATMUL_TILES, MATMUL_TILE_TO_ENUM, ENUM_TO_MATMUL_TILE = create_valid_tuples_from_enum(cublaslt.MatmulTile, "TILE_")