"""
CuPy backend for GPU computations.
This module implements the Backend interface using CuPy for
GPU-accelerated array operations with CUDA.
"""
from typing import Any, Optional, Union
import numpy as np
from .base import Backend
try:
import cupy as cp
CUPY_AVAILABLE = True
except ImportError:
CUPY_AVAILABLE = False
cp = None
[docs]
class CuPyBackend(Backend):
"""
CuPy-based backend for GPU computations.
This backend uses CuPy for GPU-accelerated array operations.
Requires CUDA and CuPy to be installed.
Parameters
----------
device_id : int, optional
CUDA device ID to use. Default is 0.
"""
[docs]
def __init__(self, device_id: int = 0):
if not CUPY_AVAILABLE:
raise RuntimeError(
"CuPy is not available. Install with: pip install cupy-cuda12x"
)
self.device_id = device_id
self.device = cp.cuda.Device(device_id)
self.device.use()
@property
def name(self) -> str:
return "cupy"
@property
def is_gpu(self) -> bool:
return True
[docs]
def zeros(self, shape: tuple[int, ...], dtype: Any = None) -> Any:
with self.device:
return cp.zeros(shape, dtype=dtype or cp.float64)
[docs]
def ones(self, shape: tuple[int, ...], dtype: Any = None) -> Any:
with self.device:
return cp.ones(shape, dtype=dtype or cp.float64)
[docs]
def empty(self, shape: tuple[int, ...], dtype: Any = None) -> Any:
with self.device:
return cp.empty(shape, dtype=dtype or cp.float64)
[docs]
def array(self, data: Any, dtype: Any = None) -> Any:
with self.device:
return cp.array(data, dtype=dtype)
[docs]
def asarray(self, data: Any, dtype: Any = None) -> Any:
with self.device:
return cp.asarray(data, dtype=dtype)
[docs]
def to_numpy(self, array: Any) -> np.ndarray:
"""Convert CuPy array to NumPy array (transfers from GPU to CPU)."""
return cp.asnumpy(array)
[docs]
def copy(self, array: Any) -> Any:
return cp.copy(array)
# Mathematical operations
[docs]
def sqrt(self, array: Any) -> Any:
return cp.sqrt(array)
[docs]
def exp(self, array: Any) -> Any:
return cp.exp(array)
[docs]
def sin(self, array: Any) -> Any:
return cp.sin(array)
[docs]
def cos(self, array: Any) -> Any:
return cp.cos(array)
[docs]
def abs(self, array: Any) -> Any:
return cp.abs(array)
[docs]
def sum(
self, array: Any, axis: Optional[Union[int, tuple[int, ...]]] = None
) -> Any:
return cp.sum(array, axis=axis)
[docs]
def max(
self, array: Any, axis: Optional[Union[int, tuple[int, ...]]] = None
) -> Any:
return cp.max(array, axis=axis)
[docs]
def min(
self, array: Any, axis: Optional[Union[int, tuple[int, ...]]] = None
) -> Any:
return cp.min(array, axis=axis)
[docs]
def mean(
self, array: Any, axis: Optional[Union[int, tuple[int, ...]]] = None
) -> Any:
return cp.mean(array, axis=axis)
# FFT operations
[docs]
def fft(self, array: Any, axis: int = -1) -> Any:
return cp.fft.fft(array, axis=axis)
[docs]
def ifft(self, array: Any, axis: int = -1) -> Any:
return cp.fft.ifft(array, axis=axis)
[docs]
def fft2(self, array: Any) -> Any:
return cp.fft.fft2(array)
[docs]
def ifft2(self, array: Any) -> Any:
return cp.fft.ifft2(array)
# Linear algebra
[docs]
def dot(self, a: Any, b: Any) -> Any:
return cp.dot(a, b)
[docs]
def matmul(self, a: Any, b: Any) -> Any:
return cp.matmul(a, b)
# Indexing and slicing helpers
[docs]
def where(self, condition: Any, x: Any, y: Any) -> Any:
return cp.where(condition, x, y)
# Memory management
[docs]
def synchronize(self) -> None:
"""Synchronize CUDA device."""
cp.cuda.Stream.null.synchronize()
[docs]
def get_memory_info(self) -> dict:
"""Get GPU memory usage information."""
mempool = cp.get_default_memory_pool()
# Try to get device name, fallback to device ID if not available
try:
device_name = self.device.name
except AttributeError:
# CuPy 12.x doesn't have device.name, use device ID instead
device_name = f"GPU:{self.device_id}"
return {
"backend": "cupy",
"device": f"GPU:{self.device_id}",
"device_name": device_name,
"used_bytes": mempool.used_bytes(),
"used_mb": mempool.used_bytes() / (1024**2),
"total_bytes": mempool.total_bytes(),
"total_mb": mempool.total_bytes() / (1024**2),
"free_bytes": self.device.mem_info[0],
"free_mb": self.device.mem_info[0] / (1024**2),
"total_device_bytes": self.device.mem_info[1],
"total_device_mb": self.device.mem_info[1] / (1024**2),
}
# Type information
@property
def float32(self) -> Any:
return cp.float32
@property
def float64(self) -> Any:
return cp.float64
@property
def complex64(self) -> Any:
return cp.complex64
@property
def complex128(self) -> Any:
return cp.complex128
@property
def int32(self) -> Any:
return cp.int32
@property
def int64(self) -> Any:
return cp.int64
@property
def pi(self) -> float:
return float(cp.pi)
[docs]
def __repr__(self) -> str:
"""String representation."""
try:
device_name = self.device.name
except AttributeError:
device_name = f"GPU:{self.device_id}"
return f"CuPyBackend(device={self.device_id}, name='{device_name}')"