#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <ctime>
#include <immintrin.h>
struct C {
double x, y;
C() : x(0), y(0) {}
C(double x, double y) : x(x), y(y) {}
inline C operator+ (const C &c) const { return C(x + c.x, y + c.y); }
inline C operator- (const C &c) const { return C(x - c.x, y - c.y); }
inline C operator* (const C &c) const { return C(x * c.x - y * c.y, x * c.y + y * c.x); }
inline C conj() const { return C(x, -y); }
};
constexpr double PI = acos(-1);
constexpr int SIZE = 100000, B = 32 - __builtin_clz(SIZE * 2 - 1), N = 1 << B;
alignas(32) C w[N >> 2], f0[N], f1[N];
void init_roots() {
w[0] = C(1, 0);
for (int len = 1; len < (N >> 2); len <<= 1) {
C wn(cos(PI / (len << 2)), sin(PI / (len << 2)));
for (int i = len; i < len << 1; i++)
w[i] = w[i - len] * wn;
}
}
__m256d avx_complex_mul(__m256d a, __m256d b) {
__m256d ax = _mm256_unpacklo_pd(a, a);
__m256d ay = _mm256_unpackhi_pd(a, a);
__m256d bs = _mm256_shuffle_pd(b, b, 5);
return _mm256_fmaddsub_pd(ax, b, _mm256_mul_pd(ay, bs));
}
void fft(C *f) {
for (int len = N >> 2; len > 1; len >>= 2) {
for (int i = 0, m = 0; i < N; i += len << 2, m++) {
double w1x = w[m].x, w1y = w[m].y;
double w2x = w1x * w1x - w1y * w1y, w2y = 2 * w1x * w1y;
double w3x = w1x * w2x - w1y * w2y, w3y = w1x * w2y + w1y * w2x;
__m256d w1 = _mm256_setr_pd(w1x, w1y, w1x, w1y);
__m256d w2 = _mm256_setr_pd(w2x, w2y, w2x, w2y);
__m256d w3 = _mm256_setr_pd(w3x, w3y, w3x, w3y);
constexpr __m256d posneg = { 0.0, -0.0, 0.0, -0.0 };
for (int j = i; j < i + len; j += 2) {
__m256d c0 = _mm256_load_pd(&f[j + len * 0].x);
__m256d c1 = avx_complex_mul(_mm256_load_pd(&f[j + len * 1].x), w1);
__m256d c2 = avx_complex_mul(_mm256_load_pd(&f[j + len * 2].x), w2);
__m256d c3 = avx_complex_mul(_mm256_load_pd(&f[j + len * 3].x), w3);
__m256d a02 = _mm256_add_pd(c0, c2);
__m256d a13 = _mm256_add_pd(c1, c3);
__m256d s02 = _mm256_sub_pd(c0, c2);
__m256d s13 = _mm256_xor_pd(posneg, _mm256_sub_pd(c1, c3));
s13 = _mm256_shuffle_pd(s13, s13, 5);
_mm256_store_pd(&f[j + len * 0].x, _mm256_add_pd(a02, a13));
_mm256_store_pd(&f[j + len * 1].x, _mm256_sub_pd(a02, a13));
_mm256_store_pd(&f[j + len * 2].x, _mm256_add_pd(s02, s13));
_mm256_store_pd(&f[j + len * 3].x, _mm256_sub_pd(s02, s13));
}
}
}
for (int i = 0, m = 0; i < N; i += 4, m++) {
C w1 = w[m], w2 = w1 * w1, w3 = w1 * w2;
C c0 = f[i + 0], c1 = f[i + 1] * w1,
c2 = f[i + 2] * w2, c3 = f[i + 3] * w3;
C a02 = c0 + c2, a13 = c1 + c3,
s02 = c0 - c2, s13 = (c1 - c3) * C(0, 1);
f[i + 0] = a02 + a13, f[i + 1] = a02 - a13;
f[i + 2] = s02 + s13, f[i + 3] = s02 - s13;
}
}
void ifft(C *f) {
for (int i = 0, m = 0; i < N; i += 4, m++) {
C w1 = w[m], w2 = w1 * w1, w3 = w1 * w2;
C c0 = f[i + 0], c1 = f[i + 1],
c2 = f[i + 2], c3 = f[i + 3];
C a01 = c0 + c1, a23 = c2 + c3,
s01 = c0 - c1, s23 = (c2 - c3) * C(0, 1);
f[i + 0] = a01 + a23, f[i + 1] = (s01 + s23) * w1;
f[i + 2] = (a01 - a23) * w2, f[i + 3] = (s01 - s23) * w3;
}
for (int len = 4; len < N; len <<= 2) {
for (int i = 0, m = 0; i < N; i += len << 2, m++) {
double w1x = w[m].x, w1y = w[m].y;
double w2x = w1x * w1x - w1y * w1y, w2y = 2 * w1x * w1y;
double w3x = w1x * w2x - w1y * w2y, w3y = w1x * w2y + w1y * w2x;
__m256d w1 = _mm256_setr_pd(w1x, w1y, w1x, w1y);
__m256d w2 = _mm256_setr_pd(w2x, w2y, w2x, w2y);
__m256d w3 = _mm256_setr_pd(w3x, w3y, w3x, w3y);
constexpr __m256d posneg = { 0.0, -0.0, 0.0, -0.0 };
for (int j = i; j < i + len; j += 2) {
__m256d c0 = _mm256_load_pd(&f[j + len * 0].x);
__m256d c1 = _mm256_load_pd(&f[j + len * 1].x);
__m256d c2 = _mm256_load_pd(&f[j + len * 2].x);
__m256d c3 = _mm256_load_pd(&f[j + len * 3].x);
__m256d a01 = _mm256_add_pd(c0, c1);
__m256d a23 = _mm256_add_pd(c2, c3);
__m256d s01 = _mm256_sub_pd(c0, c1);
__m256d s23 = _mm256_xor_pd(posneg, _mm256_sub_pd(c2, c3));
s23 = _mm256_shuffle_pd(s23, s23, 5);
_mm256_store_pd(&f[j + len * 0].x, _mm256_add_pd(a01, a23));
_mm256_store_pd(&f[j + len * 1].x, avx_complex_mul(_mm256_add_pd(s01, s23), w1));
_mm256_store_pd(&f[j + len * 2].x, avx_complex_mul(_mm256_sub_pd(a01, a23), w2));
_mm256_store_pd(&f[j + len * 3].x, avx_complex_mul(_mm256_sub_pd(s01, s23), w3));
}
}
}
}
void convolve(uint64_t *a) {
init_roots();
for (int i = 0; i < SIZE; i++)
f0[i] = C(a[i] & 4095, a[i] >> 12);
fft(f0);
double t0 = 4 * f0[0].x * f0[0].y, t1 = 4 * f0[1].x * f0[1].y;
f1[0].x = 4 * f0[0].y * f0[0].y, f1[0].y = 0;
f1[1].x = 4 * f0[1].y * f0[1].y, f1[1].y = 0;
f0[0].x = 4 * f0[0].x * f0[0].x, f0[0].y = t0;
f0[1].x = 4 * f0[1].x * f0[1].x, f0[1].y = t1;
for (int i = 2, msk = 0; i < N; i += 2) {
msk |= i >> 1;
int j = i ^ msk;
C c0 = f0[i] + f0[j].conj();
C c1 = (f0[i] - f0[j].conj()) * C(0, -1);
C c00 = c0 * c0, c01 = c0 * c1, c11 = c1 * c1;
f0[i] = c00.conj() + c01.conj() * C(0, 1), f0[j] = c00 + c01 * C(0, 1);
f1[i] = c11.conj(), f1[j] = c11;
}
ifft(f0); ifft(f1);
for (int i = 0; i < SIZE * 2 - 1; i++) {
uint64_t c0 = (uint64_t) (f0[i].x / (N << 2) + 0.5);
uint64_t c1 = (uint64_t) (f0[i].y / (N << 2) + 0.5);
uint64_t c2 = (uint64_t) (f1[i].x / (N << 2) + 0.5);
a[i] = c0 + (c1 << 13) + (c2 << 24);
}
}
uint64_t A[SIZE * 2 - 1];
int main() {
for (int i = 1; i < SIZE; i++) A[i] = 1e7;
clock_t start = clock();
convolve(A);
clock_t end = clock();
uint64_t checksum = 0;
for (int i = 0; i < SIZE * 2 - 1; i++) checksum ^= A[i];
printf("checksum: %llu\n", checksum);
printf("cpu time: %.3fs\n", (double) (end - start) / CLOCKS_PER_SEC);
return 0;
}
Try it online!
Run with -O3
or -Ofast
enabled for best performance. The -march=native
may also be necessary to run.
Uses a standard floating-point radix-4 FFT. Each input integer is split into two 12-bit chunks to avoid precision loss. Since the input comprises of only real numbers, we can get away with doing just a single FFT by using a&4095
as the real part and a>>12
as the imaginary part. Upon converting back we still need two FFTs as the output space is twice as large.
Update:
- I realized that the bit reversal was unusually time consuming, so I got rid of it using a DIF-DIT scheme. DIF takes the input in normal order and outputs in bit-reversed order, while DIT takes a bit-reversed input and outputs in normal order. Thus, they magically cancel out :)
- I made another mistake in assuming that
llround
was relatively fast. Apparently adding +0.5
and then truncating is nearly 10x faster than llround
, which led to a ~30% improvement.
- Added SIMD in the FFT loop. This one didn't help that much, presumably because GCC already auto-vectorizes it.
A = [1111111111]*100000
\$\endgroup\$