nvmath-python 0.5.0__cp313-cp313-win_amd64.whl

nvmath/linalg/_internal/algo_cap_ifc.py

@@ -0,0 +1,86 @@
+# 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 algorithm capabilities information.
+"""
+
+__all__ = ["AlgoCapInterface"]
+
+import numpy as np
+
+from nvmath.bindings import cublasLt as cublaslt
+from .enum_to_tuples import ENUM_TO_MATMUL_STAGE, ENUM_TO_MATMUL_TILE
+
+CapEnum = cublaslt.MatmulAlgoCapAttribute
+
+
+def scalar_attributes():
+    return [e.name for e in CapEnum if not e.name.endswith("_IDS")]
+
+
+CAP_ENUM_SCALAR_ATTR = scalar_attributes()
+
+
+class AlgoCapInterface:
+    def __init__(self, algorithm):
+        """ """
+        assert isinstance(algorithm, cublaslt.MatmulHeuristicResult), "Internal error."
+        self.algorithm = algorithm
+
+    def __getattr__(self, name):
+        _name = name.upper()
+        if _name not in CAP_ENUM_SCALAR_ATTR:
+            return super().__getattr__(name)
+        name = _name
+        get_dtype = cublaslt.get_matmul_algo_cap_attribute_dtype
+        attribute_buffer = np.zeros((1,), dtype=get_dtype(CapEnum[name]))
+        size_written = np.zeros((1,), dtype=np.uint64)
+        cublaslt.matmul_algo_cap_get_attribute(
+            self.algorithm["algo"].ctypes.data,
+            CapEnum[name].value,
+            attribute_buffer.ctypes.data,
+            attribute_buffer.itemsize,
+            size_written.ctypes.data,
+        )
+        return attribute_buffer[0]
+
+    def _get_array_attribute(self, name, enum_to_value_map):
+        assert name in ["TILE_IDS", "STAGES_IDS"], "Internal error."
+
+        get_dtype = cublaslt.get_matmul_algo_cap_attribute_dtype
+        dtype = get_dtype(CapEnum[name])
+        size_written = np.zeros((1,), dtype=np.uint64)
+
+        # First get the buffer size needed.
+        cublaslt.matmul_algo_cap_get_attribute(
+            self.algorithm["algo"].ctypes.data, CapEnum[name].value, 0, 0, size_written.ctypes.data
+        )
+
+        # Check if any data needs to be written.
+        size = size_written[0]
+        if size == 0:
+            return ()
+
+        # Then allocate the buffer and get the data.
+        num = int(size // dtype().itemsize)
+        attribute_buffer = np.zeros((num,), dtype=dtype)
+        cublaslt.matmul_algo_cap_get_attribute(
+            self.algorithm["algo"].ctypes.data,
+            CapEnum[name].value,
+            attribute_buffer.ctypes.data,
+            size,
+            size_written.ctypes.data,
+        )
+
+        # Convert to value.
+        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,87 @@
+# 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 algorithm configuration information.
+"""
+
+__all__ = ["AlgoConfigInterface"]
+
+from collections import namedtuple
+
+import numpy as np
+
+from nvmath.bindings import cublasLt as cublaslt
+from .enum_to_tuples import (
+    CLUSTER_SHAPE_TO_ENUM,
+    ENUM_TO_CLUSTER_SHAPE,
+    MATMUL_STAGE_TO_ENUM,
+    ENUM_TO_MATMUL_STAGE,
+    MATMUL_TILE_TO_ENUM,
+    ENUM_TO_MATMUL_TILE,
+)
+
+ConfigEnum = cublaslt.MatmulAlgoConfigAttribute
+
+
+def scalar_attributes():
+    return [e.name for e in ConfigEnum]
+
+
+CAP_ENUM_SCALAR_ATTR = scalar_attributes()
+
+# Special Handling.
+Maps = namedtuple("Maps", ["to_enumerator", "to_value"])
+SPECIAL_ATTR = {
+    "CLUSTER_SHAPE_ID": Maps(CLUSTER_SHAPE_TO_ENUM, ENUM_TO_CLUSTER_SHAPE),
+    "STAGES_ID": Maps(MATMUL_STAGE_TO_ENUM, ENUM_TO_MATMUL_STAGE),
+    "TILE_ID": Maps(MATMUL_TILE_TO_ENUM, ENUM_TO_MATMUL_TILE),
+}
+
+
+class AlgoConfigInterface:
+    def __init__(self, algorithm):
+        """ """
+        assert isinstance(algorithm, cublaslt.MatmulHeuristicResult), "Internal error."
+        self.algorithm = algorithm
+
+    def __getattr__(self, name):
+        _name = name.upper()
+        if _name not in CAP_ENUM_SCALAR_ATTR:
+            return super().__getattr__(name)
+        name = _name
+        get_dtype = cublaslt.get_matmul_algo_config_attribute_dtype
+        attribute_buffer = np.zeros((1,), dtype=get_dtype(ConfigEnum[name]))
+        size_written = np.zeros((1,), dtype=np.uint64)
+        cublaslt.matmul_algo_config_get_attribute(
+            self.algorithm["algo"].ctypes.data,
+            ConfigEnum[name].value,
+            attribute_buffer.ctypes.data,
+            attribute_buffer.itemsize,
+            size_written.ctypes.data,
+        )
+
+        if name not in SPECIAL_ATTR:
+            return attribute_buffer[0]
+
+        return SPECIAL_ATTR[name].to_value[attribute_buffer[0]]
+
+    def __setattr__(self, name, value):
+        _name = name.upper()
+        if _name not in CAP_ENUM_SCALAR_ATTR:
+            return super().__setattr__(name, value)
+        name = _name
+
+        if name in SPECIAL_ATTR:
+            value = SPECIAL_ATTR[name].to_enumerator[value]
+
+        get_dtype = cublaslt.get_matmul_algo_config_attribute_dtype
+        attribute_buffer = np.zeros((1,), dtype=get_dtype(ConfigEnum[name]))
+        attribute_buffer[0] = value
+        cublaslt.matmul_algo_config_set_attribute(
+            self.algorithm["algo"].ctypes.data,
+            ConfigEnum[name].value,
+            attribute_buffer.ctypes.data,
+            attribute_buffer.itemsize,
+        )

nvmath/linalg/_internal/enum_to_tuples.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2024-2025, 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_")