Reference r2 FFT code

gistfile1.txt

static size_t const kMaxN = 2048; // This is the largest straight FFT we support

struct complexf
{
    float re;
    float im;

    complexf() {}
    complexf(float r) : re(r), im(0.0f) {}
    complexf(float r, float i) : re(r), im(i) {}
};

complexf operator + (complexf const a, complexf const b) { return complexf(a.re + b.re, a.im + b.im); }
complexf operator - (complexf const a, complexf const b) { return complexf(a.re - b.re, a.im - b.im); }
complexf operator * (complexf const a, complexf const b) { return complexf(a.re*b.re - a.im*b.im, a.re*b.im + a.im*b.re); }
complexf &operator += (complexf &a, complexf const b) { a.re += b.re; a.im += b.im; return a; }
complexf &operator -= (complexf &a, complexf const b) { a.re -= b.re; a.im -= b.im; return a; }
complexf &operator *= (complexf &a, complexf const b) { a = a * b; return a; }

static complexf s_twiddles[(kMaxN / 4) * 2]; // FFT twiddles for reference FFT

// Reference radix-2 FFT (DIT) - recursion
static void r2_fft_rec(complexf out[], complexf const in[], size_t offs, size_t mask, size_t stride, size_t N)
{
    if (N == 1)
        out[0] = in[offs & mask];
    else if (N == 2)
    {
        complexf a = in[offs & mask];
        complexf b = in[(offs + stride) & mask];
        out[0] = a + b;
        out[1] = a - b;
    }
    else
    {
        r2_fft_rec(out + 0*N/2, in, offs + 0*stride, mask, stride * 2, N / 2);
        r2_fft_rec(out + 1*N/2, in, offs + 1*stride, mask, stride * 2, N / 2);

        // twiddles[N:2N] holds the N complex twiddles for a quarter-circle
        //   twiddles[N+k] = cexp(-2pi * i * k / 4N)
        //
        // we want to use the same quarter-circle symmetry (to reduce number
        // of twiddle factors), so even though we're doing a radix-2 step,
        // loop only up to N/4 and unroll twice (effectively).

        complexf const *twiddle = s_twiddles + N/4;
        for (size_t k = 0; k < N/4; k++)
        {
            complexf Ak  = out[k + 0*N/4];
            complexf Bk  = out[k + 1*N/4];
            complexf Ck  = out[k + 2*N/4];
            complexf Dk  = out[k + 3*N/4];
            complexf const &w = twiddle[k];

            // Twiddle
            Ck *= w; // 2 mul, 2 FMA (or 4 mul 2 add)
            Dk *= complexf(0.0, -1.0f)*w; // 2 mul, 2 FMA (or 4 mul 2 add)

            // Output butterfly
            out[k + 0*N/4] = Ak + Ck; // 2 add
            out[k + 1*N/4] = Bk + Dk; // 2 add
            out[k + 2*N/4] = Ak - Ck; // 2 add
            out[k + 3*N/4] = Bk - Dk; // 2 add
        }
    }
}

// Call this
void r2_fft(complexf out[], complexf const in[], size_t N)
{
    r2_fft_rec(out, in, 0, N - 1, 1, N);
}

void init()
{
    // init our twiddles
    for (size_t N = 1; N <= kMaxN / 4; N *= 2)
    {
        double const kPi = 3.1415926535897932384626433832795;
        double step = -2.0 * kPi / (double)(N * 4);
        complexf *twiddle = s_twiddles + N;

        for (size_t k = 0; k < N; k++)
        {
            double phase = step * (double)k;
            twiddle[k] = complexf((float)cos(phase), (float)sin(phase));
        }
    }
}