您的下一个自定义 FFT 内核#
对于实际应用场景，我们可能需要不止一个内核。一个旨在在多种架构上获得最佳性能的用例，可能需要多种不同的实现。cuFFTDx 的设计旨在自动处理这种负担，同时为用户提供对实现细节的完全控制。
使用最佳参数#
cuFFTDx 允许用户将实现某些细节（例如每个线程计算的 FFT 元素数量或每个块的 FFT 数量）的定义推迟到库。 ElementsPerThread 的默认值是建议值，它基于取决于 FFT 本身（大小、精度、GPU 架构等）的启发式方法。但是，对于选定的 ElementsPerThread（默认值或手动设置）建议的 FFTsPerBlock 可以通过 FFT::suggested_ffts_per_block 静态字段获得。 FFTsPerBlock 的默认值为 1。
#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(cudaStream_t stream = 0) {
// Base of the FFT description
using FFT_base = decltype(Size<128>() + Precision<float>() + Type<fft_type::c2c>()
+ Direction<fft_direction::forward>()
/* Notice lack of ElementsPerThread and FFTsPerBlock operators */
+ SM<700>() + Block());
// FFT description with suggested FFTs per CUDA block for the default (optimal) elements per thread
using FFT = decltype(FFT_base() + FFTsPerBlock<FFT_base::suggested_ffts_per_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, stream);
// Invokes kernel with FFT::block_dim threads in CUDA block
block_fft_kernel<FFT><<<1, FFT::block_dim, FFT::shared_memory_size, stream>>>(data, workspace);
cudaDeviceSynchronize();
cudaFree(data);
}
要检索最佳 FFTsPerBlock 值，我们需要一个完整的执行描述符（如 cufftdx::is_complete_fft_execution 所示）。这是因为某些细节仅在完全描述 FFT 操作并选择目标架构和执行模式后才可用。在主机上编译的 SM Operator 允许用户查询特定架构的启动参数。
底层发生了什么#
- 表达式模板
cuFFTDx API 正在使用一种 C++ 技术的变体，称为表达式模板。我们使用表达式模板允许用户构建编译时对象，这些对象描述要计算的 FFT 计算。编译时 C++ 机制允许 cuFFTDx 将优化的 FFT 例程附加到对象，并将它们公开为可供用户调用的计算方法。
- 仅标头
cuFFTDx FFT 例程以优化的内联 PTX 形式发布。
为什么？#
为了使库有用，它需要以面向未来的方式抽象功能。我们所说的面向未来是指现有用户代码将来不需要修改，并且新功能应该由对现有代码的简单扩展组成。在 CUDA 平台上，这需要适应快速发展的 GPU 硬件。
cuFFTDx 通过两种方式实现面向未来。一方面，API 是一种源代码级别的抽象，它将库与 ABI 更改解耦。与标头中的 PTX 代码一起，cuFFTDx 向前兼容支持 cuFFDx 发布时所支持硬件的任何 CUDA 工具包、驱动程序和编译器。 PTX 可以由 CUDA 编译器重新编译以在未来的 GPU 架构上运行。
另一方面，API 组织允许保留描述计算内容和方式的运算符。取决于类型的新功能可以自动拾取（如果代码将实现选择推迟到库），或者需要向现有表达式添加运算符。