一、Radix2的傅里叶变换的缺陷

首先,我们解释一下什么是Radix2,Radix4,Radix8等的傅里叶变换。在上一篇中,我们得出我们可以将N的傅里叶变换的输入拆为奇数和偶数,然后奇数和偶数分别进行傅里叶变换,再通过蝶形变换组合为N的傅里叶变换的结果。简单的公式形式化描述如下

$$
DFT(N) = DFT_1(\frac N 2) \oplus DFT_2(\frac N 2), \oplus表示蝶形变换
$$

我们可以知道,一次蝶形变换中所做的计算非常少,每一轮GPU派发很多任务,然后每个任务仅执行一次2个数据的蝶形变换,然后就结束写回Buffer了。这样我们的瓶颈往往会卡在DrawCall上。

其实,在拆分为奇数和偶数的过程中,我们并不一样要这样拆。我们可以拆成4组,甚至可以拆成8组,每组组内的数据的下标是同余的。
例如
$[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]$拆成四组是$[0,4,8,12][1,5,9,13][2,6,10,14][3,7,11,15]$,拆成8组就是每个2个了。

二、Radix4的快速傅里叶变换

2.1 Radix4 的蝶形变换推导

首先,假设长度为$N$的信号,按照同余分成四组,每组大小为$\frac N 4$。
令$r \in [0,1...\frac N 4]$,
我们可以得到

$$
\begin{aligned}
F[k] &= \sum_{n=0} ^{N-1} {f[n]e^{-\frac {2 \pi } {N} i k n}}
\ &= \sum_{r=0} ^{\frac N 4 - 1}f[4r]e^{-\frac {2 \pi} {N} i k (4r)} +
\sum_{r=0} ^{\frac N 4 - 1}f[4r + 1]e^{-\frac {2 \pi} {N} i k (4r + 1)} +
\sum_{r=0} ^{\frac N 4 - 1}f[4r + 2]e^{-\frac {2 \pi} {N} i k (4r + 2)} +
\sum_{r=0} ^{\frac N 4 - 1}f[4r + 3]e^{-\frac {2 \pi} {N} i k (4r + 3)}
\ &= \sum_{r=0} ^{\frac N 4 - 1}f_1[r]e^{-\frac {2 \pi} {\frac N 4} i k r} +
\sum_{r=0} ^{\frac N 4 - 1}f_2[r]e^{-\frac {2 \pi} {\frac N 4} i k r} * e^{-\frac {2 \pi} {N} i k} +
\sum_{r=0} ^{\frac N 4 - 1}f_3[r]e^{-\frac {2 \pi} {\frac N 4} i k r} * e^{-\frac {2 \pi} {N} i 2k} +
\sum_{r=0} ^{\frac N 4 - 1}f_4[r]e^{-\frac {2 \pi} {\frac N 4} i k r} * e^{-\frac {2 \pi} {N} i 3k}
\ &= \sum_{r=0} ^{\frac N 4 - 1}f_1[r]e^{-\frac {2 \pi} {\frac N 4} i k r} +
e^{-\frac {2 \pi} {N} i k} \sum_{r=0} ^{\frac N 4 - 1}f_2[r]e^{-\frac {2 \pi} {\frac N 4} i k r} +
e^{-\frac {2 \pi} {N} i 2k}\sum_{r=0} ^{\frac N 4 - 1}f_3[r]e^{-\frac {2 \pi} {\frac N 4} i k r} +
e^{-\frac {2 \pi} {N} i 3k}\sum_{r=0} ^{\frac N 4 - 1}f_4[r]e^{-\frac {2 \pi} {\frac N 4} i k r}
\end{aligned}
$$

令符号$W_N^k$ 代表 $e^{-\frac {2 \pi} {N} ik}$, 上述式子可简写为

