CUDA实现FFT并行计算

前言

  本文是个人学习心得的分享，希望大家在阅读文章后能在评论中一起学习交流！另外还可以访问我的HelloCUDA仓库查看我在学习CUDA中写的一些demo程序。

内容概要

• 复数的CUDA C++实现
• 从DFT到FFT
• FFT蝴蝶操作
• CUDA中的分治
• FFT的并行化

前置知识

• 算法基础知识
• 并行计算基础知识
• C++基础知识
• CUDA编程基础
• 离散傅里叶变换的原理

从DFT到FFT

  离散傅里叶变化(Discrete Fourier Transform)是傅里叶变换的简化版本，计算机科学家发明出离散傅里叶变换目的就是要让具有连续性质的傅里叶变换能在具有离散性质的计算机中被应用。而离散傅里叶变换算法时间复杂度高($O(n^2)$O(n2))，使其在大规模场景中的应用受到限制，从而诞生了快速傅里叶变换(Fast Fourier Transform).
  快速傅里叶变换的实现采用的是分治的思想。利用DFT的数学性质可以将一个DFT分解为多个DFT。数列$a_0, a_1, a_2, ..., a_n$a0​,a1​,a2​,...,an​的离散傅里叶变换的本质是求多项式$A(x) = a_0 + a_1x + a_2x^2 + ... + a_nx^{n-1}$A(x)=a0​+a1​x+a2​x2+...+an​xn−1在$n$n次单位复根处的取值，变换后的数列为$A(w_n^0), A(w_n^1), A(w_n^2), ..., A(w_n^{n-1})$A(wn0​),A(wn1​),A(wn2​),...,A(wnn−1​)。其中$w_n^i, i=0, 1, ..., n-1$wni​,i=0,1,...,n−1是方程$w^n = 1$wn=1的$n$n个根 (此处有疑惑请参考离散傅里叶变换的原理). 因此，对于长度为$n$n的数列进行DFT共需要进行$n$n次多项式求值，每次求值的复杂度为$O(n)$O(n)，因此总体时间复杂度为$O(n^2)$O(n2).
  了解从DFT到FFT原理的读者可以跳过以下三段内容。将上述多项式$A(x)$A(x)按照下标为奇数或偶数分解为两个多项式$A_0(x)=a_0 + a_2x + a_4x^2 + a_6x^4 + ...$A0​(x)=a0​+a2​x+a4​x2+a6​x4+...和$A_1(x) = a_1 + a_3x + a_5x^2 + a_7x^3 + ...$A1​(x)=a1​+a3​x+a5​x2+a7​x3+.... 可知$A(x)=A_0(x^2) + xA_1(x^2)$A(x)=A0​(x2)+xA1​(x2). 而$A_0$A0​和$A_1$A1​又可继续分解，形成一个递归的分解过程。由折半引理($w_n^2 = w_{n/2}$wn2​=wn/2​)可知，$A(w_n^k) = A_0(w_{n/2}^k) + w_n^kA_1(w_{n/2}^k)$A(wnk​)=A0​(wn/2k​)+wnk​A1​(wn/2k​). 因此，求原数列的DFT问题被分解为偶数下标子序列的DFT和奇数下标子序列的DFT。
  以$n=8$n=8为例，将$a_0$a0​到$a_7$a7​分解为两个子数列$a_0, a_2, a_4, a_6$a0​,a2​,a4​,a6​和$a_1, a_3, a_5, a_7$a1​,a3​,a5​,a7​，则满足$A(w_8^k) = A_0(w_4^k) + w_8^kA_1(w_4^k)$A(w8k​)=A0​(w4k​)+w8k​A1​(w4k​)，两个子数列的DFT分别为$A_0(w_4^0), A_0(w_4^1), A_0(w_4^2), A_0(w_4^3)$A0​(w40​),A0​(w41​),A0​(w42​),A0​(w43​)和$A_1(w_4^0), A_1(w_4^1), A_1(w_4^2), A_1(w_4^3)$A1​(w40​),A1​(w41​),A1​(w42​),A1​(w43​). 原数列DFT的第1项为：
