FFT#

class nvmath.fft.FFT(operand, *, axes=None, options=None, execution=None, stream=None)[source]#

创建一个有状态的对象，封装指定的 FFT 计算和所需的资源。此对象确保资源在使用期间的有效性，并在不再需要时释放它们，以防止误用。

此对象包含函数形式 API fft(), ifft(), rfft(), 和 irfft() 的所有功能，这些 API 是围绕它的便捷包装器。当需要对具有相同问题规范的多个操作数执行相同的 FFT 操作时，有状态对象还允许分摊准备成本（有关更多详细信息，请参阅 reset_operand() 和 create_key()）。

使用有状态对象通常涉及以下步骤

  1. 问题规范：使用定义的操作和选项初始化对象。

  2. 准备：使用 plan() 确定此特定 FFT 操作的最佳算法实现。

  3. 执行：使用 execute() 执行 FFT 计算，可以是正向或逆向 FFT 变换。

  4. 资源管理：通过显式调用 free() 或在上下文管理器中管理有状态对象，确保所有资源都已释放。

关于上述每个步骤的详细信息，可以通过将 logging.Logger 对象传递到 FFTOptions 或在默认使用的根 logger 对象中设置适当的选项来获得

>>> import logging
>>> logging.basicConfig(
...     level=logging.INFO,
...     format="%(asctime)s %(levelname)-8s %(message)s",
...     datefmt="%m-%d %H:%M:%S",
... )
参数:

  • operand – 张量（类似 ndarray 的对象）。当前支持的类型为 numpy.ndarray, cupy.ndarray, 和 torch.Tensor。

  • axes – 执行 FFT 的维度。 axes[-1] 是 rfft 的“最后变换”轴。目前，要求轴是连续的，并且包含第一个或最后一个维度。仅支持最多 3D FFT。

  • options – 指定 FFT 的选项，可以是一个 FFTOptions 对象。或者，也可以提供一个 dict，其中包含 FFTOptions 构造函数的参数。如果未指定，则该值将设置为默认构造的 FFTOptions 对象。

  • execution – 指定 FFT 的执行空间选项，可以是一个 ExecutionCUDA 或 ExecutionCPU 对象。或者，也可以是一个字符串（‘cuda’ 或 ‘cpu’），或者一个 dict，其中 ‘name’ 键设置为 ‘cpu’ 或 ‘cuda’，以及与给定执行空间相关的可选参数。如果未指定，将选择执行空间以匹配操作数的存储位置（在 GPU 或主机内存中），并且将默认构造相应的 ExecutionCUDA 或 ExecutionCPU 对象。

  • stream – 提供用于执行操作的 CUDA 流。可接受的输入包括 cudaStream_t (作为 Python int), cupy.cuda.Stream, 和 torch.cuda.Stream。如果未提供流，将使用来自操作数包的当前流。

参见

plan(), reset_operand(), execute(), create_key()

示例

>>> import cupy as cp
>>> import nvmath

在 GPU 上创建 3 维 complex128 ndarray

>>> shape = 128, 128, 128
>>> a = cp.random.rand(*shape) + 1j * cp.random.rand(*shape)

我们将在前两个维度上定义 2 维 C2C FFT 操作，并在最后一个维度上进行批处理

>>> axes = 0, 1

创建一个 FFT 对象，封装上述问题规范

>>> f = nvmath.fft.FFT(a, axes=axes)

上面可以提供选项，使用 options 参数来控制操作的行为（参见 FFTOptions）。 同样，可以使用 execution 参数传递执行空间（CUDA 或 CPU）和执行选项（参见 ExecutionCUDA, ExecutionCPU）。

接下来，规划 FFT。可以使用 prolog 和 epilog 选项将加载和/或存储回调函数提供给 plan()

>>> f.plan()

现在执行 FFT，并获得结果 r1 作为 CuPy ndarray。变换将在 GPU 上执行，因为 execution 未明确指定，且 a 驻留在 GPU 内存中。

>>> r1 = f.execute()

最后，释放 FFT 对象的资源。为了避免这种显式调用，建议尽可能使用 FFT 对象作为上下文管理器，如下所示。

>>> f.free()

请注意，默认情况下，所有 FFT 方法都在当前流上执行。或者，可以使用 stream 参数在指定的流上运行方法。

现在让我们看看 CPU 上 NumPy ndarray 的相同问题。

在 CPU 上创建 3 维 complex128 NumPy ndarray

>>> import numpy as np
>>> shape = 128, 128, 128
>>> a = np.random.rand(*shape) + 1j * np.random.rand(*shape)

创建一个 FFT 对象，封装前面描述的问题规范，并将其用作上下文管理器。

>>> with nvmath.fft.FFT(a, axes=axes) as f:
...     f.plan()
...
...     # Execute the FFT to get the first result.
...     r1 = f.execute()

对象使用的所有资源都在代码块结束时释放。

