C++ Fast Fourier transform attempt 2

Question

This is a follow up for C++ Fast Fourier transform . I did end up changing the algorithm slightly to reduce the memory usage.

#include <vector>
#include <iostream>
#include <complex>
#include <cmath>
#include <algorithm>

#define N 1048576 // Problem size, must be a power of 2
#define PI 3.14159265358979323846

typedef std::vector<std::complex<double>> complexVec;

/*
Calculates the N'th roots of unity.
*/
complexVec rootsOfUnityCalculator() {
    complexVec table;
    table.resize(N);

    for (size_t k = 0; k < N; k++) {
        table[k] = std::complex<double>(std::cos(-2.0 * PI * k / N), std::sin(-2.0 * PI * k / N));
    }

    return table;
}

/*
Permutes an input vector of size length - a power of 2 - according
to the bit reversal permutation.

TESTED: SUCCESS
*/
void bitReversal(complexVec& input) {
    std::vector<size_t> permutation;
    permutation.resize(N);

    permutation[0] = 0;
    size_t k = 2;

    // Calculating the permutation according to recursive relation
    while (k <= N) {
        for (size_t j = 0; j < k / 2; j++) {
            permutation[j] *= 2; // permutation_k(j) = 2 * permutation_(k/2)(j)
            permutation[j + k / 2] = permutation[j] + 1; // permutation_k(j + k/2) = 2 * permutation_(k/2)(j) + 1
        }
        k *= 2;
    }

    // Permute the input
    for (size_t j = 0; j < N; j++) {
        if (j < permutation[j]) {
            std::swap(input[j], input[ permutation[j] ]);
        }
    }
}

void unorderedFFT(complexVec& input, const complexVec& table) {
    size_t k = 2;

    while (k <= N) {
        for (size_t r = 0; r < N / k; r++) {
            for (size_t j = 0; j < k / 2; j++) {
                std::complex<double> plusMinus = input[r * k + j + k / 2] * table[N / k * j]; // omega_k^j = (omega_N^(N/k))^j = omega_N^(Nj/k)
                input[r * k + j + k / 2] = input[r * k + j] - plusMinus;
                input[r * k + j] += plusMinus;
            }
        }

        k *= 2;
    }
}

/*
Calculates the FFT

TESTED: SUCCESS
*/
void FFT(complexVec& input, const complexVec& table) {
    bitReversal(input);
    unorderedFFT(input, table);
}

// Prints array of N complex numbers
void printArray(complexVec& array) {
    for (size_t k = 0; k < N - 1; k++) {
        std::cout << array[k] << ", ";
    }

    std::cout << array[N - 1] << std::endl;
}

int main() {
    complexVec input;
    input.resize(N);

    // Initialize the input to all 0, 1, ..., N
    for (size_t k = 0; k < N; k++) {
        input[k] = std::complex<double>(k, 0.0);
    }

    const complexVec table = rootsOfUnityCalculator();

    FFT(input, table);
}

Answer 1

Profiling results

Using the profiler in Visual Studio, on my machine (Intel Haswell, 4770K) I get a breakdown roughly like this:

• unorderedFFT 107ms
• rootsOfUnityCalculator 21ms
• bitReversal 17ms

So let's improve all of those.

Roots of unity

The N'th roots of unity in the complex plane is a set of equally spaced points on the unit circle. Their coordinates can be calculated using sine and cosine, but an alternative is starting at 1 + 0i and successively applying a small vector rotation to it (alternative view: the first non-trivial root is calculated, and then iteratively raised to the powers of 0, 1, 2 .. N-1). That way, only one pair of sin/cos is necessary (by the way, std::polar can be used to replace such a sin/cos pair with simpler code). A downside is that it may be less precise. It tends to be OK, but it will perturb the results slightly.

For example, like this:

complexVec rootsOfUnityCalculator() {
    complexVec table;
    table.resize(N);

    auto step = std::polar(1.0, -2.0 * PI / N);
    std::complex<double> root = 1.0;
    for (std::size_t k = 0; k < N; k++) {
        table[k] = root;
        root *= step;
    }

    return table;
}

With this, the time it takes has gone down to 7.5ms, so a factor of 2.3, which is certainly significant. This could be done even faster than that by unrolling the loop and removing the dependency between the copies of the loop body, for example if unrolling by 2 then there would be a separate root1 and root2 that are each multiplied by stepSquared so that their calculations are independent (improving instruction-level parallelism). Doing that reduced the time to 6ms.