$A(w_8^0)=A_0(w_4^0) + w_8^0A_1(w_4^0)$A(w80​)=A0​(w40​)+w80​A1​(w40​)
即奇数下标子数列的第1项 + $w_8^0$w80​ * 偶数下标子数列的第1项。值得注意的是，与子数列第1项有关的不只有原数列DFT的第1项，还有第5项。特别注意下面出现的这个等式，原数列的DFT的第5项为（注意$k$k是从0开始的，这里$k=4$k=4）：
$A(w_8^4)=A_0(w_4^4) + w_8^4A_1(w_4^4)=A_0(w_4^0)-w_8^0A_1(w_4^0)$A(w84​)=A0​(w44​)+w84​A1​(w44​)=A0​(w40​)−w80​A1​(w40​)
即奇数下标子数列的第1项 - $w_8^0$w80​ * 偶数下标子数列的第1项，就是把前面的加变成了减。同理不难得出原数列DFT的第2项和第6项、第3项和第7项、第4项和第8项都满足这个规则。
  上述过程是建立在子数列的DFT已经求出来的假设下进行分析的，实际上还需要求解子数列的DFT才能进行原数列DFT的求解，这显然是一个分治的过程，如下图所示：

最终分解为求长度为1的数列的DFT. 数列$a_0$a0​和数列$a_4$a4​的DFT分别是它们本身，因此它们的父序列$a_0, a_4$a0​,a4​的DFT为$a_0 + w_1^0a_4, a_0 - w_1^0a_4$a0​+w10​a4​,a0​−w10​a4​，这个DFT的结果需要继续提供给$a_0, a_4$a0​,a4​的父序列使用，层层向上，如此进行递归。计算机科学家引入蝴蝶操作来描述这一过程。一个蝴蝶操作的输入是两个数，输出也是两个数，其本身还拥有一个叫旋转因子，结构如下图所示：

其中输入来自两个子数列相同位置$k$k的两个数，旋转因子中的$n$n等于原数列长度，即两个子数列长度之和。蝴蝶操作的输出可以继续作为后续蝴蝶操作的输入，用来处理上面过程所述的层层向上的过程。最终得到如下过程：（注意与分解过程的图进行对照，检验该图中进行蝴蝶操作的两个数是否来自子数列的同一位置，这样能更好地理解整个过程）

蝴蝶操作并行化

  分治过程分为分解和合并阶段，蝴蝶操作是在合并阶段进行的。第1组合并是对长度为1的数列进行两两合并，共需做$n/2$n/2次合并，每次合并需要1次蝴蝶操作，该组合并共进行了$n/2$n/2次蝴蝶操作。以此类推可知，每一组合并都需要进行$n/2$n/2次蝴蝶操作，共进行$log_2n$log2​n组合并，因此时间复杂度为$O(nlog_2n)$O(nlog2​n).
  不难看出，每组合并过程内的所有蝴蝶操作是可以同时进行的，因为它们在图中没有先后顺序关系。若分配$n/2$n/2个线程同时处理每一组合并中的$n/2$n/2次蝴蝶操作，则每一组合并的并行时间复杂度为$O(1)$O(1)，共进行$log_2n$log2​n组合并，最终的并行时间复杂度为$O(log_2n)$O(log2​n).

CUDA C++实现

前置工作：复数的C++实现

  C++自带的库无法在device上运行，因此需要自己实现复数数据结构并同时提供host和device的函数。作者在自己实现复数时加入了生成$n$n次单位复根的功能。实现复数最主要的是要实现运算符重载，即复数的加、减、乘，原理较简单，若想深究请自行查找资料。另外，本文中数列的数值都用复数数据结构来存储，如为实数则将虚部存为0，而不是直接使用浮点数存储。以下给出复数实现的代码，注意__device__标记。

class Complex {
public:
    double real;
    double imag;

    Complex() {

    }

    // Wn 获取n次单位复根中的主单位根
    __device__ static Complex W(int n) {
        Complex res = Complex(cos(2.0 * PI / n), sin(2.0 * PI / n));
        return res;
    }

    // Wn^k 获取n次单位复根中的第k个
    __device__ static Complex W(int n, int k) {
        Complex res = Complex(cos(2.0 * PI * k / n), sin(2.0 * PI * k / n));
        return res;
    }
    
    // 实例化并返回一个复数（只能在Host调用）
    static Complex GetComplex(double real, double imag) {
        Complex r;
        r.real = real;
        r.imag = imag;
        return r;
    }

