cuquantum-python-cu13 25.9.0__cp312-cp312-manylinux2014_aarch64.whl

cuquantum/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from cuquantum import bindings
+from cuquantum import densitymat
+from cuquantum import tensornet
+from cuquantum.memory import *
+from cuquantum.bindings._utils import ComputeType, cudaDataType, libraryPropertyType
+from cuquantum._version import __version__

cuquantum/__main__.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import argparse
+import os
+import site
+import sys
+import logging
+
+logger = logging.getLogger(__name__)
+
+def get_lib_path(name):
+    """Get the loaded shared library path."""
+    # Ideally we should call dl_iterate_phdr or dladdr to do the job, but this
+    # is simpler and not bad; the former two are not strictly portable anyway
+    # (not part of POSIX). Obviously this only works on Linux!
+
+    # We have switched to use dlopen, force library loading via internal API
+    if "custatevec" in name:
+        from cuquantum import bindings
+        bindings._internal.custatevec._inspect_function_pointers()
+    elif "cutensor" in name:  # cutensor or cutensornet
+        from cuquantum import bindings
+        bindings._internal.cutensornet._inspect_function_pointers()
+    elif "cudensitymat" in name:
+        from cuquantum import bindings
+        bindings._internal.cudensitymat._inspect_function_pointers()
+
+    try:
+        with open('/proc/self/maps') as f:
+            lib_map = f.read()
+    except FileNotFoundError as e:
+        raise NotImplementedError("This utility is available only on Linux.") from e
+    lib = set()
+    for line in lib_map.split('\n'):
+        if name in line:
+            fields = line.split()
+            lib.add(fields[-1])  # pathname is the last field, check "man proc"
+    if len(lib) == 0:
+        raise ValueError(f"library {name} is not loaded")
+    elif len(lib) > 1:
+        # This could happen when, e.g., a library exists in both the user env
+        # and LD_LIBRARY_PATH, and somehow both copies get loaded. This is a
+        # messy problem, but let's work around it by assuming the one in the
+        # user env is preferred.
+        lib2 = set()
+        for s in [site.getusersitepackages()] + site.getsitepackages():
+            for path in lib:
+                if path.startswith(s):
+                    lib2.add(path)
+        if len(lib2) != 1:
+            raise RuntimeError(f"cannot find the unique copy of {name}: {lib}")
+        else:
+            lib = lib2
+    return lib.pop()
+
+
+def _get_cuquantum_lib(lib):
+    try:
+        path = os.path.normpath(get_lib_path(f"lib{lib}.so"))
+        return path
+    except Exception as e:
+        logger.warning(f"library {lib} not found: {e}")
+        return None
+
+
+def _get_cuquantum_libs():
+    paths = set()
+    for lib in ('custatevec', 'cutensornet', 'cutensor', 'cudensitymat'):
+        path = _get_cuquantum_lib(lib)
+        if path is not None:
+            paths.add(path)
+    return tuple(paths)
+
+
+def _get_cuquantum_includes():
+    paths = set()
+    for path in _get_cuquantum_libs():
+        path = os.path.normpath(os.path.join(os.path.dirname(path), '..'))
+        if not os.path.isdir(os.path.join(path, 'include')):
+            path = os.path.normpath(os.path.join(path, '../include'))
+        else:
+            path = os.path.join(path, 'include')
+
+        if os.path.isdir(path):
+            paths.add(path)
+    return tuple(paths)
+
+
+def _get_cuquantum_target(target):
+    path = _get_cuquantum_lib(target)
+    if path is None:
+        # Return error as the target library was explicitly asked but not found
+        raise RuntimeError(f"library {target} not found")
+    lib = os.path.basename(path)
+    lib = '.'.join(lib.split('.')[:3])  # keep SONAME
+    flag = f"-l:{lib} "
+    return flag
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--includes', action='store_true',
+                        help='get cuQuantum include flags')
+    parser.add_argument('--libs', action='store_true',
+                        help='get cuQuantum linker flags')
+    parser.add_argument('--target', action='append', default=[],
+                        choices=('custatevec', 'cutensornet', 'cudensitymat'),
+                        help='get the linker flag for the target cuQuantum component')
+    args = parser.parse_args()
+
+    if not sys.argv[1:]:
+        parser.print_help()
+        sys.exit(1)
+    if args.includes:
+        out = ' '.join(f"-I{path}" for path in _get_cuquantum_includes())
+        print(out, end=' ')
+    if args.libs:
+        paths = set([os.path.dirname(path) for path in _get_cuquantum_libs()])
+        out = ' '.join(f"-L{path}" for path in paths)
+        print(out, end=' ')
+    flag = ''
+    for target in args.target:
+        flag += _get_cuquantum_target(target)
+        if target in ('cutensornet', 'cudensitymat') :
+            flag += _get_cuquantum_target('cutensor')
+    print(flag)

