Source code for scikits.cuda.cufft

#!/usr/bin/env python

"""
Python interface to CUFFT functions.

Note: this module does not explicitly depend on PyCUDA.
"""

import ctypes, platform, sys

# Load library:
_version_list = [6.5, 6.0, 5.5, 5.0, 4.0]
if 'linux' in sys.platform:
    _libcufft_libname_list = ['libcufft.so'] + \
                             ['libcufft.so.%s' % v for v in _version_list]
elif sys.platform == 'darwin':
    _libcufft_libname_list = ['libcufft.dylib']
elif sys.platform == 'win32':
    if platform.machine().endswith('64'):        
        _libcufft_libname_list = ['cufft.dll'] + \
                                 ['cufft64_%s.dll' % int(10*v) for v in _version_list]
    else:
        _libcufft_libname_list = ['cufft.dll'] + \
                                 ['cufft32_%s.dll' % int(10*v) for v in _version_list]

else:
    raise RuntimeError('unsupported platform')

# Print understandable error message when library cannot be found:
_libcufft = None
for _libcufft_libname in _libcufft_libname_list:
    try:
        if sys.platform == 'win32':
            _libcufft = ctypes.windll.LoadLibrary(_libcufft_libname)
        else:
            _libcufft = ctypes.cdll.LoadLibrary(_libcufft_libname)
    except OSError:
        pass
    else:
        break
if _libcufft == None:
    raise OSError('cufft library not found')

# General CUFFT error:
class cufftError(Exception):
    """CUFFT error"""
    pass

# Exceptions corresponding to different CUFFT errors:
class cufftInvalidPlan(cufftError):
    """CUFFT was passed an invalid plan handle."""
    pass

class cufftAllocFailed(cufftError):
    """CUFFT failed to allocate GPU memory."""
    pass

class cufftInvalidType(cufftError):
    """The user requested an unsupported type."""
    pass

class cufftInvalidValue(cufftError):
    """The user specified a bad memory pointer."""
    pass

class cufftInternalError(cufftError):
    """Internal driver error."""
    pass

class cufftExecFailed(cufftError):
    """CUFFT failed to execute an FFT on the GPU."""
    pass

class cufftSetupFailed(cufftError):
    """The CUFFT library failed to initialize."""
    pass

class cufftInvalidSize(cufftError):
    """The user specified an unsupported FFT size."""
    pass

class cufftUnalignedData(cufftError):
    """Input or output does not satisfy texture alignment requirements."""
    pass

cufftExceptions = {
    0x1: cufftInvalidPlan,
    0x2: cufftAllocFailed,
    0x3: cufftInvalidType,
    0x4: cufftInvalidValue,
    0x5: cufftInternalError,
    0x6: cufftExecFailed,
    0x7: cufftSetupFailed,
    0x8: cufftInvalidSize,
    0x9: cufftUnalignedData
    }


class _types:
    """Some alias types."""
    plan = ctypes.c_int
    stream = ctypes.c_void_p

