nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/tensor/_internal/cutensor_config_ifc.py

@@ -0,0 +1,247 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Interface classes to encapsulate low-level calls to get or set configuration information.
+"""
+
+__all__ = ["ContractionPlanPreference"]
+
+
+import numpy as np
+
+from nvmath.bindings import cutensor
+from nvmath.internal import utils
+
+
+class ContractionPlanPreference:
+    """
+    An interface to configure :meth:`nvmath.tensor.BinaryContraction.plan` and
+    :meth:`nvmath.tensor.TernaryContraction.plan`. The
+    current configuration can also be queried.
+    """
+
+    def __init__(self, contraction):
+        """
+        ctor for internal use only.
+        """
+        self._contraction = contraction
+        self._handle = self._contraction.handle
+
+        get_dtype = cutensor.get_plan_preference_attribute_dtype
+
+        self._autotune_mode = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.AUTOTUNE_MODE))
+        self._cache_mode = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.CACHE_MODE))
+        self._incremental_count = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.INCREMENTAL_COUNT))
+        self._algo = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.ALGO))
+        self._kernel_rank = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.KERNEL_RANK))
+        self._jit = np.zeros((1,), dtype=get_dtype(cutensor.PlanPreferenceAttribute.JIT))
+
+    def _check_valid_contraction_wrapper(self, *args, **kwargs):
+        if not self._contraction.valid_state:
+            raise RuntimeError("The ContractionPlanPreference object cannot be used after its contraction object is free'd.")
+
+    @staticmethod
+    def _get_scalar_attribute(contraction, name, attribute):
+        """
+        name      = cutensor PlanPreference enum for the attribute
+        attribute = numpy ndarray object into which the value is stored by cutensornet
+        """
+        assert contraction.plan_preference_ptr is not None, "Internal error"
+        cutensor.plan_preference_get_attribute(
+            contraction.handle, contraction.plan_preference_ptr, name, attribute.ctypes.data, attribute.dtype.itemsize
+        )
+
+    @staticmethod
+    def _set_scalar_attribute(contraction, name, attribute, value):
+        """
+        name      = cutensor PlanPreference enum for the attribute
+        attribute = numpy ndarray object into which the value is stored
+        value     = the value to set the the attribute to
+        """
+        assert contraction.plan_preference_ptr is not None, "Internal error"
+        attribute[0] = value
+        cutensor.plan_preference_set_attribute(
+            contraction.handle, contraction.plan_preference_ptr, name, attribute.ctypes.data, attribute.dtype.itemsize
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def autotune_mode(self):
+        """
+        Query the autotune mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorautotunemode-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.AUTOTUNE_MODE, self._autotune_mode
+        )
+        return self._autotune_mode.item()
+
+    @autotune_mode.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def autotune_mode(self, autotune_mode):
+        """
+        Set the autotune mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorautotunemode-t>`__
+        for more information.
+
+        Args:
+            autotune_mode: (nvmath.tensor.ContractionAutotuneMode) The autotune mode.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.AUTOTUNE_MODE, self._autotune_mode, autotune_mode
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def cache_mode(self):
+        """
+        Query the cache mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorcachemode-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.CACHE_MODE, self._cache_mode
+        )
+        return self._cache_mode.item()
+
+    @cache_mode.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def cache_mode(self, cache_mode):
+        """
+        Set the cache mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorcachemode-t>`__
+        for more information.
+
+        Args:
+            cache_mode: (nvmath.tensor.ContractionCacheMode) The cache mode.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.CACHE_MODE, self._cache_mode, cache_mode
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def incremental_count(self):
+        """
+        Query the incremental count. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorincrementalcount-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.INCREMENTAL_COUNT, self._incremental_count
+        )
+        return self._incremental_count.item()
+
+    @incremental_count.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def incremental_count(self, incremental_count):
+        """
+        Set the incremental count. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorincrementalcount-t>`__
+        for more information.
+
+        Args:
+            incremental_count: The incremental count.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.INCREMENTAL_COUNT, self._incremental_count, incremental_count
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def algo(self):
+        """
+        Query the algo. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensoralgo-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(self._contraction, cutensor.PlanPreferenceAttribute.ALGO, self._algo)
+        return self._algo.item()
+
+    @algo.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def algo(self, algo):
+        """
+        Set the algo. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensoralgo-t>`__
+        for more information.
+
+        Args:
+            algo: (nvmath.tensor.ContractionAlgo) The contraction algorithm.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.ALGO, self._algo, algo
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def kernel_rank(self):
+        """
+        Query the kernel rank. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorkernelrank-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.KERNEL_RANK, self._kernel_rank
+        )
+        return self._kernel_rank.item()
+
+    @kernel_rank.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def kernel_rank(self, kernel_rank):
+        """
+        Set the kernel rank. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorkernelrank-t>`__
+        for more information.
+
+        Args:
+            kernel_rank: The kernel rank.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(
+            self._contraction, cutensor.PlanPreferenceAttribute.KERNEL_RANK, self._kernel_rank, kernel_rank
+        )
+
+    @property
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def jit(self):
+        """
+        Query the jit compilation mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorjit-t>`__
+        for more information.
+        """
+        ContractionPlanPreference._get_scalar_attribute(self._contraction, cutensor.PlanPreferenceAttribute.JIT, self._jit)
+        return self._jit.item()
+
+    @jit.setter
+    @utils.precondition(_check_valid_contraction_wrapper)
+    def jit(self, jit):
+        """
+        Set the jit compilation mode. See the
+        `cuTensor documentation
+        <https://docs.nvidia.com/cuda/cutensor/latest/api/types.html#cutensorjit-t>`__
+        for more information.
+
+        Args:
+            jit: (nvmath.tensor.ContractionJitMode) The JIT compilation mode.
+
+        """
+        ContractionPlanPreference._set_scalar_attribute(self._contraction, cutensor.PlanPreferenceAttribute.JIT, self._jit, jit)

nvmath/tensor/_internal/cutensor_utils.py

@@ -0,0 +1,162 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import weakref
+from collections.abc import Sequence
+
+from nvmath.bindings import cutensor
+from nvmath.internal.tensor_ifc import TensorHolder
+from nvmath.internal.typemaps import NAME_TO_DATA_TYPE
+
+__all__ = [
+    "compute_pointer_alignment",
+    "compute_strides",
+    "TensorDescriptor",
+]
+
+DEFAULT_ALIGNMENT_REQUIREMENT = 256
+
+
+def compute_pointer_alignment(ptr: int) -> int:
+    """
+    Compute the pointer alignment for the given pointer.
+
+    Args:
+        ptr: Pointer address as integer
+
+    Returns:
+        The alignment value (256, 128, 64, 32, 16, 8, 4, 2, or 1)
+    """
+    return 256 if ptr == 0 else min(ptr & -ptr, 256)
+
+
+def compute_strides(shape: Sequence[int], order: str = "C") -> Sequence[int]:
+    """
+    Compute the strides for the given shape and order.
+
+    Args:
+        shape: Shape of the tensor
+
+    Returns:
+        The strides
+    """
+    assert order in {"C", "F"}, f"Invalid order: {order}"
+    strides = [0] * len(shape)
+    stride = 1
+    if order == "C":
+        for axis in range(len(shape) - 1, -1, -1):
+            strides[axis] = stride
+            stride *= shape[axis]
+    elif order == "F":
+        for axis in range(len(shape)):
+            strides[axis] = stride
+            stride *= shape[axis]
+    else:
+        raise ValueError(f"Invalid order: {order}")
+    return strides
+
+
+class TensorDescriptor:
+    """
+    A managed wrapper for cuTensor tensor descriptors with automatic resource cleanup.
+
+    This class wraps a cuTensor tensor descriptor (``cutensorTensorDescriptor_t``) and
+    automatically manages its lifetime. The underlying C descriptor is created on
+    initialization and automatically destroyed when this object is garbage collected.
+
+    Args:
+        handle: The cuTensor library handle.
+        extents: The extents (shape) of each mode of the tensor. Must have the same
+            length as ``strides``.
+        strides: The stride for each mode of the tensor, in elements of the base type.
+            Must have the same length as ``extents``.
+        dtype: The data type of the tensor as a string.
+        alignment: The alignment requirement in bytes for the base pointer that will
+            be used with this descriptor.
+
+    """
+
+    def __init__(
+        self,
+        handle: int,
+        extents: Sequence[int],
+        strides: Sequence[int],
+        dtype: str,
+        alignment: int,
+    ):
+        self._extents = extents
+        self._strides = strides
+        self._dtype = dtype
+        self._alignment = alignment
+        num_modes = len(extents)
+        assert num_modes == len(strides), "The number of modes and strides must be the same"
+        self._tensor_desc = cutensor.create_tensor_descriptor(
+            handle, num_modes, extents, strides, NAME_TO_DATA_TYPE[dtype], alignment
+        )
+        self._finalizer = weakref.finalize(self, cutensor.destroy_tensor_descriptor, self._tensor_desc)
+
+    @property
+    def ptr(self) -> int:
+        return self._tensor_desc
+
+    @property
+    def shape(self) -> Sequence[int]:
+        return self._extents
+
+    @property
+    def dtype(self) -> str:
+        return self._dtype
+
+    @property
+    def alignment(self) -> int:
+        return self._alignment
+
+    @property
+    def strides(self) -> Sequence[int]:
+        return self._strides
+
+    @classmethod
+    def from_tensor_holder(cls, handle: int, tensor: TensorHolder) -> "TensorDescriptor":
+        """
+        Create a TensorDescriptor from a TensorHolder.
+
+        Args:
+            handle: The cuTensor library handle.
+            tensor: The TensorHolder to create the descriptor from.
+
+        Returns:
+            A TensorDescriptor with properties derived from the TensorHolder.
+        """
+        return cls(
+            handle=handle,
+            extents=tensor.shape,
+            strides=tensor.strides,
+            dtype=tensor.dtype,
+            alignment=compute_pointer_alignment(tensor.data_ptr),
+        )
+
+    @classmethod
+    def from_shape_and_dtype(
+        cls, handle: int, shape: Sequence[int], dtype: str, order: str = "C", alignment: int = DEFAULT_ALIGNMENT_REQUIREMENT
+    ) -> "TensorDescriptor":
+        """
+        Create a TensorDescriptor for the given shape, data type, order, and alignment.
+
+        Args:
+            handle: The cuTensor library handle.
+            shape: The shape (extents) of the tensor.
+            dtype: The data type of the tensor as a string.
+            order: The memory layout order. Defaults to ``"C"``.
+            alignment: The alignment requirement in bytes. Defaults to 256.
+
+        Returns:
+            A TensorDescriptor with the specified properties.
+        """
+        return cls(
+            handle=handle,
+            extents=shape,
+            strides=compute_strides(shape, order),
+            dtype=dtype,
+            alignment=alignment,
+        )

nvmath/tensor/_internal/data.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Defined in CPython:
+# https://github.com/python/cpython/blob/26bc2cc06128890ac89492eca20e83abe0789c1c/Objects/unicodetype_db.h#L6311-L6349
+
+__all__ = ["WHITESPACE_UNICODE"]
+
+_WHITESPACE_UNICODE_INTS = [
+    0x0009,
+    0x000A,
+    0x000B,
+    0x000C,
+    0x000D,
+    0x001C,
+    0x001D,
+    0x001E,
+    0x001F,
+    0x0020,
+    0x0085,
+    0x00A0,
+    0x1680,
+    0x2000,
+    0x2001,
+    0x2002,
+    0x2003,
+    0x2004,
+    0x2005,
+    0x2006,
+    0x2007,
+    0x2008,
+    0x2009,
+    0x200A,
+    0x2028,
+    0x2029,
+    0x202F,
+    0x205F,
+    0x3000,
+]
+
+
+WHITESPACE_UNICODE = "".join(chr(s) for s in _WHITESPACE_UNICODE_INTS)