使用 cuFFTDx 的第一个 FFT

在本介绍中，我们将使用独立内核计算大小为 128 的 FFT。本节基于 cuFFTDx 附带的 introduction_example.cu 示例。请参阅 示例 部分以查看其他 cuFFTDx 示例。

定义基本 FFT

第一步是定义我们要执行的 FFT。这通过将 cuFFTDx 运算符加在一起来创建 FFT 描述来完成。此类型的正确性在编译时进行评估。一个定义良好的 FFT 必须包括问题大小、使用的精度（float、double 等）、运算类型（复数到复数、实数到复数等）及其方向（正向或反向）。

// cuFFTDx header
#include <cufftdx.hpp>

// FFT description:
// A 128-point single precision complex-to-complex forward FFT description
using FFT = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>() + Direction<fft_direction::forward>());

为了编码 FFT 属性，cuFFTDx 提供了运算符：大小运算符、精度运算符、类型运算符和方向运算符。如前所述，列出的运算符可以通过使用加法运算符 (+) 组合。

为了获得完全可用的 CUDA FFT 内核，我们需要提供三个附加信息。第一个是我们想要计算多少个 FFT，第二个是如何将计算映射到 CUDA 块中，最后一个是我们针对的 CUDA 架构。

注意

如果 FFT 要在单线程中执行（参见线程运算符），则 FFTsPerBlock 和 ElementsPerThread 运算符都不允许使用。在该模式下，每个线程执行一个 FFT。

在 cuFFTDx 中，我们使用每块 FFT 数运算符来指定我们想要计算多少个 FFT。它定义了在单个 CUDA 块内并行执行多少个 FFT。在本例中，我们将其设置为每个 CUDA 块 2 个 FFT（默认值为每个 CUDA 块 1 个 FFT）

// cuFFTDx header
#include <cufftdx.hpp>

// FFT description
using FFT = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>() + Direction<fft_direction::forward>()
                     + FFTsPerBlock<2>());

为了将 FFT 的计算映射到 CUDA 块，我们使用每线程元素数运算符。此运算符确定每个线程所需的寄存器数量和要使用的确切实现。它还会影响所需的 CUDA 块大小。我们将该运算符添加到描述中

#include <cufftdx.hpp>

// FFT description
using FFT = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>() + Direction<fft_direction::forward>()
                     + FFTsPerBlock<2>() + ElementsPerThread<8>());

注意

如果未设置 ElementsPerThread 和 FFTsPerBlock，则使用默认值。请参阅块配置运算符部分。

最后，我们使用 SM 运算符来指示我们要构建 FFT 描述符的目标 CUDA 架构。每个 GPU 架构可以使用不同的参数。因此，架构的选择可能会影响配置以最大化性能。在 introduction_example.cu 示例中，这作为模板参数传递，但在这里我们可以假设我们针对的是 Volta GPU (SM<700>())

#include <cufftdx.hpp>

// FFT description
using FFT = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>() + Direction<fft_direction::forward>()
                     + FFTsPerBlock<2>() + ElementsPerThread<8>()
                     + SM<700>());

一旦 FFT 描述完全形成，我们可以通过添加 块运算符来完成它。它表明我们要求由单个 CUDA 块执行集体 FFT 运算。该运算符验证描述的正确性，并且它是一种执行运算符（另一种是线程运算符）。

#include <cufftdx.hpp>

// FFT description:
// This description says that we want to execute a 128-point single precision complex-to-complex forward FFT with
// 2 batches per CUDA block and with 8 elements per thread on Volta GPU.
using FFT = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>() + Direction<fft_direction::forward>()
                     + FFTsPerBlock<2>() + ElementsPerThread<8>()
                     + SM<700>()
                     + Block());

执行 FFT

FFT 描述类型可以实例化为对象。形成对象没有计算成本，应该被视为一个句柄。FFT 描述符对象提供了一个计算方法 execute(...)，用于执行请求的 FFT。

#include <cufftdx.hpp>
using namespace cufftdx;

// FFT description
using FFT = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>()
                     + Direction<fft_direction::forward>() + FFTsPerBlock<1>()
                     + ElementsPerThread<8>() + SM<700>() + Block());
