ykcheungab
153 lines
3.8 KiB
Python
153 lines
3.8 KiB
Python
"""
|
|
Centralized backend availability checking for GPU acceleration libraries.
|
|
|
|
This module provides a single source of truth for checking whether CuPy and OpenCL
|
|
are available, avoiding duplicate import checks across the codebase.
|
|
"""
|
|
|
|
# CuPy availability and related objects
|
|
_cupy_available = False
|
|
_cupy_cp = None
|
|
_cupy_cp_signal = None
|
|
_cupy_streams_available = False
|
|
_cupy_stream = None
|
|
_cupy_torque_kernel = None
|
|
_cupy_kernel_available = False
|
|
|
|
# OpenCL availability and related objects
|
|
_opencl_available = False
|
|
_opencl_cl = None
|
|
_opencl_cla = None
|
|
_opencl_rfftn = None
|
|
_opencl_irfftn = None
|
|
|
|
# Try to import CuPy
|
|
try:
|
|
import cupy as cp
|
|
import cupyx.scipy.signal as cp_signal
|
|
_cupy_cp = cp
|
|
_cupy_cp_signal = cp_signal
|
|
_cupy_available = True
|
|
|
|
# Try to import CuPy streams for parallel execution
|
|
try:
|
|
from cupy.cuda import Stream
|
|
_cupy_stream = Stream
|
|
_cupy_streams_available = True
|
|
except ImportError:
|
|
_cupy_stream = None
|
|
_cupy_streams_available = False
|
|
|
|
# Create fused elementwise kernel for torque computation
|
|
try:
|
|
_cupy_torque_kernel = cp.ElementwiseKernel(
|
|
'float32 r_x, float32 r_y, float32 grid_0, float32 grid_1, '
|
|
'float32 max_0, float32 max_1, float32 min_0, float32 min_1',
|
|
'float32 torque',
|
|
'''
|
|
float cross_p = (grid_0 - max_0) * r_y - (grid_1 - max_1) * r_x;
|
|
float cross_n = (grid_0 - min_0) * r_y - (grid_1 - min_1) * r_x;
|
|
torque = cross_p + cross_n;
|
|
''',
|
|
'torque_kernel'
|
|
)
|
|
_cupy_kernel_available = True
|
|
except Exception:
|
|
_cupy_torque_kernel = None
|
|
_cupy_kernel_available = False
|
|
|
|
except ImportError:
|
|
pass
|
|
|
|
# Try to import OpenCL
|
|
try:
|
|
import pyopencl as cl
|
|
import pyopencl.array as cla
|
|
from pyvkfft.fft import irfftn as cl_irfftn
|
|
from pyvkfft.fft import rfftn as cl_rfftn
|
|
_opencl_cl = cl
|
|
_opencl_cla = cla
|
|
_opencl_rfftn = cl_rfftn
|
|
_opencl_irfftn = cl_irfftn
|
|
_opencl_available = True
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
def check_cupy_available():
|
|
"""Check if CuPy is available."""
|
|
return _cupy_available
|
|
|
|
|
|
def check_opencl_available():
|
|
"""Check if OpenCL/pyvkfft is available."""
|
|
return _opencl_available
|
|
|
|
|
|
def get_cupy():
|
|
"""Get CuPy module if available, None otherwise."""
|
|
return _cupy_cp
|
|
|
|
|
|
def get_cupy_signal():
|
|
"""Get cupyx.scipy.signal module if available, None otherwise."""
|
|
return _cupy_cp_signal
|
|
|
|
|
|
def get_cupy_stream():
|
|
"""Get CuPy Stream class if available, None otherwise."""
|
|
return _cupy_stream
|
|
|
|
|
|
def get_cupy_torque_kernel():
|
|
"""Get CuPy torque kernel if available, None otherwise."""
|
|
return _cupy_torque_kernel
|
|
|
|
|
|
def is_cupy_streams_available():
|
|
"""Check if CuPy streams are available."""
|
|
return _cupy_streams_available
|
|
|
|
|
|
def is_cupy_kernel_available():
|
|
"""Check if CuPy torque kernel is available."""
|
|
return _cupy_kernel_available
|
|
|
|
|
|
def get_opencl():
|
|
"""Get pyopencl module if available, None otherwise."""
|
|
return _opencl_cl
|
|
|
|
|
|
def get_opencl_array():
|
|
"""Get pyopencl.array module if available, None otherwise."""
|
|
return _opencl_cla
|
|
|
|
|
|
def get_opencl_rfftn():
|
|
"""Get pyvkfft.fft.rfftn function if available, None otherwise."""
|
|
return _opencl_rfftn
|
|
|
|
|
|
def get_opencl_irfftn():
|
|
"""Get pyvkfft.fft.irfftn function if available, None otherwise."""
|
|
return _opencl_irfftn
|
|
|
|
|
|
# For backward compatibility, export as module-level constants
|
|
CUPY_AVAILABLE = _cupy_available
|
|
OPENCL_AVAILABLE = _opencl_available
|
|
STREAMS_AVAILABLE = _cupy_streams_available
|
|
KERNEL_AVAILABLE = _cupy_kernel_available
|
|
|
|
# Export commonly used objects
|
|
cp = _cupy_cp
|
|
cp_signal = _cupy_cp_signal
|
|
Stream = _cupy_stream
|
|
_torque_kernel = _cupy_torque_kernel
|
|
cl = _opencl_cl
|
|
cla = _opencl_cla
|
|
cl_rfftn = _opencl_rfftn
|
|
cl_irfftn = _opencl_irfftn
|
|
|