$$
F[k] = \sum_{r=0} ^{\frac N 4 - 1} f_1[r]W_{\frac N 4}^{kr} +
W_{N}^{k}\sum_{r=0} ^{\frac N 4 - 1} f_2[r]W_{\frac N 4}^{kr} +
W_{N}^{2k}\sum_{r=0} ^{\frac N 4 - 1} f_3[r]W_{\frac N 4}^{kr} +
W_{N}^{3k}\sum_{r=0} ^{\frac N 4 - 1} f_4[r]W_{\frac N 4}^{kr}
$$

上述四项分别是四个长度为$\frac N 4$的傅里叶变换。
因此可化简为

$$
F[k] = {W^{0k}_N} F_1 + {W^{1k}_N} F_2 + {W^{2k}_N} F_3 + {W^{3k}_N} F_4 \tag{1}
$$

这里的$F[k]$的$k \in [0,1...\frac N 4]$,因此我们可以类似Radix2求另外4份同余的部分。

$$
F[k + \frac N 4] = {W^{0k + 0 \frac N 4}_N} F_1 + {W^{k + \frac N 4}_N} F_2 + {W^{2k + \frac {2N} 4}_N} F_3 + {W^{3k + \frac {3N} 4}_N} F_4
\ = {W^0_4}{W^{0k}_N} F_1 + {W^1_4}{W^k_N} F_2 + {W^2_4}{W^{2k}_N} F_3 + {W^3_4}{W^{3k}_N} F_4
\tag{2}
$$
$$
F[k + \frac {2N} 4] = {(W^0_4)}^2{W^{0k}_N} F_1 + {(W^1_4)}^2{W^k_N} F_2 + {(W^2_4)}^2{W^{2k}_N} F_3 + {(W^3_4)}^2{W^{3k}_N} F_4 \tag{3}
$$
$$
F[k + \frac {3N} 4] = {(W^0_4)}^3{W^{0k}_N} F_1 + {(W^1_4)}^3{W^k_N} F_2 + {(W^2_4)}^3{W^{2k}_N} F_3 + {(W^3_4)}^3{W^{3k}_N} F_4 \tag{4}
$$

我们可以求得
$$
\begin{aligned}
W^0_4 &= 1 \\
W^1_4 &= -i \\
W^2_4 &= -1 \\
W^3_4 &= i
\end{aligned}
\tag{5}
$$

在复平面上,我们可以观察到,单位复数的整数次幂,等价于原复数与实数轴的夹角的整数倍的旋转,这一点可以根据欧拉公式和辐角的几何意义推算得到。

$$
(cos\theta + isin\theta)^2 = {(e^{i\theta})}^2 = e^{i2\theta}=cos(2\theta) + isin(2\theta)
$$

式子$(5)$代入式子$(1)(2)(3)(4)$可得

$$
\begin{aligned}
F[k] &= {W^{0k}_N} F_1 + {W^{k}_N} F_2 + {W^{2k}_N} F_3 + {W^{3k}_N} F_4 \\
F[k + \frac N 4] &= {W^{0k}_N} F_1 +(-i)^1{W^k_N} F_2 + (-1)^1 {W^{2k}_N} F_3 + (i)^1{W^{3k}_N} F_4 \\
F[k + \frac {2N} 4] &= {W^{0k}_N} F_1 +(- i)^2{W^k_N} F_2 + (-1)^2 {W^{2k}_N} F_3 + (i)^2{W^{3k}_N} F_4 \\
F[k + \frac {3N} 4] &= {W^{0k}_N} F_1 +(- i)^3{W^k_N} F_2 + (-1)^3 {W^{2k}_N} F_3 + (i)^3{W^{3k}_N} F_4 \\
\end{aligned}
$$

利用旋转的思想,上述式子可得

