2维FFT算法实现——基于GPU的基2快速二维傅里叶变换

时间：2023-03-09 19:47:48
2维FFT算法实现——基于GPU的基2快速二维傅里叶变换

上篇讲述了一维FFT的GPU实现（FFT算法实现——基于GPU的基2快速傅里叶变换），后来我又由于需要做了一下二维FFT，大概思路如下。

首先看的肯定是公式：

2维FFT算法实现——基于GPU的基2快速二维傅里叶变换

如上面公式所描述的，2维FFT只需要拆分成行FFT，和列FFT就行了，其中我在下面的实现是假设原点在F(0,0)，由于我的代码需要原点在中心，所以在最后我将原点移动到了中心。

下面是原点F(0,0)的2维FFT的伪代码：

    //C2DFFT

    //被执行2DFFT的是一个N*N的矩阵，在source_2d中按行顺序储存

    //水平方向FFT

    for (int i=;i<N;i++)

    {

        fft1(&source_2d[i*N],&source_2d_1[i*N],N);

    }

    //转置列成行

    for (int i=;i<N*N;i++)

    {

        int x = i%N;

        int y = i/N;

        int index = x*N+y;

        source_2d[index] = source_2d_1[i];

    }

    //垂直FFT

    for(int i=;i<N;i++)

    {

        fft1(&source_2d[i*N],&source_2d_1[i*N],N);

    }

    //转置回来

    for (int i=;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,,GL_RGBA32F);

_origin_butterfly_2d.allocate(N,,GL_RGBA32F);

首先要做的是把长度为N的bit反转表求出来，这个只需要求一次，所以在最开始的时候就用CPU求出来：

    for(int i=;i<N;i++)

    {

          _bitrev_index_2d.setColor(i,,ofFloatColor(bit_rev(i,N-),,,));

    }

    _bitrev_index_2d.update();

    //翻转后的索引

    _fbo_bitrev_table.begin();

    _bitrev_index_2d.draw(,,N,);

    _fbo_bitrev_table.end();

然后初始化最初的蝴蝶图，这个和1维FFT是一样的，只是长度不同而已：

for(int i=;i<N;i++)

    {

        //初始化二维蝴蝶图

        if(i%==)

        {

            _data_2d.setColor(i,,ofFloatColor(.f,.f,,i+));

        }

        else

        {

            _data_2d.setColor(i,,ofFloatColor(.f,.f,,i-));

        }

    }

    _data_2d.update();

    /////////////////2D初始化/////////////////

    //初始化2D蝴蝶图

    _weight_index_2d[].begin();

    _data_2d.draw(,,N,);

    _weight_index_2d[].end();

    //备份2D初始蝴蝶图,用于下一次新的计算

    _origin_butterfly_2d.begin();

    _data_2d.draw(,,N,);

    _origin_butterfly_2d.end();