[docs]def cufftCheckStatus(status): """Raise an exception if the specified CUBLAS status is an error.""" if status != 0: try: raise cufftExceptions[status] except KeyError: raise cufftError # Data transformation types:
CUFFT_R2C = 0x2a CUFFT_C2R = 0x2c CUFFT_C2C = 0x29 CUFFT_D2Z = 0x6a CUFFT_Z2D = 0x6c CUFFT_Z2Z = 0x69 # Transformation directions: CUFFT_FORWARD = -1 CUFFT_INVERSE = 1 # FFTW compatibility modes: CUFFT_COMPATIBILITY_NATIVE = 0x00 CUFFT_COMPATIBILITY_FFTW_PADDING = 0x01 CUFFT_COMPATIBILITY_FFTW_ASYMMETRIC = 0x02 CUFFT_COMPATIBILITY_FFTW_ALL = 0x03 # FFT functions implemented by CUFFT: _libcufft.cufftPlan1d.restype = int _libcufft.cufftPlan1d.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int]
[docs]def cufftPlan1d(nx, fft_type, batch): """Create 1D FFT plan configuration.""" plan = _types.plan() status = _libcufft.cufftPlan1d(ctypes.byref(plan), nx, fft_type, batch) cufftCheckStatus(status) return plan
_libcufft.cufftPlan2d.restype = int _libcufft.cufftPlan2d.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int]
[docs]def cufftPlan2d(nx, ny, fft_type): """Create 2D FFT plan configuration.""" plan = _types.plan() status = _libcufft.cufftPlan2d(ctypes.byref(plan), nx, ny, fft_type) cufftCheckStatus(status) return plan
_libcufft.cufftPlan3d.restype = int _libcufft.cufftPlan3d.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int]
[docs]def cufftPlan3d(nx, ny, nz, fft_type): """Create 3D FFT plan configuration.""" plan = _types.plan() status = _libcufft.cufftPlan3d(ctypes.byref(plan), nx, ny, nz, fft_type) cufftCheckStatus(status) return plan
_libcufft.cufftPlanMany.restype = int _libcufft.cufftPlanMany.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int]
[docs]def cufftPlanMany(rank, n, inembed, istride, idist, onembed, ostride, odist, fft_type, batch): """Create batched FFT plan configuration.""" plan = _types.plan() status = _libcufft.cufftPlanMany(ctypes.byref(plan), rank, n, inembed, istride, idist, onembed, ostride, odist, fft_type, batch) cufftCheckStatus(status) return plan
_libcufft.cufftDestroy.restype = int _libcufft.cufftDestroy.argtypes = [_types.plan]
[docs]def cufftDestroy(plan): """Destroy FFT plan.""" status = _libcufft.cufftDestroy(plan) cufftCheckStatus(status)
_libcufft.cufftSetCompatibilityMode.restype = int _libcufft.cufftSetCompatibilityMode.argtypes = [_types.plan, ctypes.c_int]
[docs]def cufftSetCompatibilityMode(plan, mode): """Set FFTW compatibility mode.""" status = _libcufft.cufftSetCompatibilityMode(plan, mode) cufftCheckStatus(status)
_libcufft.cufftExecC2C.restype = int _libcufft.cufftExecC2C.argtypes = [_types.plan, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int]
[docs]def cufftExecC2C(plan, idata, odata, direction): """Execute single precision complex-to-complex transform plan as specified by `direction`.""" status = _libcufft.cufftExecC2C(plan, idata, odata, direction) cufftCheckStatus(status)
_libcufft.cufftExecR2C.restype = int _libcufft.cufftExecR2C.argtypes = [_types.plan, ctypes.c_void_p, ctypes.c_void_p]
[docs]def cufftExecR2C(plan, idata, odata): """Execute single precision real-to-complex forward transform plan.""" status = _libcufft.cufftExecR2C(plan, idata, odata) cufftCheckStatus(status)
_libcufft.cufftExecC2R.restype = int _libcufft.cufftExecC2R.argtypes = [_types.plan, ctypes.c_void_p, ctypes.c_void_p]
[docs]def cufftExecC2R(plan, idata, odata): """Execute single precision complex-to-real reverse transform plan.""" status = _libcufft.cufftExecC2R(plan, idata, odata) cufftCheckStatus(status)
_libcufft.cufftExecZ2Z.restype = int _libcufft.cufftExecZ2Z.argtypes = [_types.plan, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int]
[docs]def cufftExecZ2Z(plan, idata, odata, direction): """Execute double precision complex-to-complex transform plan as specified by `direction`.""" status = _libcufft.cufftExecZ2Z(plan, idata, odata, direction) cufftCheckStatus(status)
_libcufft.cufftExecD2Z.restype = int _libcufft.cufftExecD2Z.argtypes = [_types.plan, ctypes.c_void_p, ctypes.c_void_p]
[docs]def cufftExecD2Z(plan, idata, odata): """Execute double precision real-to-complex forward transform plan.""" status = _libcufft.cufftExecD2Z(plan, idata, odata) cufftCheckStatus(status)
_libcufft.cufftExecZ2D.restype = int _libcufft.cufftExecZ2D.argtypes = [_types.plan, ctypes.c_void_p, ctypes.c_void_p]
[docs]def cufftExecZ2D(plan, idata, odata): """Execute double precision complex-to-real transform plan.""" status = _libcufft.cufftExecZ2D(plan, idata, odata) cufftCheckStatus(status)
_libcufft.cufftSetStream.restype = int _libcufft.cufftSetStream.argtypes = [_types.plan, _types.stream]
[docs]def cufftSetStream(plan, stream): """Associate a CUDA stream with a CUFFT plan.""" status = _libcufft.cufftSetStream(plan, stream) cufftCheckStatus(status)