$$
\begin{aligned}
F[k] &= {W^{0k}_N} F_1 + {W^{k}_N} F_2 + {W^{2k}_N} F_3 + {W^{3k}_N} F_4 \\
F[k + \frac N 4] &= {W^{0k}_N} F_1 + (-i){W^k_N} F_2 + (-1){W^{2k}_N} F_3 + (i){W^{3k}_N} F_4 \\
F[k + \frac {2N} 4] &= {W^{0k}_N} F_1 +(-1){W^k_N} F_2 + (1) {W^{2k}_N} F_3 + (-1){W^{3k}_N} F_4 \\
F[k + \frac {3N} 4] &= {W^{0k}_N} F_1 +(i){W^k_N} F_2 + (-1) {W^{2k}_N} F_3 + (-i){W^{3k}_N} F_4 \\
\end{aligned}
\tag{6}
$$

式子$(5)$就已经求完了Radix4的傅里叶变换的全部分解式。

2.2 从Radix2的蝶形变换推导出Radix4的蝶形变换

上面虽然通过直接暴力求解出了Radix4的解析式,但是上面的式子其实在计算机算起来,效率不高,因为在这个Radix4的内部,没有利用还可以继续分治的性质。上述的方式更适合FPGA等芯片设计时使用,因为可以通过4个分别并行计算,晶体管数量会增加,功率也会增加,但是计算时间会减少。我们在GPU中,Radix4是我们一轮的Kernel函数,Kernel内部在软件层面看来没有并行,因此使用分治法更好。
下面我们通过Radix2的蝶形变换来推导一下Radix4的结果.
一次Radix4的内部,相当于N=4的数据,进行DFT,也就是相当于N=4的数据进行2轮Radix2的蝶形变换。
首先,先给出Radix2的蝶形变换公式

$$
\begin{aligned}
F[k] &= F_1[k] + W^k_N F_2[k]
\\
F[k + \frac N 2] &= F_1[k] - W^k_N F_2[k]
\end{aligned}
\tag{7}
$$

第一轮,每组Radix2蝶形变换的变换后大小是2,也就是DFT1->DFT2了。假设数据是$f_0,f_1,f_2,f_3$,则$(f_0, f_2)$与$(f_1,f_3)$分别进行一次Radix2的蝶形变换。假设这2个蝶形变换的输出分别是$F_0,F_1$。

$$
F_0[0] = f_0 + W^0_2f_2 = f_0 + f_2 \\
F_0[1] = f_0 - W^0_2f_2 = f_0 - f_2 \\
F_1[0] = f_1 + W^0_2f_3 = f_1 + f_3 \\
F_1[1] = f_1 - W^0_2f_3 = f_1 - f_3 \\
$$

第二轮,根据$(F_0, F_1)$进行蝶形变换,这次每组蝶形变换的变换后大小是4,而且也只剩下一组了。
$F_0$的第0个元素和$F_1$的第0个元素进行蝶形变换得到

$$
F[0] = F_0[0] + W^0_4F_1[0] \\
F[2] = F_0[0] - W^0_4F_1[0] \\
$$

$F_0$的第1个元素和$F_1$的第1个元素进行蝶形变换得到

$$
F[1] = F_0[1] + W^1_4 F_1[1] \\
F[3] = F_0[1] - W^1_4 F_1[1]
$$

我们代入第一轮的结果,可得

$$
\begin{aligned}
F[0] = f_0 + f_2 + W^0_4(f_1 + f_3) &= f_0 + f_2 + f_1 + f_3 \\
F[1] = f_0 - f_2 + W^1_4(f_1 - f_3) &= f_0 - f_2 + (-i)(f_1 - f_3) \\
F[2] = f_0 + f_2 - W^0_4(f_1 + f_3) &= f_0 + f_2 - (f_1 + f_3) \\
F[3] = f_0 - f_2 - W^1_4(f_1 - f_3) &= f_0 - f_2 - (-i)(f_1 - f_3) \\
\end{aligned}
\tag{8}
$$

我们可以得到跟式子$(6)$一样结果的式子啦。

2.3 Stockham的Radix4的index的映射

