nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/_internal/templates.py

@@ -0,0 +1,383 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import abc
+import contextlib
+import dataclasses
+import logging
+from collections.abc import MutableSequence
+from logging import Logger
+from typing import ClassVar, Final, Generic, Literal, TypeVar
+
+from nvmath import memory
+from nvmath.internal import utils
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class ExecutionCUDA:
+    """
+    A data class for providing GPU execution options.
+
+    Attributes:
+        device_id: CUDA device ordinal (only used if the operand resides on the CPU). The
+            default value is 0.
+
+    .. seealso::
+       :class:`ExecutionCPU`
+    """
+
+    name: ClassVar[Literal["cuda"]] = "cuda"
+    device_id: int = 0
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class ExecutionCPU:
+    """
+    A data class for providing CPU execution options.
+
+    Attributes:
+        num_threads: The number of CPU threads used to execute the operation.
+                     If not specified, defaults to the number of CPU cores available to the
+                     process.
+
+    .. seealso::
+       :class:`ExecutionCUDA`
+    """
+
+    name: ClassVar[Literal["cpu"]] = "cpu"
+    num_threads: int | None = None
+
+
+def copy_operand_perhaps(
+    internal_operand: utils.TensorHolder | None,
+    operand: utils.TensorHolder,
+    stream_holder: utils.StreamHolder | None,
+    execution_device_id: int | Literal["cpu"],
+    operands_device_id: int | Literal["cpu"],
+) -> tuple[utils.TensorHolder, utils.TensorHolder | None]:
+    """Private implementation of memory space management for tensor operands.
+
+    The `copy_operand_perhaps` function facilitates transitions of tensor operands between
+    different memory spaces, ensuring compatibility with execution requirements. Its role is
+    to determine whether a tensor operand needs to be copied to accommodate differing
+    execution and operand memory spaces, while preserving the original operand for cases
+    requiring in-place operations.
+
+    Args:
+        internal_operand: Represents an internal tensor for in-place
+            memory operations, or `None` if not applicable.
+
+        operand: Tensor to possibly copied to the execution memory space.
+
+        stream_holder: Manages the CUDA stream for device operations.
+
+        execution_device_id: Specifies the target execution space.
+
+        operands_device_id: Specifies the current operand memory space.
+
+    Returns:
+        A tuple containing:
+            - The operand copied to the execution space, or the original operand if
+              no copy is necessary.
+            - The original operand, or `None` if no copy occurred.
+
+    """
+    if execution_device_id == operands_device_id:
+        return operand, None
+    else:
+        # Copy the `operand` to memory that matches the exec space
+        # and keep the original `operand` to handle `options.inplace=True`
+        if internal_operand is None:
+            exec_space_copy = operand.to(execution_device_id, stream_holder)
+            return exec_space_copy, operand
+        else:
+            # In-place copy to existing pointer
+            internal_operand.copy_(src=operand, stream_holder=stream_holder)
+            return internal_operand, operand
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class StatefulAPIOptions:
+    """A dataclass for providing options to a :class:`StatefulAPI` object.
+
+    Attributes:
+        allocator: An object that supports the :class:`BaseCUDAMemoryManager` protocol, used
+            to draw device memory. If an allocator is not provided, a memory allocator from
+            the library package will be used (:func:`torch.cuda.caching_allocator_alloc` for
+            PyTorch operands, :func:`cupy.cuda.alloc` otherwise).
+
+        blocking: A flag specifying the behavior of the stream-ordered functions and
+            methods. When ``blocking`` is `True`, the stream-ordered methods do not return
+            until the operation is complete. When ``blocking`` is ``"auto"``, the methods
+            return immediately when the inputs are on the GPU. The stream-ordered methods
+            always block when the operands are on the CPU to ensure that the user doesn't
+            inadvertently use the result before it becomes available. The default is
+            ``"auto"``.
+
+        logger: Python Logger object. The root logger will be used if a
+            logger object is not provided.
+
+    .. seealso::
+       :class:`StatefulAPI`
+    """
+
+    allocator: memory.BaseCUDAMemoryManager | memory.BaseCUDAMemoryManagerAsync | None = None
+    blocking: Literal[True, "auto"] = "auto"
+    logger: Logger = dataclasses.field(default_factory=logging.getLogger)
+
+    def __post_init__(self):
+        if self.blocking not in (True, "auto"):
+            raise ValueError("The value specified for blocking must be either True or 'auto'.")
+
+        if self.allocator is not None and not isinstance(
+            self.allocator, memory.BaseCUDAMemoryManager | memory.BaseCUDAMemoryManagerAsync
+        ):
+            raise TypeError("The allocator must be an object of type that fulfills the BaseCUDAMemoryManager protocol.")
+
+
+OptionsPlaceholder = TypeVar("OptionsPlaceholder", bound=StatefulAPIOptions)
+
+
+class StatefulAPI(contextlib.AbstractContextManager, Generic[OptionsPlaceholder]):
+    """A base class for APIs which amortize setup costs across multiple executions.
+
+    StatefulAPIs separate planning (``plan()``) and setup (``__init__()``) actions from
+    execution (``_execute()``), so that plans may be reused with different operands. The
+    ``reset_operands()`` method allows changing the operands of the API without replanning
+    when the input and execution space do not match (the user does not have a reference to
+    the execution space buffers). If the execution and input space match, we expect the
+    user to be able update the operands by overwriting their buffers in-place.
+    """
+
+    # options is declared Final in __init__() because mypy issue #8982 is not fixed yet.
+    # options: Final[OptionsPlaceholder]
+    # See docstring for StatefulAPIOptions
+    _allocator: Final[memory.BaseCUDAMemoryManager | memory.BaseCUDAMemoryManagerAsync | None]
+    _blocking: Final[bool]
+    _logger: Final[logging.Logger]
+
+    # Metadata related to execution space
+    execution: Final[ExecutionCPU | ExecutionCUDA]
+    """A class which describes the execution space parameters."""
+    _internal_op_package: Final[str]
+    """The package of the operands in the execution space."""
+    _operands: MutableSequence[utils.TensorHolder]
+    """A copy of the operands in execution space."""
+    _result_class: Final[type[utils.TensorHolder]]
+    """The type of TensorHolder to use for the execution space result."""
+
+    # Metadata about the input/output tensors
+    _operands_backup: MutableSequence[utils.TensorHolder | None]
+    """A reference to original operands in their input space."""
+    _operands_device_id: Final[int | Literal["cpu"]]
+    """The device_id of the input space."""
+    _operands_package: Final[str]
+    """The package of the operands in the input space."""
+
+    _call_prologue: Final[str]
+    """Stores a message for logging about blocking behavior"""
+
+    _has_plan: bool
+    """True if plan has been called."""
+
+    @property
+    def options(self) -> OptionsPlaceholder:
+        """The options object used to construct this class."""
+        # This is a workaround for mypy issue #8982, where we cannot declare options as
+        # Final in the class definition, but we still want it to appear in the docs as an
+        # attribute.
+        return self._options
+
+    def __init__(
+        self,
+        operands: MutableSequence[utils.TensorHolder],
+        *,
+        options: OptionsPlaceholder,
+        execution: ExecutionCPU | ExecutionCUDA | None | Literal["cuda", "cpu"] = None,
+        stream: utils.AnyStream | int | None = None,
+    ) -> None:
+        """Copy operands to the execution space and setup options.
+
+        When inheriting from this class, you must create valid operands and options in
+        the child class before calling StatefulAPI.__init__( ... ).
+        """
+        self._options: Final[OptionsPlaceholder] = options
+        self._logger = self._options.logger
+
+        self._logger.info("= SPECIFICATION PHASE =")
+
+        operands_device_id = utils.get_operands_device_id(operands)
+
+        match execution, operands_device_id:
+            case (None | "cuda", int()):
+                execution = ExecutionCUDA(device_id=operands_device_id)
+            case ("cuda", "cpu"):
+                execution = ExecutionCUDA()
+            case (None, "cpu") | ("cpu", _):
+                execution = ExecutionCPU()
+            case (ExecutionCUDA(), int()):
+                # If operands are on a CUDA device, use the same device for execution.
+                execution = dataclasses.replace(execution, device_id=operands_device_id)
+            case (ExecutionCPU(), _) | (ExecutionCUDA(), "cpu"):
+                pass
+            case _:
+                raise ValueError(
+                    f"{self.__class__.__name__}.execution must be one of ExecutionCUDA, ExecutionCPU, None, 'cuda', or 'cpu'."
+                )
+        assert isinstance(execution, (ExecutionCPU, ExecutionCUDA))
+        self.execution = execution
+
+        self._operands_device_id = operands_device_id
+        self._operands_package = utils.get_operands_package(operands)
+        self._internal_op_package = self._internal_operand_package(self._operands_package)
+        exec_stream_holder, operand_stream_holder = self._get_or_create_stream_maybe(stream)
+
+        self._logger.info(
+            f"The input tensors are located on device {operands_device_id}, and the execution space "
+            f"is {self.execution.name}, with device {getattr(self.execution, 'device_id', 'cpu')}."
+        )
+
+        self._logger.info(
+            f"The specified stream for the {self.__class__.__name__} constructor is "
+            f"{(exec_stream_holder or operand_stream_holder) and getattr(exec_stream_holder or operand_stream_holder, 'obj', None)}."  # noqa: E501
+        )
+
+        operands_backup: list[utils.TensorHolder | None] = [None] * len(operands)
+        for i in range(len(operands)):
+            # Copy the operand to execution_space's device if needed.
+            operands[i], operands_backup[i] = copy_operand_perhaps(
+                None,
+                operands[i],
+                operand_stream_holder,
+                getattr(self.execution, "device_id", "cpu"),
+                self._operands_device_id,
+            )
+        self._operands = operands
+        self._operands_backup = operands_backup
+
+        # The result's package and device.
+        self._result_class = self._operands[0].__class__
+
+        # Set blocking or non-blocking behavior.
+        self._blocking = self._options.blocking != "auto" or self._operands_device_id == "cpu" or self.execution.name == "cpu"
+        if self._blocking:
+            call_prologue = "This call is blocking and will return only after the operation is complete."
+        else:
+            call_prologue = (
+                "This call is non-blocking and will return immediately after the operation is launched on the device."
+            )
+        self._call_prologue = call_prologue
+
+        # Set memory allocator.
+        allocator: memory.BaseCUDAMemoryManager | memory.BaseCUDAMemoryManagerAsync | None
+        match self.execution:
+            case ExecutionCUDA():
+                allocator = (
+                    memory._MEMORY_MANAGER[self._internal_op_package](self.execution.device_id, self._logger)
+                    if self._options.allocator is None
+                    else self._options.allocator
+                )
+            case ExecutionCPU() | _:
+                allocator = None  # currently, the nvpl/fftw does not support custom workspace allocation
+        self._allocator = allocator
+
+        self._has_plan = False
+
+    def _internal_operand_package(self, package_name: str) -> str:
+        if self.execution.name == "cuda":
+            return package_name if package_name != "numpy" else "cuda"
+        else:
+            return package_name if package_name != "cupy" else "cupy_host"
+
+    def _get_or_create_stream_maybe(
+        self, stream: utils.AnyStream
+    ) -> tuple[utils.StreamHolder | None, utils.StreamHolder | None]:
+        """Return a 2-tuple of Stream | None: one for execution space, one for input space.
+
+        The first stream should be used for everything in the execution space: doing work,
+        allocating workspace, allocating input/output buffers.
+
+        The second stream should be used whenever data is being moved between the input and
+        output spaces: copying data to/from the input/output tensors.
+
+        NOTE: If two streams are returned, they will be the same stream.
+        """
+        if self.execution.name == "cuda":
+            stream_holder = utils.get_or_create_stream(self.execution.device_id, stream, self._internal_op_package)
+            return stream_holder, stream_holder
+        elif isinstance(self._operands_device_id, int):
+            operand_device_steam = utils.get_or_create_stream(self._operands_device_id, stream, self._operands_package)
+            return None, operand_device_steam
+        else:
+            return None, None
+
+    # input checks
+
+    def _check_valid_operands(self, *args, **kwargs):
+        """
+        Check if the operands are available for the operation.
+        """
+        what = kwargs["what"]
+        if self._operands is None:
+            raise RuntimeError(
+                f"{what} cannot be performed if the operands have been set to None. Use reset_operands() to set the "
+                f"desired input before using performing the {what.lower()}."
+            )
+
+    def _check_planned(self, *args, **kwargs):
+        what = kwargs["what"]
+        if not self._has_plan:
+            raise RuntimeError(f"{what} cannot be performed before plan() has been called.")
+
+    # execution
+
+    @abc.abstractmethod
+    def _execute(self):
+        """Perform the main functionality of this :class:`StatefulAPI` instance without
+        safety checks."""
+        msg = f"{self.__name__}._execute() is not implemented."
+        raise NotImplementedError(msg)
+
+    @abc.abstractmethod
+    def plan(self):
+        """Plan the main functionality of this :class:`StatefulAPI` instance."""
+        msg = f"{self.__class__.__name__}.plan() is not implemented."
+        raise NotImplementedError(msg)
+
+    @abc.abstractmethod
+    def reset_operands(self):
+        """Reset the operands held by this :class:`StatefulAPI` instance."""
+        msg = f"{self.__class__.__name__}.reset_operands() is not implemented."
+        raise NotImplementedError(msg)
+
+    def free(self):
+        """Free resources held by this :class:`StatefulAPI` instance.
+
+        This method releases references to operands to ensure proper reference counting.
+        Subclasses should call super().free() if they override this method to add
+        additional cleanup logic.
+        """
+        # Release references to operands to ensure proper reference counting
+        self._operands = []
+        self._operands_backup = []
+
+
+class HasWorkspaceMemory(contextlib.AbstractContextManager):
+    """A base class for APIs which need to allocate a working buffer in memory."""
+
+    def _allocate_workspace_memory_perhaps(self, stream_holder: utils.StreamHolder):
+        """
+        Allocate workspace memory using the specified allocator, if it hasn't already been
+        done.
+        """
+        raise NotImplementedError
+
+    def _release_workspace_memory_perhaps(self, release_workspace: bool):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def free(self):
+        """Free the resources of this :class:`StatefulAPI` instance."""
+        msg = f"{self.__name__}.free() is not implemented."
+        raise NotImplementedError(msg)

