compile_epilog#

nvmath.fft.compile_epilog(
epilog_fn,
element_dtype,
user_info_dtype,
*,
compute_capability=None,
)[source]#

Compile a Python function to LTO-IR to provide as an epilog function for fft() and plan(). (将 Python 函数编译为 LTO-IR，以作为 fft() 和 plan() 的尾声函数提供。)

Parameters: (参数：)

  • epilog_fn – The epilog function to be compiled to LTO-IR. It must have the signature: epilog_fn(data_out, offset, data, user_info, reserved_for_future_use), and it essentially stores transformed data into data_out at offset. (epilog_fn – 要编译为 LTO-IR 的尾声函数。它必须具有以下签名：epilog_fn(data_out, offset, data, user_info, reserved_for_future_use)，并且它本质上将转换后的 data 存储到 data_out 的 offset 处。)

  • element_dtype – The data type of the data_in argument, one of ['float32', 'float64', 'complex64', 'complex128']. It must have the same data type as that of the FFT operand for prolog functions or the FFT result for epilog functions. (element_dtype – data_in 参数的数据类型，为 ['float32', 'float64', 'complex64', 'complex128'] 之一。它必须具有与前言函数的 FFT 操作数或尾声函数的 FFT 结果相同的数据类型。)

  • user_info_dtype – (user_info_dtype –)

    The data type of the user_info argument. It must be one of ['float32', 'float64', 'complex64', 'complex128'] or an object of type numba.types.Type. The offset is computed based on the memory layout (shape and strides) of the operand (input for prolog, output for epilog). If the user would like to pass additional tensor as user_info and access it based on the offset, it is crucial to know memory layout of the operand. Please note, the actual layout of the input tensor may differ from the layout of the tensor passed to fft call. To learn the memory layout of the input or output, please use stateful FFT API and nvmath.fft.FFT.get_input_layout() nvmath.fft.FFT.get_output_layout() respectively. ( user_info 参数的数据类型。它必须是 ['float32', 'float64', 'complex64', 'complex128'] 之一或 numba.types.Type 类型的对象。偏移量是根据操作数（前言的输入，尾声的输出）的内存布局（形状和步幅）计算的。如果用户想将额外的张量作为 user_info 传递并基于偏移量访问它，则了解操作数的内存布局至关重要。请注意，输入张量的实际布局可能与传递给 fft 调用的张量的布局不同。要了解输入或输出的内存布局，请使用有状态的 FFT API 和 nvmath.fft.FFT.get_input_layout() 和 nvmath.fft.FFT.get_output_layout()。)

    Note (注意)

    目前，在回调函数中，即使原始操作数是多维张量，输入和输出操作数中元素的位置也是用单个扁平偏移量来描述的。

  • compute_capability – 目标计算能力，指定为字符串（'80', '89', …）。 默认值为当前设备的计算能力。

返回:

编译为 LTO-IR 的函数，以 bytes 对象形式返回。

另请参阅

fft(), plan(), compile_prolog()。

示例

cuFFT 库期望最终用户管理输出的缩放，因此为了复现 其他 Python FFT 库 中的 norm 选项，我们可以定义一个执行缩放的 epilog 函数。

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

为批处理的 1-D FFT 创建数据。

>>> B, N = 256, 1024
>>> a = cp.random.rand(B, N, dtype=cp.float64) + 1j * cp.random.rand(B, N, dtype=cp.float64)

计算一个归一化因子，用于创建酉变换。

>>> norm_factor = 1.0 / math.sqrt(N)

为 FFT 定义 epilog 函数。

>>> def rescale(data_out, offset, data, user_info, unused):
...     data_out[offset] = data * norm_factor

将 epilog 编译为 LTO-IR。 在具有不同计算能力的 GPU 系统中，必须为 compile_prolog 或 compile_epilog 辅助函数指定 compute_capability 选项。 或者，可以在提供 epilog 的 FFT 执行设备的上下文中编译 epilog。 在这种情况下，我们使用当前设备上下文，其中已创建操作数。

>>> with cp.cuda.Device():
...     epilog = nvmath.fft.compile_epilog(rescale, "complex128", "complex128")

执行前向 FFT，应用重缩放作为 epilog。

>>> r = nvmath.fft.fft(a, axes=[-1], epilog=dict(ltoir=epilog))

测试融合的 FFT 运行结果是否与其他库的结果匹配。

>>> s = cp.fft.fftn(a, axes=[-1], norm="ortho")
>>> assert cp.allclose(r, s)

注释

  • 用户必须确保指定的参数类型满足上面列出的要求。

本页目录