Fast Fourier Transform

原始题目：LeetGPU - Fast Fourier Transform

题目描述

在 GPU 上实现一维复信号的快速傅里叶变换（FFT）。给定包含 $N$ 个复数的输入信号数组（以交织实部/虚部对存储），计算离散傅里叶变换：

$$X_k = \sum_{n=0}^{N-1} x_n \cdot e^{-j 2\pi k n / N}, \quad k = 0, 1, \dots, N-1$$

FFT 算法利用旋转因子的对称性将复杂度从 $O(N^2)$ 降到 $O(N \log N)$。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。

约束条件

• $1 \le N \le 1{,}000{,}000$，$N$ 为 2 的幂。
• 信号值为 32 位浮点（实部和虚部）。

解题思路

GPU 上高效的 FFT 通常使用 Cooley-Tukey 算法。对于 2 的幂大小，可以分解为 $\log_2 N$ 级蝶形运算。每个线程处理一个蝶形对，级间需要同步。共享内存可用于缓存当前级的数据。Stockham 自排序 FFT 避免了 bit-reversal 重排的额外开销。实际应用中通常使用 cuFFT，但手写实现是理解 GPU 上频域计算的好练习。

代码实现

CUDA

#include <cuda_runtime.h>
#include <math.h>
__global__ void fft_kernel(float* re, float* im, int N, int step) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    int half = N / (2 * step);
    if(tid < half) {
        int i = tid * 2 * step + (blockIdx.x % step);
        int j = i + step;
        float wr = cosf(-M_PI * (tid * step) / (half * step));
        float wi = sinf(-M_PI * (tid * step) / (half * step));
        float tr = re[j] * wr - im[j] * wi;
        float ti = re[j] * wi + im[j] * wr;
        re[j] = re[i] - tr; im[j] = im[i] - ti;
        re[i] = re[i] + tr; im[i] = im[i] + ti;
    }
}
extern "C" void solve(float* re, float* im, int N) {
    for(int step = 1; step < N; step <<= 1) {
        int half = N / (2 * step);
        fft_kernel<<<(half+255)/256, 256>>>(re, im, N, step);
        cudaDeviceSynchronize();
    }
}

Triton

import torch
def solve(signal):
    return torch.fft.fft(signal)