cuquantum/_internal/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause

cuquantum/_internal/package_ifc.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+An abstract interface to certain package-provided operations.
+"""
+
+__all__ = ['Package']
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from contextlib import nullcontext
+from typing import Any
+
+
+class Package(ABC):
+
+    @staticmethod
+    @abstractmethod
+    def get_current_stream(device_id):
+        """
+        Obtain the current stream on the device.
+
+        Args:
+            device_id: The id (ordinal) of the device.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def to_stream_pointer(stream):
+        """
+        Obtain the stream pointer.
+
+        Args:
+            stream: The stream object.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def to_stream_context(stream):
+        """
+        Create a context manager from the stream.
+
+        Args:
+            stream: The stream object.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def create_external_stream(device_id, stream_ptr):
+        """
+        Wrap a stream pointer into an external stream object.
+
+        Args:
+            device_id: The id (ordinal) of the device.
+            stream: The stream pointer (int) to be wrapped.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def create_stream(device_id):
+        """
+        Create a new stream on the specified device.
+
+        Args:
+            device_id: The id (ordinal) of the device.
+        """
+        raise NotImplementedError
+
+
+@dataclass
+class StreamHolder:
+    """A data class for easing CUDA stream manipulation.
+
+    Attributes:
+        ctx: A context manager for using the specified stream.
+        device_id (int): The device ID where the encapsulated stream locates.
+        obj: A foreign object that holds the stream alive.
+        package (str):
+        ptr (int): The address of the underlying ``cudaStream_t`` object.
+    """
+    ctx: Any = nullcontext()
+    device_id: int = -2
+    obj: Any = None
+    package: str = ""
+    ptr: int = 0

cuquantum/_internal/tensor_ifc.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Interface to seamlessly use tensors (or ndarray-like objects) from different libraries.
+"""
+
+from abc import ABC, abstractmethod
+
+from nvmath.internal import typemaps
+
+class Tensor(ABC):
+    """
+    A simple wrapper type for tensors to make the API package-agnostic.
+    """
+
+    def __init__(self, tensor):
+        self.tensor = tensor
+
+    @property
+    @abstractmethod
+    def data_ptr(self):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def device(self):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def device_id(self):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def dtype(self):
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def empty(cls, shape, **context):
+        raise NotImplementedError
+
+    @abstractmethod
+    def numpy(self, stream_holder):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def shape(self):
+        raise NotImplementedError
+    
+    @property
+    def ndim(self):
+        return len(self.shape)
+    
+    @property
+    def T(self):
+        return self.__class__(self.tensor.T)
+
+    @property
+    @abstractmethod
+    def strides(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def to(self, device='cpu', stream_holder=None):
+        raise NotImplementedError
+
+    @abstractmethod
+    def copy_(self, src, stream_holder=None):
+        raise NotImplementedError
+
+    @staticmethod
+    def create_name_dtype_map(conversion_function, exception_type):
+        """
+        Create a map between CUDA data type names and the corresponding package dtypes for supported data types.
+        """
+        names = typemaps.NAME_TO_DATA_TYPE.keys()
+        name_to_dtype = dict()
+        for name in names:
+            try:
+                name_to_dtype[name] = conversion_function(name)
+            except exception_type:
+                pass
+        return name_to_dtype

cuquantum/_internal/tensor_ifc_cupy.py

@@ -0,0 +1,126 @@
+# Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Interface to seamlessly use Cupy ndarray objects.
+"""
+
+__all__ = ['CupyTensor']
+
+import cupy
+
+from nvmath.internal import utils
+from .package_ifc import StreamHolder
+from .tensor_ifc import Tensor
+
+
+class CupyTensor(Tensor):
+    """
+    Tensor wrapper for cupy ndarrays.
+    """
+    name = 'cupy'
+    module = cupy
+    name_to_dtype = Tensor.create_name_dtype_map(conversion_function=lambda name: cupy.dtype(name), exception_type=TypeError)
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    @property
+    def data_ptr(self):
+        return self.tensor.data.ptr
+
+    @property
+    def device(self):
+        return 'cuda'
+
+    @property
+    def device_id(self):
+        return self.tensor.device.id
+
+    @property
+    def dtype(self):
+        """Name of the data type"""
+        return self.tensor.dtype.name
+
+    @property
+    def shape(self):
+        return tuple(self.tensor.shape)
+
+    @property
+    def strides(self):
+        return tuple(stride_in_bytes // self.tensor.itemsize for stride_in_bytes in self.tensor.strides)
+
+    def numpy(self, stream_holder=StreamHolder()):
+        stream = stream_holder.obj
+        # cupy/cupy#7820
+        if stream is not None:
+            with stream:
+                out = self.tensor.get(stream=stream)
+            stream.synchronize()
+        else:
+            out = self.tensor.get(stream=stream)
+            # note that synchronization semantics are inconsistent for stream_holder 
+            # TODO: add synchronization on current stream here (better to recompose the if/else branch then)
+        return out
+
+    @classmethod
+    def empty(cls, shape, **context):
+        """
+        Create an empty tensor of the specified shape and data type.
+        """
+        name = context.get('dtype', 'float32')
+        dtype = CupyTensor.name_to_dtype[name]
+        device = context.get('device', None)
+        strides = context.get('strides', None)
+
+        if isinstance(device, cupy.cuda.Device):
+           device_id = device.id
+        elif isinstance(device, int):
+           device_id = device
+        else:
+            raise ValueError(f"The device must be specified as an integer or cupy.cuda.Device instance, not '{device}'.")
+
+        with utils.device_ctx(device_id):
+            if strides:
+                # need an explicit allocation due to cupy/cupy#7818
+                size = dtype.itemsize
+                for s in shape:
+                    size = size * s
+                ptr = cupy.cuda.alloc(size)
+                # when strides is not None, it should be of unit counts not bytes
+                strides = tuple(s * dtype.itemsize for s in strides)
+                tensor = cupy.ndarray(shape, dtype=dtype, strides=strides, memptr=ptr)
+            else:
+                tensor = cupy.ndarray(shape, dtype=dtype)
+
+        return tensor
+
+    def to(self, device='cpu', stream_holder=StreamHolder()):
+        """
+        Create a copy of the tensor on the specified device (integer or 
+          'cpu'). Copy to  Numpy ndarray if CPU, otherwise return Cupy type.
+        """
+        if device == 'cpu':
+            return self.numpy(stream_holder=stream_holder)
+
+        if not isinstance(device, int):
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device}'.")
+
+        with utils.device_ctx(device), stream_holder.ctx:
+            tensor_device = cupy.asarray(self.tensor)
+
+        return tensor_device
+
+    def copy_(self, src, stream_holder=StreamHolder()):
+        """
+        Inplace copy of src (copy the data from src into self).
+        """
+        with stream_holder.ctx:
+            cupy.copyto(self.tensor, src)
+
+    def istensor(self):
+        """
+        Check if the object is ndarray-like.
+        """
+        return isinstance(self.tensor, cupy.ndarray)

cuquantum/_internal/tensor_ifc_numpy.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Interface to seamlessly use Numpy ndarray objects.
+"""
+
+__all__ = ['NumpyTensor']
+
+import cupy
+import numpy
+
+from nvmath.internal import utils
+from .package_ifc import StreamHolder
+from .tensor_ifc import Tensor
+
+
+class NumpyTensor(Tensor):
+    """
+    Tensor wrapper for numpy ndarrays.
+    """
+    name = 'numpy'
+    module = numpy
+    name_to_dtype = Tensor.create_name_dtype_map(conversion_function=lambda name: numpy.dtype(name), exception_type=TypeError)
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    @property
+    def data_ptr(self):
+        return self.tensor.ctypes.data
+
+    @property
+    def device(self):
+        return 'cpu'
+
+    @property
+    def device_id(self):
+        return None
+
+    @property
+    def dtype(self):
+        """Name of the data type"""
+        return self.tensor.dtype.name
+
+    @property
+    def shape(self):
+        return tuple(self.tensor.shape)
+
+    @property
+    def strides(self):
+        return tuple(stride_in_bytes // self.tensor.itemsize for stride_in_bytes in self.tensor.strides)
+
+    def numpy(self, stream_holder=StreamHolder()):
+        return self.tensor
+
+    @classmethod
+    def empty(cls, shape, **context):
+        """
+        Create an empty tensor of the specified shape and data type.
+        """
+        name = context.get('dtype', 'float32')
+        dtype = NumpyTensor.name_to_dtype[name]
+        strides = context.get('strides', None)
+        # when strides is not None, it should be of unit counts not bytes
+        return cls(numpy.ndarray(shape, dtype=dtype, strides=(tuple(s * dtype.itemsize for s in strides) if strides else None)))
+
+    def to(self, device='cpu', stream_holder=StreamHolder()):
+        """
+        Create a copy of the tensor on the specified device (integer or 
+          'cpu'). Copy to  Cupy ndarray on the specified device if it 
+          is not CPU. Otherwise, return self.
+        """
+        if device == 'cpu':
+            return self.tensor
+
+        if not isinstance(device, int):
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device}'.")
+
+        with utils.device_ctx(device), stream_holder.ctx:
+            tensor_device = cupy.asarray(self.tensor)
+
+        return tensor_device
+
+    def copy_(self, src, stream_holder=StreamHolder()):
+        package = utils.infer_object_package(src)
+        # Handle NumPy <=> CuPy CPU-GPU ndarray asymmetry.
+        if package == 'cupy':
+            stream = stream_holder.obj
+            with stream:
+                out = src.get(stream=stream, out=self.tensor)
+            # cupy/cupy#7820
+            if stream is not None:
+                stream.synchronize()
+
+            return out
+        else:
+            raise NotImplementedError
+
+    def istensor(self):
+        """
+        Check if the object is ndarray-like.
+        """
+        return isinstance(self.tensor, numpy.ndarray)

cuquantum/_internal/tensor_ifc_torch.py

@@ -0,0 +1,104 @@
+# Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Interface to seamlessly use Torch tensor objects.
+"""
+
+__all__ = ['TorchTensor']
+
+import torch
+
+from .package_ifc import StreamHolder
+from .tensor_ifc import Tensor
+
+
+class TorchTensor(Tensor):
+    """
+    Tensor wrapper for Torch Tensors.
+    """
+    name = 'torch'
+    module = torch
+    name_to_dtype = Tensor.create_name_dtype_map(conversion_function=lambda name: getattr(torch, name), exception_type=AttributeError)
+
+    def __init__(self, tensor):
+        super().__init__(tensor)
+
+    @property
+    def data_ptr(self):
+        return self.tensor.data_ptr()
+
+    @property
+    def device(self):
+        return str(self.tensor.device).split(':')[0]
+
+    @property
+    def device_id(self):
+        return self.tensor.device.index
+
+    @property
+    def dtype(self):
+        """Name of the data type"""
+        return str(self.tensor.dtype).split('.')[-1]
+
+    @property
+    def shape(self):
+        return tuple(self.tensor.shape)
+
+    @property
+    def strides(self):
+        return self.tensor.stride()
+    
+    @property
+    def T(self):
+        return self.__class__(self.tensor.permute(*torch.arange(self.tensor.ndim - 1, -1, -1)))
+
+    def numpy(self, stream_holder=StreamHolder()):
+        # We currently do not use this.
+        raise NotImplementedError
+
+    @classmethod
+    def empty(cls, shape, **context):
+        """
+        Create an empty tensor of the specified shape and data type on the specified device (None, 'cpu', or device id).
+        """
+        name = context.get('dtype', 'float32')
+        dtype = TorchTensor.name_to_dtype[name]
+        device = context.get('device', None)
+        strides = context.get('strides', None)
+        if strides:
+            # note: torch strides is not scaled by bytes
+            tensor = torch.empty_strided(shape, strides, dtype=dtype, device=device)
+        else:
+            tensor = torch.empty(shape, dtype=dtype, device=device)
+
+        return tensor
+
+    def to(self, device='cpu', stream_holder=StreamHolder()):
+        """
+        Create a copy of the tensor on the specified device (integer or 
+          'cpu'). Copy to  Numpy ndarray if CPU, otherwise return Cupy type.
+        """
+        if not (device == 'cpu' or isinstance(device, int)):
+            raise ValueError(f"The device must be specified as an integer or 'cpu', not '{device}'.")
+
+        non_blocking = False if device == 'cpu' else True
+
+        with stream_holder.ctx:
+            tensor_device = self.tensor.to(device=device, non_blocking=non_blocking)
+
+        return tensor_device
+
+    def copy_(self, src, stream_holder=StreamHolder()):
+        """
+        Inplace copy of src (copy the data from src into self).
+        """
+        with stream_holder.ctx:
+            self.tensor.copy_(src)
+
+    def istensor(self):
+        """
+        Check if the object is ndarray-like.
+        """
+        return isinstance(self.tensor, torch.Tensor)