操作在 CPU 上执行，因为 a 驻留在主机内存中。如果将 execution 指定为 ‘cuda’，NumPy 数组将被临时复制到设备内存并在 GPU 上进行变换

>>> with nvmath.fft.FFT(a, axes=axes, execution="cuda") as f:
...     f.plan()
...
...     # Execute the FFT to get the first result.
...     r1 = f.execute()

更多示例可以在 nvmath/examples/fft 目录中找到。

注解

  • 当 FFTOptions.fft_type 为 ‘C2R’ 时，输入必须是 Hermitian 对称的，否则结果是未定义的。作为一个具体的例子，如果 C2R FFT 的输入是使用 R2C FFT 生成的，且最后一个轴的大小为奇数，则必须将 FFTOptions.last_axis_size 设置为 odd 才能恢复原始信号。

方法

__init__(
operand,
*,
axes=None,
options=None,
execution=None,
stream=None,
)[source]#
static create_key(
operand,
*,
axes=None,
options=None,
execution=None,
prolog=None,
epilog=None,
)[source]#

基于给定的操作数、轴和 FFT 选项，创建一个键，作为 FFT 问题规范的紧凑表示。请注意，操作数布局、轴和选项的不同组合可能对应于相同的底层问题规范（键）。当不同的输入问题映射到相同的键时，用户可以重用 FFT 对象。

参数:

  • operand – 张量（类似 ndarray 的对象）。当前支持的类型为 numpy.ndarray, cupy.ndarray, 和 torch.Tensor。

  • axes – 执行 FFT 的维度。 axes[-1] 是 rfft 的“最后变换”轴。目前，要求轴是连续的，并且包含第一个或最后一个维度。仅支持最多 3D FFT。

  • options – 指定 FFT 的选项，可以是一个 FFTOptions 对象。或者，也可以提供一个 dict，其中包含 FFTOptions 构造函数的参数。如果未指定，则该值将设置为默认构造的 FFTOptions 对象。

  • execution – 指定 FFT 的执行空间选项，可以是一个 ExecutionCUDA 或 ExecutionCPU 对象。或者，也可以是一个字符串（‘cuda’ 或 ‘cpu’），或者一个 dict，其中 ‘name’ 键设置为 ‘cpu’ 或 ‘cuda’，以及与给定执行空间相关的可选参数。如果未指定，将选择执行空间以匹配操作数的存储位置（在 GPU 或主机内存中），并且将默认构造相应的 ExecutionCUDA 或 ExecutionCPU 对象。

  • prolog – 提供 LTO-IR 格式的可设备调用函数，用作加载回调，类型为 DeviceCallable 的对象。或者，也可以提供一个 dict，其中包含 DeviceCallable 构造函数的参数。默认情况下没有 prolog。目前，回调仅在 CUDA 执行中受支持。

  • epilog – 提供 LTO-IR 格式的可设备调用函数，用作存储回调，类型为 DeviceCallable 的对象。或者，也可以提供一个 dict，其中包含 DeviceCallable 构造函数的参数。默认情况下没有 epilog。目前，回调仅在 CUDA 执行中受支持。

返回值:

表示输入 FFT 问题的元组键。

注解

  • 用户可以利用此方法基于有状态对象 API 创建 fft() 的缓存版本（有关示例实现，请参见 caching.py）。

  • 此键仅供运行时使用，不适用于序列化或在不同的机器上使用。

  • 如果用户使用流序内存池，则用户有责任使用流来扩充此键。

execute(direction=None, stream=None, release_workspace=False)[source]#

执行 FFT 操作。

参数:

  • direction – 指定执行正向还是反向 FFT (FFTDirection 对象，或作为来自 [‘forward’, ‘inverse’] 的字符串，或作为来自 [-1, 1] 的整数，分别表示正向和反向)。

  • stream – 提供用于执行操作的 CUDA 流。可接受的输入包括 cudaStream_t (作为 Python int), cupy.cuda.Stream, 和 torch.cuda.Stream。如果未提供流，将使用来自操作数包的当前流。

  • release_workspace – 值 True 指定有状态对象应在函数返回时将工作区内存释放回包内存池，而值 False 指定对象应保留内存。 如果应用程序在连续调用（相同或不同的）execute() API 之间执行其他消耗大量内存的操作，则可以将此选项设置为 True，但这会因每次调用时从包内存池获取和释放工作区内存而产生少量开销。 默认为 False。

返回值:

变换后的操作数，它保留在与输入操作数相同的设备上，并使用相同的包。变换后操作数的数据类型和形状取决于输入操作数的类型

  • 对于 C2C FFT，数据类型和形状与输入保持相同。

  • 对于 R2C 和 C2R FFT，数据类型和形状都与输入不同。

free()[source]#

释放 FFT 资源。

建议在上下文中使用 FFT 对象，但如果不可能，则必须显式调用此方法，以确保正确清理 FFT 资源（尤其是内部库对象）。

get_input_layout()[source]#

返回一对元组：FFT 输入的形状和步幅。

注意