stockham_fft16_radix4.svg
如上图所示,是Stockham FFT的16个数据时候的例子。它与Stockham FFT Radix2的时候映射方式基本一致。

令$x$表示第$x$个蝶形变换,则$x \in [0,1,...\frac N 4 - 1]$. $N$是DFT的总长度。
令$m$表示第$m$轮DFT,则本次每组子DFT长度是$4^m$,$m \in [0,1,...log_4{N}-1]$。

例如总长度为16的DFT,需要执行两轮Radix4的蝶形变换。
第一轮,每组1个,每同余的4组合并进行蝶形变换。得到4组 长度为4的DFT的结果。
第二轮,这4组,合并进行蝶形变换,得到1组 长度为16的DFT的结果,这就是最终结果。

令$p$表示本轮每组子DFT长度, $p = 4^m$。
令$k$表示第$x$个蝶形变换,在本轮所属的组内的DFT的数据序号,则$k = x mod p$.
则,

  1. 每一轮第$x$个蝶形变换的四个输入的下标分别为 $x$, $x + \frac {N} 4$, $x + \frac {2N} 4$, $x + \frac {3N} 4$.
  2. 第$m$轮 第$x$个蝶形变换的第一个输出下标是 $(x - k) * 4 + k$
  3. 第$m$轮 第$x$个蝶形变换的每个输出下标的间隔是 $p$

2.4 Unity中的实现

根据上面的2.3和2.2的推导,可以很轻松的写出下面的实现。
注意,为了提高运算效率,我把twiddle操作挪到上一层去了。例如第二轮的$F=u_i +W^k_4 u_j$的$W^k_8$的系数,实际上放到了第一轮的$u_j$的变量上。
容易出错的是,一些笔误或者对应关系搞错了。出错了,不容易查,最好用很小的数据集手算一下函数的流程。

// compute shader, x direction of radix4
inline float2 twiddle_w1_4(float2 m) {
    return float2(m.y, -m.x);
}

[numthreads(NTHREADS, NTHREADS, 1)]
void dit_x_radix4(uint3 id: SV_DispatchThreadID) {
    int offset_in = n >> 2;
    int k = id.x & (p - 1);  // id.x mod p
#ifdef IDFT
    float theta = float(k)/ (p * 2);
#else
    float theta = -float(k)/ (p * 2);
#endif

    int2 in_0 = int2(id.x + offset_in * 0, id.y);
    int2 in_1 = int2(id.x + offset_in * 1, id.y);
    int2 in_2 = int2(id.x + offset_in * 2, id.y);
    int2 in_3 = int2(id.x + offset_in * 3, id.y);

    float2 f0 = fftin[in_0];
    float2 f1 = mul(fftin[in_1], exp_theta(theta));
    float2 f2 = mul(fftin[in_2], exp_theta(2.0f * theta));
    float2 f3 = mul(fftin[in_3], exp_theta(3.0f * theta));

    float2 u0 = f0 + f2;
    float2 u1 = f0 - f2;
    float2 u2 = f1 + f3;
    float2 u3 = twiddle_w1_4(f1 - f3);

    float2 F0 = u0 + u2;
    float2 F1 = u1 + u3;
    float2 F2 = u0 - u2;
    float2 F3 = u1 - u3;

    int out_index = ((id.x - k) << 2) + k;

    fftout[int2(out_index, id.y)] = F0;
    fftout[int2(out_index + p, id.y)] = F1;
    fftout[int2(out_index + 2 * p, id.y)] = F2;
    fftout[int2(out_index + 3 * p, id.y)] = F3;
}

三、Radix8的快速傅里叶变换