辅助函数：

    static unsigned int bit_rev(unsigned int v, unsigned int maxv)

    {

        unsigned int t = log(maxv + )/log();

        unsigned int ret = ;

        unsigned int s = 0x80000000>>();

        for (unsigned int i = ; i < t; ++i)

        {

            unsigned int r = v&(s << i);

            ret |= (r << (t-i-)) >> (i);

        }

        return ret;

    }

    static void bit_reverse_copy(RBVector2 src[], RBVector2 des[], int len)

    {

        for (int i = ; i < len;i++)

        {

            des[bit_rev(i, len-)] = 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(,,N,);

    _weight_index_2d[_cur_2d].end();

    _fbo_in_bitreved_2d.begin();

    _bit_reverse_shader_2d.begin();

    _bit_reverse_shader_2d.setUniform3f("iResolution",N,N,);

    _bit_reverse_shader_2d.setUniform1i("N",N);

    _bit_reverse_shader_2d.setUniform1i("dir",);

    _bit_reverse_shader_2d.setUniformTexture("tex_origin",in.getTextureReference(),);

    _bit_reverse_shader_2d.setUniformTexture("tex_bitreverse_table",_fbo_bitrev_table.getTextureReference(),);

    ofRect(,,N,N);

    _bit_reverse_shader_2d.end();

    _fbo_in_bitreved_2d.end();

    //翻转后的数据

    _res_back_2d[_cur_2d].begin();

    _fbo_in_bitreved_2d.draw(,,N,N);

    _res_back_2d[_cur_2d].end();

    for(int i = ;i<N;i*=)

    {

        _res_back_2d[-_cur_2d].begin();

        ofClear(,,,);

        _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,);

        _gpu_fft_shader_2d.setUniform1i("dir",);

        _gpu_fft_shader_2d.setUniformTexture("tex_index_weight",_weight_index_2d[_cur_2d].getTextureReference(),);

        _gpu_fft_shader_2d.setUniformTexture("tex_res_back",_res_back_2d[_cur_2d].getTextureReference(),);

        //_gpu_fft_shader_2d.setUniformTexture("test",imag_test.getTextureReference(),4);

        ofRect(,,N,N);

        _gpu_fft_shader_2d.end();

        _res_back_2d[-_cur_2d].end();

        _weight_index_2d[-_cur_2d].begin();

        ofClear(,,,);

        _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,,);

        _weight_index_shader_2d.setUniform1i("dir",);

        _weight_index_shader_2d.setUniformTexture("tex_input",_weight_index_2d[_cur_2d].getTextureReference(),);

        ofRect(,,N,);

        _weight_index_shader_2d.end();

        _weight_index_2d[-_cur_2d].end();

        _cur_2d =  - _cur_2d;

    }

    //for ifft

    _res_back_2d[-_cur_2d].begin();

    _res_back_2d[_cur_2d].draw(,,N,N);

    _res_back_2d[-_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,);

    _ifft_div_shader_2d.setUniformTexture("tex_rgb",_res_back_2d[-_cur_2d].getTextureReference(),);

    ofRect(,,N,N);

    _ifft_div_shader_2d.end();

    _res_back_2d[_cur_2d].end();

    //垂直FFT

    //垂直方向的所有都是计算都按照垂直方向来

    _weight_index_2d[_cur_2d].begin();

    _origin_butterfly_2d.draw(,,N,);

    _weight_index_2d[_cur_2d].end();

    //这一步不会将垂直水平化

    _fbo_in_bitreved_2d.begin();

    _bit_reverse_shader_2d.begin();

    _bit_reverse_shader_2d.setUniform3f("iResolution",N,N,);

    _bit_reverse_shader_2d.setUniform1i("N",N);

    _bit_reverse_shader_2d.setUniform1i("dir",);

    _bit_reverse_shader_2d.setUniformTexture("tex_origin",_res_back_2d[_cur_2d].getTextureReference(),);

    _bit_reverse_shader_2d.setUniformTexture("tex_bitreverse_table",_fbo_bitrev_table.getTextureReference(),);

    ofRect(,,N,N);

    _bit_reverse_shader_2d.end();

    _fbo_in_bitreved_2d.end();

    //翻转后的数据

    _res_back_2d[_cur_2d].begin();

    _fbo_in_bitreved_2d.draw(,,N,N);

    _res_back_2d[_cur_2d].end();

    for(int i = ;i<N;i*=)

    {

        _res_back_2d[-_cur_2d].begin();

        ofClear(,,,);

        _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,);

        _gpu_fft_shader_2d.setUniform1i("dir",);

        _gpu_fft_shader_2d.setUniformTexture("tex_index_weight",_weight_index_2d[_cur_2d].getTextureReference(),);

        _gpu_fft_shader_2d.setUniformTexture("tex_res_back",_res_back_2d[_cur_2d].getTextureReference(),);

        //_gpu_fft_shader_2d.setUniformTexture("test",imag_test.getTextureReference(),4);

        ofRect(,,N,N);

        _gpu_fft_shader_2d.end();

        _res_back_2d[-_cur_2d].end();

        _weight_index_2d[-_cur_2d].begin();

        ofClear(,,,);

        _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,,);

        _weight_index_shader_2d.setUniform1i("dir",);

        _weight_index_shader_2d.setUniformTexture("tex_input",_weight_index_2d[_cur_2d].getTextureReference(),);

        ofRect(,,N,);

        _weight_index_shader_2d.end();

        _weight_index_2d[-_cur_2d].end();

        _cur_2d =  - _cur_2d;

    }

    //for ifft

    _res_back_2d[-_cur_2d].begin();

    _res_back_2d[_cur_2d].draw(,,N,N);

    _res_back_2d[-_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,);

    _ifft_div_shader_2d.setUniformTexture("tex_rgb",_res_back_2d[-_cur_2d].getTextureReference(),);

    ofRect(,,N,N);

    _ifft_div_shader_2d.end();

    _res_back_2d[_cur_2d].end();

    out.begin();

    _res_back_2d[_cur_2d].draw(,,N,N);

    out.end();

    //恢复Alpha混合

    //ofEnableAlphaBlending();

}

现在来看shader内容：

_bit_reverse_shader_2d

这个shader用于将整个N*N的数据全部按照行或者按照列进行翻装，使之满足执行fft的条件：

#version

uniform sampler2D tex_origin;

//1xN查找表，用于查找索引对应的bitreverse数

uniform sampler2D tex_bitreverse_table;

//1 for x direction,2 for y direction

uniform 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==)

        table_index.x = tex_coord.x;

    else

        table_index.x = tex_coord.y;

    float bitreverse = texture(tex_bitreverse_table,table_index).r;

    vec2 origin_index;

    if(dir==)

    {

        //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,,);

}

_gpu_fft_shader_2d

这是fft执行计算的部分，同样分为按行和按列：

#version

//NX1

uniform sampler2D tex_index_weight;

//NXN

uniform sampler2D tex_res_back;

uniform sampler2D test;

uniform int size;

uniform int n_step;

//1 for x direction,2 for y direction

uniform int dir;

uniform vec3 iResolution;

out vec4 outColor;

void main()

{

    vec2 tex_coord = gl_FragCoord.xy/iResolution.xy;

    vec2 first_index;

    if(dir==)

    {

        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);

    //ifft

    vec2 weight = vec2(cos(temp.r/temp.g**3.141592653),-sin(temp.r/temp.g**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==)

    {

        _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==)

    {

        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,,);

}

_weight_index_shader_2d

更新蝴蝶图索引：

#version 

uniform sampler2D tex_input;

uniform int size;

uniform int n_total;

//start with 2

uniform int n_step;

//1 for x direction,2 for y direction

uniform 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==)

    {

        //x dir

        tex_coord_n = cur_x;

    }

    else

    {

        //y dir

        tex_coord_n = cur_x;

    }

    //updata weight

    vec2 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*;

    if(int(tex_coord_n)%(n_step*) > n_step*-)

    {

        new_i += n_step*;

    }

    outv.r = new_i;

    outv.g = new_n;

    //outv.rg = tex_coord;

    //updata index

    vec2 pre_index = fetch.ba;

    int x = int(pre_index.x);

    int y = int(pre_index.y);

    int ni = n_step*;

    float new_tex_coord_n = tex_coord_n;

    if((int(tex_coord_n)/ni)%==)

    {

        new_tex_coord_n += ni;

    }

    else

    {

        new_tex_coord_n -= ni;

    }

    outv.b = ;

    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

uniform 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,,);

}

最后，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;