Bit-reversal permutation

This is a well-studied permutation, because it is of practical interest, namely this exact application. Various more efficient techniques are known, including some variants of "increment a reversed integer". If an integer i and its reverse rev are both available, calculating the next values for both of them is efficiently possible.

Of course i can just be incremented. rev needs some more work, but there is hope: the XOR of an integer and its successor is a "nice mask" consisting of a contiguous range of ones and a contiguous range of zeroes, with no "mess". This results from how, in the increment, the carry goes through the trailing ones, flipping each of them, then finally the first zero that is found is also flipped and the carry is "absorbed" by that zero - so the bits the flipped start at the least significant bit and form a contiguous range. That is relevant because a "nice mask" can be reversed by shifting it, and then XORing rev by it results in a "bit-reversed increment". A variant of that trick is using a trailing-zero-count instruction to find how long the contiguous range of flipped bits is, and then basing a mask directly on that. Let's go with that:

void bitReversal(complexVec& input) {
    std::size_t maskLen = std::countr_zero(unsigned(N));

    // Permute the input
    for (std::size_t j = 0, rev = 0; j < N; j++) {
        if (j < rev)
            std::swap(input[j], input[rev]);
        std::size_t maskLen = std::countr_zero(j + 1) + 1;
        rev ^= N - (N >> maskLen);
    }
}

This uses the very new std::countr_zero from the <bit> header, a C++20 feature. There are other ways to write that too, for example _tzcnt_u64. unsigned(N) is a bit unfortunate, but required for countr_zero, which refuses to work on signed types. Preferably this would be avoided by making N a constant of type size_t instead of #define-ing it, which I recommend anyway.

Using this trick, bitReversal has gone down from taking 17ms to 14ms (additionally, the memory used by the permutation vector is saved). It's a bit better, but nothing as great as improving the calculation of the roots of unity.

This is not the best way to do it. The if corresponds to a badly-predicted branch, and the memory access pattern is semi-random. There are various useful papers about the technique I used here and further improvements, for example practically efficient methods for performing bit-reversed permutation in C++11 on the x86-64 architecture.

The actual FFT

The real meat of the algorithm. Since you asked in the previous question not to bother with suggestions for different algorithms, I won't. However, I can still suggest a performance improvement: use SSE3. SSE2 is actually sufficient, but SSE3 adds the ADDSUBPD instruction which is handy for complex multiplication:

__m128d multiplyComplex(__m128d A, __m128d B)
{
    __m128d ARealReal = _mm_shuffle_pd(A, A, 0);
    __m128d AImagImag = _mm_shuffle_pd(A, A, 3);
    __m128d BRealImag = B;
    __m128d BImagReal = _mm_shuffle_pd(B, B, 1);
    return _mm_addsub_pd(_mm_mul_pd(ARealReal, BRealImag), _mm_mul_pd(AImagImag, BImagReal));
}

Which could be used like this:

void unorderedFFT2(complexVec& input, const complexVec& table) {
    std::size_t k = 2;

    while (k <= N) {
        for (std::size_t r = 0; r < N / k; r++) {
            for (std::size_t j = 0; j < k / 2; j++) {
                __m128d input0 = _mm_loadu_pd((double*)&input[r * k + j + k / 2]);
                __m128d input1 = _mm_loadu_pd((double*)&input[r * k + j]);
                __m128d twiddle = _mm_loadu_pd((double*)&table[N / k * j]);
                __m128d product = multiplyComplex(input0, twiddle);
                _mm_storeu_pd((double*)&input[r * k + j + k / 2], _mm_sub_pd(input1, product));
                _mm_storeu_pd((double*)&input[r * k + j], _mm_add_pd(input1, product));
            }
        }

        k *= 2;
    }
}

Casting those pointers looks scary, but it seems to be allowed according to the paragraph about Array Oriented Access of std::complex. reinterpret_cast could be used if you dislike C-style casts.

Admittedly, using SSE is not very beginner-friendly. But it does help significantly, the time of unorderedFFT went down to 70ms, the biggest improvement in absolute terms.

Using intrinsic functions like this requires including the relevant header, for example #include <tmmintrin.h> (or newer) in this case. Depending on the compiler, it may also require passing certain command line options (GCC and Clang need that, MSVC does not).