#pragma GCC target("arch=skylake")
#include <immintrin.h>
#include <stdint.h>
#include <stdlib.h>
#ifdef LOCAL_PROFILE
#include <stdio.h>
#include <time.h>
#endif
// 2,000,000 = 8 * 2 * 5^6. A cyclic transform of exactly this length is
// sufficient: only the coefficient of x^2,000,000 wraps, and it is restored
// separately at the end.
static const uint32_t MOD = 22000001u;
static const uint32_t TWO_MOD = 44000002u;
static const uint32_t FOUR_MOD = 88000004u;
static const uint32_t EIGHT_MOD = 176000008u;
static const uint32_t SIXTEEN_MOD = 352000016u;
static const uint32_t G = 3u;
static const uint32_t QINV = 1352495487u;
static const uint32_t R2 = 15222743u;
static const uint32_t R = 4967101u;
static const int TRANSFORM_N = 2000000;
static const int HALF_BLOCK = 125000;
#ifndef FA_BYTE_OFFSET
#define FA_BYTE_OFFSET 98304u
#endif
#ifndef FB_WORD_GAP
#define FB_WORD_GAP 8192
#endif
static inline uint32_t mont_reduce(uint64_t x) {
uint32_t y = (uint32_t)x * QINV;
return (uint32_t)((x + (uint64_t)y * MOD) >> 32);
}
static inline uint32_t mont_mul(uint32_t a, uint32_t b) {
return mont_reduce((uint64_t)a * b);
}
static inline uint32_t to_mont(uint32_t x) {
return mont_mul(x, R2);
}
static uint32_t pow_mod(uint32_t a, uint32_t e) {
uint32_t r = 1;
while (e) {
if (e & 1) r = (uint32_t)((uint64_t)r * a % MOD);
a = (uint32_t)((uint64_t)a * a % MOD);
e >>= 1;
}
return r;
}
static inline uint32_t add_mod(uint32_t a, uint32_t b) {
uint32_t s = a + b;
return s >= TWO_MOD ? s - TWO_MOD : s;
}
static inline uint32_t sub_mod(uint32_t a, uint32_t b) {
uint32_t d = a + TWO_MOD - b;
return d >= TWO_MOD ? d - TWO_MOD : d;
}
static inline __m256i add_vec(__m256i a, __m256i b) {
const __m256i mod = _mm256_set1_epi32((int)TWO_MOD);
const __m256i hi = _mm256_set1_epi32((int)(TWO_MOD - 1));
__m256i s = _mm256_add_epi32(a, b);
__m256i m = _mm256_cmpgt_epi32(s, hi);
return _mm256_sub_epi32(s, _mm256_and_si256(m, mod));
}
static inline __m256i sub_vec(__m256i a, __m256i b) {
const __m256i mod = _mm256_set1_epi32((int)TWO_MOD);
const __m256i hi = _mm256_set1_epi32((int)(TWO_MOD - 1));
__m256i d = _mm256_sub_epi32(_mm256_add_epi32(a, mod), b);
__m256i m = _mm256_cmpgt_epi32(d, hi);
return _mm256_sub_epi32(d, _mm256_and_si256(m, mod));
}
static inline __m256i add_raw(__m256i a, __m256i b) {
return _mm256_add_epi32(a, b);
}
static inline __m256i sub2_raw(__m256i a, __m256i b) {
return _mm256_sub_epi32(_mm256_add_epi32(a, _mm256_set1_epi32((int)TWO_MOD)), b);
}
static inline __m256i sub4_raw(__m256i a, __m256i b) {
return _mm256_sub_epi32(_mm256_add_epi32(a, _mm256_set1_epi32((int)FOUR_MOD)), b);
}
static inline __m256i sub8_raw(__m256i a, __m256i b) {
return _mm256_sub_epi32(_mm256_add_epi32(a, _mm256_set1_epi32((int)EIGHT_MOD)), b);
}
static inline __m256i reduce16_to2_vec(__m256i x) {
const __m256i two = _mm256_set1_epi32((int)TWO_MOD);
const __m256i four = _mm256_set1_epi32((int)FOUR_MOD);
const __m256i eight = _mm256_set1_epi32((int)EIGHT_MOD);
__m256i m = _mm256_cmpgt_epi32(x, _mm256_set1_epi32((int)(EIGHT_MOD - 1)));
x = _mm256_sub_epi32(x, _mm256_and_si256(m, eight));
m = _mm256_cmpgt_epi32(x, _mm256_set1_epi32((int)(FOUR_MOD - 1)));
x = _mm256_sub_epi32(x, _mm256_and_si256(m, four));
m = _mm256_cmpgt_epi32(x, _mm256_set1_epi32((int)(TWO_MOD - 1)));
return _mm256_sub_epi32(x, _mm256_and_si256(m, two));
}
static inline __m256i reduce_mod_vec(__m256i x) {
const __m256i mod = _mm256_set1_epi32((int)MOD);
__m256i m = _mm256_cmpgt_epi32(x, _mm256_set1_epi32((int)(MOD - 1)));
return _mm256_sub_epi32(x, _mm256_and_si256(m, mod));
}
static inline __m256i mont_mul_vec(__m256i a, __m256i b) {
const __m256i mod = _mm256_set1_epi32((int)MOD);
const __m256i qinv = _mm256_set1_epi32((int)QINV);
__m256i bq = _mm256_mullo_epi32(b, qinv);
__m256i m = _mm256_mullo_epi32(a, bq);
__m256i t0 = _mm256_mul_epu32(a, b);
__m256i mp0 = _mm256_mul_epu32(m, mod);
__m256i r0 = _mm256_srli_epi64(_mm256_add_epi64(t0, mp0), 32);
__m256i a1 = _mm256_srli_epi64(a, 32);
__m256i b1 = _mm256_srli_epi64(b, 32);
__m256i m1 = _mm256_srli_epi64(m, 32);
__m256i t1 = _mm256_mul_epu32(a1, b1);
__m256i mp1 = _mm256_mul_epu32(m1, mod);
__m256i r1 = _mm256_srli_epi64(_mm256_add_epi64(t1, mp1), 32);
return _mm256_or_si256(r0, _mm256_slli_epi64(r1, 32));
}
static inline __m256i mulhi_epu32_vec(__m256i a, __m256i b) {
__m256i e = _mm256_srli_epi64(_mm256_mul_epu32(a, b), 32);
__m256i o = _mm256_srli_epi64(
_mm256_mul_epu32(_mm256_srli_epi64(a, 32), _mm256_srli_epi64(b, 32)), 32);
return _mm256_or_si256(e, _mm256_slli_epi64(o, 32));
}
// Shoup multiplication by a precomputed ordinary-domain constant. All callers
// keep a < 2p, hence the unreduced result is also in [0, 2p).
static inline __m256i shoup_mul_vec(__m256i a, __m256i b, __m256i bp) {
__m256i q = mulhi_epu32_vec(a, bp);
return _mm256_sub_epi32(_mm256_mullo_epi32(a, b),
_mm256_mullo_epi32(q, _mm256_set1_epi32((int)MOD)));
}
static inline uint32_t shoup_mul(uint32_t a, uint32_t b, uint32_t bp) {
uint32_t q = (uint32_t)(((uint64_t)a * bp) >> 32);
return (uint32_t)((uint64_t)a * b - (uint64_t)q * MOD);
}
static inline uint32_t shoup_quot(uint32_t x) {
return (uint32_t)(((uint64_t)x << 32) / MOD);
}
static inline void fill_mont_digits(uint32_t *dst, const unsigned *src, int count) {
const __m256i vr = _mm256_set1_epi32((int)R);
const __m256i two = _mm256_set1_epi32((int)TWO_MOD);
const __m256i hi = _mm256_set1_epi32((int)(TWO_MOD - 1));
int i = 0;
for (; i + 8 <= count; i += 8) {
__m256i x = _mm256_mullo_epi32(
_mm256_loadu_si256((const __m256i *)(src + i)), vr);
__m256i m = _mm256_cmpgt_epi32(x, hi);
x = _mm256_sub_epi32(x, _mm256_and_si256(m, two));
_mm256_store_si256((__m256i *)(dst + i), x);
}
for (; i < count; ++i) {
uint32_t x = src[i] * R;
dst[i] = x >= TWO_MOD ? x - TWO_MOD : x;
}
}
static inline __m256i make_wvec(uint32_t step) {
uint32_t t[8];
t[0] = R;
for (int i = 1; i < 8; ++i) t[i] = mont_mul(t[i - 1], step);
return _mm256_loadu_si256((const __m256i *)t);
}
static inline uint32_t step_pow8(uint32_t step) {
uint32_t x = R;
for (int i = 0; i < 8; ++i) x = mont_mul(x, step);
return x;
}
static inline void transpose8_epi32(__m256i &r0, __m256i &r1, __m256i &r2, __m256i &r3,
__m256i &r4, __m256i &r5, __m256i &r6, __m256i &r7) {
__m256i t0 = _mm256_unpacklo_epi32(r0, r1);
__m256i t1 = _mm256_unpackhi_epi32(r0, r1);
__m256i t2 = _mm256_unpacklo_epi32(r2, r3);
__m256i t3 = _mm256_unpackhi_epi32(r2, r3);
__m256i t4 = _mm256_unpacklo_epi32(r4, r5);
__m256i t5 = _mm256_unpackhi_epi32(r4, r5);
__m256i t6 = _mm256_unpacklo_epi32(r6, r7);
__m256i t7 = _mm256_unpackhi_epi32(r6, r7);
__m256i u0 = _mm256_unpacklo_epi64(t0, t2);
__m256i u1 = _mm256_unpackhi_epi64(t0, t2);
__m256i u2 = _mm256_unpacklo_epi64(t1, t3);
__m256i u3 = _mm256_unpackhi_epi64(t1, t3);
__m256i u4 = _mm256_unpacklo_epi64(t4, t6);
__m256i u5 = _mm256_unpackhi_epi64(t4, t6);
__m256i u6 = _mm256_unpacklo_epi64(t5, t7);
__m256i u7 = _mm256_unpackhi_epi64(t5, t7);
r0 = _mm256_permute2x128_si256(u0, u4, 0x20);
r1 = _mm256_permute2x128_si256(u1, u5, 0x20);
r2 = _mm256_permute2x128_si256(u2, u6, 0x20);
r3 = _mm256_permute2x128_si256(u3, u7, 0x20);
r4 = _mm256_permute2x128_si256(u0, u4, 0x31);
r5 = _mm256_permute2x128_si256(u1, u5, 0x31);
r6 = _mm256_permute2x128_si256(u2, u6, 0x31);
r7 = _mm256_permute2x128_si256(u3, u7, 0x31);
}
static inline void dft8_vec(__m256i &a0, __m256i &a1, __m256i &a2, __m256i &a3,
__m256i &a4, __m256i &a5, __m256i &a6, __m256i &a7,
__m256i vz, __m256i vi, __m256i vz3) {
__m256i e04 = add_raw(a0, a4);
__m256i f04 = sub2_raw(a0, a4);
__m256i e26 = add_raw(a2, a6);
__m256i f26 = mont_mul_vec(sub2_raw(a2, a6), vi);
__m256i E0 = add_raw(e04, e26);
__m256i E1 = add_raw(f04, f26);
__m256i E2 = sub4_raw(e04, e26);
__m256i E3 = sub4_raw(f04, f26);
__m256i o15 = add_raw(a1, a5);
__m256i g15 = sub2_raw(a1, a5);
__m256i o37 = add_raw(a3, a7);
__m256i g37 = mont_mul_vec(sub2_raw(a3, a7), vi);
__m256i O0 = add_raw(o15, o37);
__m256i O1 = mont_mul_vec(add_raw(g15, g37), vz);
__m256i O2 = mont_mul_vec(sub4_raw(o15, o37), vi);
__m256i O3 = mont_mul_vec(sub4_raw(g15, g37), vz3);
a0 = reduce16_to2_vec(add_raw(E0, O0));
a1 = reduce16_to2_vec(add_raw(E1, O1));
a2 = reduce16_to2_vec(add_raw(E2, O2));
a3 = reduce16_to2_vec(add_raw(E3, O3));
a4 = reduce16_to2_vec(sub8_raw(E0, O0));
a5 = reduce16_to2_vec(sub8_raw(E1, O1));
a6 = reduce16_to2_vec(sub8_raw(E2, O2));
a7 = reduce16_to2_vec(sub8_raw(E3, O3));
}
// First-stage specialization for the two degree-1,000,000 inputs. Except for
// one lane handled separately, the high four limbs of every radix-8 butterfly
// are zero.
static inline void dft8_half_vec(__m256i &a0, __m256i &a1, __m256i &a2, __m256i &a3,
__m256i &a4, __m256i &a5, __m256i &a6, __m256i &a7,
__m256i vz, __m256i vi, __m256i vz3) {
__m256i E0 = add_raw(a0, a2);
__m256i E1 = add_raw(a0, mont_mul_vec(a2, vi));
__m256i E2 = sub2_raw(a0, a2);
__m256i E3 = sub2_raw(a0, mont_mul_vec(a2, vi));
__m256i O0 = add_raw(a1, a3);
__m256i ia3 = mont_mul_vec(a3, vi);
__m256i O1 = mont_mul_vec(add_raw(a1, ia3), vz);
__m256i O2 = mont_mul_vec(sub2_raw(a1, a3), vi);
__m256i O3 = mont_mul_vec(sub2_raw(a1, ia3), vz3);
a0 = reduce16_to2_vec(add_raw(E0, O0));
a1 = reduce16_to2_vec(add_raw(E1, O1));
a2 = reduce16_to2_vec(add_raw(E2, O2));
a3 = reduce16_to2_vec(add_raw(E3, O3));
a4 = reduce16_to2_vec(sub8_raw(E0, O0));
a5 = reduce16_to2_vec(sub8_raw(E1, O1));
a6 = reduce16_to2_vec(sub8_raw(E2, O2));
a7 = reduce16_to2_vec(sub8_raw(E3, O3));
}
static inline void dft8_scalar(uint32_t *x, uint32_t z, uint32_t im, uint32_t z3) {
uint32_t e04 = add_mod(x[0], x[4]), f04 = sub_mod(x[0], x[4]);
uint32_t e26 = add_mod(x[2], x[6]);
uint32_t f26 = mont_mul(sub_mod(x[2], x[6]), im);
uint32_t E0 = add_mod(e04, e26), E1 = add_mod(f04, f26);
uint32_t E2 = sub_mod(e04, e26), E3 = sub_mod(f04, f26);
uint32_t o15 = add_mod(x[1], x[5]), g15 = sub_mod(x[1], x[5]);
uint32_t o37 = add_mod(x[3], x[7]);
uint32_t g37 = mont_mul(sub_mod(x[3], x[7]), im);
uint32_t O0 = add_mod(o15, o37);
uint32_t O1 = mont_mul(add_mod(g15, g37), z);
uint32_t O2 = mont_mul(sub_mod(o15, o37), im);
uint32_t O3 = mont_mul(sub_mod(g15, g37), z3);
x[0] = add_mod(E0, O0); x[1] = add_mod(E1, O1);
x[2] = add_mod(E2, O2); x[3] = add_mod(E3, O3);
x[4] = sub_mod(E0, O0); x[5] = sub_mod(E1, O1);
x[6] = sub_mod(E2, O2); x[7] = sub_mod(E3, O3);
}
// Roots for the six radix-5 levels inside a 125,000-point block. Each level
// stores w^j, w^(2j), w^(3j), w^(4j), then their Shoup quotients.
static const int ROOT_COUNT = 124992;
static const int STAGE_OFF[6] = {0, 100000, 120000, 124000, 124800, 124960};
alignas(32) static uint32_t root_value[ROOT_COUNT];
alignas(32) static uint32_t root_shoup[ROOT_COUNT];
static void build_radix5_roots(bool inverse) {
int len = HALF_BLOCK;
for (int level = 0; level < 6; ++level, len /= 5) {
int q = len / 5;
uint32_t step = pow_mod(G, (MOD - 1) / (uint32_t)len);
if (inverse) step = pow_mod(step, MOD - 2);
uint32_t stepk = 1;
int off = STAGE_OFF[level];
for (int k = 1; k <= 4; ++k) {
stepk = (uint32_t)((uint64_t)stepk * step % MOD);
uint32_t sm = to_mont(stepk), w = R;
uint32_t *v = root_value + off + (k - 1) * q;
uint32_t *s = root_shoup + off + (k - 1) * q;
for (int j = 0; j < q; ++j) {
uint32_t plain = mont_reduce(w);
v[j] = plain;
s[j] = shoup_quot(plain);
w = mont_mul(w, sm);
}
}
}
}
template <bool INVERSE>
static inline void dft5_vec(__m256i &x0, __m256i &x1, __m256i &x2,
__m256i &x3, __m256i &x4) {
const uint32_t km = INVERSE ? 615934u : 21384067u;
const uint32_t kp = INVERSE ? 7267153u : 14732848u;
const uint32_t k2 = INVERSE ? 14325610u : 7674391u;
const __m256i vu = _mm256_set1_epi32(5500000);
const __m256i vup = _mm256_set1_epi32(1073741775);
const __m256i vv = _mm256_set1_epi32(11587557);
const __m256i vvp = _mm256_set1_epi32((int)2262189822u);
const __m256i vkm = _mm256_set1_epi32((int)km);
const __m256i vkmp = _mm256_set1_epi32((int)shoup_quot(km));
const __m256i vkp = _mm256_set1_epi32((int)kp);
const __m256i vkpp = _mm256_set1_epi32((int)shoup_quot(kp));
const __m256i vk2 = _mm256_set1_epi32((int)k2);
const __m256i vk2p = _mm256_set1_epi32((int)shoup_quot(k2));
__m256i s1 = add_vec(x1, x4), s2 = add_vec(x2, x3);
__m256i d1 = sub_vec(x1, x4), d2 = sub_vec(x2, x3);
__m256i s = add_vec(s1, s2), d = sub_vec(s1, s2);
__m256i common = add_vec(x0, shoup_mul_vec(s, vu, vup));
__m256i vd = shoup_mul_vec(d, vv, vvp);
__m256i A1 = add_vec(common, vd), A2 = sub_vec(common, vd);
__m256i P = shoup_mul_vec(d1, vkm, vkmp);
__m256i Q = shoup_mul_vec(d2, vkp, vkpp);
__m256i RR = shoup_mul_vec(add_vec(d1, d2), vk2, vk2p);
__m256i B1 = add_vec(P, RR), B2 = sub_vec(RR, Q);
x0 = add_vec(x0, s);
x1 = add_vec(A1, B1); x4 = sub_vec(A1, B1);
x2 = add_vec(A2, B2); x3 = sub_vec(A2, B2);
}
static inline void dft5_scalar(uint32_t *x, bool inverse) {
uint32_t km = inverse ? 615934u : 21384067u;
uint32_t kp = inverse ? 7267153u : 14732848u;
uint32_t k2 = inverse ? 14325610u : 7674391u;
uint32_t s1 = add_mod(x[1], x[4]), s2 = add_mod(x[2], x[3]);
uint32_t d1 = sub_mod(x[1], x[4]), d2 = sub_mod(x[2], x[3]);
uint32_t s = add_mod(s1, s2), d = sub_mod(s1, s2);
uint32_t common = add_mod(x[0], shoup_mul(s, 5500000u, 1073741775u));
uint32_t vd = shoup_mul(d, 11587557u, 2262189822u);
uint32_t A1 = add_mod(common, vd), A2 = sub_mod(common, vd);
uint32_t P = shoup_mul(d1, km, shoup_quot(km));
uint32_t Q = shoup_mul(d2, kp, shoup_quot(kp));
uint32_t RR = shoup_mul(add_mod(d1, d2), k2, shoup_quot(k2));
uint32_t B1 = add_mod(P, RR), B2 = sub_mod(RR, Q);
x[0] = add_mod(x[0], s);
x[1] = add_mod(A1, B1); x[4] = sub_mod(A1, B1);
x[2] = add_mod(A2, B2); x[3] = sub_mod(A2, B2);
}
static void top8_forward_input(uint32_t *a) {
const int q = 250000;
uint32_t z = to_mont(pow_mod(G, (MOD - 1) / 8));
uint32_t im = mont_mul(z, z), z3 = mont_mul(im, z);
__m256i vz = _mm256_set1_epi32((int)z);
__m256i vi = _mm256_set1_epi32((int)im);
__m256i vz3 = _mm256_set1_epi32((int)z3);
uint32_t step = to_mont(pow_mod(G, (MOD - 1) / TRANSFORM_N));
__m256i w1 = make_wvec(step);
__m256i w2 = mont_mul_vec(w1, w1), w3 = mont_mul_vec(w2, w1);
__m256i w4 = mont_mul_vec(w2, w2), w5 = mont_mul_vec(w4, w1);
__m256i w6 = mont_mul_vec(w4, w2), w7 = mont_mul_vec(w4, w3);
uint32_t s1 = step_pow8(step), s2 = mont_mul(s1, s1), s3 = mont_mul(s2, s1);
uint32_t s4 = mont_mul(s2, s2), s5 = mont_mul(s4, s1);
uint32_t s6 = mont_mul(s4, s2), s7 = mont_mul(s4, s3);
__m256i vs1 = _mm256_set1_epi32((int)s1), vs2 = _mm256_set1_epi32((int)s2);
__m256i vs3 = _mm256_set1_epi32((int)s3), vs4 = _mm256_set1_epi32((int)s4);
__m256i vs5 = _mm256_set1_epi32((int)s5), vs6 = _mm256_set1_epi32((int)s6);
__m256i vs7 = _mm256_set1_epi32((int)s7);
for (int j = 0; j < q; j += 8) {
__m256i x0 = _mm256_load_si256((const __m256i *)(a + j));
__m256i x1 = _mm256_load_si256((const __m256i *)(a + q + j));
__m256i x2 = _mm256_load_si256((const __m256i *)(a + q * 2 + j));
__m256i x3 = _mm256_load_si256((const __m256i *)(a + q * 3 + j));
__m256i x4, x5, x6, x7;
if (j == 0) {
x4 = _mm256_set_epi32(0, 0, 0, 0, 0, 0, 0, (int)a[1000000]);
x5 = x6 = x7 = _mm256_setzero_si256();
dft8_vec(x0, x1, x2, x3, x4, x5, x6, x7, vz, vi, vz3);
} else {
dft8_half_vec(x0, x1, x2, x3, x4, x5, x6, x7, vz, vi, vz3);
}
_mm256_store_si256((__m256i *)(a + j), x0);
_mm256_store_si256((__m256i *)(a + q + j), mont_mul_vec(x1, w1));
_mm256_store_si256((__m256i *)(a + q * 2 + j), mont_mul_vec(x2, w2));
_mm256_store_si256((__m256i *)(a + q * 3 + j), mont_mul_vec(x3, w3));
_mm256_store_si256((__m256i *)(a + q * 4 + j), mont_mul_vec(x4, w4));
_mm256_store_si256((__m256i *)(a + q * 5 + j), mont_mul_vec(x5, w5));
_mm256_store_si256((__m256i *)(a + q * 6 + j), mont_mul_vec(x6, w6));
_mm256_store_si256((__m256i *)(a + q * 7 + j), mont_mul_vec(x7, w7));
w1 = mont_mul_vec(w1, vs1); w2 = mont_mul_vec(w2, vs2);
w3 = mont_mul_vec(w3, vs3); w4 = mont_mul_vec(w4, vs4);
w5 = mont_mul_vec(w5, vs5); w6 = mont_mul_vec(w6, vs6);
w7 = mont_mul_vec(w7, vs7);
}
}
static void radix2_forward(uint32_t *a) {
const int q = HALF_BLOCK;
uint32_t step = to_mont(pow_mod(G, (MOD - 1) / 250000));
__m256i w = make_wvec(step);
__m256i ws = _mm256_set1_epi32((int)step_pow8(step));
for (int j = 0; j < q; j += 8) {
__m256i x0 = _mm256_load_si256((const __m256i *)(a + j));
__m256i x1 = _mm256_load_si256((const __m256i *)(a + q + j));
_mm256_store_si256((__m256i *)(a + j), add_vec(x0, x1));
_mm256_store_si256((__m256i *)(a + q + j), mont_mul_vec(sub_vec(x0, x1), w));
w = mont_mul_vec(w, ws);
}
}
static void radix5_forward_block(uint32_t *a) {
int len = HALF_BLOCK;
for (int level = 0; level < 6; ++level, len /= 5) {
int q = len / 5, off = STAGE_OFF[level];
const uint32_t *v1 = root_value + off, *v2 = v1 + q;
const uint32_t *v3 = v2 + q, *v4 = v3 + q;
const uint32_t *s1 = root_shoup + off, *s2 = s1 + q;
const uint32_t *s3 = s2 + q, *s4 = s3 + q;
for (int base = 0; base < HALF_BLOCK; base += len) {
uint32_t *p = a + base;
for (int j = 0; j < q; j += 8) {
__m256i x0 = _mm256_load_si256((const __m256i *)(p + j));
__m256i x1 = _mm256_load_si256((const __m256i *)(p + q + j));
__m256i x2 = _mm256_load_si256((const __m256i *)(p + q * 2 + j));
__m256i x3 = _mm256_load_si256((const __m256i *)(p + q * 3 + j));
__m256i x4 = _mm256_load_si256((const __m256i *)(p + q * 4 + j));
dft5_vec<false>(x0, x1, x2, x3, x4);
x1 = shoup_mul_vec(x1, _mm256_load_si256((const __m256i *)(v1 + j)),
_mm256_load_si256((const __m256i *)(s1 + j)));
x2 = shoup_mul_vec(x2, _mm256_load_si256((const __m256i *)(v2 + j)),
_mm256_load_si256((const __m256i *)(s2 + j)));
x3 = shoup_mul_vec(x3, _mm256_load_si256((const __m256i *)(v3 + j)),
_mm256_load_si256((const __m256i *)(s3 + j)));
x4 = shoup_mul_vec(x4, _mm256_load_si256((const __m256i *)(v4 + j)),
_mm256_load_si256((const __m256i *)(s4 + j)));
_mm256_store_si256((__m256i *)(p + j), x0);
_mm256_store_si256((__m256i *)(p + q + j), x1);
_mm256_store_si256((__m256i *)(p + q * 2 + j), x2);
_mm256_store_si256((__m256i *)(p + q * 3 + j), x3);
_mm256_store_si256((__m256i *)(p + q * 4 + j), x4);
}
}
}
}
static void base8_forward(uint32_t *a) {
uint32_t z = to_mont(pow_mod(G, (MOD - 1) / 8));
uint32_t im = mont_mul(z, z), z3 = mont_mul(im, z);
__m256i vz = _mm256_set1_epi32((int)z), vi = _mm256_set1_epi32((int)im);
__m256i vz3 = _mm256_set1_epi32((int)z3);
int base = 0;
for (; base + 64 <= HALF_BLOCK; base += 64) {
__m256i x0 = _mm256_load_si256((const __m256i *)(a + base));
__m256i x1 = _mm256_load_si256((const __m256i *)(a + base + 8));
__m256i x2 = _mm256_load_si256((const __m256i *)(a + base + 16));
__m256i x3 = _mm256_load_si256((const __m256i *)(a + base + 24));
__m256i x4 = _mm256_load_si256((const __m256i *)(a + base + 32));
__m256i x5 = _mm256_load_si256((const __m256i *)(a + base + 40));
__m256i x6 = _mm256_load_si256((const __m256i *)(a + base + 48));
__m256i x7 = _mm256_load_si256((const __m256i *)(a + base + 56));
transpose8_epi32(x0, x1, x2, x3, x4, x5, x6, x7);
dft8_vec(x0, x1, x2, x3, x4, x5, x6, x7, vz, vi, vz3);
transpose8_epi32(x0, x1, x2, x3, x4, x5, x6, x7);
_mm256_store_si256((__m256i *)(a + base), x0);
_mm256_store_si256((__m256i *)(a + base + 8), x1);
_mm256_store_si256((__m256i *)(a + base + 16), x2);
_mm256_store_si256((__m256i *)(a + base + 24), x3);
_mm256_store_si256((__m256i *)(a + base + 32), x4);
_mm256_store_si256((__m256i *)(a + base + 40), x5);
_mm256_store_si256((__m256i *)(a + base + 48), x6);
_mm256_store_si256((__m256i *)(a + base + 56), x7);
}
uint32_t x[8];
for (int k = 0; k < 8; ++k) x[k] = a[base + k];
dft8_scalar(x, z, im, z3);
for (int k = 0; k < 8; ++k) a[base + k] = x[k];
}
static void forward_all(uint32_t *a) {
top8_forward_input(a);
for (int k = 0; k < 8; ++k) {
uint32_t *p = a + k * 250000;
radix2_forward(p);
radix5_forward_block(p);
base8_forward(p);
radix5_forward_block(p + HALF_BLOCK);
base8_forward(p + HALF_BLOCK);
}
}
static void base8_inverse_mul(uint32_t *a, const uint32_t *b) {
uint32_t z0 = pow_mod(G, (MOD - 1) / 8);
uint32_t z = to_mont(pow_mod(z0, MOD - 2));
uint32_t im = mont_mul(z, z), z3 = mont_mul(im, z);
__m256i vz = _mm256_set1_epi32((int)z), vi = _mm256_set1_epi32((int)im);
__m256i vz3 = _mm256_set1_epi32((int)z3);
int base = 0;
for (; base + 64 <= HALF_BLOCK; base += 64) {
__m256i x0 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + base)),
_mm256_load_si256((const __m256i *)(b + base)));
__m256i x1 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + base + 8)),
_mm256_load_si256((const __m256i *)(b + base + 8)));
__m256i x2 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + base + 16)),
_mm256_load_si256((const __m256i *)(b + base + 16)));
__m256i x3 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + base + 24)),
_mm256_load_si256((const __m256i *)(b + base + 24)));
__m256i x4 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + base + 32)),
_mm256_load_si256((const __m256i *)(b + base + 32)));
__m256i x5 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + base + 40)),
_mm256_load_si256((const __m256i *)(b + base + 40)));
__m256i x6 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + base + 48)),
_mm256_load_si256((const __m256i *)(b + base + 48)));
__m256i x7 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + base + 56)),
_mm256_load_si256((const __m256i *)(b + base + 56)));
transpose8_epi32(x0, x1, x2, x3, x4, x5, x6, x7);
dft8_vec(x0, x1, x2, x3, x4, x5, x6, x7, vz, vi, vz3);
transpose8_epi32(x0, x1, x2, x3, x4, x5, x6, x7);
_mm256_store_si256((__m256i *)(a + base), x0);
_mm256_store_si256((__m256i *)(a + base + 8), x1);
_mm256_store_si256((__m256i *)(a + base + 16), x2);
_mm256_store_si256((__m256i *)(a + base + 24), x3);
_mm256_store_si256((__m256i *)(a + base + 32), x4);
_mm256_store_si256((__m256i *)(a + base + 40), x5);
_mm256_store_si256((__m256i *)(a + base + 48), x6);
_mm256_store_si256((__m256i *)(a + base + 56), x7);
}
uint32_t x[8];
for (int k = 0; k < 8; ++k) x[k] = mont_mul(a[base + k], b[base + k]);
dft8_scalar(x, z, im, z3);
for (int k = 0; k < 8; ++k) a[base + k] = x[k];
}
static void radix5_inverse_block(uint32_t *a, const uint32_t *b) {
base8_inverse_mul(a, b);
int len = 40;
for (int level = 5; level >= 0; --level, len *= 5) {
int q = len / 5, off = STAGE_OFF[level];
const uint32_t *v1 = root_value + off, *v2 = v1 + q;
const uint32_t *v3 = v2 + q, *v4 = v3 + q;
const uint32_t *s1 = root_shoup + off, *s2 = s1 + q;
const uint32_t *s3 = s2 + q, *s4 = s3 + q;
for (int base = 0; base < HALF_BLOCK; base += len) {
uint32_t *p = a + base;
for (int j = 0; j < q; j += 8) {
__m256i x0 = _mm256_load_si256((const __m256i *)(p + j));
__m256i x1 = shoup_mul_vec(
_mm256_load_si256((const __m256i *)(p + q + j)),
_mm256_load_si256((const __m256i *)(v1 + j)),
_mm256_load_si256((const __m256i *)(s1 + j)));
__m256i x2 = shoup_mul_vec(
_mm256_load_si256((const __m256i *)(p + q * 2 + j)),
_mm256_load_si256((const __m256i *)(v2 + j)),
_mm256_load_si256((const __m256i *)(s2 + j)));
__m256i x3 = shoup_mul_vec(
_mm256_load_si256((const __m256i *)(p + q * 3 + j)),
_mm256_load_si256((const __m256i *)(v3 + j)),
_mm256_load_si256((const __m256i *)(s3 + j)));
__m256i x4 = shoup_mul_vec(
_mm256_load_si256((const __m256i *)(p + q * 4 + j)),
_mm256_load_si256((const __m256i *)(v4 + j)),
_mm256_load_si256((const __m256i *)(s4 + j)));
dft5_vec<true>(x0, x1, x2, x3, x4);
_mm256_store_si256((__m256i *)(p + j), x0);
_mm256_store_si256((__m256i *)(p + q + j), x1);
_mm256_store_si256((__m256i *)(p + q * 2 + j), x2);
_mm256_store_si256((__m256i *)(p + q * 3 + j), x3);
_mm256_store_si256((__m256i *)(p + q * 4 + j), x4);
}
}
}
}
static void radix2_inverse(uint32_t *a) {
const int q = HALF_BLOCK;
uint32_t r0 = pow_mod(G, (MOD - 1) / 250000);
uint32_t step = to_mont(pow_mod(r0, MOD - 2));
__m256i w = make_wvec(step);
__m256i ws = _mm256_set1_epi32((int)step_pow8(step));
for (int j = 0; j < q; j += 8) {
__m256i x0 = _mm256_load_si256((const __m256i *)(a + j));
__m256i x1 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + q + j)), w);
_mm256_store_si256((__m256i *)(a + j), add_vec(x0, x1));
_mm256_store_si256((__m256i *)(a + q + j), sub_vec(x0, x1));
w = mont_mul_vec(w, ws);
}
}
static void top8_inverse(uint32_t *a) {
const int q = 250000;
uint32_t z0 = pow_mod(G, (MOD - 1) / 8);
uint32_t z = to_mont(pow_mod(z0, MOD - 2));
uint32_t im = mont_mul(z, z), z3 = mont_mul(im, z);
__m256i vz = _mm256_set1_epi32((int)z), vi = _mm256_set1_epi32((int)im);
__m256i vz3 = _mm256_set1_epi32((int)z3);
uint32_t r0 = pow_mod(G, (MOD - 1) / TRANSFORM_N);
uint32_t step = to_mont(pow_mod(r0, MOD - 2));
__m256i w1 = make_wvec(step);
__m256i w2 = mont_mul_vec(w1, w1), w3 = mont_mul_vec(w2, w1);
__m256i w4 = mont_mul_vec(w2, w2), w5 = mont_mul_vec(w4, w1);
__m256i w6 = mont_mul_vec(w4, w2), w7 = mont_mul_vec(w4, w3);
uint32_t s1 = step_pow8(step), s2 = mont_mul(s1, s1), s3 = mont_mul(s2, s1);
uint32_t s4 = mont_mul(s2, s2), s5 = mont_mul(s4, s1);
uint32_t s6 = mont_mul(s4, s2), s7 = mont_mul(s4, s3);
__m256i vs1 = _mm256_set1_epi32((int)s1), vs2 = _mm256_set1_epi32((int)s2);
__m256i vs3 = _mm256_set1_epi32((int)s3), vs4 = _mm256_set1_epi32((int)s4);
__m256i vs5 = _mm256_set1_epi32((int)s5), vs6 = _mm256_set1_epi32((int)s6);
__m256i vs7 = _mm256_set1_epi32((int)s7);
for (int j = 0; j < q; j += 8) {
__m256i x0 = _mm256_load_si256((const __m256i *)(a + j));
__m256i x1 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + q + j)), w1);
__m256i x2 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + q * 2 + j)), w2);
__m256i x3 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + q * 3 + j)), w3);
__m256i x4 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + q * 4 + j)), w4);
__m256i x5 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + q * 5 + j)), w5);
__m256i x6 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + q * 6 + j)), w6);
__m256i x7 = mont_mul_vec(_mm256_load_si256((const __m256i *)(a + q * 7 + j)), w7);
dft8_vec(x0, x1, x2, x3, x4, x5, x6, x7, vz, vi, vz3);
_mm256_store_si256((__m256i *)(a + j), x0);
_mm256_store_si256((__m256i *)(a + q + j), x1);
_mm256_store_si256((__m256i *)(a + q * 2 + j), x2);
_mm256_store_si256((__m256i *)(a + q * 3 + j), x3);
_mm256_store_si256((__m256i *)(a + q * 4 + j), x4);
_mm256_store_si256((__m256i *)(a + q * 5 + j), x5);
_mm256_store_si256((__m256i *)(a + q * 6 + j), x6);
_mm256_store_si256((__m256i *)(a + q * 7 + j), x7);
w1 = mont_mul_vec(w1, vs1); w2 = mont_mul_vec(w2, vs2);
w3 = mont_mul_vec(w3, vs3); w4 = mont_mul_vec(w4, vs4);
w5 = mont_mul_vec(w5, vs5); w6 = mont_mul_vec(w6, vs6);
w7 = mont_mul_vec(w7, vs7);
}
}
static void inverse_all(uint32_t *a, const uint32_t *b) {
for (int k = 0; k < 8; ++k) {
uint32_t *pa = a + k * 250000;
const uint32_t *pb = b + k * 250000;
radix5_inverse_block(pa, pb);
radix5_inverse_block(pa + HALF_BLOCK, pb + HALF_BLOCK);
radix2_inverse(pa);
}
top8_inverse(a);
}
void poly_multiply(unsigned *a, int n, unsigned *b, int m, unsigned *c) {
#ifdef LOCAL_PROFILE
clock_t t0 = clock();
#endif
uint32_t wrapped = a[n] * b[m];
size_t extra = (size_t)(FA_BYTE_OFFSET / 4u) + FB_WORD_GAP + 16;
uint32_t *raw = (uint32_t *)malloc(((size_t)TRANSFORM_N * 2 + extra) * sizeof(uint32_t));
uint32_t *fa = (uint32_t *)((((uintptr_t)raw + 31u) & ~(uintptr_t)31u) + FA_BYTE_OFFSET);
uint32_t *fb = fa + TRANSFORM_N + FB_WORD_GAP;
fill_mont_digits(fa, a, n + 1);
fill_mont_digits(fb, b, m + 1);
#ifdef LOCAL_PROFILE
clock_t t1 = clock();
#endif
build_radix5_roots(false);
#ifdef LOCAL_PROFILE
clock_t t2 = clock();
#endif
forward_all(fa);
#ifdef LOCAL_PROFILE
clock_t t3 = clock();
#endif
forward_all(fb);
#ifdef LOCAL_PROFILE
clock_t t4 = clock();
#endif
build_radix5_roots(true);
#ifdef LOCAL_PROFILE
clock_t t5 = clock();
#endif
inverse_all(fa, fb);
#ifdef LOCAL_PROFILE
clock_t t6 = clock();
#endif
const __m256i vk = _mm256_set1_epi32(11793621);
const __m256i vkp = _mm256_set1_epi32((int)2302418826u);
int i = 0;
for (; i + 8 <= TRANSFORM_N; i += 8) {
__m256i x = shoup_mul_vec(_mm256_load_si256((const __m256i *)(fa + i)), vk, vkp);
_mm256_storeu_si256((__m256i *)(c + i), reduce_mod_vec(x));
}
c[0] = c[0] >= wrapped ? c[0] - wrapped : c[0] + MOD - wrapped;
c[TRANSFORM_N] = wrapped;
#ifdef LOCAL_PROFILE
clock_t t7 = clock();
fprintf(stderr, "fill %.3f rootsF %.3f fwdA %.3f fwdB %.3f rootsI %.3f inv %.3f out %.3f total %.3f\n",
1000.0 * (double)(t1-t0)/CLOCKS_PER_SEC,
1000.0 * (double)(t2-t1)/CLOCKS_PER_SEC,
1000.0 * (double)(t3-t2)/CLOCKS_PER_SEC,
1000.0 * (double)(t4-t3)/CLOCKS_PER_SEC,
1000.0 * (double)(t5-t4)/CLOCKS_PER_SEC,
1000.0 * (double)(t6-t5)/CLOCKS_PER_SEC,
1000.0 * (double)(t7-t6)/CLOCKS_PER_SEC,
1000.0 * (double)(t7-t0)/CLOCKS_PER_SEC);
#endif
}
#ifdef LOCAL_BENCH
#include <stdio.h>
#include <time.h>
static unsigned aa[1000001], bb[1000001], cc[2000001];
int main() {
for (int i = 0; i <= 1000000; ++i) {
aa[i] = (unsigned)((i * 7 + 3) % 10);
bb[i] = (unsigned)((i * 5 + 1) % 10);
}
clock_t st = clock();
poly_multiply(aa, 1000000, bb, 1000000, cc);
clock_t ed = clock();
unsigned long long sample = 0;
for (int i = 0; i <= 2000000; i += 137) sample += cc[i];
printf("%.3f ms sample=%llu edge=%u,%u,%u,%u\n",
1000.0 * (double)(ed-st)/CLOCKS_PER_SEC, sample,
cc[0], cc[1], cc[1999999], cc[2000000]);
return 0;
}
#endif
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 39.698 ms | 23 MB + 884 KB | Accepted | Score: 100 | 显示更多 |