cupy.fft.config.set_cufft_callbacks#

class cupy.fft.config.set_cufft_callbacks(unicode cb_load=u'', unicode cb_store=u'', ndarray cb_load_aux_arr=None, *, ndarray cb_store_aux_arr=None)[source]#

一个用于设置加载和/或存储回调的上下文管理器。

参数:
  • cb_load (str) – 包含加载回调的设备内核的字符串。它必须定义 d_loadCallbackPtr

  • cb_store (str) – 包含存储回调的设备内核的字符串。它必须定义 d_storeCallbackPtr

  • cb_load_aux_arr (cupy.ndarray, 可选) – 一个 CuPy 数组,包含用于加载回调的数据。

  • cb_store_aux_arr (cupy.ndarray, 可选) – 一个 CuPy 数组,包含用于存储回调的数据。

注意

在此上下文中进行的任何 FFT 调用都将设置回调。下面显示了加载回调的一个示例:

code = r'''
__device__ cufftComplex CB_ConvertInputC(
    void *dataIn,
    size_t offset,
    void *callerInfo,
    void *sharedPtr) {
  // implementation
}

__device__ cufftCallbackLoadC d_loadCallbackPtr = CB_ConvertInputC;
'''

with cp.fft.config.set_cufft_callbacks(cb_load=code):
    out_arr = cp.fft.fft(in_arr, ...)

注意

以下是使用此功能的运行时要求

  • cython >= 0.29.0

  • 支持 C++11 及更高版本的主机编译器;可能需要设置 CXX 环境变量。

  • nvcc 和完整的 CUDA Toolkit。请注意,Conda-Forge 中的 cudatoolkit 包不足够,因为它不包含静态库。

注意

回调仅适用于连续轴上的变换;对于非连续变换的行为通常是未定义的。

警告

使用 cuFFT 回调需要在运行时编译和加载 Python 模块,并对每个不同的变换和回调进行静态链接,因此每次组合的首次调用将非常慢。这是 cuFFT 的一个限制,所以仅当已知启用回调的变换性能更优并且可以重复使用以分摊成本时才使用此功能。

警告

生成的 Python 模块默认缓存到 ~/.cupy/callback_cache 中,以便可能重用(使用相同的加载/存储回调集)。然而,由于静态链接,文件大小可能过大!缓存位置可以通过设置 CUPY_CACHE_DIR 来更改。

另请参阅

cuFFT 回调例程

方法

__enter__(self)#
__exit__(self, exc_type, exc_value, traceback)#
__eq__(value, /)#

返回 self==value。

__ne__(value, /)#

返回 self!=value。

__lt__(value, /)#

返回 self<value。

__le__(value, /)#

返回 self<=value。

__gt__(value, /)#

返回 self>value。

__ge__(value, /)#

返回 self>=value。