#pragma GCC optimize("O3","unroll-loops","inline-functions","ipa-cp-clone","ipa-sra","rename-registers")
#pragma GCC target("arch=skylake")
// Blocked radix-4 DIF/DIT NTT mod 2013265921 using Shoup modular multiplication.
// Per stage: 6 tables (t1,t2,t3 = twiddle values; s1,s2,s3 = Shoup constants), contiguous.
// DRAM levels: block sizes n, n/4, ...; then in-cache base blocks of size ~2^15..2^19.
// DIF forward: natural in -> digit-reversed out. DIT inverse: digit-reversed in -> natural out.
#pragma once
#include <immintrin.h>
#include <cstdint>
#include <cstring>
#include <cstdlib>
// AVX-256 split-access fixes. g++-9 -O2 lowers an unaligned 32-byte access into a 128-bit
// access plus vinserti128 (load) or vextracti128 (store): two uops, one of them on p5.
// These emit the single real instruction instead. The store's memory is an OUTPUT operand
// so gcc knows the location is written and no "memory" clobber is needed; do NOT use the
// aligned intrinsic, which lowers to vmovdqa with no alignment proof and would fault.
__attribute__((target("avx2"))) static inline __m256i ldu256(const void *p) {
__m256i r; __asm__("vmovdqu %1, %0" : "=x"(r) : "m"(*(const __m256i *)p)); return r; }
__attribute__((target("avx2"))) static inline void stu256(void *p, __m256i v) {
__asm__("vmovdqu %1, %0" : "=m"(*(__m256i *)p) : "x"(v)); }
// Non-temporal 32-byte store: the line is written without a read-for-ownership.
// PRECONDITION: the destination is 32-byte aligned (vmovntdq faults otherwise).
__attribute__((target("avx2"))) static inline void stu256_nt(void *p, __m256i v) {
__asm__("vmovntdq %1, %0" : "=m"(*(__m256i *)p) : "x"(v)); }
// NT stores are weakly ordered: the write-combining buffers are NOT snooped by
// ordinary loads, so every reader of the tables must sit behind this fence.
static inline void nt_fence(void) { __asm__ __volatile__("sfence" ::: "memory"); }
typedef uint32_t u32;
typedef uint16_t u16;
typedef uint64_t u64;
#define NTP 2013265921u
#define TGT __attribute__((target("avx2")))
#ifndef NTT_MAXN
#define NTT_MAXN (1u << 21)
#endif
static u32 g_n, g_Sb;
static u32 g_P = 2013265921u, g_Ml = 572662301u, g_Msh = 1;
static u32 g_pinv;
static int g_log2n, g_L, g_odd, g_nb, g_BT = 15;
static u32 g_ciF, g_cisF, g_ciI, g_cisI;
// 64-byte aligned so that every sub-series the walk in build_tabs lands on is
// 32-byte aligned whenever its size M >= 8 u32 (see NT_OK).
static u32 g_tabF[5 * NTT_MAXN / 2 + 4096] __attribute__((aligned(64)));
static u32 g_tabI[5 * NTT_MAXN / 2 + 4096] __attribute__((aligned(64)));
// A series gets non-temporal stores only when its start is 32-byte aligned (the
// loop steps a whole 32 bytes, so every store in the series then stays aligned)
// and it is big enough that it cannot be re-read out of L3 (6 MB, 12-way).
#ifndef NT_MIN
#define NT_MIN 4096u
#endif
#define NT_OK(P, CNT) ((u32)(CNT) >= NT_MIN && (((uintptr_t)(P) & 31u) == 0u))
struct Tabs {
const u32 *l1[40], *l2[40], *l3[40], *m1[40], *m2[40], *m3[40];
const u32 *b2, *b2s;
const u32 *h1[40], *h2[40], *h3[40], *q1[40], *q2[40], *q3[40];
};
static Tabs g_TF, g_TI;
static inline u32 mulmod(u32 a, u32 b) { return (u32)((u64)a * b % g_P); }
static u32 powmod(u32 a, u32 e) { u32 r = 1; while (e) { if (e & 1) r = mulmod(r, a); a = mulmod(a, a); e >>= 1; } return r; }
static inline u32 shoup1s(u32 a, u32 w, u32 ws) {
u32 t = (u32)((u64)a * w);
u32 q = (u32)(((u64)a * ws) >> 32);
u32 r = t - (u32)((u64)q * g_P);
return r >= g_P ? r - g_P : r;
}
static inline u32 shoup_c(u32 w) { return (u32)(((u64)w << 32) / g_P); }
TGT static inline __m256i vmulhi(__m256i a, __m256i b) {
__m256i e = _mm256_mul_epu32(a, b);
__m256i o = _mm256_mul_epu32(_mm256_srli_epi64(a, 32), _mm256_srli_epi64(b, 32));
e = _mm256_srli_epi64(e, 32);
return _mm256_blend_epi32(e, o, 0xAA);
}
TGT static inline __m256i vmont(__m256i a, __m256i b, __m256i pv) {
const __m256i pinv = _mm256_set1_epi32((int)g_pinv);
__m256i tlo = _mm256_mullo_epi32(a, b);
__m256i thi = vmulhi(a, b);
__m256i m = _mm256_mullo_epi32(tlo, pinv);
__m256i mhi = vmulhi(m, pv);
__m256i nz = _mm256_andnot_si256(_mm256_cmpeq_epi32(tlo, _mm256_setzero_si256()), _mm256_set1_epi32(1));
__m256i u = _mm256_add_epi32(_mm256_add_epi32(thi, mhi), nz);
return _mm256_min_epu32(u, _mm256_sub_epi32(u, pv));
}
// M = floor(2^64/p) = 2^33 + MLO ; ws ~= 2*w + mulhi(w, MLO) (error <= 1, valid for Shoup)
#define MLO_V 572662301u
TGT static inline __m256i der_ws(__m256i w) {
__m256i mlo = _mm256_set1_epi32((int)g_Ml);
__m256i t = _mm256_add_epi32(vmulhi(w, mlo), _mm256_sll_epi32(w, _mm_cvtsi32_si128(g_Msh)));
return t;
}
TGT static inline __m256i vshoup(__m256i a, __m256i w, __m256i ws, __m256i pv) {
__m256i t = _mm256_mullo_epi32(a, w);
__m256i q = vmulhi(a, ws);
__m256i r = _mm256_sub_epi32(t, _mm256_mullo_epi32(q, pv));
return _mm256_min_epu32(r, _mm256_sub_epi32(r, pv));
}
TGT static inline __m256i vaddm(__m256i a, __m256i b, __m256i pv) {
__m256i s = _mm256_add_epi32(a, b);
return _mm256_min_epu32(s, _mm256_sub_epi32(s, pv));
}
TGT static inline __m256i vsubm(__m256i a, __m256i b, __m256i pv) {
__m256i d = _mm256_sub_epi32(a, b);
return _mm256_min_epu32(d, _mm256_add_epi32(d, pv));
}
TGT static inline __m128i vmulhi4(__m128i a, __m128i b) {
__m128i e = _mm_mul_epu32(a, b);
__m128i o = _mm_mul_epu32(_mm_srli_epi64(a, 32), _mm_srli_epi64(b, 32));
e = _mm_srli_epi64(e, 32);
return _mm_blend_epi32(e, o, 0xAA);
}
TGT static inline __m128i sh4(__m128i a, __m128i w, __m128i ws, __m128i pv) {
__m128i t = _mm_mullo_epi32(a, w);
__m128i q = vmulhi4(a, ws);
__m128i r = _mm_sub_epi32(t, _mm_mullo_epi32(q, pv));
return _mm_min_epu32(r, _mm_sub_epi32(r, pv));
}
TGT static inline __m128i add4(__m128i a, __m128i b, __m128i pv) {
__m128i s = _mm_add_epi32(a, b);
return _mm_min_epu32(s, _mm_sub_epi32(s, pv));
}
TGT static inline __m128i sub4(__m128i a, __m128i b, __m128i pv) {
__m128i d = _mm_sub_epi32(a, b);
return _mm_min_epu32(d, _mm_add_epi32(d, pv));
}
TGT static inline __m256i ld2(const __m128i *lo, const __m128i *hi) {
return _mm256_inserti128_si256(_mm256_castsi128_si256(_mm_loadu_si128(lo)), _mm_loadu_si128(hi), 1);
}
TGT static inline void st2(__m128i *lo, __m128i *hi, __m256i v) {
_mm_storeu_si128(lo, _mm256_castsi256_si128(v));
_mm_storeu_si128(hi, _mm256_extracti128_si256(v, 1));
}
TGT static inline __m256i bc128(__m128i v) {
return _mm256_inserti128_si256(_mm256_castsi128_si256(v), v, 1);
}
// fused level-0 DIF stage: reads coefficients directly from src ([0,len] valid, zero beyond).
// Requires len >= h and len < 2h+... (x2 = x3 = 0 assumed).
TGT static void stage0_dif_src(u32 *a, const u32 *src, u32 len, u32 h,
const u32 *t1, const u32 *t2, const u32 *t3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
const __m256i zero = _mm256_setzero_si256();
u32 j = 0;
u32 lim1 = len - h;
for (; j + 8 <= lim1 + 1; j += 8) {
__builtin_prefetch(src + j + 256);
__builtin_prefetch(src + j + h + 256);
__builtin_prefetch(a + j + 256);
__builtin_prefetch(a + j + h + 256);
__builtin_prefetch(a + j + 2 * h + 256);
__builtin_prefetch(a + j + 3 * h + 256);
__builtin_prefetch(t1 + j + 256);
__m256i x0 = ldu256((src + j));
__m256i x1 = ldu256((src + j + h));
__m256i u0 = vaddm(x0, x1, pv), u2 = vsubm(x0, x1, pv);
__m256i E = vshoup(x1, cI, cIs, pv);
__m256i u1 = vaddm(x0, E, pv), u3 = vsubm(x0, E, pv);
{ __m256i w = ldu256((t1 + j)); u1 = vshoup(u1, w, der_ws(w), pv); }
{ __m256i w = ldu256((t2 + j)); u2 = vshoup(u2, w, der_ws(w), pv); }
{ __m256i w = ldu256((t3 + j)); u3 = vshoup(u3, w, der_ws(w), pv); }
stu256((a + j), u0);
stu256((a + j + h), u1);
stu256((a + j + 2 * h), u2);
stu256((a + j + 3 * h), u3);
}
for (; j <= lim1 && j < h; j++) {
u32 x0 = (j <= len) ? src[j] : 0, x1 = (j + h <= len) ? src[j + h] : 0;
u32 u0 = (u32)(((u64)x0 + x1) % g_P);
u32 u2 = (u32)(((u64)x0 + g_P - x1) % g_P);
u32 E = mulmod(x1, ci);
u32 u1 = (u32)(((u64)x0 + E) % g_P), u3 = (u32)(((u64)x0 + g_P - E) % g_P);
u1 = shoup1s(u1, t1[j], shoup_c(t1[j]));
u2 = shoup1s(u2, t2[j], shoup_c(t2[j]));
u3 = shoup1s(u3, t3[j], shoup_c(t3[j]));
a[j] = u0; a[j + h] = u1; a[j + 2 * h] = u2; a[j + 3 * h] = u3;
}
for (; j + 8 <= h; j += 8) {
__m256i x0 = (j + 8 <= len + 1) ? ldu256((src + j)) : zero;
__m256i u0 = x0, u1 = x0, u2 = x0, u3 = x0;
{ __m256i w = ldu256((t1 + j)); u1 = vshoup(u1, w, der_ws(w), pv); }
{ __m256i w = ldu256((t2 + j)); u2 = vshoup(u2, w, der_ws(w), pv); }
{ __m256i w = ldu256((t3 + j)); u3 = vshoup(u3, w, der_ws(w), pv); }
stu256((a + j), u0);
stu256((a + j + h), u1);
stu256((a + j + 2 * h), u2);
stu256((a + j + 3 * h), u3);
}
for (; j < h; j++) {
u32 x0 = (j <= len) ? src[j] : 0;
u32 u1 = shoup1s(x0, t1[j], shoup_c(t1[j]));
u32 u2 = shoup1s(x0, t2[j], shoup_c(t2[j]));
u32 u3 = shoup1s(x0, t3[j], shoup_c(t3[j]));
a[j] = x0; a[j + h] = u1; a[j + 2 * h] = u2; a[j + 3 * h] = u3;
}
}
// ================= forward (DIF) =================
TGT static void stage4_dif(u32 *a, u32 blk, u32 h, const u32 *t1, const u32 *t2, const u32 *t3,
const u32 *s1, const u32 *s2, const u32 *s3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
for (u32 i = 0; i < blk; i += 4 * h) {
u32 *p = a + i;
u32 j = 0;
for (; j + 16 <= h; j += 16) {
__builtin_prefetch(p + j + 256);
__builtin_prefetch(p + j + h + 256);
__builtin_prefetch(p + j + 2 * h + 256);
__builtin_prefetch(p + j + 3 * h + 256);
__builtin_prefetch(t1 + j + 256);
__builtin_prefetch(t2 + j + 256);
__builtin_prefetch(t3 + j + 256);
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + h));
__m256i x2 = ldu256((p + j + 2 * h));
__m256i x3 = ldu256((p + j + 3 * h));
__m256i y0 = ldu256((p + j + 8));
__m256i y1 = ldu256((p + j + h + 8));
__m256i y2 = ldu256((p + j + 2 * h + 8));
__m256i y3 = ldu256((p + j + 3 * h + 8));
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i A2 = vaddm(y0, y2, pv), B2 = vsubm(y0, y2, pv);
__m256i C2 = vaddm(y1, y3, pv), D2 = vsubm(y1, y3, pv);
__m256i u0 = vaddm(A, C, pv), u2 = vsubm(A, C, pv);
__m256i v0 = vaddm(A2, C2, pv), v2 = vsubm(A2, C2, pv);
__m256i E = vshoup(D, cI, cIs, pv);
__m256i F = vshoup(D2, cI, cIs, pv);
__m256i u1 = vaddm(B, E, pv), u3 = vsubm(B, E, pv);
__m256i v1 = vaddm(B2, F, pv), v3 = vsubm(B2, F, pv);
{ __m256i w = ldu256((t1 + j)); u1 = vshoup(u1, w, der_ws(w), pv); }
{ __m256i w = ldu256((t2 + j)); u2 = vshoup(u2, w, der_ws(w), pv); }
{ __m256i w = ldu256((t3 + j)); u3 = vshoup(u3, w, der_ws(w), pv); }
{ __m256i w = ldu256((t1 + j + 8)); v1 = vshoup(v1, w, der_ws(w), pv); }
{ __m256i w = ldu256((t2 + j + 8)); v2 = vshoup(v2, w, der_ws(w), pv); }
{ __m256i w = ldu256((t3 + j + 8)); v3 = vshoup(v3, w, der_ws(w), pv); }
stu256((p + j), u0);
stu256((p + j + h), u1);
stu256((p + j + 2 * h), u2);
stu256((p + j + 3 * h), u3);
stu256((p + j + 8), v0);
stu256((p + j + h + 8), v1);
stu256((p + j + 2 * h + 8), v2);
stu256((p + j + 3 * h + 8), v3);
}
for (; j + 8 <= h; j += 8) {
__builtin_prefetch(p + j + 256);
__builtin_prefetch(p + j + h + 256);
__builtin_prefetch(p + j + 2 * h + 256);
__builtin_prefetch(p + j + 3 * h + 256);
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + h));
__m256i x2 = ldu256((p + j + 2 * h));
__m256i x3 = ldu256((p + j + 3 * h));
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i u0 = vaddm(A, C, pv), u2 = vsubm(A, C, pv);
__m256i E = vshoup(D, cI, cIs, pv);
__m256i u1 = vaddm(B, E, pv), u3 = vsubm(B, E, pv);
{ __m256i w = ldu256((t1 + j)); u1 = vshoup(u1, w, der_ws(w), pv); }
{ __m256i w = ldu256((t2 + j)); u2 = vshoup(u2, w, der_ws(w), pv); }
{ __m256i w = ldu256((t3 + j)); u3 = vshoup(u3, w, der_ws(w), pv); }
stu256((p + j), u0);
stu256((p + j + h), u1);
stu256((p + j + 2 * h), u2);
stu256((p + j + 3 * h), u3);
}
}
}
TGT static void stage4_dif_s(u32 *a, u32 blk, u32 h, const u32 *t1, const u32 *t2, const u32 *t3,
const u32 *s1, const u32 *s2, const u32 *s3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
for (u32 i = 0; i < blk; i += 4 * h) {
u32 *p = a + i;
u32 j = 0;
for (; j + 16 <= h; j += 16) {
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + h));
__m256i x2 = ldu256((p + j + 2 * h));
__m256i x3 = ldu256((p + j + 3 * h));
__m256i y0 = ldu256((p + j + 8));
__m256i y1 = ldu256((p + j + h + 8));
__m256i y2 = ldu256((p + j + 2 * h + 8));
__m256i y3 = ldu256((p + j + 3 * h + 8));
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i A2 = vaddm(y0, y2, pv), B2 = vsubm(y0, y2, pv);
__m256i C2 = vaddm(y1, y3, pv), D2 = vsubm(y1, y3, pv);
__m256i u0 = vaddm(A, C, pv), u2 = vsubm(A, C, pv);
__m256i v0 = vaddm(A2, C2, pv), v2 = vsubm(A2, C2, pv);
__m256i E = vshoup(D, cI, cIs, pv);
__m256i F = vshoup(D2, cI, cIs, pv);
__m256i u1 = vaddm(B, E, pv), u3 = vsubm(B, E, pv);
__m256i v1 = vaddm(B2, F, pv), v3 = vsubm(B2, F, pv);
{ __m256i w = ldu256((t1 + j)); u1 = vshoup(u1, w, ldu256((s1 + j)), pv); }
{ __m256i w = ldu256((t2 + j)); u2 = vshoup(u2, w, ldu256((s2 + j)), pv); }
{ __m256i w = ldu256((t3 + j)); u3 = vshoup(u3, w, ldu256((s3 + j)), pv); }
{ __m256i w = ldu256((t1 + j + 8)); v1 = vshoup(v1, w, ldu256((s1 + j + 8)), pv); }
{ __m256i w = ldu256((t2 + j + 8)); v2 = vshoup(v2, w, ldu256((s2 + j + 8)), pv); }
{ __m256i w = ldu256((t3 + j + 8)); v3 = vshoup(v3, w, ldu256((s3 + j + 8)), pv); }
stu256((p + j), u0);
stu256((p + j + h), u1);
stu256((p + j + 2 * h), u2);
stu256((p + j + 3 * h), u3);
stu256((p + j + 8), v0);
stu256((p + j + h + 8), v1);
stu256((p + j + 2 * h + 8), v2);
stu256((p + j + 3 * h + 8), v3);
}
for (; j + 8 <= h; j += 8) {
__builtin_prefetch(p + j + 256);
__builtin_prefetch(p + j + h + 256);
__builtin_prefetch(p + j + 2 * h + 256);
__builtin_prefetch(p + j + 3 * h + 256);
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + h));
__m256i x2 = ldu256((p + j + 2 * h));
__m256i x3 = ldu256((p + j + 3 * h));
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i u0 = vaddm(A, C, pv), u2 = vsubm(A, C, pv);
__m256i E = vshoup(D, cI, cIs, pv);
__m256i u1 = vaddm(B, E, pv), u3 = vsubm(B, E, pv);
{ __m256i w = ldu256((t1 + j)); u1 = vshoup(u1, w, ldu256((s1 + j)), pv); }
{ __m256i w = ldu256((t2 + j)); u2 = vshoup(u2, w, ldu256((s2 + j)), pv); }
{ __m256i w = ldu256((t3 + j)); u3 = vshoup(u3, w, ldu256((s3 + j)), pv); }
stu256((p + j), u0);
stu256((p + j + h), u1);
stu256((p + j + 2 * h), u2);
stu256((p + j + 3 * h), u3);
}
}
}
template<int H>
TGT static void stage4_dif_h(u32 *a, u32 blk, u32 _h, const u32 *t1, const u32 *t2, const u32 *t3,
const u32 *s1, const u32 *s2, const u32 *s3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
for (u32 i = 0; i < blk; i += 4 * H) {
u32 *p = a + i;
u32 j = 0;
for (; j + 16 <= H; j += 16) {
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + H));
__m256i x2 = ldu256((p + j + 2 * H));
__m256i x3 = ldu256((p + j + 3 * H));
__m256i y0 = ldu256((p + j + 8));
__m256i y1 = ldu256((p + j + H + 8));
__m256i y2 = ldu256((p + j + 2 * H + 8));
__m256i y3 = ldu256((p + j + 3 * H + 8));
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i A2 = vaddm(y0, y2, pv), B2 = vsubm(y0, y2, pv);
__m256i C2 = vaddm(y1, y3, pv), D2 = vsubm(y1, y3, pv);
__m256i u0 = vaddm(A, C, pv), u2 = vsubm(A, C, pv);
__m256i v0 = vaddm(A2, C2, pv), v2 = vsubm(A2, C2, pv);
__m256i E = vshoup(D, cI, cIs, pv);
__m256i F = vshoup(D2, cI, cIs, pv);
__m256i u1 = vaddm(B, E, pv), u3 = vsubm(B, E, pv);
__m256i v1 = vaddm(B2, F, pv), v3 = vsubm(B2, F, pv);
{ __m256i w = ldu256((t1 + j)); u1 = vshoup(u1, w, ldu256((s1 + j)), pv); }
{ __m256i w = ldu256((t2 + j)); u2 = vshoup(u2, w, ldu256((s2 + j)), pv); }
{ __m256i w = ldu256((t3 + j)); u3 = vshoup(u3, w, ldu256((s3 + j)), pv); }
{ __m256i w = ldu256((t1 + j + 8)); v1 = vshoup(v1, w, ldu256((s1 + j + 8)), pv); }
{ __m256i w = ldu256((t2 + j + 8)); v2 = vshoup(v2, w, ldu256((s2 + j + 8)), pv); }
{ __m256i w = ldu256((t3 + j + 8)); v3 = vshoup(v3, w, ldu256((s3 + j + 8)), pv); }
stu256((p + j), u0);
stu256((p + j + H), u1);
stu256((p + j + 2 * H), u2);
stu256((p + j + 3 * H), u3);
stu256((p + j + 8), v0);
stu256((p + j + H + 8), v1);
stu256((p + j + 2 * H + 8), v2);
stu256((p + j + 3 * H + 8), v3);
}
for (; j + 8 <= H; j += 8) {
__builtin_prefetch(p + j + 256);
__builtin_prefetch(p + j + H + 256);
__builtin_prefetch(p + j + 2 * H + 256);
__builtin_prefetch(p + j + 3 * H + 256);
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + H));
__m256i x2 = ldu256((p + j + 2 * H));
__m256i x3 = ldu256((p + j + 3 * H));
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i u0 = vaddm(A, C, pv), u2 = vsubm(A, C, pv);
__m256i E = vshoup(D, cI, cIs, pv);
__m256i u1 = vaddm(B, E, pv), u3 = vsubm(B, E, pv);
{ __m256i w = ldu256((t1 + j)); u1 = vshoup(u1, w, ldu256((s1 + j)), pv); }
{ __m256i w = ldu256((t2 + j)); u2 = vshoup(u2, w, ldu256((s2 + j)), pv); }
{ __m256i w = ldu256((t3 + j)); u3 = vshoup(u3, w, ldu256((s3 + j)), pv); }
stu256((p + j), u0);
stu256((p + j + H), u1);
stu256((p + j + 2 * H), u2);
stu256((p + j + 3 * H), u3);
}
}
}
TGT static void stage4_dif_nt(u32 *a, u32 blk, u32 h, const u32 *t1, const u32 *t2, const u32 *t3,
const u32 *s1, const u32 *s2, const u32 *s3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
for (u32 i = 0; i < blk; i += 4 * h) {
u32 *p = a + i;
u32 j = 0;
for (; j + 16 <= h; j += 16) {
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + h));
__m256i x2 = ldu256((p + j + 2 * h));
__m256i x3 = ldu256((p + j + 3 * h));
__m256i y0 = ldu256((p + j + 8));
__m256i y1 = ldu256((p + j + h + 8));
__m256i y2 = ldu256((p + j + 2 * h + 8));
__m256i y3 = ldu256((p + j + 3 * h + 8));
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i A2 = vaddm(y0, y2, pv), B2 = vsubm(y0, y2, pv);
__m256i C2 = vaddm(y1, y3, pv), D2 = vsubm(y1, y3, pv);
__m256i u0 = vaddm(A, C, pv), u2 = vsubm(A, C, pv);
__m256i v0 = vaddm(A2, C2, pv), v2 = vsubm(A2, C2, pv);
__m256i E = vshoup(D, cI, cIs, pv);
__m256i F = vshoup(D2, cI, cIs, pv);
__m256i u1 = vaddm(B, E, pv), u3 = vsubm(B, E, pv);
__m256i v1 = vaddm(B2, F, pv), v3 = vsubm(B2, F, pv);
u1 = vshoup(u1, ldu256((t1 + j)), der_ws(ldu256((t1 + j))), pv);
u2 = vshoup(u2, ldu256((t2 + j)), der_ws(ldu256((t2 + j))), pv);
u3 = vshoup(u3, ldu256((t3 + j)), der_ws(ldu256((t3 + j))), pv);
v1 = vshoup(v1, ldu256((t1 + j + 8)), der_ws(ldu256((t1 + j + 8))), pv);
v2 = vshoup(v2, ldu256((t2 + j + 8)), der_ws(ldu256((t2 + j + 8))), pv);
v3 = vshoup(v3, ldu256((t3 + j + 8)), der_ws(ldu256((t3 + j + 8))), pv);
_mm256_stream_si256((__m256i *)(p + j), u0);
_mm256_stream_si256((__m256i *)(p + j + h), u1);
_mm256_stream_si256((__m256i *)(p + j + 2 * h), u2);
_mm256_stream_si256((__m256i *)(p + j + 3 * h), u3);
_mm256_stream_si256((__m256i *)(p + j + 8), v0);
_mm256_stream_si256((__m256i *)(p + j + h + 8), v1);
_mm256_stream_si256((__m256i *)(p + j + 2 * h + 8), v2);
_mm256_stream_si256((__m256i *)(p + j + 3 * h + 8), v3);
}
for (; j + 8 <= h; j += 8) {
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + h));
__m256i x2 = ldu256((p + j + 2 * h));
__m256i x3 = ldu256((p + j + 3 * h));
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i u0 = vaddm(A, C, pv), u2 = vsubm(A, C, pv);
__m256i E = vshoup(D, cI, cIs, pv);
__m256i u1 = vaddm(B, E, pv), u3 = vsubm(B, E, pv);
u1 = vshoup(u1, ldu256((t1 + j)), der_ws(ldu256((t1 + j))), pv);
u2 = vshoup(u2, ldu256((t2 + j)), der_ws(ldu256((t2 + j))), pv);
u3 = vshoup(u3, ldu256((t3 + j)), der_ws(ldu256((t3 + j))), pv);
_mm256_stream_si256((__m256i *)(p + j), u0);
_mm256_stream_si256((__m256i *)(p + j + h), u1);
_mm256_stream_si256((__m256i *)(p + j + 2 * h), u2);
_mm256_stream_si256((__m256i *)(p + j + 3 * h), u3);
}
}
}
TGT static void stage4_dif_h4(u32 *a, u32 nsub, const u32 *t1, const u32 *t2, const u32 *t3,
const u32 *s1, const u32 *s2, const u32 *s3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
if (nsub < 32) {
const __m128i pv4 = _mm256_castsi256_si128(pv), cI4 = _mm256_castsi256_si128(cI), cIs4 = _mm256_castsi256_si128(cIs);
__m128i W1 = _mm_loadu_si128((const __m128i *)t1), X1 = _mm_loadu_si128((const __m128i *)s1);
__m128i W2 = _mm_loadu_si128((const __m128i *)t2), X2 = _mm_loadu_si128((const __m128i *)s2);
__m128i W3 = _mm_loadu_si128((const __m128i *)t3), X3 = _mm_loadu_si128((const __m128i *)s3);
for (u32 i = 0; i < nsub; i += 16) {
u32 *p = a + i;
__m128i x0 = _mm_loadu_si128((const __m128i *)(p));
__m128i x1 = _mm_loadu_si128((const __m128i *)(p + 4));
__m128i x2 = _mm_loadu_si128((const __m128i *)(p + 8));
__m128i x3 = _mm_loadu_si128((const __m128i *)(p + 12));
__m128i A = add4(x0, x2, pv4), B = sub4(x0, x2, pv4);
__m128i C = add4(x1, x3, pv4), D = sub4(x1, x3, pv4);
__m128i u0 = add4(A, C, pv4), u2 = sub4(A, C, pv4);
__m128i E = sh4(D, cI4, cIs4, pv4);
__m128i u1 = add4(B, E, pv4), u3 = sub4(B, E, pv4);
u1 = sh4(u1, W1, X1, pv4); u2 = sh4(u2, W2, X2, pv4); u3 = sh4(u3, W3, X3, pv4);
_mm_storeu_si128((__m128i *)(p), u0);
_mm_storeu_si128((__m128i *)(p + 4), u1);
_mm_storeu_si128((__m128i *)(p + 8), u2);
_mm_storeu_si128((__m128i *)(p + 12), u3);
}
return;
}
__m256i W1 = bc128(_mm_loadu_si128((const __m128i *)t1));
__m256i W2 = bc128(_mm_loadu_si128((const __m128i *)t2));
__m256i W3 = bc128(_mm_loadu_si128((const __m128i *)t3));
__m256i X1 = bc128(_mm_loadu_si128((const __m128i *)s1));
__m256i X2 = bc128(_mm_loadu_si128((const __m128i *)s2));
__m256i X3 = bc128(_mm_loadu_si128((const __m128i *)s3));
for (u32 i = 0; i < nsub; i += 32) {
u32 *p = a + i;
__m256i x0 = ld2((const __m128i *)(p), (const __m128i *)(p + 16));
__m256i x1 = ld2((const __m128i *)(p + 4), (const __m128i *)(p + 20));
__m256i x2 = ld2((const __m128i *)(p + 8), (const __m128i *)(p + 24));
__m256i x3 = ld2((const __m128i *)(p + 12), (const __m128i *)(p + 28));
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i u0 = vaddm(A, C, pv), u2 = vsubm(A, C, pv);
__m256i E = vshoup(D, cI, cIs, pv);
__m256i u1 = vaddm(B, E, pv), u3 = vsubm(B, E, pv);
u1 = vshoup(u1, W1, X1, pv); u2 = vshoup(u2, W2, X2, pv); u3 = vshoup(u3, W3, X3, pv);
st2((__m128i *)(p), (__m128i *)(p + 16), u0);
st2((__m128i *)(p + 4), (__m128i *)(p + 20), u1);
st2((__m128i *)(p + 8), (__m128i *)(p + 24), u2);
st2((__m128i *)(p + 12), (__m128i *)(p + 28), u3);
}
}
// h==1 radix-4 (every stage twiddle at this level is w^0 = 1). 32 points per
// iteration: a 4x4 32-bit-lane transpose turns four 8-point shuffle groups
// (32 insns each) into one 67-instruction group.
TGT static void h1_rad4(u32 *a, u32 blk, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
const __m256i SORT = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7);
const __m256i RS = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
u32 i = 0;
for (; i + 32 <= blk; i += 32) {
u32 *p = a + i;
__m256i Y0 = ldu256((p));
__m256i Y1 = ldu256((p + 8));
__m256i Y2 = ldu256((p + 16));
__m256i Y3 = ldu256((p + 24));
__m256i T0 = _mm256_unpacklo_epi32(Y0, Y1);
__m256i T1 = _mm256_unpackhi_epi32(Y0, Y1);
__m256i T2 = _mm256_unpacklo_epi32(Y2, Y3);
__m256i T3 = _mm256_unpackhi_epi32(Y2, Y3);
__m256i Q0 = _mm256_permutevar8x32_epi32(T0, SORT);
__m256i Q1 = _mm256_permutevar8x32_epi32(T1, SORT);
__m256i Q2 = _mm256_permutevar8x32_epi32(T2, SORT);
__m256i Q3 = _mm256_permutevar8x32_epi32(T3, SORT);
__m256i X0 = _mm256_permute2x128_si256(Q0, Q2, 0x20);
__m256i X1 = _mm256_permute2x128_si256(Q0, Q2, 0x31);
__m256i X2 = _mm256_permute2x128_si256(Q1, Q3, 0x20);
__m256i X3 = _mm256_permute2x128_si256(Q1, Q3, 0x31);
__m256i A = vaddm(X0, X2, pv), B = vsubm(X0, X2, pv);
__m256i C = vaddm(X1, X3, pv), D = vsubm(X1, X3, pv);
__m256i E = vshoup(D, cI, cIs, pv);
__m256i U0 = vaddm(A, C, pv), U1 = vaddm(B, E, pv);
__m256i U2 = vsubm(A, C, pv), U3 = vsubm(B, E, pv);
__m256i PA = _mm256_unpacklo_epi32(U0, U1);
__m256i QA = _mm256_unpacklo_epi32(U2, U3);
__m256i RA = _mm256_unpackhi_epi32(U0, U1);
__m256i SA = _mm256_unpackhi_epi32(U2, U3);
__m256i W0 = _mm256_permute2x128_si256(PA, QA, 0x20);
__m256i W1 = _mm256_permute2x128_si256(PA, QA, 0x31);
__m256i W2 = _mm256_permute2x128_si256(RA, SA, 0x20);
__m256i W3 = _mm256_permute2x128_si256(RA, SA, 0x31);
stu256((p), _mm256_permutevar8x32_epi32(W0, RS));
stu256((p + 8), _mm256_permutevar8x32_epi32(W1, RS));
stu256((p + 16), _mm256_permutevar8x32_epi32(W2, RS));
stu256((p + 24), _mm256_permutevar8x32_epi32(W3, RS));
}
for (; i + 8 <= blk; i += 8) {
u32 *p = a + i;
__m256i x = ldu256(p);
__m256i t = _mm256_shuffle_epi32(x, _MM_SHUFFLE(1, 0, 3, 2));
__m256i s = vaddm(x, t, pv);
__m256i d = vsubm(x, t, pv);
__m256i Z = _mm256_shuffle_epi32(d, 0x00);
__m256i P = _mm256_blend_epi32(s, Z, 0xAA);
__m256i E = vshoup(d, cI, cIs, pv);
__m256i E2 = _mm256_shuffle_epi32(E, 0x55);
__m256i W = _mm256_shuffle_epi32(s, 0x55);
__m256i Q = _mm256_blend_epi32(W, E2, 0xAA);
__m256i sum = vaddm(P, Q, pv);
__m256i dif = vsubm(P, Q, pv);
stu256(p, _mm256_blend_epi32(sum, dif, 0xCC));
}
}
TGT static void stage4_dif_h1(u32 *a, u32 blk, u32 ci, u32 cis) { h1_rad4(a, blk, ci, cis); }
TGT static void stage2_dif(u32 *a, u32 blk, const u32 *t, const u32 *s) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
u32 h = blk >> 1;
for (u32 j = 0; j < h; j += 8) {
__m256i x0 = ldu256((a + j));
__m256i x1 = ldu256((a + j + h));
__m256i u0 = vaddm(x0, x1, pv);
__m256i u1 = vshoup(vsubm(x0, x1, pv), ldu256((t + j)),
ldu256((s + j)), pv);
stu256((a + j), u0);
stu256((a + j + h), u1);
}
}
// ================= inverse (DIT) =================
TGT static void merge4_dit(u32 *a, u32 blk, u32 h, const u32 *t1, const u32 *t2, const u32 *t3,
const u32 *s1, const u32 *s2, const u32 *s3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
for (u32 i = 0; i < blk; i += 4 * h) {
u32 *p = a + i;
for (u32 j = 0; j < h; j += 8) {
__builtin_prefetch(p + j + 256);
__builtin_prefetch(p + j + h + 256);
__builtin_prefetch(p + j + 2 * h + 256);
__builtin_prefetch(p + j + 3 * h + 256);
__builtin_prefetch(t1 + j + 256);
__builtin_prefetch(t2 + j + 256);
__builtin_prefetch(t3 + j + 256);
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + h));
__m256i x2 = ldu256((p + j + 2 * h));
__m256i x3 = ldu256((p + j + 3 * h));
{ __m256i w = ldu256((t1 + j)); x1 = vshoup(x1, w, der_ws(w), pv); }
{ __m256i w = ldu256((t2 + j)); x2 = vshoup(x2, w, der_ws(w), pv); }
{ __m256i w = ldu256((t3 + j)); x3 = vshoup(x3, w, der_ws(w), pv); }
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i E = vshoup(D, cI, cIs, pv);
stu256((p + j), vaddm(A, C, pv));
stu256((p + j + h), vaddm(B, E, pv));
stu256((p + j + 2 * h), vsubm(A, C, pv));
stu256((p + j + 3 * h), vsubm(B, E, pv));
}
}
}
TGT static void merge4_dit_s(u32 *a, u32 blk, u32 h, const u32 *t1, const u32 *t2, const u32 *t3,
const u32 *s1, const u32 *s2, const u32 *s3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
for (u32 i = 0; i < blk; i += 4 * h) {
u32 *p = a + i;
for (u32 j = 0; j < h; j += 8) {
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + h));
__m256i x2 = ldu256((p + j + 2 * h));
__m256i x3 = ldu256((p + j + 3 * h));
{ __m256i w = ldu256((t1 + j)); x1 = vshoup(x1, w, ldu256((s1 + j)), pv); }
{ __m256i w = ldu256((t2 + j)); x2 = vshoup(x2, w, ldu256((s2 + j)), pv); }
{ __m256i w = ldu256((t3 + j)); x3 = vshoup(x3, w, ldu256((s3 + j)), pv); }
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i E = vshoup(D, cI, cIs, pv);
stu256((p + j), vaddm(A, C, pv));
stu256((p + j + h), vaddm(B, E, pv));
stu256((p + j + 2 * h), vsubm(A, C, pv));
stu256((p + j + 3 * h), vsubm(B, E, pv));
}
}
}
template<int H>
TGT static void merge4_dit_h(u32 *a, u32 blk, u32 _h, const u32 *t1, const u32 *t2, const u32 *t3,
const u32 *s1, const u32 *s2, const u32 *s3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
for (u32 i = 0; i < blk; i += 4 * H) {
u32 *p = a + i;
for (u32 j = 0; j < H; j += 8) {
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + H));
__m256i x2 = ldu256((p + j + 2 * H));
__m256i x3 = ldu256((p + j + 3 * H));
{ __m256i w = ldu256((t1 + j)); x1 = vshoup(x1, w, ldu256((s1 + j)), pv); }
{ __m256i w = ldu256((t2 + j)); x2 = vshoup(x2, w, ldu256((s2 + j)), pv); }
{ __m256i w = ldu256((t3 + j)); x3 = vshoup(x3, w, ldu256((s3 + j)), pv); }
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i E = vshoup(D, cI, cIs, pv);
stu256((p + j), vaddm(A, C, pv));
stu256((p + j + H), vaddm(B, E, pv));
stu256((p + j + 2 * H), vsubm(A, C, pv));
stu256((p + j + 3 * H), vsubm(B, E, pv));
}
}
}
TGT static void merge4_dit_nt(u32 *a, u32 blk, u32 h, const u32 *t1, const u32 *t2, const u32 *t3,
const u32 *s1, const u32 *s2, const u32 *s3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
for (u32 i = 0; i < blk; i += 4 * h) {
u32 *p = a + i;
for (u32 j = 0; j < h; j += 8) {
__m256i x0 = ldu256((p + j));
__m256i x1 = ldu256((p + j + h));
__m256i x2 = ldu256((p + j + 2 * h));
__m256i x3 = ldu256((p + j + 3 * h));
x1 = vshoup(x1, ldu256((t1 + j)), der_ws(ldu256((t1 + j))), pv);
x2 = vshoup(x2, ldu256((t2 + j)), der_ws(ldu256((t2 + j))), pv);
x3 = vshoup(x3, ldu256((t3 + j)), der_ws(ldu256((t3 + j))), pv);
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i E = vshoup(D, cI, cIs, pv);
_mm256_stream_si256((__m256i *)(p + j), vaddm(A, C, pv));
_mm256_stream_si256((__m256i *)(p + j + h), vaddm(B, E, pv));
_mm256_stream_si256((__m256i *)(p + j + 2 * h), vsubm(A, C, pv));
_mm256_stream_si256((__m256i *)(p + j + 3 * h), vsubm(B, E, pv));
}
}
}
TGT static void merge4_dit_h4(u32 *a, u32 nsub, const u32 *t1, const u32 *t2, const u32 *t3,
const u32 *s1, const u32 *s2, const u32 *s3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
if (nsub < 32) {
const __m128i pv4 = _mm256_castsi256_si128(pv), cI4 = _mm256_castsi256_si128(cI), cIs4 = _mm256_castsi256_si128(cIs);
__m128i W1 = _mm_loadu_si128((const __m128i *)t1), X1 = _mm_loadu_si128((const __m128i *)s1);
__m128i W2 = _mm_loadu_si128((const __m128i *)t2), X2 = _mm_loadu_si128((const __m128i *)s2);
__m128i W3 = _mm_loadu_si128((const __m128i *)t3), X3 = _mm_loadu_si128((const __m128i *)s3);
for (u32 i = 0; i < nsub; i += 16) {
u32 *p = a + i;
__m128i x0 = _mm_loadu_si128((const __m128i *)(p));
__m128i x1 = sh4(_mm_loadu_si128((const __m128i *)(p + 4)), W1, X1, pv4);
__m128i x2 = sh4(_mm_loadu_si128((const __m128i *)(p + 8)), W2, X2, pv4);
__m128i x3 = sh4(_mm_loadu_si128((const __m128i *)(p + 12)), W3, X3, pv4);
__m128i A = add4(x0, x2, pv4), B = sub4(x0, x2, pv4);
__m128i C = add4(x1, x3, pv4), D = sub4(x1, x3, pv4);
__m128i E = sh4(D, cI4, cIs4, pv4);
_mm_storeu_si128((__m128i *)(p), add4(A, C, pv4));
_mm_storeu_si128((__m128i *)(p + 4), add4(B, E, pv4));
_mm_storeu_si128((__m128i *)(p + 8), sub4(A, C, pv4));
_mm_storeu_si128((__m128i *)(p + 12), sub4(B, E, pv4));
}
return;
}
__m256i W1 = bc128(_mm_loadu_si128((const __m128i *)t1));
__m256i W2 = bc128(_mm_loadu_si128((const __m128i *)t2));
__m256i W3 = bc128(_mm_loadu_si128((const __m128i *)t3));
__m256i X1 = bc128(_mm_loadu_si128((const __m128i *)s1));
__m256i X2 = bc128(_mm_loadu_si128((const __m128i *)s2));
__m256i X3 = bc128(_mm_loadu_si128((const __m128i *)s3));
for (u32 i = 0; i < nsub; i += 32) {
u32 *p = a + i;
__m256i x0 = ld2((const __m128i *)(p), (const __m128i *)(p + 16));
__m256i x1 = ld2((const __m128i *)(p + 4), (const __m128i *)(p + 20));
__m256i x2 = ld2((const __m128i *)(p + 8), (const __m128i *)(p + 24));
__m256i x3 = ld2((const __m128i *)(p + 12), (const __m128i *)(p + 28));
x1 = vshoup(x1, W1, X1, pv); x2 = vshoup(x2, W2, X2, pv); x3 = vshoup(x3, W3, X3, pv);
__m256i A = vaddm(x0, x2, pv), B = vsubm(x0, x2, pv);
__m256i C = vaddm(x1, x3, pv), D = vsubm(x1, x3, pv);
__m256i E = vshoup(D, cI, cIs, pv);
st2((__m128i *)(p), (__m128i *)(p + 16), vaddm(A, C, pv));
st2((__m128i *)(p + 4), (__m128i *)(p + 20), vaddm(B, E, pv));
st2((__m128i *)(p + 8), (__m128i *)(p + 24), vsubm(A, C, pv));
st2((__m128i *)(p + 12), (__m128i *)(p + 28), vsubm(B, E, pv));
}
}
TGT static void merge4_dit_h1(u32 *a, u32 blk, u32 ci, u32 cis) { h1_rad4(a, blk, ci, cis); }
TGT static void merge2_dit(u32 *a, u32 blk, const u32 *t, const u32 *s) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
u32 h = blk >> 1;
for (u32 j = 0; j < h; j += 8) {
__m256i x0 = ldu256((a + j));
__m256i x1 = vshoup(ldu256((a + j + h)),
ldu256((t + j)),
ldu256((s + j)), pv);
stu256((a + j), vaddm(x0, x1, pv));
stu256((a + j + h), vsubm(x0, x1, pv));
}
}
// ================= table generation =================
// dst[i] = v0*r^i ; sc[i] = floor(dst[i]*2^32/p) (within 1, sufficient for Shoup)
TGT static void gen_geo(u32 *dst, u32 *sc, u32 cnt, u32 v0, u32 r) {
if (!cnt) return;
const __m256i pv = _mm256_set1_epi32((int)g_P);
u32 seq[8], cur = v0 % g_P;
for (int i = 0; i < 8; i++) { seq[i] = cur; cur = mulmod(cur, r); }
u32 r8 = 1; for (int t = 0; t < 8; t++) r8 = mulmod(r8, r);
u32 rs8 = shoup_c(r8);
__m256i v = ldu256(seq);
__m256i r8v = _mm256_set1_epi32((int)r8), rs8v = _mm256_set1_epi32((int)rs8);
u32 i = 0;
if (NT_OK(dst, cnt) && NT_OK(sc, cnt)) {
for (; i + 8 <= cnt; i += 8) {
stu256_nt((dst + i), v);
stu256_nt((sc + i), der_ws(v));
v = vshoup(v, r8v, rs8v, pv);
}
if (i < cnt) {
u32 x = v0 % g_P;
for (u32 t = 0; t < i; t++) x = mulmod(x, r);
for (; i < cnt; i++) { dst[i] = x; sc[i] = shoup_c(x); x = mulmod(x, r); }
}
return;
}
for (; i + 8 <= cnt; i += 8) {
stu256((dst + i), v);
stu256((sc + i), der_ws(v));
v = vshoup(v, r8v, rs8v, pv);
}
if (i < cnt) {
u32 x = v0 % g_P;
for (u32 t = 0; t < i; t++) x = mulmod(x, r);
for (; i < cnt; i++) { dst[i] = x; sc[i] = shoup_c(x); x = mulmod(x, r); }
}
}
TGT static void gen_geo_val(u32 *dst, u32 cnt, u32 v0, u32 r) {
if (!cnt) return;
const __m256i pv = _mm256_set1_epi32((int)g_P);
u32 seq[8], cur = v0 % g_P;
for (int i = 0; i < 8; i++) { seq[i] = cur; cur = mulmod(cur, r); }
u32 r8 = 1; for (int t = 0; t < 8; t++) r8 = mulmod(r8, r);
u32 rs8 = shoup_c(r8);
__m256i v = ldu256(seq);
__m256i r8v = _mm256_set1_epi32((int)r8), rs8v = _mm256_set1_epi32((int)rs8);
u32 i = 0;
if (NT_OK(dst, cnt)) {
if (cnt >= 32) {
__m256i w1 = vshoup(v, r8v, rs8v, pv);
__m256i w2 = vshoup(w1, r8v, rs8v, pv);
__m256i w3 = vshoup(w2, r8v, rs8v, pv);
u32 r32 = r8; for (int t = 0; t < 3; t++) r32 = mulmod(r32, r8);
u32 rs32 = shoup_c(r32);
__m256i r32v = _mm256_set1_epi32((int)r32), rs32v = _mm256_set1_epi32((int)rs32);
for (; i + 32 <= cnt; i += 32) {
stu256_nt((dst + i), v); v = vshoup(v, r32v, rs32v, pv);
stu256_nt((dst + i + 8), w1); w1 = vshoup(w1, r32v, rs32v, pv);
stu256_nt((dst + i + 16), w2); w2 = vshoup(w2, r32v, rs32v, pv);
stu256_nt((dst + i + 24), w3); w3 = vshoup(w3, r32v, rs32v, pv);
}
}
for (; i + 8 <= cnt; i += 8) { stu256_nt((dst + i), v); v = vshoup(v, r8v, rs8v, pv); }
if (i < cnt) { u32 x = v0 % g_P; for (u32 t = 0; t < i; t++) x = mulmod(x, r); for (; i < cnt; i++) { dst[i] = x; x = mulmod(x, r); } }
return;
}
if (cnt >= 32) {
__m256i w1 = vshoup(v, r8v, rs8v, pv);
__m256i w2 = vshoup(w1, r8v, rs8v, pv);
__m256i w3 = vshoup(w2, r8v, rs8v, pv);
u32 r32 = r8; for (int t = 0; t < 3; t++) r32 = mulmod(r32, r8);
u32 rs32 = shoup_c(r32);
__m256i r32v = _mm256_set1_epi32((int)r32), rs32v = _mm256_set1_epi32((int)rs32);
for (; i + 32 <= cnt; i += 32) {
stu256((dst + i), v); v = vshoup(v, r32v, rs32v, pv);
stu256((dst + i + 8), w1); w1 = vshoup(w1, r32v, rs32v, pv);
stu256((dst + i + 16), w2); w2 = vshoup(w2, r32v, rs32v, pv);
stu256((dst + i + 24), w3); w3 = vshoup(w3, r32v, rs32v, pv);
}
}
for (; i + 8 <= cnt; i += 8) { stu256((dst + i), v); v = vshoup(v, r8v, rs8v, pv); }
if (i < cnt) { u32 x = v0 % g_P; for (u32 t = 0; t < i; t++) x = mulmod(x, r); for (; i < cnt; i++) { dst[i] = x; x = mulmod(x, r); } }
}
// Round a table cursor up to the next 32-byte boundary. Applied before every
// series group in build_tabs so that each sub-series start is 32-byte aligned
// whenever its size M is a multiple of 8 u32; series smaller than that have no
// vector loop at all, so no store in them needs the precondition.
static inline u32 *pad8(u32 *p) {
return (u32 *)(((uintptr_t)p + 31u) & ~(uintptr_t)31u);
}
static void build_tabs(Tabs &T, u32 *tab, u32 n, u32 Sb, int L, int odd, u32 root) {
u32 *p = tab;
u32 S = n;
for (int l = 0; l < L; l++) {
p = pad8(p);
u32 M = S >> 2;
u32 St = M + 32; /* L2 is 4-way with a 4 KB set period; M u32 is a multiple of
1024 u32 = 4 KB, so t1/t2/t3 would share one set. */
T.l1[l] = p; T.l2[l] = p + St; T.l3[l] = p + 2 * St;
T.m1[l] = p; T.m2[l] = p + St; T.m3[l] = p + 2 * St;
u32 base = powmod(root, n / S);
u32 b1 = base, b2 = mulmod(base, base), b3 = mulmod(b2, base);
gen_geo_val(p, M, 1, b1);
gen_geo_val(p + St, M, 1, b2);
gen_geo_val(p + 2 * St, M, 1, b3);
p += 3 * St;
S >>= 2;
}
T.b2 = 0; T.b2s = 0;
if (odd) {
p = pad8(p);
u32 M = Sb >> 1;
T.b2 = p; T.b2s = p + M;
gen_geo(p, p + M, M, 1, powmod(root, n / Sb));
p += 2 * M;
}
int s = 0;
u32 S2v = 4;
while (s < g_nb) {
p = pad8(p);
u32 M = S2v >> 2;
// The six base-case tables are M u32 each; at M = 4096 (16 KB) they would
// land on IDENTICAL L1/L2 sets (both spans are multiples of the 4 KB set
// period). One 64-byte line of gap per table rotates the set index, and
// keeps every start 32-byte aligned.
u32 St = M + 32;
T.h1[s] = p; T.h2[s] = p + St; T.h3[s] = p + 2 * St;
T.q1[s] = p + 3 * St; T.q2[s] = p + 4 * St; T.q3[s] = p + 5 * St;
u32 base = powmod(root, n / S2v);
u32 b1 = base, b2 = mulmod(base, base), b3 = mulmod(b2, base);
gen_geo(p, p + 3 * St, M, 1, b1);
gen_geo(p + St, p + 4 * St, M, 1, b2);
gen_geo(p + 2 * St, p + 5 * St, M, 1, b3);
p += 6 * St;
S2v <<= 2;
s++;
}
}
static u32 g_tabwords;
static u32 g_initP = 2013265921u;
static void ntt_init(u32 n) {
g_P = g_initP; g_pinv = 0;
{ u32 iv = 1; for (int q = 0; q < 5; q++) iv *= 2u - g_P * iv; g_pinv = 0u - iv;
u64 M = (u64)(~0ull) / g_P; u32 mh = (u32)(M >> 32); g_Msh = 0; while ((1u << g_Msh) < mh) g_Msh++; g_Ml = (u32)(M & 0xffffffffu); }
g_n = n;
int lg = 0; while ((1u << lg) < n) lg++;
g_log2n = lg;
const int BT = g_BT;
int L = 0;
while (lg - 2 * (L + 1) >= BT) L++;
g_L = L;
g_Sb = n >> (2 * L);
g_odd = ((lg - 2 * L) & 1) ? 1 : 0;
u32 S2 = g_Sb; if (g_odd) S2 >>= 1;
g_nb = 0; while (S2 >= 4) { S2 >>= 2; g_nb++; }
u32 w = 0;
{ u32 e = (g_P - 1) / n;
for (u32 g = 2; g < 400 && !w; g++) {
u32 c = powmod(g, e);
if (c != 1 && powmod(c, n >> 1) != 1) w = c; /* order exactly n (n = 2^k) */
} }
u32 wi = powmod(w, g_P - 2);
u32 iF = powmod(w, n >> 2), iI = powmod(wi, n >> 2);
g_ciF = iF; g_cisF = shoup_c(iF);
g_ciI = iI; g_cisI = shoup_c(iI);
{ u32 S=n; u32 tot=0; for (int l=0;l<L;l++){tot+=3*((S>>2)+16);S>>=2;} u32 S2v=4; for(int s=0;s<g_nb;s++){tot+=6*(S2v>>2);S2v<<=2;} if(g_odd) tot+=2*(g_Sb>>1); g_tabwords=tot; }
build_tabs(g_TF, g_tabF, n, g_Sb, L, g_odd, w);
build_tabs(g_TI, g_tabI, n, g_Sb, L, g_odd, wi);
// Every reader of the tables (ntt_fwd / ntt_inv) runs after this call, and the
// build above may have used non-temporal stores whose WC buffers ordinary loads
// do not snoop. One fence retires the whole build.
nt_fence();
}
TGT static void ntt_fwd(u32 *a, const u32 *src = 0, u32 len = 0) {
u32 S = g_n;
int l0 = 0;
if (src && g_L >= 1 && len >= (g_n >> 2) && len <= (g_n >> 1)) {
stage0_dif_src(a, src, len, g_n >> 2, g_TF.l1[0], g_TF.l2[0], g_TF.l3[0], g_ciF, g_cisF);
S = g_n >> 2;
l0 = 1;
}
for (int l = l0; l < g_L; l++) {
u32 h = S >> 2;
for (u32 b = 0; b < g_n; b += S)
stage4_dif(a + b, S, h, g_TF.l1[l], g_TF.l2[l], g_TF.l3[l], g_TF.m1[l], g_TF.m2[l], g_TF.m3[l], g_ciF, g_cisF);
S = h;
}
for (u32 b = 0; b < g_n; b += g_Sb) {
u32 *p = a + b;
if (g_odd) stage2_dif(p, g_Sb, g_TF.b2, g_TF.b2s);
for (int s = g_nb - 1; s >= 0; s--) {
u32 S2 = 4u << (2 * s);
if (S2 >= 32) {
for (u32 i = 0; i < g_Sb; i += S2)
{ u32 *pp = p + i;
switch (S2) {
case 16384: stage4_dif_h<4096>(pp, 16384, 4096, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF); break;
case 4096: stage4_dif_h<1024>(pp, 4096, 1024, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF); break;
case 1024: stage4_dif_h<256>(pp, 1024, 256, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF); break;
case 256: stage4_dif_h<64>(pp, 256, 64, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF); break;
case 64: stage4_dif_h<16>(pp, 64, 16, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF); break;
default: stage4_dif_s(pp, S2, S2 >> 2, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF);
} }
} else if (S2 == 16) {
stage4_dif_h4(p, g_Sb, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF);
} else {
stage4_dif_h1(p, g_Sb, g_ciF, g_cisF);
}
}
}
}
TGT static void ntt_inv(u32 *a, const u32 *Bmul = 0, u32 ninvarg = 1) {
for (u32 b = 0; b < g_n; b += g_Sb) {
u32 *p = a + b;
if (Bmul) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
u32 R2v = (u32)(((u64)1 << 32) % g_P * (((u64)1 << 32) % g_P) % g_P);
R2v = (u32)((u64)R2v * ninvarg % g_P);
u32 Cv = (u32)((u64)ninvarg * (((u64)1 << 32) % g_P) % g_P);
__m256i R2 = _mm256_set1_epi32((int)Cv);
const __m256i R2s = der_ws(R2);
const u32 *q = Bmul + b;
for (u32 i = 0; i < g_Sb; i += 8) {
__m256i x = ldu256((p + i));
__m256i y = ldu256((q + i));
stu256((p + i), vshoup(vmont(x, y, pv), R2, R2s, pv));
}
}
for (int s = 0; s < g_nb; s++) {
u32 S = 4u << (2 * s);
if (S >= 32) {
for (u32 i = 0; i < g_Sb; i += S)
{ u32 *pp = p + i;
switch (S) {
case 16384: merge4_dit_h<4096>(pp, 16384, 4096, g_TI.h1[s], g_TI.h2[s], g_TI.h3[s], g_TI.q1[s], g_TI.q2[s], g_TI.q3[s], g_ciI, g_cisI); break;
case 4096: merge4_dit_h<1024>(pp, 4096, 1024, g_TI.h1[s], g_TI.h2[s], g_TI.h3[s], g_TI.q1[s], g_TI.q2[s], g_TI.q3[s], g_ciI, g_cisI); break;
case 1024: merge4_dit_h<256>(pp, 1024, 256, g_TI.h1[s], g_TI.h2[s], g_TI.h3[s], g_TI.q1[s], g_TI.q2[s], g_TI.q3[s], g_ciI, g_cisI); break;
case 256: merge4_dit_h<64>(pp, 256, 64, g_TI.h1[s], g_TI.h2[s], g_TI.h3[s], g_TI.q1[s], g_TI.q2[s], g_TI.q3[s], g_ciI, g_cisI); break;
case 64: merge4_dit_h<16>(pp, 64, 16, g_TI.h1[s], g_TI.h2[s], g_TI.h3[s], g_TI.q1[s], g_TI.q2[s], g_TI.q3[s], g_ciI, g_cisI); break;
default: merge4_dit_s(pp, S, S >> 2, g_TI.h1[s], g_TI.h2[s], g_TI.h3[s], g_TI.q1[s], g_TI.q2[s], g_TI.q3[s], g_ciI, g_cisI);
} }
} else if (S == 16) {
merge4_dit_h4(p, g_Sb, g_TI.h1[s], g_TI.h2[s], g_TI.h3[s], g_TI.q1[s], g_TI.q2[s], g_TI.q3[s], g_ciI, g_cisI);
} else {
merge4_dit_h1(p, g_Sb, g_ciI, g_cisI);
}
}
if (g_odd) merge2_dit(p, g_Sb, g_TI.b2, g_TI.b2s);
}
for (int l = g_L - 1; l >= 0; l--) {
u32 SB = g_n >> (2 * l);
u32 h = SB >> 2;
for (u32 b = 0; b < g_n; b += SB)
merge4_dit(a + b, SB, h, g_TI.l1[l], g_TI.l2[l], g_TI.l3[l], g_TI.m1[l], g_TI.m2[l], g_TI.m3[l], g_ciI, g_cisI);
}
}
// ==================== 1004e7: decimal multiply, n = m = 1e7 digits ====================
struct DI {
unsigned long abi;
const char *s; unsigned long sn;
char *o; unsigned long ol; unsigned long os;
char *e; unsigned long el; unsigned long es;
const char *IB; unsigned long IBl;
char *OB; unsigned long OBl;
unsigned long tsc;
} __attribute__((packed));
static const u32 P1 = 1998585857u, P2 = 2013265921u, P3 = 2088763393u;
#define MAXLIMB 1001024
static u64 LB1[MAXLIMB], LB2[MAXLIMB];
static u32 RED1[MAXLIMB] __attribute__((aligned(64))), RED2[MAXLIMB] __attribute__((aligned(64)));
// Aligned so the final scale loop's 32-byte stores are 32-byte aligned: RA/RB/RC
// are written once and read once (by the CRT, long after eviction), so the RFO
// they cost under WB stores buys nothing.
static u32 RA[NTT_MAXN + 64] __attribute__((aligned(64)));
static u32 RB[NTT_MAXN + 64] __attribute__((aligned(64)));
static u32 RC[NTT_MAXN + 64] __attribute__((aligned(64)));
static u64 RO[MAXLIMB * 2 + 8];
static u32 PA_[NTT_MAXN], PB_[NTT_MAXN] __attribute__((aligned(64)));
static const char DIG2[201] =
"00010203040506070809101112131415161718192021222324252627282930313233343536373839"
"40414243444546474849505152535455565758596061626364656667686970717273747576777879"
"8081828384858687888990919293949596979899";
// little-endian limbs: out[0] = least significant group of up to 6 digits
// 8 ASCII digits -> numeric value (SWAR, 4 multiplies)
// 8 ASCII digits -> numeric value (SWAR, 4 multiplies)
// AVX2 scan: end of a run of ASCII digits starting at `from`.
TGT static long scan_digits(const char *s, long from, long sn) {
long i = from;
const __m256i lo = _mm256_set1_epi8((char)('0' - 1));
const __m256i hi = _mm256_set1_epi8((char)('9' + 1));
for (; i + 64 <= sn; i += 64) {
__m256i v = ldu256((s + i));
__m256i w = ldu256((s + i + 32));
__m256i d = _mm256_and_si256(_mm256_cmpgt_epi8(v, lo), _mm256_cmpgt_epi8(hi, v));
__m256i e = _mm256_and_si256(_mm256_cmpgt_epi8(w, lo), _mm256_cmpgt_epi8(hi, w));
unsigned m = (unsigned)_mm256_movemask_epi8(d) & (unsigned)_mm256_movemask_epi8(e);
if (m != 0xFFFFFFFFu) {
unsigned md = (unsigned)_mm256_movemask_epi8(d);
if (md != 0xFFFFFFFFu) return i + __builtin_ctz(~md);
unsigned me = (unsigned)_mm256_movemask_epi8(e);
return i + 32 + __builtin_ctz(~me);
}
}
for (; i + 32 <= sn; i += 32) {
__m256i v = ldu256((s + i));
__m256i d = _mm256_and_si256(_mm256_cmpgt_epi8(v, lo), _mm256_cmpgt_epi8(hi, v));
unsigned m = (unsigned)_mm256_movemask_epi8(d);
if (m != 0xFFFFFFFFu) return i + __builtin_ctz(~m);
}
while (i < sn && s[i] >= '0' && s[i] <= '9') i++;
return i;
}
static inline u64 swar8(const char *s) {
u64 x;
memcpy(&x, s, 8);
x -= 0x3030303030303030ull;
x = (x * 10 + (x >> 8)) & 0x00FF00FF00FF00FFull;
x = (x * 100 + (x >> 16)) & 0x0000FFFF0000FFFFull;
return (x * 10000 + (x >> 32)) & 0xFFFFFFFFull;
}
static inline u32 ld2d(const char *s) {
u32 v = (u32)(unsigned char)s[0] - '0';
return (u32)((u32)(unsigned char)s[1] - '0') + v * 10u;
}
// base 10^10 limbs, little-endian; emits nlm limbs
static u32 parse_limbs(const char *s, long len, u64 *out) {
long nlm = (len + 9) / 10;
if (nlm <= 0) { out[0] = 0; return 1; }
int rem = (int)(len - (nlm - 1) * 10);
u32 k = (u32)nlm;
long i;
{ u64 v = 0; for (int j = 0; j < rem; j++) v = v * 10 + (u32)(s[j] - '0'); out[--k] = v; i = rem; }
if (i + 40 <= len) {
const __m256i z0 = _mm256_set1_epi8((char)'0');
const __m256i p1010 = _mm256_set1_epi16(0x010A);
const __m256i p100_1 = _mm256_set1_epi32(0x00010064);
const __m128i z0h = _mm_set1_epi8((char)'0');
const __m128i p1010h = _mm_set1_epi16(0x010A);
const __m128i p100_1h = _mm_set1_epi32(0x00010064);
u32 G[12];
while (k >= 4 && i + 40 <= len) {
const char *q = s + i;
__m256i v0 = ldu256(q);
__m256i t0 = _mm256_maddubs_epi16(_mm256_sub_epi8(v0, z0), p1010);
stu256(G, _mm256_madd_epi16(t0, p100_1));
__m128i v1 = _mm_loadu_si128((const __m128i *)(q + 24));
__m128i t1 = _mm_maddubs_epi16(_mm_sub_epi8(v1, z0h), p1010h);
_mm_storeu_si128((__m128i *)(G + 8), _mm_madd_epi16(t1, p100_1h));
u32 g0 = G[0], g1 = G[1], g2 = G[2], g3 = G[3], g4 = G[4];
u32 g5 = G[5], g6 = G[6], g7 = G[7], g8 = G[10], g9 = G[11];
_mm_stream_si64((long long *)(out + k - 1), (long long)((u64)g0 * 1000000ull + (u64)g1 * 100ull + (u32)(g2 / 100u)));
_mm_stream_si64((long long *)(out + k - 2), (long long)((u64)(g2 % 100u) * 100000000ull + (u64)g3 * 10000ull + (u64)g4));
_mm_stream_si64((long long *)(out + k - 3), (long long)((u64)g5 * 1000000ull + (u64)g6 * 100ull + (u32)(g7 / 100u)));
_mm_stream_si64((long long *)(out + k - 4), (long long)((u64)(g7 % 100u) * 100000000ull + (u64)g8 * 10000ull + (u64)g9));
k -= 4; i += 40;
}
}
while (i + 10 <= len) { out[--k] = swar8(s + i) * 100ull + ld2d(s + i + 8); i += 10; }
while (k > 0) { u64 v = 0; for (long j = i; j < len; j++) v = v * 10 + (u64)(s[j] - '0'); out[--k] = v; i = len; }
return (u32)nlm;
}
static inline u32 mulmodp(u32 a, u32 b, u32 p) { return (u32)((u64)a * b % p); }
static u32 powmodp(u32 a, u32 e, u32 p) { u32 r = 1; a %= p; while (e) { if (e & 1) r = mulmodp(r, a, p); a = mulmodp(a, a, p); e >>= 1; } return r; }
#define RED1E(I) { u64 v = in[z+(I)]; u64 q = (v * mn) >> 40; u64 r = v - q * p; \
u64 t = r + p; r = (r > v) ? t : r; t = r - p; r = (r >= p) ? t : r; \
out[z+(I)] = (u32)r; }
// Exact reduction of v < 2^34 mod p (2^30 < p < 2^31) via 2^31 = (2^31-p) mod p.
// v = q*2^31 + s, q = v>>31 <= 4 ; t = q*(2^31-p) + s < 2p ; r = t - (t>=p ? p : 0).
TGT static inline __m256i red4(__m256i v, __m256i cV, __m256i pV, __m256i pm1V, __m256i m31V) {
__m256i q = _mm256_srli_epi64(v, 31);
__m256i s = _mm256_and_si256(v, m31V);
__m256i t = _mm256_add_epi64(_mm256_mul_epu32(q, cV), s);
__m256i ge = _mm256_cmpgt_epi64(t, pm1V);
return _mm256_sub_epi64(t, _mm256_and_si256(ge, pV));
}
static void redn(const u64 *in, u32 *out, u32 n, u32 p, u32 mn) {
const int nt_hint = NT_OK(out, n);
u32 z = 0;
const u64 c64 = (u64)2147483648u - p;
const __m256i cV = _mm256_set1_epi64x((long long)c64);
const __m256i pV = _mm256_set1_epi64x((long long)p);
const __m256i pm1V = _mm256_set1_epi64x((long long)(p - 1));
const __m256i m31V = _mm256_set1_epi64x(0x7fffffffLL);
const __m256i shuf = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
for (; z + 8 <= n; z += 8) {
__m256i v0 = ldu256((in + z));
__m256i v1 = ldu256((in + z + 4));
__m256i t0 = _mm256_permutevar8x32_epi32(red4(v0, cV, pV, pm1V, m31V), shuf);
__m256i t1 = _mm256_permutevar8x32_epi32(red4(v1, cV, pV, pm1V, m31V), shuf);
__m256i m = _mm256_inserti128_si256(t0, _mm256_castsi256_si128(t1), 1);
if (nt_hint) stu256_nt((out + z), m); else stu256((out + z), m);
}
for (; z < n; z++) {
u64 v = in[z];
u64 t = ((v >> 31) * c64) + (v & 0x7fffffff);
out[z] = (u32)(t >= p ? t - p : t);
}
}
__attribute__((target("avx2")))
// ==========================================================================
// FUSED level-0 DIF stage: reads the base-10^10 limbs DIRECTLY and reduces them
// in-register, so the 4 MB RED1/RED2 arrays and their write+read round trip do
// not exist. Byte-exact by construction: the reduction below is redn's red4 /
// scalar tail applied to the same limbs at the same indices.
// ==========================================================================
static inline __m256i ldred8(const u64 *lb, __m256i cV, __m256i pV, __m256i pm1V,
__m256i m31V, __m256i shuf) {
__m256i v0 = ldu256((const void *)lb);
__m256i v1 = ldu256((const void *)(lb + 4));
__m256i t0 = _mm256_permutevar8x32_epi32(red4(v0, cV, pV, pm1V, m31V), shuf);
__m256i t1 = _mm256_permutevar8x32_epi32(red4(v1, cV, pV, pm1V, m31V), shuf);
return _mm256_inserti128_si256(t0, _mm256_castsi256_si128(t1), 1);
}
static inline u32 red1s(u64 v, u64 c64, u32 p) {
u64 t = ((v >> 31) * c64) + (v & 0x7fffffff);
return (u32)(t >= p ? t - p : t);
}
TGT static void stage0_dif_u64(u32 *a, const u64 *lb, u32 len, u32 h,
const u32 *t1, const u32 *t2, const u32 *t3, u32 ci, u32 cis) {
const __m256i pv = _mm256_set1_epi32((int)g_P);
const __m256i cI = _mm256_set1_epi32((int)ci), cIs = _mm256_set1_epi32((int)cis);
const __m256i zero = _mm256_setzero_si256();
const u64 c64 = (u64)2147483648u - (u64)g_P;
const __m256i cV = _mm256_set1_epi64x((long long)c64);
const __m256i pV = _mm256_set1_epi64x((long long)g_P);
const __m256i pm1V = _mm256_set1_epi64x((long long)(g_P - 1));
const __m256i m31V = _mm256_set1_epi64x(0x7fffffffLL);
const __m256i shuf = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
u32 j = 0;
u32 lim1 = len - h;
for (; j + 8 <= lim1 + 1; j += 8) {
__m256i x0 = ldred8(lb + j, cV, pV, pm1V, m31V, shuf);
__m256i x1 = ldred8(lb + j + h, cV, pV, pm1V, m31V, shuf);
__m256i u0 = vaddm(x0, x1, pv), u2 = vsubm(x0, x1, pv);
__m256i E = vshoup(x1, cI, cIs, pv);
__m256i u1 = vaddm(x0, E, pv), u3 = vsubm(x0, E, pv);
{ __m256i w = ldu256((t1 + j)); u1 = vshoup(u1, w, der_ws(w), pv); }
{ __m256i w = ldu256((t2 + j)); u2 = vshoup(u2, w, der_ws(w), pv); }
{ __m256i w = ldu256((t3 + j)); u3 = vshoup(u3, w, der_ws(w), pv); }
stu256((a + j), u0);
stu256((a + j + h), u1);
stu256((a + j + 2 * h), u2);
stu256((a + j + 3 * h), u3);
}
for (; j <= lim1 && j < h; j++) {
u32 x0 = (j <= len) ? red1s(lb[j], c64, g_P) : 0;
u32 x1 = (j + h <= len) ? red1s(lb[j + h], c64, g_P) : 0;
u32 u0 = (u32)(((u64)x0 + x1) % g_P);
u32 u2 = (u32)(((u64)x0 + g_P - x1) % g_P);
u32 E = mulmod(x1, ci);
u32 u1 = (u32)(((u64)x0 + E) % g_P), u3 = (u32)(((u64)x0 + g_P - E) % g_P);
u1 = shoup1s(u1, t1[j], shoup_c(t1[j]));
u2 = shoup1s(u2, t2[j], shoup_c(t2[j]));
u3 = shoup1s(u3, t3[j], shoup_c(t3[j]));
a[j] = u0; a[j + h] = u1; a[j + 2 * h] = u2; a[j + 3 * h] = u3;
}
for (; j + 8 <= h; j += 8) {
__m256i x0 = (j + 8 <= len + 1) ? ldred8(lb + j, cV, pV, pm1V, m31V, shuf) : zero;
__m256i u0 = x0, u1 = x0, u2 = x0, u3 = x0;
{ __m256i w = ldu256((t1 + j)); u1 = vshoup(u1, w, der_ws(w), pv); }
{ __m256i w = ldu256((t2 + j)); u2 = vshoup(u2, w, der_ws(w), pv); }
{ __m256i w = ldu256((t3 + j)); u3 = vshoup(u3, w, der_ws(w), pv); }
stu256((a + j), u0);
stu256((a + j + h), u1);
stu256((a + j + 2 * h), u2);
stu256((a + j + 3 * h), u3);
}
for (; j < h; j++) {
u32 x0 = (j <= len) ? red1s(lb[j], c64, g_P) : 0;
u32 u1 = shoup1s(x0, t1[j], shoup_c(t1[j]));
u32 u2 = shoup1s(x0, t2[j], shoup_c(t2[j]));
u32 u3 = shoup1s(x0, t3[j], shoup_c(t3[j]));
a[j] = x0; a[j + h] = u1; a[j + 2 * h] = u2; a[j + 3 * h] = u3;
}
}
// ntt_fwd with the level-0 stage reading u64 limbs directly.
TGT static void ntt_fwd_u64(u32 *a, const u64 *lb, u32 len) {
stage0_dif_u64(a, lb, len, g_n >> 2, g_TF.l1[0], g_TF.l2[0], g_TF.l3[0], g_ciF, g_cisF);
u32 S = g_n >> 2;
int l0 = 1;
for (int l = l0; l < g_L; l++) {
u32 h = S >> 2;
for (u32 b = 0; b < g_n; b += S)
stage4_dif(a + b, S, h, g_TF.l1[l], g_TF.l2[l], g_TF.l3[l], g_TF.m1[l], g_TF.m2[l], g_TF.m3[l], g_ciF, g_cisF);
S = h;
}
for (u32 b = 0; b < g_n; b += g_Sb) {
u32 *p = a + b;
if (g_odd) stage2_dif(p, g_Sb, g_TF.b2, g_TF.b2s);
for (int s = g_nb - 1; s >= 0; s--) {
u32 S2 = 4u << (2 * s);
if (S2 >= 32) {
for (u32 i = 0; i < g_Sb; i += S2)
{ u32 *pp = p + i;
switch (S2) {
case 16384: stage4_dif_h<4096>(pp, 16384, 4096, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF); break;
case 4096: stage4_dif_h<1024>(pp, 4096, 1024, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF); break;
case 1024: stage4_dif_h<256>(pp, 1024, 256, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF); break;
case 256: stage4_dif_h<64>(pp, 256, 64, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF); break;
case 64: stage4_dif_h<16>(pp, 64, 16, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF); break;
default: stage4_dif_s(pp, S2, S2 >> 2, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF);
} }
} else if (S2 == 16) {
stage4_dif_h4(p, g_Sb, g_TF.h1[s], g_TF.h2[s], g_TF.h3[s], g_TF.q1[s], g_TF.q2[s], g_TF.q3[s], g_ciF, g_cisF);
} else {
stage4_dif_h1(p, g_Sb, g_ciF, g_cisF);
}
}
}
}
static void mul_run_job(DI *d) {
const char *s = d->s;
long sn = (long)d->sn;
long l1 = scan_digits(s, 0, sn);
long b0 = l1;
while (b0 < sn && (s[b0] < '0' || s[b0] > '9')) b0++;
// The second operand's digit run ends where the trailing non-digits begin, so
// it is reachable in O(trailing non-digits) instead of a second 10 MB pass.
long e2 = sn;
while (e2 > b0 && (s[e2 - 1] < '0' || s[e2 - 1] > '9')) e2--;
long l2 = e2 - b0;
if (l1 <= 0 || l2 <= 0) return;
u32 na = parse_limbs(s, l1, LB1);
u32 nb = parse_limbs(s + b0, l2, LB2);
// parse_limbs wrote LB1/LB2 with movnti; redn is their first reader.
nt_fence();
u32 need = na + nb;
u32 n = 1; while (n < need) n <<= 1;
u32 *RES[3]; RES[0] = RA; RES[1] = RB; RES[2] = RC;
const u32 PR[3] = {P1, P2, P3};
for (int pr = 0; pr < 3; pr++) {
g_initP = PR[pr];
ntt_init(n);
u32 *dst = RES[pr];
u32 *A = dst, *B = PB_;
ntt_fwd_u64(A, LB1, na);
ntt_fwd_u64(B, LB2, nb);
u32 ninv = powmodp(n % g_P, g_P - 2, g_P);
ntt_inv(A, B, ninv);
nt_fence();
}
// RC (the last RES array) was written with non-temporal stores above; the CRT
// below is the first reader of RA/RB/RC, so the WC buffers must be retired first.
nt_fence();
// ---------------- 3-prime CRT, base 10^10, u64-only carry ----------------
const u64 p1 = P1, p2 = P2, p3 = P3;
const u64 M2 = p1 * p2; // < 2^63
u32 i12 = powmodp((u32)(p1 % p2), P2 - 2, P2);
u32 Kinv2 = (u32)((u64)i12 % p2 * (((u64)1 << 32) % p2) % p2);
u32 pinv2 = 0; { u32 iv = 1; for (int q = 0; q < 5; q++) iv *= 2u - P2 * iv; pinv2 = 0u - iv; }
u32 iM2 = powmodp((u32)(M2 % p3), P3 - 2, P3);
u32 Kinv3 = (u32)((u64)iM2 % p3 * (((u64)1 << 32) % p3) % p3);
u32 pinv3 = 0; { u32 iv = 1; for (int q = 0; q < 5; q++) iv *= 2u - P3 * iv; pinv3 = 0u - iv; }
const u64 B10 = 10000000000ull;
const u64 Q64 = 1844674407ull; // floor(2^64 / 10^10)
const u64 R64 = 3709551616ull; // 2^64 mod 10^10
u64 carry = 0;
u64 *res = RO;
u32 j = 0, top = 0;
#define CRTRECON(J, cLo, cHi) do { \
u32 r1 = RA[J], r2 = RB[J], r3 = RC[J]; \
u32 dd2 = r2 - r1; if (dd2 >= P2) dd2 += P2; \
u64 x2 = (u64)dd2 * Kinv2; \
u32 m2 = (u32)x2 * pinv2; \
u64 u2 = (x2 + (u64)m2 * P2) >> 32; if (u2 >= P2) u2 -= P2; \
u64 c12 = (u64)r1 + p1 * u2; \
u32 dd3 = r3 - (u32)(c12 % P3); if (dd3 >= P3) dd3 += P3; \
u64 x3 = (u64)dd3 * Kinv3; \
u32 m3 = (u32)x3 * pinv3; \
u64 u3 = (x3 + (u64)m3 * P3) >> 32; if (u3 >= P3) u3 -= P3; \
unsigned __int128 t = (unsigned __int128)M2 * (u32)u3 + c12; \
cLo = (u64)t; cHi = (u64)(t >> 64); \
} while (0)
#define CRTEX(J, cLo, cHi) do { \
u64 cLm = (cLo) % B10, cLd = (cLo) / B10; \
u64 hi = (cHi) * R64 + cLm; \
u64 v = (cHi) * Q64 + cLd + hi / B10; \
u64 dg = hi % B10 + carry; \
res[J] = dg % B10; \
carry = v + dg / B10; \
} while (0)
for (; j + 8 <= need; j += 8) {
u64 L0,H0; u64 L1,H1; u64 L2,H2; u64 L3,H3; u64 L4,H4; u64 L5,H5; u64 L6,H6; u64 L7,H7;
CRTRECON(j+0,L0,H0); CRTRECON(j+1,L1,H1); CRTRECON(j+2,L2,H2); CRTRECON(j+3,L3,H3); CRTRECON(j+4,L4,H4); CRTRECON(j+5,L5,H5); CRTRECON(j+6,L6,H6); CRTRECON(j+7,L7,H7);
CRTEX(j+0,L0,H0); CRTEX(j+1,L1,H1); CRTEX(j+2,L2,H2); CRTEX(j+3,L3,H3); CRTEX(j+4,L4,H4); CRTEX(j+5,L5,H5); CRTEX(j+6,L6,H6); CRTEX(j+7,L7,H7);
}
for (; j < need; j++) { u64 L,H; CRTRECON(j,L,H); CRTEX(j,L,H); }
#undef CRTRECON
#undef CRTEX
top = need; while (top > 0 && res[top-1] == 0) top--;
while (carry) { res[top++] = carry % B10; carry /= B10; }
if (top == 0) top = 1;
// ---------------- output: 10 digits per limb ----------------
char *o = d->o;
char *olim = d->o + d->ol - 64;
{
u64 v = res[top - 1];
char tmp[24]; int k = 0;
if (v == 0) tmp[k++] = '0';
while (v) { tmp[k++] = (char)('0' + (v % 10)); v /= 10; }
while (k) { if (o > olim) { d->os = (unsigned long)(o - d->o); return; } *o++ = tmp[--k]; }
}
static u32 t2d[100], t4d[10000];
for (u32 j = 0; j < 100; j++) t2d[j] = (u32)('0' + j / 10) | ((u32)('0' + j % 10) << 8);
for (u32 j = 0; j < 10000; j++)
t4d[j] = (u32)('0' + j / 1000) | ((u32)('0' + (j / 100) % 10) << 8)
| ((u32)('0' + (j / 10) % 10) << 16) | ((u32)('0' + j % 10) << 24);
for (int i = (int)top - 2; i >= 0; ) {
int lim = i - 1024; if (lim < -1) lim = -1;
if (o + 10240 > olim) { d->os = (unsigned long)(o - d->o); return; }
for (; i > lim; i--) {
u64 v = res[i];
u64 h10 = v / 100000000ull;
u32 lo = (u32)(v - h10 * 100000000ull);
u32 B4 = lo / 10000u;
u32 C4 = lo - B4 * 10000u;
u64 w = (u64)t4d[B4] | ((u64)t4d[C4] << 32);
memcpy(o + 2, &w, 8);
memcpy(o, &t2d[h10], 2);
o += 10;
}
}
if (o <= olim) *o++ = '\n';
d->os = (unsigned long)(o - d->o);
}
#ifndef NO_LOCAL_TEST
extern "C" void __libc_start_main(void *m, int argc, char **argv) {
(void)m;
unsigned long *p = (unsigned long *)(argv + argc + 1);
while (*p) p++;
p++;
DI *d = 0;
for (int i = 0; i < 32 && p[0]; i++, p += 2)
if (p[0] == 0x6b637564UL) { d = (DI *)p[1]; break; }
if (d) mul_run_job(d);
__asm__ volatile("syscall" ::"a"(60), "D"(0) : "rcx", "r11", "memory");
for (;;);
}
int main() { return 0; }
#endif
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 120.447 ms | 97 MB + 952 KB | Accepted | Score: 100 | 显示更多 |