    // 随机返回一个复数
    static Complex GetRandomComplex() {
        Complex r;
        r.real = (double)rand() / rand();
        r.imag = (double)rand() / rand();
        return r;
    }

    // 随即返回一个实数
    static Complex GetRandomReal() {
        Complex r;
        r.real = (double)rand() / rand();
        r.imag = 0;
        return r;
    }

    // 随即返回一个纯虚数
    static Complex GetRandomPureImag() {
        Complex r;
        r.real = 0;
        r.imag = (double)rand() / rand();
        return r;
    }

    // 构造函数（只能在Device上调用）
    __device__ Complex(double real, double imag) {
        this->real = real;
        this->imag = imag;
    }
    
    // 运算符重载
    __device__ Complex operator+(const Complex &other) {
        Complex res(this->real + other.real, this->imag + other.imag);
        return res;
    }

    __device__ Complex operator-(const Complex &other) {
        Complex res(this->real - other.real, this->imag - other.imag);
        return res;
    }

    __device__ Complex operator*(const Complex &other) {
        Complex res(this->real * other.real - this->imag * other.imag, this->imag * other.real + this->real * other.imag);
        return res;
    }
};

分治的并行实现

  在使用串行化的方法实现分治时，可以借助递归进行问题的分解。然而CUDA是基于内存共享的并行计算框架，递归不利于程序的并行化，因此需要使用循环代替递归。

__global__ void Reduce(int nums[], int n) {
    int tid = threadIdx.x + blockDim.x * blockIdx.x;  // 线程号, 每个数对应一个线程
    if (tid >= n) return;  // 线程号超出数列范围, 返回
    for (int i = 2; i < 2 * n; i *= 2) {
        if (tid % i == 0) {
            nums[tid] += nums[tid + i / 2];  // 基本代码
        }
        __syncthreads();  // 同步线程
    }
}

  上述代码展示的是使用分治法进行数列求和的过程，最终的运行结果是数列的和被求出并存储在nums[0]中。对每次循环进行分析。第一次循环中，能进入基本代码的进程的tid为0, 2, 4, 6, …，本次循环执行完基本代码后，nums中1, 3, 5, 7, … 位置上的数都被加到了0, 2, 4, 6, … 位置上。第二次循环中，能进入基本代码的进程的tid为0, 4, 8, 12, …，本次循环执行完成后，nums中2, 6, 10, 14, … 位置上的数（存储的是第一次相加的结果）都被加到了0, 4, 8, 12, …  位置上。往后的循环都同理，其中最后一次循环，只有tid为0的进程能进入基本代码，将nums[n/2]加到nums[0]上，完成所有数的求和。

FFT分治面临的问题

  FFT分治的最大问题在于每次合并的两项并不来自相邻的进程。以$n=8$n=8为例，在数列求和问题中，第一次循环进行合并操作的两项分别为：0和1、2和3、4和5、6和7，但在FFT中，第一次循环进行合并操作的两项分别为：0和4、2和6、1和5、3和7. 往后的循环同样存在这个问题。解决方案是将数列按照0, 4, 2, 6, 1, 5, 3, 7的顺序重新排序。
  那么0, 4, 2, 6, 1, 5, 3, 7和原数列之间有什么关系呢？该数列的二进制数列为000, 100, 010, 110, 001, 101, 011, 111，将每个二进制数逆转，得到数列000, 001, 010, 011, 100, 101, 110, 111即原数列0, 1, 2, 3, 4, 5, 6, 7. 引入二进制逆转操作，对原数列重新排序可解决此问题。以下给出一个完整的程序来解释其C++实现：

// 根据数列长度n计算二进制最高位数
int GetBits(int n) {
    int bits = 0;
    while (n >>= 1) {
        bits++;
    }
    return bits;
}

// 在二进制位数为bits的前提下求数值i的二进制逆转
int BinaryReverse(int i, int bits) {
    int r = 0;
    do {
        r += i % 2 << --bits;
    } while (i /= 2);
    return r;
}

int main() {
    const int n = 8;
    const int bits = GetBits(n);
    int nums[n] = { 3, 8, 2, 4, 1, 6, 7, 9 };
    int j = BinaryReverse(1, bits);  // 排序后的数列的第1个位置对应原数列中的第j个位置
    printf("排序后的数列第1个位置上的数为%d", nums[j]);
}

