2维FFT算法实现——基于GPU的基2快速二维傅里叶变换
上篇讲述了一维FFT的GPU实现（FFT算法实现——基于GPU的基2快速傅里叶变换），后来我又由于需要做了一下二维FFT，大概思路如下。
首先看的肯定是公式：
如上面公式所描述的，2维FFT只需要拆分成行FFT，和列FFT就行了，其中我在下面的实现是假设原点在F(0,0)，由于我的代码需要原点在中心，所以在最后我将原点移动到了中心。
下面是原点F(0,0)的2维FFT的伪代码：
//C2DFFT//被执行2DFFT的是一个N*N的矩阵，在source_2d中按行顺序储存//水平方向FFTfor (int i=0;i<N;i++){fft1(&source_2d[i*N],&source_2d_1[i*N],N);}//转置列成行for (int i=0;i<N*N;i++){int x = i%N;int y = i/N;int index = x*N+y;source_2d[index] = source_2d_1[i];}//垂直FFTfor(int i=0;i<N;i++){fft1(&source_2d[i*N],&source_2d_1[i*N],N);}//转置回来for (int i=0;i<N*N;i++){int x = i%N;int y = i/N;int index = x*N+y;source_2d[index] = source_2d_1[i];}
GPU实现无非把这些东西转换到GPU上。
我基于OpenGL的fragment shader来计算fft；数据都存放在纹理或者FBO里面。和1维fft不同的是，NXN的数据里面，只是对当前列或者当前排做一维FFT，所以bit反转表只需要一个1*N的buffer就可以了。对应的蝴蝶图数据也只需要1*N即可。所以我们有如下的分配：
static ofFbo _fbo_bitrev_table;static ofFbo _origin_butterfly_2d;_fbo_bitrev_table.allocate(N,1,GL_RGBA32F);_origin_butterfly_2d.allocate(N,1,GL_RGBA32F);
首先要做的是把长度为N的bit反转表求出来，这个只需要求一次，所以在最开始的时候就用CPU求出来：
for(int i=0;i<N;i++){_bitrev_index_2d.setColor(i,0,ofFloatColor(bit_rev(i,N-1),0,0,0));}_bitrev_index_2d.update();//翻转后的索引 _fbo_bitrev_table.begin();_bitrev_index_2d.draw(0,0,N,1);_fbo_bitrev_table.end();
然后初始化最初的蝴蝶图，这个和1维FFT是一样的，只是长度不同而已：
for(int i=0;i<N;i++){//初始化二维蝴蝶图if(i%2==0){_data_2d.setColor(i,0,ofFloatColor(0.f,2.f,0,i+1));}else{_data_2d.setColor(i,0,ofFloatColor(1.f,2.f,0,i-1));}}_data_2d.update();/////2D初始化///////初始化2D蝴蝶图_weight_index_2d[0].begin();_data_2d.draw(0,0,N,1);_weight_index_2d[0].end();//备份2D初始蝴蝶图,用于下一次新的计算 _origin_butterfly_2d.begin();_data_2d.draw(0,0,N,1);_origin_butterfly_2d.end();
辅助函数：
static unsigned int bit_rev(unsigned int v, unsigned int maxv){unsigned int t = log(maxv + 1)/log(2);unsigned int ret = 0;unsigned int s = 0x80000000>>(31);for (unsigned int i = 0; i < t; ++i){unsigned int r = v&(s << i);ret |= (r << (t-i-1)) >> (i); }return ret;}static void bit_reverse_copy(RBVector2 src[], RBVector2 des[], int len){for (int i = 0; i < len;i++){des[bit_rev(i, len-1)] = src[i];}}
下面定义计算2维IFFT的函数：
void GPUFFT::ifft_2d(ofFbo& in,ofFbo& out,int size);
其中in是输入，out是输出，size就是N，由初始化的时候传入了一次，在这里写是为了方便调试的时候临时改变尺寸。
Ifft本身的代码和上面形式一样，内容变成了各种shader计算：
void GPUFFT::ifft_2d(ofFbo& in,ofFbo& out,int size){//禁用Alpha混合，否则绘制到FBO会混合Alpha,造成数据丢失 ofDisableAlphaBlending();//水平FFT _weight_index_2d[_cur_2d].begin();_origin_butterfly_2d.draw(0,0,N,1);_weight_index_2d[_cur_2d].end();_fbo_in_bitreved_2d.begin();_bit_reverse_shader_2d.begin();_bit_reverse_shader_2d.setUniform3f("iResolution",N,N,0);_bit_reverse_shader_2d.setUniform1i("N",N);_bit_reverse_shader_2d.setUniform1i("dir",1);_bit_reverse_shader_2d.setUniformTexture("tex_origin",in.getTextureReference(),1);_bit_reverse_shader_2d.setUniformTexture("tex_bitreverse_table",_fbo_bitrev_table.getTextureReference(),2);ofRect(0,0,N,N);_bit_reverse_shader_2d.end();_fbo_in_bitreved_2d.end();//翻转后的数据 _res_back_2d[_cur_2d].begin();_fbo_in_bitreved_2d.draw(0,0,N,N);_res_back_2d[_cur_2d].end();for(int i = 1;i<N;i*=2){_res_back_2d[1-_cur_2d].begin();ofClear(0,0,0,0);_gpu_fft_shader_2d.begin();_gpu_fft_shader_2d.setUniform1i("size",N);_gpu_fft_shader_2d.setUniform1i("n_step",i);_gpu_fft_shader_2d.setUniform3f("iResolution",N,N,0);_gpu_fft_shader_2d.setUniform1i("dir",1);_gpu_fft_shader_2d.setUniformTexture("tex_index_weight",_weight_index_2d[_cur_2d].getTextureReference(),1);_gpu_fft_shader_2d.setUniformTexture("tex_res_back",_res_back_2d[_cur_2d].getTextureReference(),2);//_gpu_fft_shader_2d.setUniformTexture("test",imag_test.getTextureReference(),4);ofRect(0,0,N,N);_gpu_fft_shader_2d.end();_res_back_2d[1-_cur_2d].end();_weight_index_2d[1-_cur_2d].begin();ofClear(0,0,0,0);_weight_index_shader_2d.begin();_weight_index_shader_2d.setUniform1i("size",N);_weight_index_shader_2d.setUniform1i("n_step",i);_weight_index_shader_2d.setUniform3f("iResolution",N,1,0);_weight_index_shader_2d.setUniform1i("dir",1);_weight_index_shader_2d.setUniformTexture("tex_input",_weight_index_2d[_cur_2d].getTextureReference(),1);ofRect(0,0,N,1);_weight_index_shader_2d.end();_weight_index_2d[1-_cur_2d].end();_cur_2d = 1 - _cur_2d;}//for ifft_res_back_2d[1-_cur_2d].begin();_res_back_2d[_cur_2d].draw(0,0,N,N);_res_back_2d[1-_cur_2d].end();_res_back_2d[_cur_2d].begin();_ifft_div_shader_2d.begin();_ifft_div_shader_2d.setUniform1i("N",N);_ifft_div_shader_2d.setUniform3f("iResolution",N,N,0);_ifft_div_shader_2d.setUniformTexture("tex_rgb",_res_back_2d[1-_cur_2d].getTextureReference(),1);ofRect(0,0,N,N);_ifft_div_shader_2d.end();_res_back_2d[_cur_2d].end();//垂直FFT//垂直方向的所有都是计算都按照垂直方向来 _weight_index_2d[_cur_2d].begin();_origin_butterfly_2d.draw(0,0,N,1);_weight_index_2d[_cur_2d].end();//这一步不会将垂直水平化 _fbo_in_bitreved_2d.begin();_bit_reverse_shader_2d.begin();_bit_reverse_shader_2d.setUniform3f("iResolution",N,N,0);_bit_reverse_shader_2d.setUniform1i("N",N);_bit_reverse_shader_2d.setUniform1i("dir",2);_bit_reverse_shader_2d.setUniformTexture("tex_origin",_res_back_2d[_cur_2d].getTextureReference(),1);_bit_reverse_shader_2d.setUniformTexture("tex_bitreverse_table",_fbo_bitrev_table.getTextureReference(),2);ofRect(0,0,N,N);_bit_reverse_shader_2d.end();_fbo_in_bitreved_2d.end();//翻转后的数据 _res_back_2d[_cur_2d].begin();_fbo_in_bitreved_2d.draw(0,0,N,N);_res_back_2d[_cur_2d].end();for(int i = 1;i<N;i*=2){_res_back_2d[1-_cur_2d].begin();ofClear(0,0,0,0);_gpu_fft_shader_2d.begin();_gpu_fft_shader_2d.setUniform1i("size",N);_gpu_fft_shader_2d.setUniform1i("n_step",i);_gpu_fft_shader_2d.setUniform3f("iResolution",N,N,0);_gpu_fft_shader_2d.setUniform1i("dir",2);_gpu_fft_shader_2d.setUniformTexture("tex_index_weight",_weight_index_2d[_cur_2d].getTextureReference(),1);_gpu_fft_shader_2d.setUniformTexture("tex_res_back",_res_back_2d[_cur_2d].getTextureReference(),2);//_gpu_fft_shader_2d.setUniformTexture("test",imag_test.getTextureReference(),4);ofRect(0,0,N,N);_gpu_fft_shader_2d.end();_res_back_2d[1-_cur_2d].end();_weight_index_2d[1-_cur_2d].begin();ofClear(0,0,0,0);_weight_index_shader_2d.begin();_weight_index_shader_2d.setUniform1i("size",N);_weight_index_shader_2d.setUniform1i("n_step",i);_weight_index_shader_2d.setUniform3f("iResolution",N,1,0);_weight_index_shader_2d.setUniform1i("dir",2);_weight_index_shader_2d.setUniformTexture("tex_input",_weight_index_2d[_cur_2d].getTextureReference(),1);ofRect(0,0,N,1);_weight_index_shader_2d.end();_weight_index_2d[1-_cur_2d].end();_cur_2d = 1 - _cur_2d;}//for ifft_res_back_2d[1-_cur_2d].begin();_res_back_2d[_cur_2d].draw(0,0,N,N);_res_back_2d[1-_cur_2d].end();_res_back_2d[_cur_2d].begin();_ifft_div_shader_2d.begin();_ifft_div_shader_2d.setUniform1i("N",N);_ifft_div_shader_2d.setUniform3f("iResolution",N,N,0);_ifft_div_shader_2d.setUniformTexture("tex_rgb",_res_back_2d[1-_cur_2d].getTextureReference(),1);ofRect(0,0,N,N);_ifft_div_shader_2d.end();_res_back_2d[_cur_2d].end();out.begin();_res_back_2d[_cur_2d].draw(0,0,N,N);out.end();//恢复Alpha混合//ofEnableAlphaBlending();}
现在来看shader内容：
_bit_reverse_shader_2d
这个shader用于将整个N*N的数据全部按照行或者按照列进行翻装，使之满足执行fft的条件：
#version 130uniform sampler2D tex_origin;//1xN查找表，用于查找索引对应的bitreverse数uniform sampler2D tex_bitreverse_table;//1 for x direction,2 for y directionuniform int dir;uniform int N;uniform vec3 iResolution;out vec4 outColor;void main() { vec2 tex_coord = gl_FragCoord.xy/iResolution.xy;vec2 table_index;table_index.y = 0.5;if(dir==1)table_index.x = tex_coord.x;elsetable_index.x = tex_coord.y;float bitreverse = texture(tex_bitreverse_table,table_index).r;vec2 origin_index;if(dir==1){//x方向origin_index.y = tex_coord.y;origin_index.x = (bitreverse+0.5)/N;}else{//y方向origin_index.x = tex_coord.x;origin_index.y = (bitreverse+0.5)/N;}vec2 param = texture(tex_origin,origin_index).xy;outColor = vec4(param,0,1);}
_gpu_fft_shader_2d
这是fft执行计算的部分，同样分为按行和按列：
#version 130//NX1uniform sampler2D tex_index_weight;//NXNuniform sampler2D tex_res_back;uniform sampler2D test;uniform int size;uniform int n_step;//1 for x direction,2 for y directionuniform int dir;uniform vec3 iResolution;out vec4 outColor;void main() { vec2 tex_coord = gl_FragCoord.xy/iResolution.xy;vec2 first_index;if(dir==1){first_index.y = 0.5;first_index.x = tex_coord.x;}else{first_index.y = 0.5;first_index.x = tex_coord.y;}float cur_x = gl_FragCoord.x - 0.5;float cur_y = gl_FragCoord.y - 0.5;vec2 outv;vec4 temp = texture(tex_index_weight,first_index);//ifftvec2 weight = vec2(cos(temp.r/temp.g*2*3.141592653),-sin(temp.r/temp.g*2*3.141592653));//fft//vec2 weight = vec2(cos(temp.r/temp.g*2*3.141592653),sin(temp.r/temp.g*2*3.141592653)); vec2 _param2_index;if(dir==1){_param2_index.x = (temp.a + 0.5)/size;_param2_index.y = tex_coord.y;}else{_param2_index.x = tex_coord.x;_param2_index.y = (temp.a + 0.5)/size;}vec2 param1 = texture(tex_res_back,tex_coord).rg;vec2 param2 = texture(tex_res_back,_param2_index).rg;float tex_coord_n1;float tex_coord_n2;if(dir==1){tex_coord_n1 = cur_x;}else{tex_coord_n1 = cur_y;}tex_coord_n2 = temp.a;if(tex_coord_n1<tex_coord_n2){outv.r = param1.r + param2.r*weight.r-weight.g*param2.g;outv.g = param1.g +weight.r*param2.g + weight.g*param2.r;}else{outv.r = param2.r + param1.r*weight.r-weight.g*param1.g;outv.g = param2.g +weight.r*param1.g + weight.g*param1.r;}outColor = vec4(outv,0,1);}
_weight_index_shader_2d
更新蝴蝶图索引：
#version 130uniform sampler2D tex_input;uniform int size;uniform int n_total;//start with 2uniform int n_step;//1 for x direction,2 for y directionuniform int dir;uniform vec3 iResolution;out vec4 outColor;void main() {vec2 tex_coord = gl_FragCoord.xy/iResolution.xy;vec4 fetch = texture(tex_input,tex_coord);float cur_x = gl_FragCoord.x - 0.5;float cur_y = gl_FragCoord.y - 0.5;vec4 outv;float tex_coord_n;if(dir==1){//x dirtex_coord_n = cur_x;}else{//y dirtex_coord_n = cur_x;}//updata weightvec2 pre_w = fetch.rg;float i = pre_w.r;float n = pre_w.g;float new_i;float new_n;new_i = i;new_n = n*2;if(int(tex_coord_n)%(n_step*4) > n_step*2-1){new_i += n_step*2;}outv.r = new_i;outv.g = new_n;//outv.rg = tex_coord;//updata indexvec2 pre_index = fetch.ba;int x = int(pre_index.x);int y = int(pre_index.y);int ni = n_step*2;float new_tex_coord_n = tex_coord_n;if((int(tex_coord_n)/ni)%2==0){new_tex_coord_n += ni;}else{new_tex_coord_n -= ni;}outv.b = 0;outv.a = new_tex_coord_n;outColor = outv;//outColor = vec4(tex_coord_n,tex_coord_n%n_step,tex_coord_n%n_step,tex_coord_n%n_step);}
最后的
_ifft_div_shader_2d
是为了计算ifft，将每个计算结果除以一个N：
#version 130uniform sampler2D tex_rgb;uniform int N;uniform vec3 iResolution;out vec4 outColor;void main() { vec2 tex_coord = gl_FragCoord.xy/iResolution.xy;vec2 outv;vec4 c = texture(tex_rgb,tex_coord);outv.r = c.r/N;outv.g = c.g/N;outColor = vec4(outv,0,1);}
最后，out里面就是结果了。
对于将原点移动到中心多了以下shader:
vec4 c;if(tex_coord.x>0.5&&tex_coord.y>0.5){c = texture(tex_rgb,tex_coord-vec2(0.5,0.5));}if(tex_coord.x>0.5&&tex_coord.y<0.5){c = texture(tex_rgb,tex_coord+vec2(-0.5,0.5));}if(tex_coord.x<0.5&&tex_coord.y>0.5){c = texture(tex_rgb,tex_coord+vec2(0.5,-0.5));}if(tex_coord.x<0.5&&tex_coord.y<0.5){c = texture(tex_rgb,tex_coord+vec2(0.5,0.5));}outColor = c;
