上一篇
琪露诺:详尽的快速傅里叶变换推导与Unity中的实现zhuanlan.zhihu.com,我推导了最基本最简单的快速傅里叶变换,以及描述了市面上一下不同的变种的实现的区别,最后在Unity Compute Shader中实现了FFT。
本篇,我主要讲傅里叶变换的一些性能优化,以及更加复杂的变种实现。
不想细看,想直接跑工程的,可以点这里:
https://github.com/MioMioReimu/RadixNFFTgithub.com一、Radix2的傅里叶变换的缺陷
首先,我们解释一下什么是Radix2,Radix4,Radix8等的傅里叶变换。在上一篇中,我们得出我们可以将N的傅里叶变换的输入拆为奇数和偶数,然后奇数和偶数分别进行傅里叶变换,再通过蝶形变换组合为N的傅里叶变换的结果。简单的公式形式化描述如下
我们可以知道,一次蝶形变换中所做的计算非常少,每一轮GPU派发很多任务,然后每个任务仅执行一次2个数据的蝶形变换,然后就结束写回Buffer了。这样我们的瓶颈往往会卡在DrawCall上。
其实,在拆分为奇数和偶数的过程中,我们并不一样要这样拆。我们可以拆成4组,甚至可以拆成8组,每组组内的数据的下标是同余的。
例如
当我们拆成2组,就是Radix2;当拆成4组就是Radix4,以此类推。
二、Radix4的快速傅里叶变换
2.1 Radix4 的蝶形变换推导
首先,假设长度为
令
我们可以得到
令符号
上述四项分别是四个长度为
因此可化简为
这里的
我们可以求得
在复平面上,我们可以观察到,单位复数的整数次幂,等价于原复数与实数轴的夹角的整数倍的旋转,这一点可以根据欧拉公式和辐角的几何意义推算得到。
式子
利用旋转的思想,上述式子可得
式子
2.2 从Radix2的蝶形变换推导出Radix4的蝶形变换
上面虽然通过直接暴力求解出了Radix4的解析式,但是上面的式子其实在计算机算起来,效率不高,因为在这个Radix4的内部,没有利用还可以继续分治的性质。上述的方式更适合FPGA等芯片设计时使用,因为可以通过4个分别并行计算,晶体管数量会增加,功率也会增加,但是计算时间会减少。我们在GPU中,Radix4是我们一轮的Kernel函数,Kernel内部在软件层面看来没有并行,因此使用分治法更好。
下面我们通过Radix2的蝶形变换来推导一下Radix4的结果.
一次Radix4的内部,相当于N=4的数据,进行DFT,也就是相当于N=4的数据进行2轮Radix2的蝶形变换。
首先,先给出Radix2的蝶形变换公式
第一轮,每组Radix2蝶形变换的变换后大小是2,也就是DFT1->DFT2了。假设数据是
第二轮,根据
我们代入第一轮的结果,可得
稍微化简一下式子
2.3 Stockham的Radix4的index的映射
如上图所示,是Stockham FFT的16个数据时候的例子。它与Stockham FFT Radix2的时候映射方式基本一致。
令
令
例如总长度为16的DFT,需要执行两轮Radix4的蝶形变换。
第一轮,每组1个,每同余的4组合并进行蝶形变换。得到4组 长度为4的DFT的结果。
第二轮,这4组,合并进行蝶形变换,得到1组 长度为16的DFT的结果,这就是最终结果。
令
令
则,
1. 每一轮第个蝶形变换的四个输入的下标分别为,,,.
2. 第轮 第个蝶形变换的第一个输出下标是
3. 第轮 第个蝶形变换的每个输出下标的间隔是
2.4 Unity中的实现
根据上面的2.3和2.2的推导,可以很轻松的写出下面的实现。 注意,为了提高运算效率,我把twiddle操作挪到上一层去了。例如第二轮的
// 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的分解式
令
下面这个是Radix8的因子查找表。
3.1 Radix2推导Radix8
由于上述式子基本不能用于GPU的实现啦,我们推导一下如何将一个Radix8的内部逻辑分解为3轮Radix2的逻辑。
我们假设数据是
为了避免复杂的id映射导致混乱,先给出8个数据的Radix2的蝶形变换图。我们要做的就是,将它变成一次Radix8。
首先,第一轮蝶形变换
第二轮
第三轮
将上式子中的
3.2 Unity中的实现
Unity中实现,几乎也是照搬上面的公式。注意,为了提高运算效率,我把twiddle操作挪到上一层去了。例如第三轮的
// 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个长度为
// 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的测试。
经过测试发现,对于1024的图来说,以32作为Radix是最佳的效率,刚好一个方向做2次DFT32就可以完成了,加上一个Normalize,总共才5个Drawcall。
另外,有个性能急剧下降的点,在最大Radix为256的时候,性能产生了急剧的下降。我认为,在DFT256的情况下,显卡的SM单元里的寄存器可能用完了。根据NVIDIA的白皮书,Turing架构的显卡,一个SM单元有16384个Register,每个Register是32bit。多个线程同时在SM单元里并行运行,数据太大就产生了这种寄存器不够的情况。
总结下来,我觉得对于size为128以下的FFT,可能直接一轮就搞完最快,对于更大的,以三轮或者2轮较为合适。
我暂时没有在手机上测试过,但是很明显,最优的Radix大小 跟GPU的寄存器大小和L1 Cache的大小有很大关系,应该会有不同的结果。