各线程蝴蝶操作的时机

  设计CUDA程序需要明确数据在kernel中的位置，以及每个线程执行操作的时机，有时还需要考虑线程之间的同步。
  前面已经讲述了并行分治的每次循环分别有哪些进程可以进入基本代码。以$n=8$n=8为例，用以下表格对比数列求和和FFT操作的第一次循环中需要进行合并操作的两项：

线程	数列求和操作	FFT操作
0	nums[0]和nums[1]	nums[BinaryReverse(0)]和nums[BinaryReverse(1)]
2	nums[2]和nums[3]	nums[BinaryReverse(2)]和nums[BinaryReverse(3)]
4	nums[4]和nums[5]	nums[BinaryReverse(4)]和nums[BinaryReverse(5)]
6	nums[6]和nums[7]	nums[BinaryReverse(6)]和nums[BinaryReverse(7)]
tid	nums[tid]和nums[tid+1]	nums[BinaryReverse(tid)]和nums[BinaryReverse(tid+1)]

通过对比第一次循环可以看出，定义了BinaryReverse操作之后，确认合并对象在数列中的位置这个问题已经解决，只需用BinaryReverse(tid)替代tid即可确认数的位置。
  数列求和每次合并操作合并的是两个数，但FFT每次合并操作合并的都是两个数列。为了便于理解，仍然以$n=8$n=8为例，用以下表格对比数列求和和FFT操作的第二次和第三次循环中需要进行合并操作的两项：

线程	数列求和操作	FFT操作
0	nums[0]和nums[2]	1. nums[BinaryReverse(0)]和nums[BinaryReverse(2)] 2. nums[BinaryReverse(1)]和nums[BinaryReverse(3)]
4	nums[4]和nums[6]	1. nums[BinaryReverse(4)]和nums[BinaryReverse(6)] 2. nums[BinaryReverse(5)]和nums[BinaryReverse(7)]
tid	nums[tid]和nums[tid+2]	1. nums[BinaryReverse(tid)]和nums[BinaryReverse(tid+2)] 2. nums[BinaryReverse(tid+1)]和nums[BinaryReverse(tid+3)]

线程	数列求和操作	FFT操作
0	nums[0]和nums[4]	1. nums[BinaryReverse(0)]和nums[BinaryReverse(4)] 2. nums[BinaryReverse(1)]和nums[BinaryReverse(5)] 3. nums[BinaryReverse(2)]和nums[BinaryReverse(6)] 4. nums[BinaryReverse(3)]和nums[BinaryReverse(7)]
tid	nums[tid]和nums[tid+4]	1. nums[BinaryReverse(tid)]和nums[BinaryReverse(tid+4)] 2. nums[BinaryReverse(tid+1)]和nums[BinaryReverse(tid+5)] 3. nums[BinaryReverse(tid + 2)]和nums[BinaryReverse(tid+6)] 4. nums[BinaryReverse(tid + 3)]和nums[BinaryReverse(tid+7)]

通过以上对比可以看出，在FFT操作中，tid和tid+2 (或tid+4) 的作用是进行定位，找到需要合并的两个数列的开头。在本文第一份代码的基础上引入一层内部循环，可以达到这个要求：

// 蝴蝶操作, 输出结果直接覆盖原存储单元的数据, factor是旋转因子
__device__ void Bufferfly(Complex *a, Complex *b, Complex factor) {
    Complex a1 = (*a) + factor * (*b);
    Complex b1 = (*a) - factor * (*b);
    *a = a1;
    *b = b1;
}

__global__ void FFT(Complex nums[], int n, int bits) {
    int tid = threadIdx.x + blockDim.x * blockIdx.x;
    if (tid >= n) return; 
    for (int i = 2; i < 2 * n; i *= 2) {
        if (tid % i == 0) {
            // 新引入的循环
            for (int j = 0; j < k / 2; ++j) {
                Bufferfly(&nums[BinaryReverse(tid + j, bits)], &nums[BinaryReverse(tid + j + k / 2, bits)], Complex::W(k, j));
            }
        }
        __syncthreads();
    }
}

上述代码已经完成了FFT运算，但FFT结果数列的顺序是按照逆转二进制数的顺序排列的，因此在拷贝结果时应按逆转二进制数的顺序来拷贝。为此，引入一个数组result来存储最终结果：