可类似上述的Radix4的分解式$(1)$,轻松推导出Radix8的分解式。
令$k \in [0,1,...\frac N 8 - 1], \omega \in [0,1,2...7]$
$$
F[k + \frac {\omega N} 4] =
{(W^0_8)}^{\omega}{W^{0k}_N} F_1 + {(W^1_8)}^{\omega}{W^k_N} F_2 +
{(W^2_8)}^{\omega}{W^{2k}_N} F_3 + {(W^3_8)}^{\omega}{W^{3k}_N} F_4 + \\
{(W^4_8)}^{\omega}{W^{2k}_N} F_5 + {(W^5_8)}^{\omega}{W^{3k}_N} F_6 +
{(W^6_8)}^{\omega}{W^{2k}_N} F_7 + {(W^7_8)}^{\omega}{W^{3k}_N} F_8
\tag{9}
$$

下面这个是Radix8的因子查找表。

$W^0_8$ $W^1_8$ $W^2_8$ $W^3_8$ $W^4_8$ $W^5_8$ $W^6_8$ $W^7_8$
0 1 $W^0_8$ $W^0_8$ $W^0_8$ $W^0_8$ $W^0_8$ $W^0_8$ $W^0_8$
1 1 $W^1_8$ $W^2_8$ $W^3_8$ $W^4_8$ $W^5_8$ $W^6_8$ $W^7_8$
2 1 $W^2_8$ $W^4_8$ $W^6_8$ $W^0_8$ $W^2_8$ $W^4_8$ $W^6_8$
3 1 $W^3_8$ $W^6_8$ $W^1_8$ $W^3_8$ $W^6_8$ $W^1_8$ $W^3_8$
4 1 $W^4_8$ $W^0_8$ $W^4_8$ $W^0_8$ $W^4_8$ $W^0_8$ $W^4_8$
5 1 $W^5_8$ $W^2_8$ $W^7_8$ $W^4_8$ $W^1_8$ $W^6_8$ $W^3_8$
6 1 $W^6_8$ $W^4_8$ $W^2_8$ $W^0_8$ $W^6_8$ $W^4_8$ $W^2_8$
7 1 $W^7_8$ $W^6_8$ $W^5_8$ $W^4_8$ $W^3_8$ $W^2_8$ $W^1_8$

3.1 Radix2推导Radix8

由于上述式子基本不能用于GPU的实现啦,我们推导一下如何将一个Radix8的内部逻辑分解为3轮Radix2的逻辑。
我们假设数据是$f_0,f_1,f_2,f_3,f_4,f_5,f_6,f_7$.
为了避免复杂的id映射导致混乱,先给出8个数据的Radix2的蝶形变换图。我们要做的就是,将它变成一次Radix8。
StockhamFFT.svg
首先,第一轮蝶形变换

$$
u_0 = f_0 + f_4 \\
u_4 = f_0 - f_4 \\
u_1 = f_1 + f_5 \\
u_5 = f_1 - f_5 \\
u_2 = f_2 + f_6 \\
u_6 = f_2 - f_6 \\
u_3 = f_3 + f_7 \\
u_7 = f_3 - f_7 \\
$$

第二轮
$$
v_0 = u_0 + W^0_4 u_2 = f_0 + f_4 + W^0_8 f_2 + W^0_8 f_6 \\
v_2 = u_0 - W^0_4 u_2 = f_0 + f_4 - W^0_8 f_2 - W^0_8 f_6 \\
v_4 = u_4 + W^1_4 u_6 = f_0 - f_4 + W^2_8 f_2 - W^2_8 f_6 \\
v_6 = u_4 - W^1_4 u_6 = f_0 - f_4 - W^2_8 f_2 + W^2_8 f_6 \\
v_1 = u_1 + W^0_4 u_3 = f_1 + f_5 + W^0_8 f_3 + W^0_8 f_7 \\
v_3 = u_1 - W^0_4 u_3 = f_1 + f_5 - W^0_8 f_3 - W^0_8 f_7 \\
v_5 = u_5 + W^1_4 u_7 = f_1 - f_5 + W^2_8 f_3 - W^2_8 f_7 \\
v_7 = u_5 - W^1_4 u_7 = f_1 - f_5 - W^2_8 f_3 + W^2_8 f_7 \\
$$

