快速傅里叶变换

概述

nvmath-python 中的快速傅里叶变换 (FFT) 模块 nvmath.fft 利用了 NVIDIA cuFFT 库，并提供了一套强大的 API，可以直接从主机调用以高效地执行离散傅里叶变换。提供了无状态函数形式的 API 和有状态类形式的 API，以支持各种 N 维 FFT 操作。这些操作包括正向和反向变换，以及复数到复数 (C2C)、复数到实数 (C2R) 和实数到复数 (R2C) 变换。

• 通过 nvmath.fft.fft() 进行 N 维正向 C2C FFT 变换。

• 通过 nvmath.fft.ifft() 进行 N 维反向 C2C FFT 变换。

• 通过 nvmath.fft.rfft() 进行 N 维正向 R2C FFT 变换。

• 通过 nvmath.fft.irfft() 进行 N 维反向 C2R FFT 变换。

• 通过有状态的 nvmath.fft.FFT 进行所有类型的 N 维 FFT。

此外，nvmath.fft.FFT 类包含实用程序 API，旨在帮助用户缓存 FFT 计划，从而促进跨各种计算任务高效执行重复计算（请参阅 create_key()）。

在 GPU 上执行的 FFT 变换可以与使用 FFT 回调 的其他操作融合。这使得用户可以使用 Python 编写自定义函数进行预处理或后处理，同时利用即时 (JIT) 和链接时优化 (LTO)。

用户还可以选择 CPU 执行 以利用所有可用的计算资源。

注意

API fft() 和相关函数形式的 API 执行 N 维 FFT 操作，类似于 numpy.fft.fftn()。没有特殊的 1 维 (numpy.fft.fft()) 或 2 维 FFT (numpy.fft.fft2()) API。这不仅减少了 API 表面，而且还避免了潜在的错误使用，因为对于 numpy.fft.fft()，批处理维度的数量为 \(N - 1\)，对于 numpy.fft.fft2()，批处理维度的数量为 \(N - 2\)，其中 \(N\) 是操作数维度。

FFT 回调

用户定义的函数可以编译为 LTO-IR 格式，并作为 FFT 操作的尾声或序言提供，从而实现链接时优化和融合。例如，这可以用于实现基于 DFT 的卷积或缩放 FFT 输出。

FFT 模块附带了便捷的辅助函数 nvmath.fft.compile_prolog() 和 nvmath.fft.compile_epilog()，它们将用 Python 编写的函数编译为 LTO-IR 格式。在底层，这些辅助函数依赖于 Numba 作为编译器。编译后的回调可以作为 DeviceCallable 传递给函数式或有状态的 FFT API。或者，用户可以使用他们选择的编译器将回调编译为 LTO-IR 格式，并将它们作为 DeviceCallable 传递给 FFT 调用。

有关序言和尾声函数用法的示例，请参见 FFT 示例目录。

注意

Windows 上目前不支持 FFT 回调。

设置

开始将 cuFFT LTO 与 nvmath 结合使用的最快方法是安装它以及设备 API 依赖项。Pip 用户应运行以下命令

pip install nvmath-python[cu12,dx]

必需的依赖项

对于那些需要手动收集所需依赖项的用户

• cuFFT 11.3 支持 LTO 回调，该版本随 CUDA Toolkit 12.6 Update 2 及更高版本一起发布。

• 使用 cuFFT LTO 回调需要来自同一 CUDA 工具包或更新版本的 nvJitLink（在同一 CUDA 主要版本内，例如版本 12）。

• 使用 nvmath.fft.compile_prolog() 和 nvmath.fft.compile_epilog() 辅助函数编译回调需要 Numba 0.59+ 和来自与 nvJitLink 相同的 CUDA 工具包或更旧版本的 nvcc/nvvm（在同一 CUDA 主要版本内）。这些辅助函数要求目标设备具有 7.0 或更高的计算能力。

有关更多详细信息，请参阅 cuFFT LTO 文档。

较旧的 CTK

