首次使用cuFFTDx进行FFT

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

定义基础FFT

第一步是定义我们要执行的FFT。这通过将cuFFTDx算子组合在一起来创建FFT描述来实现。该类型的正确性会在编译时进行评估。一个定义良好的FFT必须包含问题规模、使用的精度(float、double等)、操作类型(complex-to-complex、real-to-complex等)及其方向(正向或逆向)。

// 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提供了Size OperatorPrecision OperatorType OperatorDirection Operator等运算符。如前所述,列出的运算符可以通过使用加法运算符(+)进行组合。

要获得一个完全可用的CUDA FFT内核,我们需要提供三个额外信息。第一个是我们想要计算多少个FFT,第二个是如何将计算映射到CUDA块中,最后一个是我们目标使用的CUDA架构。

注意

如果FFT要在单线程中执行(参见Thread Operator),则不允许使用FFTsPerBlockElementsPerThread运算符。在该模式下,每个线程执行一个FFT。

在cuFFTDx中,我们使用FFTs Per Block Operator来指定要计算的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块,我们使用Elements Per Thread Operator。该运算符决定了每个线程所需的寄存器数量以及要使用的具体实现。它还会影响所需的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>());

注意

如果未设置ElementsPerThreadFFTsPerBlock,则将使用默认值。请参阅Block Configuration Operators部分。

最后,我们使用SM Operator来指定构建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,更多详情请参阅Block Execute Method章节。

某些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返回的错误代码。请查看完整的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进行编译。详情请参阅需求与功能章节。