using complex_type = typename FFT::value_type;

__global__ void block_fft_kernel(complex_type* data) {
  // Execute FFT
  FFT().execute(/* What are the arguments? */);
}

cuFFTDx 运算需要寄存器和共享内存才能运行。用户可以查询 FFT 描述符以获取所需的资源。

#include <cufftdx.hpp>
using namespace cufftdx;

// FFT description
using FFT = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>()
                     + Direction<fft_direction::forward>() + FFTsPerBlock<1>()
                     + ElementsPerThread<8>() + SM<700>() + Block());
using complex_type = typename FFT::value_type;

__global__ void block_fft_kernel(complex_type* data) {
  // Registers
  complex_type thread_data[FFT::storage_size];

  // FFT::shared_memory_size bytes of shared memory
  extern __shared__ __align__(alignof(float4)) complex_type shared_mem[];

  // Execute FFT
  FFT().execute(thread_data, shared_mem);
}

在上面介绍的 execute(...) 方法中，cuFFTDx 要求输入数据位于 thread_data 寄存器中，并将 FFT 结果存储在那里。用户还可以使用仅接受指向共享内存的指针的 API，并假定所有数据都以自然顺序排列在那里，有关更多详细信息，请参阅块执行方法部分。

某些 FFT，根据选择的大小，可能还需要额外的全局内存工作区，这需要在主机上分配并传递给内核。您可以使用 FFT::requires_workspace 特性检查是否必须创建工作区。

#include <cufftdx.hpp>
using namespace cufftdx;

using FFT = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>()
                     + Direction<fft_direction::forward>() + FFTsPerBlock<1>()
                     + ElementsPerThread<8>() + SM<700>() + Block());
using complex_type = typename FFT::value_type;

__global__ void block_fft_kernel(complex_type* data, typename FFT::workspace_type workspace) {
  // Registers
  complex_type thread_data[FFT::storage_size];

  // FFT::shared_memory_size bytes of shared memory
  extern __shared__ __align__(alignof(float4)) complex_type shared_mem[];

  // Execute FFT
  FFT().execute(thread_data, shared_mem, workspace /* additional workspace */);
}

启动 FFT 内核

要启动内核，我们需要知道执行 FFT 运算所需的块大小和共享内存量。两者都是固定的，由 FFT 描述确定。

由于我们在设备代码中定义了 FFT 描述，因此有关块大小的信息需要传播到主机。当所有参数都完全指定时，所有 GPU 架构都使用相同的块大小，因此内核可以以相同的方式在所有架构上启动。

#include <cufftdx.hpp>
using namespace cufftdx;

template<class FFT>
__global__ void block_fft_kernel(FFT::value_type* data, typename FFT::workspace_type workspace) {
  using complex_type = typename FFT::value_type;
  // Registers
  complex_type thread_data[FFT::storage_size];

  // FFT::shared_memory_size bytes of shared memory
  extern __shared__ __align__(alignof(float4)) complex_type shared_mem[];

  // Execute FFT
  FFT().execute(thread_data, shared_mem, workspace);
}

// CUDA_CHECK_AND_EXIT - marco checks if function returns cudaSuccess; if not it prints
// the error code and exits the program
void introduction_example(FFT::value_type* data /* data is a manage memory pointer*/) {
  // FFT description
  using FFT = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>()
                       + Direction<fft_direction::forward>() + FFTsPerBlock<1>()
                       + ElementsPerThread<8>() + SM<700>() + Block());

  cudaError_t error_code = cudaSuccess;
  auto workspace = make_workspace<FFT>(error_code);
  CUDA_CHECK_AND_EXIT(error_code);

  // Invokes kernel with FFT::block_dim threads in CUDA block
  block_fft_kernel<FFT><<<1, FFT::block_dim, FFT::shared_memory_size>>>(data, workspace);
  CUDA_CHECK_AND_EXIT(cudaPeekAtLastError());
  CUDA_CHECK_AND_EXIT(cudaDeviceSynchronize());
}