第三轮
$$
F[0] = v_0 + W^0_8 v_1 = f_0 + f_4 + W^0_8 f_2 + W^0_8 f_6 + W^0_8(f_1 + f_5 + W^0_8 f_3 + W^0_8 f_7) \\
F[1] = v_4 + W^1_8 v_5 = f_0 - f_4 + W^2_8 f_2 - W^2_8 f_6 + W^1_8(f_1 - f_5 + W^2_8 f_3 - W^2_8 f_7) \\
F[2] = v_2 + W^2_8 v_3 = f_0 + f_4 - W^0_8 f_2 - W^0_8 f_6 + W^2_8(f_1 + f_5 - W^0_8 f_3 - W^0_8 f_7) \\
F[3] = v_6 + W^3_8 v_7 = f_0 - f_4 - W^2_8 f_2 + W^2_8 f_6 + W^3_8(f_1 - f_5 - W^2_8 f_3 + W^2_8 f_7) \\
F[4] = v_0 - W^0_8 v_1 = f_0 + f_4 + W^0_8 f_2 + W^0_8 f_6 - W^0_8(f_1 + f_5 + W^0_8 f_3 + W^0_8 f_7) \\
F[5] = v_4 - W^1_8 v_5 = f_0 - f_4 + W^2_8 f_2 - W^2_8 f_6 - W^1_8(f_1 - f_5 + W^2_8 f_3 - W^2_8 f_7) \\
F[6] = v_2 - W^2_8 v_3 = f_0 + f_4 - W^0_8 f_2 - W^0_8 f_6 - W^2_8(f_1 + f_5 - W^0_8 f_3 - W^0_8 f_7) \\
F[7] = v_6 - W^3_8 v_7 = f_0 - f_4 - W^2_8 f_2 + W^2_8 f_6 - W^3_8(f_1 - f_5 - W^2_8 f_3 + W^2_8 f_7) \\
$$

将上式子中的$-1$系数替换为$W^4_8$, 化简可得与式子 $(9)$一样的结果。

3.2 Unity中的实现

Unity中实现,几乎也是照搬上面的公式。注意,为了提高运算效率,我把twiddle操作挪到上一层去了。例如第三轮的$F=v_i +W^k_8 v_j$的$W^k_8$的系数,实际上放到了第二轮的$v_j$的变量上。

// compute shader, x direction of radix8
inline float2 twiddle_w1_4(float2 m) {
    return float2(m.y, -m.x);
}

inline float2 twiddle_w1_8(float2 m) {
    return float2(COS_PI_4, COS_PI_4) * float2(m.x + m.y, -m.x + m.y);
}

inline float2 twiddle_w2_8(float2 m) {
    return twiddle_w1_4(m);
}

inline float2 twiddle_w3_8(float2 m) {
    return float2(COS_PI_4, COS_PI_4) * float2(-m.x + m.y, -m.x - m.y);
}