nvmath/_utils.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import re
+import sys
+from enum import IntEnum
+from functools import cache
+
+from cuda import pathfinder
+
+logger = logging.getLogger()
+
+#
+# Note: This module should not depend on anything from the nvmath namespace!
+#
+
+
+# The (subset of) compute types below are shared by cuStateVec and cuTensorNet
+class ComputeType(IntEnum):
+    """An enumeration of CUDA compute types."""
+
+    COMPUTE_DEFAULT = 0
+    COMPUTE_16F = 1 << 0
+    COMPUTE_32F = 1 << 2
+    COMPUTE_64F = 1 << 4
+    COMPUTE_8U = 1 << 6
+    COMPUTE_8I = 1 << 8
+    COMPUTE_32U = 1 << 7
+    COMPUTE_32I = 1 << 9
+    COMPUTE_16BF = 1 << 10
+    COMPUTE_TF32 = 1 << 12
+
+
+# TODO: use those exposed by CUDA Python instead, but before removing these
+# duplicates, check if they are fixed to inherit IntEnum instead of Enum.
+class CudaDataType(IntEnum):
+    """An enumeration of `cudaDataType_t`."""
+
+    CUDA_R_16F = 2
+    CUDA_C_16F = 6
+    CUDA_R_16BF = 14
+    CUDA_C_16BF = 15
+    CUDA_R_32F = 0
+    CUDA_C_32F = 4
+    CUDA_R_64F = 1
+    CUDA_C_64F = 5
+    CUDA_R_4I = 16
+    CUDA_C_4I = 17
+    CUDA_R_4U = 18
+    CUDA_C_4U = 19
+    CUDA_R_8I = 3
+    CUDA_C_8I = 7
+    CUDA_R_8U = 8
+    CUDA_C_8U = 9
+    CUDA_R_16I = 20
+    CUDA_C_16I = 21
+    CUDA_R_16U = 22
+    CUDA_C_16U = 23
+    CUDA_R_32I = 10
+    CUDA_C_32I = 11
+    CUDA_R_32U = 12
+    CUDA_C_32U = 13
+    CUDA_R_64I = 24
+    CUDA_C_64I = 25
+    CUDA_R_64U = 26
+    CUDA_C_64U = 27
+    CUDA_R_8F_E4M3 = 28
+    CUDA_R_8F_E5M2 = 29
+    CUDA_R_4F_E2M1 = 33  # FP4 (supported on SM >= 10.0)
+
+
+class LibraryPropertyType(IntEnum):
+    """An enumeration of library version information."""
+
+    MAJOR_VERSION = 0
+    MINOR_VERSION = 1
+    PATCH_LEVEL = 2
+
+
+del IntEnum
+
+
+PLATFORM_LINUX = sys.platform.startswith("linux")
+PLATFORM_WIN = sys.platform.startswith("win32")
+
+
+def module_init_force_cupy_lib_load():
+    """
+    Attempt to preload libraries at module import time. We want to do it before
+    cupy, since it does not know how to properly search for libraries:
+    https://github.com/cupy/cupy/issues/9127
+    Fail silently if preload fails.
+    """
+    from nvmath.bindings import _internal
+
+    libs = ("cublas", "cufft", "curand", "cusolverDn", "cusparse", "cutensor")
+    for lib in libs:
+        try:
+            mod = getattr(_internal, lib)
+            mod._inspect_function_pointers()
+        except (_internal.utils.NotSupportedError, RuntimeError):
+            pass
+
+    try:
+        pathfinder.load_nvidia_dynamic_lib("nvrtc")
+    except pathfinder.DynamicLibNotFoundError:
+        pass
+
+
+@cache
+def get_nvrtc_build_id(minimal=True) -> int:
+    try:
+        from cuda.core import ObjectCode, Program, ProgramOptions
+    except ImportError:
+        from cuda.core.experimental import ObjectCode, Program, ProgramOptions
+
+    code = r"""
+    extern "C" __global__ void get_build_id(int* build_id) {
+
+        *build_id = __CUDACC_VER_BUILD__;
+    }
+    """
+
+    prog = Program(code, "c++", ProgramOptions(std="c++17", minimal=minimal, arch="compute_75"))
+    obj = prog.compile("ptx")
+    assert isinstance(obj, ObjectCode)
+
+    pattern = re.compile(r"mov\.u32\s+%\w+,\s+(\d+)")
+    m = pattern.search(obj.code.decode())
+    assert m is not None
+
+    return int(m.group(1))
+
+
+@cache
+def get_nvrtc_version() -> tuple[int, int, int]:
+    """
+    Returns the NVRTC version as a tuple of (major, minor, build).
+    """
+    from cuda.bindings import nvrtc
+
+    err, major, minor = nvrtc.nvrtcVersion()
+    if err != nvrtc.nvrtcResult.NVRTC_SUCCESS:
+        raise RuntimeError(f"nvrtcVersion error: {err}")
+    # minimal support was added in CUDA 12.4
+    # https://docs.nvidia.com/cuda/archive/12.4.0/cuda-toolkit-release-notes/index.html#cuda-compilers
+    build = get_nvrtc_build_id(minimal=major >= 12 and minor >= 4)
+    return major, minor, build

nvmath/bindings/__init__.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# type: ignore
+
+from nvmath.bindings import cublas
+from nvmath.bindings import cublasLt
+from nvmath.bindings import cudss
+from nvmath.bindings import cufft
+from nvmath.bindings import curand
+from nvmath.bindings import cusolver
+from nvmath.bindings import cusolverDn
+from nvmath.bindings import cusparse
+from nvmath.bindings import cusparseLt
+from nvmath.bindings import cutensor
+
+try:
+    # cufftMp is Linux-only.
+    from nvmath.bindings import cufftMp
+except ImportError:
+    cufftMp = None
+
+try:
+    # nvshmem is Linux-only.
+    from nvmath.bindings import nvshmem
+except ImportError:
+    nvshmem = None
+
+try:
+    # cublasMp is Linux-only.
+    from nvmath.bindings import cublasMp
+except ImportError:
+    cublasMp = None
+
+__all__ = [
+    "cublas",
+    "cublasLt",
+    "cublasMp",
+    "cudss",
+    "cufft",
+    "cufftMp",
+    "curand",
+    "cusolver",
+    "cusolverDn",
+    "cusparse",
+    "cusparseLt",
+    "cutensor",
+    "nvpl",
+    "nvshmem",
+]