__global__ void FFT(Complex nums[], Complex result[], int n, int bits) {
    int tid = threadIdx.x + blockDim.x * blockIdx.x;
    if (tid >= n) return;
    for (int i = 2; i < 2 * n; i *= 2) {
        if (tid % i == 0) {
            for (int j = 0; j < k / 2; ++j) {
                Bufferfly(&nums[BinaryReverse(tid + j, bits)], &nums[BinaryReverse(tid + j + k / 2, bits)], Complex::W(k, j));
            }
        }
        __syncthreads(); 
    }
    result[tid] = nums[BinaryReverse(tid, bits)];  // 拷贝到result中的对应地址
}

完整代码

// 一维FFT
#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include "../utils/Complex.cu"  // 自定义的复数数据结构
#include <iostream>
#include <string>
#include <stdlib.h>
#include <time.h>
#include <Windows.h>

// 根据数列长度n获取二进制位数
int GetBits(int n) {
    int bits = 0;
    while (n >>= 1) {
        bits++;
    }
    return bits;
}

// 在二进制位数为bits的前提下求数值i的二进制逆转
__device__ int BinaryReverse(int i, int bits) {
    int r = 0;
    do {
        r += i % 2 << --bits;
    } while (i /= 2);
    return r;
}

// 蝴蝶操作, 输出结果直接覆盖原存储单元的数据, factor是旋转因子
__device__ void Bufferfly(Complex *a, Complex *b, Complex factor) {
    Complex a1 = (*a) + factor * (*b);
    Complex b1 = (*a) - factor * (*b);
    *a = a1;
    *b = b1;
}

__global__ void FFT(Complex nums[], Complex result[], int n, int bits) {
    int tid = threadIdx.x + blockDim.x * blockIdx.x;
    if (tid >= n) return;
    for (int i = 2; i < 2 * n; i *= 2) {
        if (tid % i == 0) {
            int k = i;
            if (n - tid < k) k = n - tid;
            for (int j = 0; j < k / 2; ++j) {
                Bufferfly(&nums[BinaryReverse(tid + j, bits)], &nums[BinaryReverse(tid + j + k / 2, bits)], Complex::W(k, j));
            }
        }
        __syncthreads();
    }
    result[tid] = nums[BinaryReverse(tid, bits)];
}

// 打印数列
void printSequence(Complex nums[], const int N) {
    printf("[");
    for (int i = 0; i < N; ++i) {
        double real = nums[i].real, imag = nums[i].imag;
        if (imag == 0) printf("%.16f", real);
        else {
            if (imag > 0) printf("%.16f+%.16fi", real, imag);
            else printf("%.16f%.16fi", real, imag);
        }
        if (i != N - 1) printf(", ");
    }
    printf("]\n");
}

int main() {
    srand(time(0));  // 设置随机数种子
    const int TPB = 1024;  // 每个Block的线程数，即blockDim.x
    const int N = 1024 * 32;  // 数列大小
    const int bits = GetBits(N);
    
    // 随机生成实数数列
    Complex *nums = (Complex*)malloc(sizeof(Complex) * N), *dNums, *dResult;
    for (int i = 0; i < N; ++i) {
        nums[i] = Complex::GetRandomReal();
    }
    printf("Length of Sequence: %d\n", N);
    // printf("Before FFT: \n");
    // printSequence(nums, N);
    
    // 保存开始时间
    float s = GetTickCount();
    
    // 分配device内存，拷贝数据到device
    cudaMalloc((void**)&dNums, sizeof(Complex) * N);
    cudaMalloc((void**)&dResult, sizeof(Complex) * N);
    cudaMemcpy(dNums, nums, sizeof(Complex) * N, cudaMemcpyHostToDevice);
    
    // 调用kernel
    dim3 threadPerBlock = dim3(TPB);
    dim3 blockNum = dim3((N + threadPerBlock.x - 1) / threadPerBlock.x);
    FFT<<<blockNum, threadPerBlock>>>(dNums, dResult, N, bits);

    // 拷贝回结果
    cudaMemcpy(nums, dResult, sizeof(Complex) * N, cudaMemcpyDeviceToHost);
    
    // 计算用时
    float cost = GetTickCount() - s;
    // printf("After FFT: \n");
    // printSequence(nums, N);
    printf("Time of Transfromation: %fms", cost);
  
    // 释放内存
    free(nums);
    cudaFree(dNums);
    cudaFree(dResult);
}