[numthreads(NTHREADS, NTHREADS, 1)]
void dit_x_radix8(uint3 id: SV_DispatchThreadID) {
    int offset_in = n >> 3;
    int k = id.x & (p - 1);  // id.x mod p
#ifdef IDFT
    float theta = float(k) / (p * 4);
#else
    float theta = -float(k) / (p * 4);
#endif

    int2 in_0 = int2(id.x + offset_in * 0, id.y);
    int2 in_1 = int2(id.x + offset_in * 1, id.y);
    int2 in_2 = int2(id.x + offset_in * 2, id.y);
    int2 in_3 = int2(id.x + offset_in * 3, id.y);
    int2 in_4 = int2(id.x + offset_in * 4, id.y);
    int2 in_5 = int2(id.x + offset_in * 5, id.y);
    int2 in_6 = int2(id.x + offset_in * 6, id.y);
    int2 in_7 = int2(id.x + offset_in * 7, id.y);

    float2 f0 = fftin[in_0];
    float2 f1 = mul(fftin[in_1], exp_theta(theta));
    float2 f2 = mul(fftin[in_2], exp_theta(2.0f * theta));
    float2 f3 = mul(fftin[in_3], exp_theta(3.0f * theta));
    float2 f4 = mul(fftin[in_4], exp_theta(4.0f * theta));
    float2 f5 = mul(fftin[in_5], exp_theta(5.0f * theta));
    float2 f6 = mul(fftin[in_6], exp_theta(6.0f * theta));
    float2 f7 = mul(fftin[in_7], exp_theta(7.0f * theta));

    float2 u0 =              f0 + f4;
    float2 u4 =              f0 - f4;
    float2 u1 =              f1 + f5;
    float2 u5 =              f1 - f5;
    float2 u2 =              f2 + f6;
    float2 u6 = twiddle_w1_4(f2 - f6);
    float2 u3 =              f3 + f7;
    float2 u7 = twiddle_w1_4(f3 - f7);

    float2 v0 =              u0 + u2;
    float2 v2 =              u0 - u2;
    float2 v4 =              u4 + u6;
    float2 v6 =              u4 - u6;
    float2 v1 =              u1 + u3;
    float2 v3 = twiddle_w2_8(u1 - u3);
    float2 v5 = twiddle_w1_8(u5 + u7);
    float2 v7 = twiddle_w3_8(u5 - u7);

    int out_index = ((id.x - k) << 3) + k;

    fftout[int2(out_index, id.y)]         = v0 + v1;
    fftout[int2(out_index + p, id.y)]     = v4 + v5;
    fftout[int2(out_index + 2 * p, id.y)] = v2 + v3;
    fftout[int2(out_index + 3 * p, id.y)] = v6 + v7;
    fftout[int2(out_index + 4 * p, id.y)] = v0 - v1;
    fftout[int2(out_index + 5 * p, id.y)] = v4 - v5;
    fftout[int2(out_index + 6 * p, id.y)] = v2 - v3;
    fftout[int2(out_index + 7 * p, id.y)] = v6 - v7;
}

四、RadixR的快速傅里叶变换

Radix16以后,直接手写代码非常复杂,但是实际推算也和上面的Radix8类似,主要根据Stockham FFT或者Cooley-Tukey FFT来实现。
注意到,Stockham FFT每轮输入和输出的下标都会变化,实现起来,会需要2个长度为$2^R$的临时变量数组,类似上面Radix8的推算里的$u_i$和$v_i$一样。
临时变量太多,会大大降低Shader效率,这一点在后面测试可以看到。
因此,我在实现RadixR的快速傅里叶变换时,使用的是Cooley-Tukey FFT,这样可以用一个数组即可实现In-place FFT.

// compute shader, x direction of radixR
#ifdef RADIX_16
        #define RADIX_SIZE 16U
        #define RADIX_POW 4U
#else
    #ifdef RADIX_32
        #define RADIX_SIZE 32U
        #define RADIX_POW 5U
    #else
        #ifdef RADIX_64
            #define RADIX_SIZE 64U
            #define RADIX_POW 6U
        #else
            #ifdef RADIX_128
                #define RADIX_SIZE 128U
                #define RADIX_POW 7U
            #else
                #ifdef RADIX_256
                    #define RADIX_SIZE 256U
                    #define RADIX_POW 8U
                #else
                    #ifdef RADIX_512
                        #define RADIX_SIZE 512U
                        #define RADIX_POW 9U
                    #else
                        #define RADIX_SIZE 1024U
                        #define RADIX_POW 10U
                    #endif
                #endif
            #endif
        #endif
    #endif
#endif

#define twiddle(w_frac, w_base, m) mul(\
    float2(cos(-2 * PI * (w_frac) / float(w_base)),\
    sin(-2 * PI * (w_frac) / float(w_base))), m)