想要尝试回调功能但无法将 CUDA Toolkit 升级到 12.6U2 的冒险用户，可以从此处下载并安装较旧的预览版本 cuFFT LTO EA 版本 11.1.3.0，这至少需要 CUDA Toolkit 12.2。使用 LTO EA 时，可能需要设置环境变量，以便 nvmath 选择所需的 cuFFT 版本。用户应调整 LD_PRELOAD 变量，以便使用正确的 cuFFT 共享库

export LD_PRELOAD="/path_to_cufft_lto_ea/libcufft.so"

执行空间

FFT 变换可以在 NVIDIA GPU 或 CPU 上执行。默认情况下，执行空间是根据传递给 FFT 调用的操作数的内存空间选择的，但可以使用 ExecutionCUDA 和 ExecutionCPU 作为 execution 选项传递给调用（例如 FFT 或 fft()）来显式控制。

注意

Windows 上目前不支持 CPU 执行。

必需的依赖项

对于 ARM CPU（例如 NVIDIA Grace），nvmath-python 可以利用 NVPL (Nvidia Performance Libraries) FFT 来运行变换。在 x86_64 架构上，可以使用 MKL 库。

对于 pip 用户，获取所需依赖项的最快方法是使用 'cu12' / 'cu11' 和 'cpu' 扩展

# for CPU-only dependencies
pip install nvmath-python[cpu]

# for CUDA-only dependencies (assuming CUDA 12)
pip install nvmath-python[cu12]

# for CUDA 12 and CPU dependencies
pip install nvmath-python[cu12,cpu]

自定义 CPU 库

其他符合 FFTW3 API 并在单个 so 文件中提供单精度和双精度符号的库，可用于支持 CPU FFT 执行。想要对 CPU FFT 使用不同库，或指向 NVPL 或 MKL 库的自定义安装的用户，可以通过在 LD_LIBRARY_PATH 中包含库路径并使用 NVMATH_FFT_CPU_LIBRARY 指定库名称来完成。例如

# nvpl
export LD_LIBRARY_PATH=/path/to/nvpl/:$LD_LIBRARY_PATH
export NVMATH_FFT_CPU_LIBRARY=libnvpl_fftw.so.0

# mkl
export LD_LIBRARY_PATH=/path/to/mkl/:$LD_LIBRARY_PATH
export NVMATH_FFT_CPU_LIBRARY=libmkl_rt.so.2

主机 API 参考

FFT 支持 (nvmath.fft)

fft(x, *[, axes, direction, options, ...])	fft(operand, axes=None, direction=None, options=None, execution=None, prolog=None, epilog=None, stream=None)
ifft(x, *[, axes, direction, options, ...])	ifft(operand, axes=None, options=None, execution=None, prolog=None, epilog=None, stream=None)
rfft(operand, *[, axes, options, execution, ...])	rfft(operand, axes=None, options=None, execution=None, prolog=None, epilog=None, stream=None)
irfft(x, *[, axes, options, execution, ...])	irfft(operand, axes=None, options=None, execution=None, prolog=None, epilog=None, stream=None)
FFT(operand, *[, axes, options, execution, ...])	创建一个状态对象，该对象封装指定的 FFT 计算和所需资源。
compile_prolog(prolog_fn, element_dtype, ...)	将 Python 函数编译为 LTO-IR，以作为 fft() 和 plan() 的序言函数提供。
compile_epilog(epilog_fn, element_dtype, ...)	将 Python 函数编译为 LTO-IR，以作为 fft() 和 plan() 的尾声函数提供。
UnsupportedLayoutError(message, permutation, ...)	库不支持的布局的错误类型。
FFTOptions([fft_type, inplace, ...])	用于为 FFT 对象和包装函数系列 fft()、ifft()、rfft() 和 irfft() 提供选项的数据类。
FFTDirection(value[, names, module, ...])	一个 IntEnum 类，用于指定变换的方向。
ExecutionCUDA([device_id])	用于为 FFT 对象和包装函数系列 fft()、ifft()、rfft() 和 irfft() 提供 GPU 执行选项的数据类。
ExecutionCPU([num_threads])	用于为 FFT 对象和包装函数系列 fft()、ifft()、rfft() 和 irfft() 提供 CPU 执行选项的数据类。
DeviceCallable([ltoir, size, data])	一个捕获 LTO-IR 可调用对象的数据类。