如果我们还添加从/到全局内存的输入/输出操作，我们将获得一个在功能上等同于 cuFFT 大小为 128 和单精度复数到复数内核的内核。数据从全局内存加载并存储到寄存器中，如输入/输出数据格式部分所述，类似地，结果保存回全局内存。

为了简单起见，在本示例中，我们为设备输入/输出数组分配托管内存，假设使用 Volta 架构，并且不检查 CUDA API 函数和 cufftdx::make_workspace 返回的 CUDA 错误代码。请查看完整的 introduction_example.cu 示例以及 cuFFTDx 附带的其他示例，以获取更详细的代码。

#include <cufftdx.hpp>
using namespace cufftdx;

template<class FFT>
__global__ void block_fft_kernel(FFT::value_type* data, typename FFT::workspace_type workspace) {
  using complex_type = typename FFT::value_type;

  // Registers
  complex_type thread_data[FFT::storage_size];

  // Local batch id of this FFT in CUDA block, in range [0; FFT::ffts_per_block)
  const unsigned int local_fft_id = threadIdx.y;
  // Global batch id of this FFT in CUDA grid is equal to number of batches per CUDA block (ffts_per_block)
  // times CUDA block id, plus local batch id.
  const unsigned int global_fft_id = (blockIdx.x * FFT::ffts_per_block) + local_fft_id;

  // Load data from global memory to registers
  const unsigned int offset = cufftdx::size_of<FFT>::value * global_fft_id;
  const unsigned int stride = FFT::stride;
  unsigned int       index  = offset + threadIdx.x;
  for (unsigned int i = 0; i < FFT::elements_per_thread; i++) {
      // Make sure not to go out-of-bounds
      if ((i * stride + threadIdx.x) < cufftdx::size_of<FFT>::value) {
          thread_data[i] = data[index];
          index += stride;
      }
  }

  // FFT::shared_memory_size bytes of shared memory
  extern __shared__ __align__(alignof(float4)) complex_type shared_mem[];

  // Execute FFT
  FFT().execute(thread_data, shared_mem, workspace);

  // Save results
  index = offset + threadIdx.x;
  for (unsigned int i = 0; i < FFT::elements_per_thread; i++) {
      if ((i * stride + threadIdx.x) < cufftdx::size_of<FFT>::value) {
          data[index] = thread_data[i];
          index += stride;
      }
  }
}

void introduction_example() {
  // FFT description
  using FFT = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>()
                       + Direction<fft_direction::forward>() + FFTsPerBlock<1>()
                       + ElementsPerThread<8>() + SM<700>() + Block());

  // Allocate managed memory for input/output
  complex_type* data;
  auto          size       = FFT::ffts_per_block * cufftdx::size_of<FFT>::value;
  auto          size_bytes = size * sizeof(complex_type);
  cudaMallocManaged(&data, size_bytes);
  // Generate data
  for (size_t i = 0; i < size; i++) {
      data[i] = complex_type {float(i), -float(i)};
  }

  cudaError_t error_code = cudaSuccess;
  auto workspace = make_workspace<FFT>(error_code);

  // Invokes kernel with FFT::block_dim threads in CUDA block
  block_fft_kernel<FFT><<<1, FFT::block_dim, FFT::shared_memory_size>>>(data, workspace);
  cudaDeviceSynchronize();

  cudaFree(data);
}

重要的是要注意，与 cuFFT 不同，cuFFTDx 在执行 FFT 运算后不需要将数据移回全局内存。这可能是一个主要的性能优势，因为 FFT 计算可以与自定义预处理和后处理操作融合在一起。

编译

为了编译包含 cufftdx.hpp 的程序，用户只需传递 cuFFTDx 库的位置（包含 cufftdx.hpp 文件的目录）。有关如何在您的项目中使用 cuFFTDx 的更多信息，请参阅快速安装指南。

nvcc -std=c++17 -arch sm_70 -O3 -I<mathdx_include_dir> introduction_example.cu -o introduction_example

注意

自 0.3.0 版本以来，cuFFTDx 已实验性地支持使用 NVRTC 进行编译。请参阅要求和功能部分。