inline int bitreverse(int i, int max_bit) {
    int res = 0;
    int count = 0;
    while (count < max_bit) {
        res <<= 1;
        res |= (i & 1);
        i >>= 1;
        count++;
    }
    return res;
}

[numthreads(NTHREADS, NTHREADS, 1)]
void dit_x_radixq(uint3 id: SV_DispatchThreadID) {
    int offset_in = n >> RADIX_POW;
    int k = id.x & (p - 1);  // id.x mod p
#ifdef IDFT
    float theta = float(k) / (p * (1 << (RADIX_POW - 1)));
#else
    float theta = -float(k) / (p * (1 << (RADIX_POW - 1)));
#endif

    float2 u[RADIX_SIZE];
    for (uint i = 0; i < RADIX_SIZE; ++i) {
        int2 in_index = int2(id.x + offset_in * i, id.y);
        u[bitreverse(i, RADIX_POW)] = mul(fftin[in_index], exp_theta(float(i) * theta));
    } 

    for (uint j = 0; j < RADIX_POW; ++j) {
        for (uint i = 0; i < (RADIX_SIZE / 2); ++i) {
            uint u_p = 1 << j;
            uint u_2p = 1 << (j + 1);
            uint u_i = u_2p * (i / u_p) + (i % u_p);
            float2 a = u[u_i];
            float2 b = u[u_i + u_p];
            float2 twiddle_b = twiddle(i % u_p, u_2p, (b));
            float2 tmp = a - twiddle_b;
            u[u_i] = a + twiddle_b;
            u[u_i + u_p] = tmp;
        }
    }

    int out_index = ((id.x - k) << RADIX_POW) + k;
    for (uint i1 = 0; i1 < RADIX_SIZE / 2; ++i1) {
        fftout[int2(out_index + i1 * p, id.y)] = u[i1];
        fftout[int2(out_index + (i1 + RADIX_SIZE / 2) * p, id.y)] = u[i1 + RADIX_SIZE / 2];
    }
}

五、 Unity中FFT的性能分析

通过上面的分析可知,Radix越大,一次做的计算就越多,DrawCall越少,但是Dispatch的并行程度越小,需要的临时变量也越多。
我在自己电脑上进行了性能测试。我自己电脑的配置是i7 6800k + 48GRAM + GTX 2080Ti + 11G显存。
我用了一张1024x1024xR8的图,做了针对不同Radix的测试。

最大的Radix 运行时间(ms) DrawCall 一次X方向的DFT组合
2 6.257 21 DFT2 * 10次
4 2.149 11 DFT4 * 5次
8 1.763 9 DFT8 * 3次, DFT2
16 1.308 7 DFT16 * 2次, DFT4
32 1.001 5 DFT32 * 2次
64 1.074 5 DFT64, DFT16
128 1.34 5 DFT128, DFT8
256 15.4 5 DFT256, DFT4
512 17.2 5 DFT512, DFT2
1024 17.85 3 DFT1024

经过测试发现,对于1024的图来说,以32作为Radix是最佳的效率,刚好一个方向做2次DFT32就可以完成了,加上一个Normalize,总共才5个Drawcall。
另外,有个性能急剧下降的点,在最大Radix为256的时候,性能产生了急剧的下降。我认为,在DFT256的情况下,显卡的SM单元里的寄存器可能用完了,根据NVIDIA的白皮书,一个SM单元有16384个Register,每个Register是32bit。多个线程同时在SM单元里并行运行,数据太大就产生了这种寄存器不够的情况。

总结下来,我觉得对于size为128以下的FFT,可能直接一轮就搞完最快,对于更大的,以三轮或者2轮较为合适。
很明显,这个性能应该跟GPU的寄存器和Cache的大小有很大关系,我暂时在手机上测试过,应该会有不同的结果。