compile_prolog#

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

将 Python 函数编译为 LTO-IR，以作为 prolog 函数提供给 fft() 和 plan()。

参数:

  • prolog_fn – 要编译为 LTO-IR 的 prolog 函数。它必须具有以下签名： prolog_fn(data_in, offset, user_info, reserved_for_future_use)，并且它本质上返回在 offset 处转换后的 data_in。

  • element_dtype – data_in 参数的数据类型，为 ['float32', 'float64', 'complex64', 'complex128'] 之一。它必须具有与 prolog 函数的 FFT 操作数或 epilog 函数的 FFT 结果相同的数据类型。

  • user_info_dtype –

    user_info 参数的数据类型。它必须是 ['float32', 'float64', 'complex64', 'complex128'] 之一或 numba.types.Type 类型的对象。偏移量是根据操作数的内存布局（形状和步幅）计算的（prolog 的输入，epilog 的输出）。如果用户希望将额外的张量作为 user_info 传递并根据偏移量访问它，则了解操作数的内存布局至关重要。请注意，输入张量的实际布局可能与传递给 fft 调用的张量的布局不同。要了解输入或输出的内存布局，请使用有状态 FFT API 和 nvmath.fft.FFT.get_input_layout() nvmath.fft.FFT.get_output_layout() 分别。

    注意

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

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

返回:

编译为 LTO-IR 的函数，作为 bytes 对象。

另请参阅

fft(), plan(), compile_epilog()。

注意

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

本页内容