在某些情况下，FFT 操作需要复制输入张量（例如 C2R cuFFT，或者提供的张量驻留在 CPU 上，但 FFT 在 GPU 上执行）。 如果原始张量的步幅不符合密集的 C 样式布局，则复制张量的步幅可能与用户传递的输入张量不同。

get_key(*, prolog=None, epilog=None)[source]#

获取此对象数据的键，并辅以回调。

参数:

  • prolog – 提供 LTO-IR 格式的可设备调用函数，用作加载回调，类型为 DeviceCallable 的对象。或者，也可以提供一个 dict，其中包含 DeviceCallable 构造函数的参数。默认情况下没有 prolog。目前，回调仅在 CUDA 执行中受支持。

  • epilog – 提供 LTO-IR 格式的可设备调用函数，用作存储回调，类型为 DeviceCallable 的对象。或者，也可以提供一个 dict，其中包含 DeviceCallable 构造函数的参数。默认情况下没有 epilog。目前，回调仅在 CUDA 执行中受支持。

返回值:

表示输入 FFT 问题的元组键。

参见

create_key()

get_output_layout()[source]#

返回一对元组：FFT 输出的形状和步幅。

plan(*, prolog=None, epilog=None, stream=None, direction=None)[source]#

规划 FFT。

参数:

  • prolog – 提供 LTO-IR 格式的可设备调用函数，用作加载回调，类型为 DeviceCallable 的对象。或者，也可以提供一个 dict，其中包含 DeviceCallable 构造函数的参数。默认情况下没有 prolog。目前，回调仅在 CUDA 执行中受支持。

  • epilog – 提供 LTO-IR 格式的可设备调用函数，用作存储回调，类型为 DeviceCallable 的对象。或者，也可以提供一个 dict，其中包含 DeviceCallable 构造函数的参数。默认情况下没有 epilog。目前，回调仅在 CUDA 执行中受支持。

  • stream – 提供用于执行操作的 CUDA 流。可接受的输入包括 cudaStream_t (作为 Python int), cupy.cuda.Stream, 和 torch.cuda.Stream。如果未提供流，将使用来自操作数包的当前流。

  • direction – 如果指定，则必须将相同的方向传递给后续的 execute() 调用。 它可以作为优化 CPU FFT 调用的 C2C 规划的提示。

reset_operand(operand=None, *, stream=None)[source]#
重置此 FFT 实例持有的操作数。 此方法有两种用例

  1. 它可以用于为执行提供新的操作数

  2. 它可以用于释放对先前操作数的内部引用，并通过传递 operand=None 使其内存可用于其他用途。

参数:

  • operand –

    与前一个操作数兼容的张量（类似 ndarray 的对象）或 None (默认值)。 值 None 将释放对先前操作数的内部引用，并且用户应在再次调用 execute() 之前设置新的操作数。 如果新操作数的所有以下属性与前一个操作数匹配，则认为新操作数是兼容的：

    • 新操作数的问题规范键。 通常，如果操作数共享相同的布局（形状、步幅和数据类型），则键将匹配。 对于某些布局不同的操作数，键可能仍然匹配，有关详细信息，请参见 create_key()。

    • 新操作数所属的包。

    • 新操作数的内存空间（CPU 或 GPU）。

    • 新操作数所属的设备（如果在 GPU 上）。

  • stream – 提供用于执行操作的 CUDA 流。 可接受的输入包括 cudaStream_t (作为 Python int), cupy.cuda.Stream, 和 torch.cuda.Stream。 如果未提供流，将使用来自操作数包的当前流。

示例

>>> import cupy as cp
>>> import nvmath

在 GPU 上创建 3 维 complex128 ndarray

>>> shape = 128, 128, 128
>>> a = cp.random.rand(*shape) + 1j * cp.random.rand(*shape)

将 FFT 对象创建为上下文管理器

>>> axes = 0, 1
>>> with nvmath.fft.FFT(a, axes=axes) as f:
...     # Plan the FFT
...     f.plan()
...
...     # Execute the FFT to get the first result.
...     r1 = f.execute()
...
...     # Reset the operand to a new CuPy ndarray.
...     b = cp.random.rand(*shape) + 1j * cp.random.rand(*shape)
...     f.reset_operand(b)
...
...     # Execute to get the new result corresponding to the updated operand.
...     r2 = f.execute()

使用 reset_operand()，可以实现最小的开销，因为问题规范和规划仅执行一次。

对于上面的特定示例，显式调用 reset_operand() 等效于就地更新操作数，即，将 f.reset_operand(b) 替换为 a[:]=b。 请注意，应谨慎采用就地更新操作数，因为它只能在以下附加约束条件下产生预期的结果并且不会产生额外的副本

  • 操作不是复数到实数 (C2R) FFT。

  • 操作数的内存与 FFT 执行空间匹配。 更准确地说，操作数内存空间应可从执行空间（CPU 或 CUDA）访问。

有关更多详细信息，请参阅 inplace update example。