// 1002e8 poly_multiply (cyclic NTT mod 2013265921 + exact wrap correction)
#ifndef LGR_
#define LGR_ 12
#endif
#ifndef M3L_
#define M3L_ 14
#endif
#ifndef LGR2_
#define LGR2_ 13
#endif
#ifndef LGC2_
#define LGC2_ 13
#endif
#define EXACT 1
// pad (in u32) added to the Aarr/Barr row stride to break the
// power-of-two set conflicts of the 64x64 corner transposes.
#define PADC_ 16
#define PADC2_ 16
// C = 3 * 2^M3L_ (radix-3 column dimension), R = 2^LGR_ (pure power of two,
// so the corner's bit-reversal algebra over the R index stays exact).
// N1 = R*C = 3*2^26 = 201 326 592 >= n+m+1 = 200 000 001 => M < 0: the wrap
// region is EMPTY and the wrap-correction engine is never used.
#define N3_ 3
#define M3_ (1u << M3L_)
#define NR_ (1u << LGR_)
#define NC_ (N3_ * M3_)
#define N1_ ((u32)(NR_ * NC_))
#define CSX_ (NC_ + PADC_)
#define AOFF(k) (((long)(k) / (long)NC_) * CSX_ + ((long)(k) % (long)NC_))
#define NR2_ (1u << LGR2_)
#define NC2_ (1u << LGC2_)
#define N2_ (1u << (LGR2_ + LGC2_))
#define CSX2_ (NC2_ + PADC2_)
// ============================================================================
// 1002e8 core: NTT mod p = 2013265921 = 15*2^27+1 (primitive root 31)
//
// Transform of length N = R*C ("four step"):
// input (natural flat): x[j1*C + j0], j1 in [0,R), j0 in [0,C)
// scratch (transposed) : S[j0*R + j1]
// spectral output : X[k0*C + k1] holds A_hat[k1*R + k0]
// inverse output : S[j0*R + j1] holds x[j1*C + j0]
// ============================================================================
#include <cstdint>
#define NO_AVX2_UNUSED 0
#include <cstring>
#include <cstdio>
#include <ctime>
static double LMD_PT[24]; static long LMD_PC[24];
static inline double LMD_now(){struct timespec t;clock_gettime(CLOCK_MONOTONIC,&t);return t.tv_sec+1e-9*t.tv_nsec;}
static inline void LMD_acc(int i,double t0){LMD_PT[i]+=LMD_now()-t0;LMD_PC[i]++;}
static void LMD_dump(){const char*nm[24]={"E2_fwd","E2_inv","E2_untr","E2_pt","E1_fwd","E1_inv","E1_untr","E1_pt","p3","p1copy","loadblk","rowdif","corner","colfwd","x","x","x","x","x","x","x","x","x","x"};for(int i=0;i<24;i++)if(LMD_PC[i])fprintf(stderr,"LMD %-9s %8.4f s n=%ld\n",nm[i],LMD_PT[i],LMD_PC[i]);}
#ifndef NO_AVX2
#include <immintrin.h>
#endif
typedef uint16_t u16;
typedef uint32_t u32;
typedef uint64_t u64;
typedef int64_t i64;
typedef int32_t i32;
typedef unsigned __int128 u128;
// ===========================================================================
// SPLIT STORAGE FOR THE B/S ARRAY.
//
// The b spectrum B and the inverse's transposed scratch S are the same array at
// different times (every use of B ends at the pointwise; S is written only
// after it). In the radix-3 layout that array is 201.4 M u32 = 805.6 MB, and
// Aarr + it = 1.611 GB does not fit the 2 GiB limit together with the
// harness's own `c` (800 MB, which the mem_kb measurement DOES count).
//
// So the array is stored as the harness's `c` for its first g_BS elements and
// in a static tail for the rest. g_BS is a multiple of BOTH row strides
// (CSX_ for B, R for S), so no row ever straddles the split, and every access
// to a spectrum array is one row or one 64-row tile whose base is computed per
// row -- which is why one base pointer per row is enough.
//
// lcm(CSX_, R) = lcm(49168, 4096) = 12 587 008
// g_BS = 15 * 12 587 008 = 188 805 120 (<= n+m+1 = 200 000 001) [720 MB of c]
// tail = 201 392 128 - 188 805 120 = 12 587 008 u32 = 48.0 MB static
// ===========================================================================
#ifndef BSTAIL_SPLIT
#define BSTAIL_SPLIT 188805120L
#endif
static u32 *g_B = 0; // element 0 of the c-resident part
static u32 *g_BT = 0; // element BSTAIL_SPLIT
static inline u32 *BROW(long el) {
return (el < BSTAIL_SPLIT) ? g_B + el : g_BT + (el - BSTAIL_SPLIT);
}
// ------------------- forced single-instruction 256-bit memory ops ----------
// Ubuntu g++-9 defaults to -mavx256-split-unaligned-load/store: every
// _mm256_loadu_si256 becomes vmovdqu xmm + vinserti128 (2 insns + extra p5
// uop) and every store becomes vmovups xmm + vextracti128. The target pragma
// cannot turn those off, so emit the single vmovdqu ourselves. No alignment
// requirement -> safe on the harness-owned arrays.
#pragma GCC push_options
#pragma GCC target("avx2,tune=skylake")
static inline __attribute__((always_inline)) __m256i ldu256(const void *p) {
__m256i r; __asm__("vmovdqu %1, %0" : "=x"(r) : "m"(*(const __m256i *)p)); return r;
}
static inline __attribute__((always_inline)) void stu256(void *p, __m256i v) {
__asm__ volatile("vmovdqu %1, %0" : "=m"(*(__m256i *)p) : "x"(v));
}
// (a+b) mod P and (a-b) mod P for a,b < P, in 2 vector ops (P < 2^31, so a+b
// and (a-b) mod 2^32 both stay below 2P and unsigned-min selects the residue).
static inline __attribute__((always_inline)) __m256i vaddm256(__m256i a, __m256i b, __m256i Pv) {
__m256i s = _mm256_add_epi32(a, b);
return _mm256_min_epu32(s, _mm256_sub_epi32(s, Pv));
}
static inline __attribute__((always_inline)) __m256i vsubm256(__m256i a, __m256i b, __m256i Pv) {
__m256i d = _mm256_sub_epi32(a, b);
return _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
}
#undef CMBSTEP
#pragma GCC pop_options
#pragma GCC optimize("O3","rename-registers","live-range-shrinkage","ira-loop-pressure","no-stack-protector")
#define TS 64
// ---- streamed forward: one 64-column block of the transposed scratch lives in
// L1/L2 instead of in the 671 MB S array (which cost a full write + a full read
// per forward). The block holds TS rows of the corner's A index by R columns.
#define PADB_ 16
#define NRMAX_ ((NR_ > NR2_) ? NR_ : NR2_)
static const u32 P = 2013265921u;
static const u32 G = 31u;
static u64 IM;
static u32 TBUF[TS * TS];
static u32 TW[TS * TS];
static inline u32 mm(u32 a, u32 b) {
u64 z = (u64)a * b;
u64 q = (u64)(((u128)z * IM) >> 64);
u64 r = z - q * (u64)P;
return (u32)(r >= P ? r - P : r);
}
static inline u32 shoup_c(u32 w) { return (u32)(((u64)w << 32) / (u64)P); }
static u32 pw(u32 a, u32 e) {
u32 r = 1;
while (e) { if (e & 1) r = (u32)((u64)r * a % P); a = (u32)((u64)a * a % P); e >>= 1; }
return r;
}
// ---- staged twiddle table for a DIT row transform (bitrev in, natural out) ----
static void build_tab(u32 *tw, int L, u32 w) {
int off = 0;
for (int len = 2; len <= L; len <<= 1) {
u32 wl = pw(w, (u32)(L / len));
u32 c = 1;
int h = len >> 1;
for (int j = 0; j < h; j++) { tw[off + j] = c; c = mm(c, wl); }
off += h;
}
}
static void build_br(u16 *br, int L) {
int bits = 0; while ((1 << bits) < L) bits++;
for (int i = 0; i < L; i++) {
int r = 0;
for (int b = 0; b < bits; b++) if (i & (1 << b)) r |= 1 << (bits - 1 - b);
br[i] = (u16)r;
}
}
// ---------------- AVX2 kernels (Shoup modmul) ----------------
#ifdef NO_AVX2
#define HAVE_AVX2 0
#else
#define HAVE_AVX2 1
#pragma GCC push_options
#pragma GCC target("avx2,tune=skylake")
// 32x32 -> high 32, 8 lanes. AVX2 has no vpmulhu32, so build it from the two
// half-products; 5 instructions instead of the 6 the previous form paid, and it
// lets the whole Shoup run in 32-bit lanes (vpmulld) instead of 64-bit lanes
// carrying two packed 32-bit strands.
static inline __m256i vmulhi8(__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);
}
// Harvey-Shoup: q = floor(v*ws/2^32), r = v*w - q*P, reduced to [0,P).
// Requires v < 2^32 and w < P; ws = floor(w*2^32/P), which is what build_tab2's
// shoup_c() already stores.
static inline __attribute__((always_inline)) __m256i mullo256(__m256i a, __m256i b) {
__m256i r; __asm__("vpmulld %2, %1, %0" : "=x"(r) : "x"(a), "x"(b)); return r;
}
static inline __m256i shoup8(__m256i v, __m256i w, __m256i ws) {
const __m256i Pv = _mm256_set1_epi32((int)P);
__m256i t = mullo256(v, w);
__m256i q = vmulhi8(v, ws);
__m256i r = _mm256_sub_epi32(t, mullo256(q, Pv));
return _mm256_min_epu32(r, _mm256_sub_epi32(r, Pv));
}
#define SP(w, ws) _mm256_and_si256(w,M32), _mm256_srli_epi64(w,32), _mm256_and_si256(ws,M32), _mm256_srli_epi64(ws,32)
// ---------------------------------------------------------------------------
// len=2,4,8 stages as ONE in-register 8-point codelet. DIT with bit-reversed
// input: the three smallest stages act entirely inside each aligned 8-block, so
// one YMM register holds a whole 8-block and the stages are lane shuffles plus
// three Shoup multiplies with pre-broadcast twiddle vectors. The twiddles are
// the first 7 entries of the staged table (len=2: 1 entry, len=4: 2, len=8: 4),
// so this must run before the len>=16 stages.
static void row_dit_avx(u32 *x, int L, const u32 *tw, const u32 *tws) {
__m256i Pv = _mm256_set1_epi32((int)P), M1 = _mm256_set1_epi32(-1), Z = _mm256_setzero_si256();
int off = 0;
if (L >= 8) {
const __m256i V4 = _mm256_setr_epi32((int)tw[1], (int)tw[2], 1, 1, (int)tw[1], (int)tw[2], 1, 1);
const __m256i V4s = _mm256_setr_epi32((int)tws[1], (int)tws[2], 2, 2, (int)tws[1], (int)tws[2], 2, 2);
const __m256i V8 = _mm256_setr_epi32((int)tw[3], (int)tw[4], (int)tw[5], (int)tw[6], 1, 1, 1, 1);
const __m256i V8s = _mm256_setr_epi32((int)tws[3], (int)tws[4], (int)tws[5], (int)tws[6], 2, 2, 2, 2);
const __m256i IDX8 = _mm256_setr_epi32(4, 5, 6, 7, 0, 1, 2, 3);
const __m256i IDXQ4 = _mm256_setr_epi32(0, 1, 0, 1, 4, 5, 4, 5);
const __m256i IDXQ8 = _mm256_setr_epi32(0, 1, 2, 3, 0, 1, 2, 3);
int i4 = 0;
for (; i4 + 64 <= L; i4 += 64) {
__m256i vr0 = ldu256((const void *)((x + i4 + 0)));
__m256i vr1 = ldu256((const void *)((x + i4 + 8)));
__m256i vr2 = ldu256((const void *)((x + i4 + 16)));
__m256i vr3 = ldu256((const void *)((x + i4 + 24)));
__m256i vr4 = ldu256((const void *)((x + i4 + 32)));
__m256i vr5 = ldu256((const void *)((x + i4 + 40)));
__m256i vr6 = ldu256((const void *)((x + i4 + 48)));
__m256i vr7 = ldu256((const void *)((x + i4 + 56)));
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr0, 0xB1); // [1,0,3,2,5,4,7,6]
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr0); // d[2j+1] = vr0[2j]-vr0[2j+1]
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr0 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr1, 0xB1); // [1,0,3,2,5,4,7,6]
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr1); // d[2j+1] = vr1[2j]-vr1[2j+1]
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr1 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr2, 0xB1); // [1,0,3,2,5,4,7,6]
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr2); // d[2j+1] = vr2[2j]-vr2[2j+1]
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr2 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr3, 0xB1); // [1,0,3,2,5,4,7,6]
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr3); // d[2j+1] = vr3[2j]-vr3[2j+1]
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr3 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr4, 0xB1); // [1,0,3,2,5,4,7,6]
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr4); // d[2j+1] = vr4[2j]-vr4[2j+1]
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr4 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr5, 0xB1); // [1,0,3,2,5,4,7,6]
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr5); // d[2j+1] = vr5[2j]-vr5[2j+1]
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr5 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr6, 0xB1); // [1,0,3,2,5,4,7,6]
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr6); // d[2j+1] = vr6[2j]-vr6[2j+1]
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr6 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr7, 0xB1); // [1,0,3,2,5,4,7,6]
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr7); // d[2j+1] = vr7[2j]-vr7[2j+1]
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr7 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=4: pairs (0,2),(1,3) and (4,6),(5,7), twiddles [1,w4]
__m256i y = _mm256_shuffle_epi32(vr0, 0x4E); // [2,3,0,1,6,7,4,5]
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr0, y), _mm256_add_epi32(_mm256_sub_epi32(vr0, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4); // [_,_,d0,d1,_,_,d4,d5]
vr0 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (0,2),(1,3) and (4,6),(5,7), twiddles [1,w4]
__m256i y = _mm256_shuffle_epi32(vr1, 0x4E); // [2,3,0,1,6,7,4,5]
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr1, y), _mm256_add_epi32(_mm256_sub_epi32(vr1, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4); // [_,_,d0,d1,_,_,d4,d5]
vr1 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (0,2),(1,3) and (4,6),(5,7), twiddles [1,w4]
__m256i y = _mm256_shuffle_epi32(vr2, 0x4E); // [2,3,0,1,6,7,4,5]
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr2, y), _mm256_add_epi32(_mm256_sub_epi32(vr2, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4); // [_,_,d0,d1,_,_,d4,d5]
vr2 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (0,2),(1,3) and (4,6),(5,7), twiddles [1,w4]
__m256i y = _mm256_shuffle_epi32(vr3, 0x4E); // [2,3,0,1,6,7,4,5]
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr3, y), _mm256_add_epi32(_mm256_sub_epi32(vr3, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4); // [_,_,d0,d1,_,_,d4,d5]
vr3 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (0,2),(1,3) and (4,6),(5,7), twiddles [1,w4]
__m256i y = _mm256_shuffle_epi32(vr4, 0x4E); // [2,3,0,1,6,7,4,5]
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr4, y), _mm256_add_epi32(_mm256_sub_epi32(vr4, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4); // [_,_,d0,d1,_,_,d4,d5]
vr4 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (0,2),(1,3) and (4,6),(5,7), twiddles [1,w4]
__m256i y = _mm256_shuffle_epi32(vr5, 0x4E); // [2,3,0,1,6,7,4,5]
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr5, y), _mm256_add_epi32(_mm256_sub_epi32(vr5, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4); // [_,_,d0,d1,_,_,d4,d5]
vr5 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (0,2),(1,3) and (4,6),(5,7), twiddles [1,w4]
__m256i y = _mm256_shuffle_epi32(vr6, 0x4E); // [2,3,0,1,6,7,4,5]
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr6, y), _mm256_add_epi32(_mm256_sub_epi32(vr6, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4); // [_,_,d0,d1,_,_,d4,d5]
vr6 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (0,2),(1,3) and (4,6),(5,7), twiddles [1,w4]
__m256i y = _mm256_shuffle_epi32(vr7, 0x4E); // [2,3,0,1,6,7,4,5]
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr7, y), _mm256_add_epi32(_mm256_sub_epi32(vr7, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4); // [_,_,d0,d1,_,_,d4,d5]
vr7 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=8: pairs (0,4),(1,5),(2,6),(3,7), twiddles [1,w8,w8^2,w8^3]
__m256i y = _mm256_permutevar8x32_epi32(vr0, IDX8); // [4,5,6,7,0,1,2,3]
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr0, y), _mm256_add_epi32(_mm256_sub_epi32(vr0, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8); // [d0,d1,d2,d3,d0,d1,d2,d3]
vr0 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (0,4),(1,5),(2,6),(3,7), twiddles [1,w8,w8^2,w8^3]
__m256i y = _mm256_permutevar8x32_epi32(vr1, IDX8); // [4,5,6,7,0,1,2,3]
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr1, y), _mm256_add_epi32(_mm256_sub_epi32(vr1, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8); // [d0,d1,d2,d3,d0,d1,d2,d3]
vr1 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (0,4),(1,5),(2,6),(3,7), twiddles [1,w8,w8^2,w8^3]
__m256i y = _mm256_permutevar8x32_epi32(vr2, IDX8); // [4,5,6,7,0,1,2,3]
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr2, y), _mm256_add_epi32(_mm256_sub_epi32(vr2, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8); // [d0,d1,d2,d3,d0,d1,d2,d3]
vr2 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (0,4),(1,5),(2,6),(3,7), twiddles [1,w8,w8^2,w8^3]
__m256i y = _mm256_permutevar8x32_epi32(vr3, IDX8); // [4,5,6,7,0,1,2,3]
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr3, y), _mm256_add_epi32(_mm256_sub_epi32(vr3, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8); // [d0,d1,d2,d3,d0,d1,d2,d3]
vr3 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (0,4),(1,5),(2,6),(3,7), twiddles [1,w8,w8^2,w8^3]
__m256i y = _mm256_permutevar8x32_epi32(vr4, IDX8); // [4,5,6,7,0,1,2,3]
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr4, y), _mm256_add_epi32(_mm256_sub_epi32(vr4, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8); // [d0,d1,d2,d3,d0,d1,d2,d3]
vr4 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (0,4),(1,5),(2,6),(3,7), twiddles [1,w8,w8^2,w8^3]
__m256i y = _mm256_permutevar8x32_epi32(vr5, IDX8); // [4,5,6,7,0,1,2,3]
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr5, y), _mm256_add_epi32(_mm256_sub_epi32(vr5, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8); // [d0,d1,d2,d3,d0,d1,d2,d3]
vr5 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (0,4),(1,5),(2,6),(3,7), twiddles [1,w8,w8^2,w8^3]
__m256i y = _mm256_permutevar8x32_epi32(vr6, IDX8); // [4,5,6,7,0,1,2,3]
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr6, y), _mm256_add_epi32(_mm256_sub_epi32(vr6, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8); // [d0,d1,d2,d3,d0,d1,d2,d3]
vr6 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (0,4),(1,5),(2,6),(3,7), twiddles [1,w8,w8^2,w8^3]
__m256i y = _mm256_permutevar8x32_epi32(vr7, IDX8); // [4,5,6,7,0,1,2,3]
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(vr7, y), _mm256_add_epi32(_mm256_sub_epi32(vr7, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8); // [d0,d1,d2,d3,d0,d1,d2,d3]
vr7 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
stu256((void *)((x + i4 + 0)), vr0);
stu256((void *)((x + i4 + 8)), vr1);
stu256((void *)((x + i4 + 16)), vr2);
stu256((void *)((x + i4 + 24)), vr3);
stu256((void *)((x + i4 + 32)), vr4);
stu256((void *)((x + i4 + 40)), vr5);
stu256((void *)((x + i4 + 48)), vr6);
stu256((void *)((x + i4 + 56)), vr7);
}
for (; i4 < L; i4 += 8) {
__m256i v = ldu256((const void *)((x + i4)));
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(v, 0xB1); // [1,0,3,2,5,4,7,6]
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, v); // d[2j+1] = v[2j]-v[2j+1]
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
v = _mm256_blend_epi32(sa, d, 0xAA);
}
{ // len=4: pairs (0,2),(1,3) and (4,6),(5,7), twiddles [1,w4]
__m256i y = _mm256_shuffle_epi32(v, 0x4E); // [2,3,0,1,6,7,4,5]
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(v, y), _mm256_add_epi32(_mm256_sub_epi32(v, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4); // [_,_,d0,d1,_,_,d4,d5]
v = _mm256_blend_epi32(sa, d, 0xCC);
}
{ // len=8: pairs (0,4),(1,5),(2,6),(3,7), twiddles [1,w8,w8^2,w8^3]
__m256i y = _mm256_permutevar8x32_epi32(v, IDX8); // [4,5,6,7,0,1,2,3]
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(v, y), _mm256_add_epi32(_mm256_sub_epi32(v, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8); // [d0,d1,d2,d3,d0,d1,d2,d3]
v = _mm256_blend_epi32(sa, d, 0xF0);
}
stu256((void *)((x + i4)), v);
}
off = 7;
// Fused radix-4 DIT: stages (len, 2len) in one pass over four h=len/2
// quarters -- half the loads/stores of two separate radix-2 passes, same
// modmul count and the same twiddle-table segments as the radix-2 form.
int len = 16;
for (; (len << 1) <= L; len <<= 2) {
int h = len >> 1; // quarter size
const u32 *t1 = tw + off, *ts1 = tws + off; // w_len^j
const u32 *t2 = t1 + h, *ts2 = ts1 + h; // w_{2len}^j
const u32 *t3 = t2 + h, *ts3 = ts2 + h; // w_{2len}^{h+j}
off += 3 * h;
for (int i = 0; i < L; i += (len << 1)) {
u32 *p = x + i;
for (int j = 0; j < h; j += 8) {
__m256i WV1 = ldu256((const void *)((t1 + j))), WS1 = ldu256((const void *)((ts1 + j)));
__m256i WV2 = ldu256((const void *)((t2 + j))), WS2 = ldu256((const void *)((ts2 + j)));
__m256i WV3 = ldu256((const void *)((t3 + j))), WS3 = ldu256((const void *)((ts3 + j)));
__m256i A = ldu256((const void *)(p + j));
__m256i B = ldu256((const void *)(p + j + h));
__m256i C = ldu256((const void *)(p + j + 2 * h));
__m256i D = ldu256((const void *)(p + j + 3 * h));
__m256i Bm = shoup8(B, WV1, WS1), Dm = shoup8(D, WV1, WS1);
__m256i Ap = vaddm256(A, Bm, Pv), Bp = vsubm256(A, Bm, Pv);
__m256i Cp = vaddm256(C, Dm, Pv), Dp = vsubm256(C, Dm, Pv);
__m256i Cm = shoup8(Cp, WV2, WS2), Dn = shoup8(Dp, WV3, WS3);
stu256((void *)(p + j), vaddm256(Ap, Cm, Pv));
stu256((void *)(p + j + 2 * h), vsubm256(Ap, Cm, Pv));
stu256((void *)(p + j + h), vaddm256(Bp, Dn, Pv));
stu256((void *)(p + j + 3 * h), vsubm256(Bp, Dn, Pv));
}
}
}
if (len <= L) { // odd stage count
int h = len >> 1;
const u32 *t = tw + off, *ts = tws + off; off += h;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j += 8) {
__m256i u = ldu256((const void *)(p + j));
__m256i v = shoup8(ldu256((const void *)(q + j)),
ldu256((const void *)(t + j)), ldu256((const void *)(ts + j)));
stu256((void *)(p + j), vaddm256(u, v, Pv));
stu256((void *)(q + j), vsubm256(u, v, Pv));
}
}
}
} else {
for (int len = 2; len <= L; len <<= 1) {
int h = len >> 1;
const u32 *t = tw + off; off += h;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j++) {
u32 u = p[j], v = mm(q[j], t[j]);
u32 sa = u + v; if (sa >= P) sa -= P;
u32 d = u + P - v; if (d >= P) d -= P;
p[j] = sa; q[j] = d;
}
}
}
}
}
// ---------------------------------------------------------------------------
// DIF row transform: natural input, output at position m holds the value of
// index rev(m) (length preserved). It is the exact mirror of bitrev+row_dit,
// so a DIF forward pairs with a bitrev-free DIT inverse -- no bitrev passes.
// Stage order is len = L .. 16; the last three stages (8,4,2) act inside each
// aligned 8-block and are done in registers, exactly as the DIT codelet does
// for the first three.
static void row_dif_avx(u32 *x, int L, const u32 *tw, const u32 *tws) {
__m256i Pv = _mm256_set1_epi32((int)P), M1 = _mm256_set1_epi32(-1), Z = _mm256_setzero_si256();
if (L >= 8) {
// Fused radix-4: the stages len and len/2 in one pass over four h2-blocks.
// 4 loads / 4 stores / 4 shoup8 per 32 points per 2 stages (half the memory
// ops and ~14 % fewer uops than two separate radix-2 passes).
int len = L;
for (; len >= 32; len >>= 2) {
int h = len >> 1, h2 = len >> 2;
const u32 *t1 = tw + h - 1, *ts1 = tws + h - 1; // w_len^j (j < h2)
const u32 *t2 = t1 + h2, *ts2 = ts1 + h2; // w_len^{j+h2}
const u32 *t3 = tw + h2 - 1, *ts3 = tws + h2 - 1; // w_{len/2}^j = w_len^{2j}
for (int i = 0; i < L; i += len) {
u32 *p = x + i;
for (int j = 0; j < h2; j += 8) {
__m256i WV1 = ldu256((const void *)((t1 + j))), WS1 = ldu256((const void *)((ts1 + j)));
__m256i WV2 = ldu256((const void *)((t2 + j))), WS2 = ldu256((const void *)((ts2 + j)));
__m256i WV3 = ldu256((const void *)((t3 + j))), WS3 = ldu256((const void *)((ts3 + j)));
__m256i v0 = ldu256((const void *)((p + j)));
__m256i v1 = ldu256((const void *)((p + j + h2)));
__m256i v2 = ldu256((const void *)((p + j + h)));
__m256i v3 = ldu256((const void *)((p + j + h + h2)));
__m256i s02 = _mm256_min_epu32(_mm256_add_epi32(v0, v2), _mm256_sub_epi32(_mm256_add_epi32(v0, v2), Pv));
__m256i d02 = _mm256_sub_epi32(v0, v2);
__m256i s13 = _mm256_min_epu32(_mm256_add_epi32(v1, v3), _mm256_sub_epi32(_mm256_add_epi32(v1, v3), Pv));
__m256i d13 = _mm256_sub_epi32(v1, v3);
d02 = _mm256_min_epu32(d02, _mm256_add_epi32(d02, Pv));
d13 = _mm256_min_epu32(d13, _mm256_add_epi32(d13, Pv));
d02 = shoup8(d02, WV1, WS1);
d13 = shoup8(d13, WV2, WS2);
__m256i ra = _mm256_min_epu32(_mm256_add_epi32(s02, s13), _mm256_sub_epi32(_mm256_add_epi32(s02, s13), Pv));
__m256i rb = _mm256_min_epu32(_mm256_sub_epi32(s02, s13), _mm256_add_epi32(_mm256_sub_epi32(s02, s13), Pv));
rb = shoup8(rb, WV3, WS3);
__m256i rc = _mm256_min_epu32(_mm256_add_epi32(d02, d13), _mm256_sub_epi32(_mm256_add_epi32(d02, d13), Pv));
__m256i rd = _mm256_min_epu32(_mm256_sub_epi32(d02, d13), _mm256_add_epi32(_mm256_sub_epi32(d02, d13), Pv));
rd = shoup8(rd, WV3, WS3);
stu256((void *)((p + j)), ra);
stu256((void *)((p + j + h2)), rb);
stu256((void *)((p + j + h)), rc);
stu256((void *)((p + j + h + h2)), rd);
}
}
}
if (len >= 16) { // odd stage count: one radix-2 DIF
int h = len >> 1;
const u32 *t = tw + h - 1, *ts = tws + h - 1;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j += 8) {
__m256i u = ldu256((const void *)((p + j)));
__m256i v = ldu256((const void *)((q + j)));
__m256i w = ldu256((const void *)((t + j)));
__m256i ws = ldu256((const void *)((ts + j)));
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(u, v), _mm256_sub_epi32(_mm256_add_epi32(u, v), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(u, v), _mm256_add_epi32(_mm256_sub_epi32(u, v), Pv));
d = shoup8(d, w, ws);
stu256((void *)((p + j)), sa);
stu256((void *)((q + j)), d);
}
}
}
const __m256i V8 = _mm256_setr_epi32(1, 1, 1, 1, (int)tw[3], (int)tw[4], (int)tw[5], (int)tw[6]);
const __m256i V8s = _mm256_setr_epi32(2, 2, 2, 2, (int)tws[3], (int)tws[4], (int)tws[5], (int)tws[6]);
const __m256i V4 = _mm256_setr_epi32(1, 1, (int)tw[1], (int)tw[2], 1, 1, (int)tw[1], (int)tw[2]);
const __m256i V4s = _mm256_setr_epi32(2, 2, (int)tws[1], (int)tws[2], 2, 2, (int)tws[1], (int)tws[2]);
const __m256i IDX8 = _mm256_setr_epi32(4, 5, 6, 7, 0, 1, 2, 3);
int i4 = 0;
for (; i4 + 64 <= L; i4 += 64) {
__m256i vr0 = ldu256((const void *)((x + i4 + 0)));
__m256i vr1 = ldu256((const void *)((x + i4 + 8)));
__m256i vr2 = ldu256((const void *)((x + i4 + 16)));
__m256i vr3 = ldu256((const void *)((x + i4 + 24)));
__m256i vr4 = ldu256((const void *)((x + i4 + 32)));
__m256i vr5 = ldu256((const void *)((x + i4 + 40)));
__m256i vr6 = ldu256((const void *)((x + i4 + 48)));
__m256i vr7 = ldu256((const void *)((x + i4 + 56)));
{
{ // len=8: pairs (j, j+4); the difference gets w8^0..w8^3
__m256i y = _mm256_permutevar8x32_epi32(vr0, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr0); // lanes 4..7 = x_j - x_{j+4}
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr0 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (j, j+4); the difference gets w8^0..w8^3
__m256i y = _mm256_permutevar8x32_epi32(vr1, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr1); // lanes 4..7 = x_j - x_{j+4}
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr1 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (j, j+4); the difference gets w8^0..w8^3
__m256i y = _mm256_permutevar8x32_epi32(vr2, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr2); // lanes 4..7 = x_j - x_{j+4}
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr2 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (j, j+4); the difference gets w8^0..w8^3
__m256i y = _mm256_permutevar8x32_epi32(vr3, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr3); // lanes 4..7 = x_j - x_{j+4}
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr3 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (j, j+4); the difference gets w8^0..w8^3
__m256i y = _mm256_permutevar8x32_epi32(vr4, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr4); // lanes 4..7 = x_j - x_{j+4}
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr4 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (j, j+4); the difference gets w8^0..w8^3
__m256i y = _mm256_permutevar8x32_epi32(vr5, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr5); // lanes 4..7 = x_j - x_{j+4}
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr5 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (j, j+4); the difference gets w8^0..w8^3
__m256i y = _mm256_permutevar8x32_epi32(vr6, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr6); // lanes 4..7 = x_j - x_{j+4}
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr6 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=8: pairs (j, j+4); the difference gets w8^0..w8^3
__m256i y = _mm256_permutevar8x32_epi32(vr7, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, vr7); // lanes 4..7 = x_j - x_{j+4}
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr7 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{ // len=4: pairs (j, j+2) per half; differences get [1, w4]
__m256i y = _mm256_shuffle_epi32(vr0, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr0); // lanes 2,3,6,7 = diffs
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr0 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (j, j+2) per half; differences get [1, w4]
__m256i y = _mm256_shuffle_epi32(vr1, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr1); // lanes 2,3,6,7 = diffs
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr1 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (j, j+2) per half; differences get [1, w4]
__m256i y = _mm256_shuffle_epi32(vr2, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr2); // lanes 2,3,6,7 = diffs
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr2 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (j, j+2) per half; differences get [1, w4]
__m256i y = _mm256_shuffle_epi32(vr3, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr3); // lanes 2,3,6,7 = diffs
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr3 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (j, j+2) per half; differences get [1, w4]
__m256i y = _mm256_shuffle_epi32(vr4, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr4); // lanes 2,3,6,7 = diffs
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr4 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (j, j+2) per half; differences get [1, w4]
__m256i y = _mm256_shuffle_epi32(vr5, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr5); // lanes 2,3,6,7 = diffs
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr5 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (j, j+2) per half; differences get [1, w4]
__m256i y = _mm256_shuffle_epi32(vr6, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr6); // lanes 2,3,6,7 = diffs
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr6 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=4: pairs (j, j+2) per half; differences get [1, w4]
__m256i y = _mm256_shuffle_epi32(vr7, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d = _mm256_sub_epi32(y, vr7); // lanes 2,3,6,7 = diffs
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr7 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr0, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr0, y), _mm256_sub_epi32(_mm256_add_epi32(vr0, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr0), _mm256_add_epi32(_mm256_sub_epi32(y, vr0), Pv));
vr0 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr1, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr1, y), _mm256_sub_epi32(_mm256_add_epi32(vr1, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr1), _mm256_add_epi32(_mm256_sub_epi32(y, vr1), Pv));
vr1 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr2, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr2, y), _mm256_sub_epi32(_mm256_add_epi32(vr2, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr2), _mm256_add_epi32(_mm256_sub_epi32(y, vr2), Pv));
vr2 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr3, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr3, y), _mm256_sub_epi32(_mm256_add_epi32(vr3, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr3), _mm256_add_epi32(_mm256_sub_epi32(y, vr3), Pv));
vr3 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr4, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr4, y), _mm256_sub_epi32(_mm256_add_epi32(vr4, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr4), _mm256_add_epi32(_mm256_sub_epi32(y, vr4), Pv));
vr4 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr5, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr5, y), _mm256_sub_epi32(_mm256_add_epi32(vr5, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr5), _mm256_add_epi32(_mm256_sub_epi32(y, vr5), Pv));
vr5 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr6, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr6, y), _mm256_sub_epi32(_mm256_add_epi32(vr6, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr6), _mm256_add_epi32(_mm256_sub_epi32(y, vr6), Pv));
vr6 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(vr7, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(vr7, y), _mm256_sub_epi32(_mm256_add_epi32(vr7, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, vr7), _mm256_add_epi32(_mm256_sub_epi32(y, vr7), Pv));
vr7 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
stu256((void *)((x + i4 + 0)), vr0);
stu256((void *)((x + i4 + 8)), vr1);
stu256((void *)((x + i4 + 16)), vr2);
stu256((void *)((x + i4 + 24)), vr3);
stu256((void *)((x + i4 + 32)), vr4);
stu256((void *)((x + i4 + 40)), vr5);
stu256((void *)((x + i4 + 48)), vr6);
stu256((void *)((x + i4 + 56)), vr7);
}
for (; i4 < L; i4 += 8) {
__m256i v = ldu256((const void *)((x + i4)));
{ // len=8: pairs (j, j+4); the difference gets w8^0..w8^3
__m256i y = _mm256_permutevar8x32_epi32(v, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d0 = _mm256_sub_epi32(y, v); // lanes 4..7 = x_j - x_{j+4}
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
v = _mm256_blend_epi32(sa, d, 0xF0);
}
{ // len=4: pairs (j, j+2) per half; differences get [1, w4]
__m256i y = _mm256_shuffle_epi32(v, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_sub_epi32(y, v); // lanes 2,3,6,7 = diffs
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
v = _mm256_blend_epi32(sa, d, 0xCC);
}
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(v, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, v), _mm256_add_epi32(_mm256_sub_epi32(y, v), Pv));
v = _mm256_blend_epi32(sa, d, 0xAA);
}
stu256((void *)((x + i4)), v);
}
} else {
for (int len = L; len >= 2; len >>= 1) {
int h = len >> 1;
const u32 *t = tw + h - 1;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j++) {
u32 u = p[j], v = q[j];
u32 sa = u + v; if (sa >= P) sa -= P;
u32 d = u + P - v; if (d >= P) d -= P;
p[j] = sa; q[j] = mm(d, t[j]);
}
}
}
}
}
// Shoup constant recovered on the fly: ws = floor(w*2^32/P) == 2w + mulhi32(w,MLO)
// (error <= 1, absorbed by shoup8's single conditional subtract).
static inline __m256i shoup_of(__m256i w) {
const __m256i ML = _mm256_set1_epi32(572662301);
__m256i lo = _mm256_mul_epu32(w, ML);
__m256i hi = _mm256_mul_epu32(_mm256_srli_epi64(w, 32), ML);
__m256i m = _mm256_blend_epi32(_mm256_srli_epi64(lo, 32), _mm256_slli_epi64(_mm256_srli_epi64(hi, 32), 32), 0xAA);
return _mm256_add_epi32(_mm256_add_epi32(w, w), m);
}
// 8x8 in-register transpose of u32 lanes
static inline void t8x8(__m256i *r) {
__m256i t0 = _mm256_unpacklo_epi32(r[0], r[1]), t1 = _mm256_unpackhi_epi32(r[0], r[1]);
__m256i t2 = _mm256_unpacklo_epi32(r[2], r[3]), t3 = _mm256_unpackhi_epi32(r[2], r[3]);
__m256i t4 = _mm256_unpacklo_epi32(r[4], r[5]), t5 = _mm256_unpackhi_epi32(r[4], r[5]);
__m256i t6 = _mm256_unpacklo_epi32(r[6], r[7]), t7 = _mm256_unpackhi_epi32(r[6], r[7]);
__m256i u0 = _mm256_unpacklo_epi64(t0, t2), u1 = _mm256_unpackhi_epi64(t0, t2);
__m256i u2 = _mm256_unpacklo_epi64(t1, t3), u3 = _mm256_unpackhi_epi64(t1, t3);
__m256i u4 = _mm256_unpacklo_epi64(t4, t6), u5 = _mm256_unpackhi_epi64(t4, t6);
__m256i u6 = _mm256_unpacklo_epi64(t5, t7), u7 = _mm256_unpackhi_epi64(t5, t7);
r[0] = _mm256_permute2x128_si256(u0, u4, 0x20); r[1] = _mm256_permute2x128_si256(u1, u5, 0x20);
r[2] = _mm256_permute2x128_si256(u2, u6, 0x20); r[3] = _mm256_permute2x128_si256(u3, u7, 0x20);
r[4] = _mm256_permute2x128_si256(u0, u4, 0x31); r[5] = _mm256_permute2x128_si256(u1, u5, 0x31);
r[6] = _mm256_permute2x128_si256(u2, u6, 0x31); r[7] = _mm256_permute2x128_si256(u3, u7, 0x31);
}
// A[i] = A[i]*B[i]*sc mod P, whole-array, AVX2
static void pointwise_mul_avx(u32 *A, const u32 *B, u32 n, u32 sc, u32 sx) {
// `sc` (= 1/N1) is no longer applied here: it is folded into corner_inv's
// twiddle-table SEED, where it costs ZERO extra modmuls instead of one per element.
// (void)sc keeps the signature so the call sites are unchanged.
(void)sc;
long k = 0;
for (int r = 0; k < (long)n; r++) {
long lim = (long)sx - PADC_;
if (lim > (long)n - k) lim = (long)n - k;
u32 *a = A + (long)r * sx;
const u32 *b = BROW((long)r * sx);
long i = 0;
for (; i + 8 <= lim; i += 8) {
__m256i x = ldu256((const void *)(a + i));
__m256i y = ldu256((const void *)(b + i));
stu256((void *)(a + i), shoup8(x, y, shoup_of(y)));
}
for (; i < lim; i++) a[i] = mm(a[i], b[i]);
k += lim;
}
}
// ---- Montgomery arithmetic (corner only): mont_mul(a,b) = a*b*R^-1 mod P, R = 2^32.
// The corner's twiddle tables are kept in MONTGOMERY FORM (TW_M = TW*R) and the data
// stays in normal form, so mont8(x, TW_M) = x*TW -- i.e. the corner gets its modmul with
// NO per-element Shoup reciprocal (shoup_of costs 9 uops + 1 load per 8 modmuls, and
// mont8 is 16 uops against shoup8's 20).
static u32 NPI = 0; // -P^-1 mod 2^32
static const u32 R32 = (u32)((1ull << 32) % P);
static inline u32 mont1(u32 a, u32 b) {
u64 z = (u64)a * (u64)b;
u32 m = (u32)((u64)(u32)z * (u64)NPI);
u64 t = (z + (u64)m * (u64)P) >> 32;
return (u32)(t >= P ? t - P : t);
}
static inline __m256i mont8(__m256i a, __m256i b) {
__m256i P64 = _mm256_set1_epi64x((long long)P), NP = _mm256_set1_epi64x((long long)NPI);
__m256i Pv = _mm256_set1_epi32((int)P);
__m256i ahi = _mm256_srli_epi64(a, 32), bhi = _mm256_srli_epi64(b, 32);
__m256i tlo = _mm256_mul_epu32(a, b), thi = _mm256_mul_epu32(ahi, bhi);
__m256i e2 = _mm256_mul_epu32(tlo, NP), o2 = _mm256_mul_epu32(thi, NP);
__m256i e4 = _mm256_mul_epu32(e2, P64), o4 = _mm256_mul_epu32(o2, P64);
__m256i e5 = _mm256_add_epi64(tlo, e4), o5 = _mm256_add_epi64(thi, o4);
__m256i ue = _mm256_srli_epi64(e5, 32), uo = _mm256_srli_epi64(o5, 32);
__m256i u = _mm256_blend_epi32(ue, _mm256_slli_epi64(uo, 32), 0xAA);
return _mm256_min_epu32(u, _mm256_sub_epi32(u, Pv));
}
// dst[v] = mont(src[v], kM) over one TS-row (both operands Montgomery-form)
static void scale_row_mont(u32 *dst, const u32 *src, u32 kM) {
const __m256i kv = _mm256_set1_epi32((int)kM); // full 8-lane broadcast (NOT set1_epi64x:
// that would zero the odd u32 lanes)
for (int v = 0; v < TS; v += 8)
stu256((void *)((dst + v)), mont8(ldu256((const void *)((src + v))), kv));
}
// TW[u] = mont(TW[u-1], RVM)
static void twiddle_rows_mont(u32 *TW, const u32 *RVM) {
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v += 8)
stu256((void *)((cur + v)), mont8(ldu256((const void *)((prv + v))), ldu256((const void *)((RVM + v)))));
}
}
// dst[v] = src[v]*k mod P over one TS-row
static void scale_row_avx(u32 *dst, const u32 *src, u32 k) {
const __m256i kv = _mm256_set1_epi32((int)k), ks = _mm256_set1_epi32((int)shoup_c(k));
for (int v = 0; v < TS; v += 8)
stu256((void *)((dst + v)), shoup8(ldu256((const void *)((src + v))), kv, ks));
}
// row recurrence TW[u][v] = TW[u-1][v]*RV[v] mod P, 512 modmuls per call -> 64 shoup8
static void twiddle_rows_avx(u32 *TW, const u32 *RV) {
__m256i RVS[TS / 8];
for (int v = 0; v < TS; v += 8) RVS[v >> 3] = shoup_of(ldu256((const void *)((RV + v))));
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v += 8) {
__m256i r = ldu256((const void *)((RV + v)));
stu256((void *)((cur + v)), shoup8(ldu256((const void *)((prv + v))), r, RVS[v >> 3]));
}
}
}
// corner tile data path: whole 64x64 tile through registers (no TBUF round trip)
static void corner_tile_avx(const u32 *in, u32 *out, int A0, int B0, int strIn, int strOut, int split) {
for (int u0 = 0; u0 < TS; u0 += 8) {
for (int v0 = 0; v0 < TS; v0 += 8) {
__m256i r[8];
for (int k = 0; k < 8; k++) {
const u32 *s = in + (u64)(A0 + u0 + k) * strIn + B0 + v0;
const u32 *tw = TW + (u0 + k) * TS + v0;
__m256i x = ldu256((const void *)(s));
__m256i w = ldu256((const void *)(tw)); // TW is in Montgomery form
r[k] = mont8(x, w);
}
t8x8(r);
for (int k = 0; k < 8; k++)
stu256((void *)(TBUF + (v0 + k) * TS + u0), r[k]);
}
}
for (int v = 0; v < TS; v++) { // 256 B contiguous per destination row
u32 *d = split ? BROW((long)(B0 + v) * strOut + A0) : out + (u64)(B0 + v) * strOut + A0;
const u32 *s = TBUF + v * TS;
for (int u = 0; u < TS; u += 8) stu256((void *)(d + u), ldu256((const void *)(s + u)));
}
}
// Corner tile for the streamed forward: the A index is block-local (0..TS-1).
static void corner_tile_avx_blk(const u32 *in, u32 *out, int A0c, int B0, int strIn, int strOut, int split) {
for (int u0 = 0; u0 < TS; u0 += 8) {
for (int v0 = 0; v0 < TS; v0 += 8) {
__m256i r[8];
for (int k = 0; k < 8; k++) {
const u32 *s = in + (u64)(u0 + k) * strIn + B0 + v0;
const u32 *tw = TW + (u0 + k) * TS + v0;
r[k] = mont8(ldu256((const void *)(s)), ldu256((const void *)(tw)));
}
t8x8(r);
for (int k = 0; k < 8; k++)
stu256((void *)(TBUF + (v0 + k) * TS + u0), r[k]);
}
}
for (int v = 0; v < TS; v++) {
u32 *d = split ? BROW((long)(B0 + v) * strOut + A0c) : out + (u64)(B0 + v) * strOut + A0c;
const u32 *s = TBUF + v * TS;
for (int u = 0; u < TS; u += 8) stu256((void *)(d + u), ldu256((const void *)(s + u)));
}
}
// load_T fast tile: src tile -> transposed into S, with the -4 shift; allzero=1
// when the whole tile lies past nsrc (copy the zeros straight through).
static void load_tile_avx(const u32 *src, u32 *S, int r0, int c0, int R, int C, int ctr, int allzero) {
for (int rb0 = 0; rb0 < TS; rb0 += 8) {
for (int cb0 = 0; cb0 < TS; cb0 += 16) {
__m256i r[8];
if (allzero) {
for (int k = 0; k < 8; k++) r[k] = _mm256_setzero_si256();
} else {
for (int k = 0; k < 8; k++) {
const u32 *s = src + (u64)(r0 + rb0 + k) * C + c0 + cb0;
__m256i x = ldu256((const void *)(s));
if (ctr) {
__m256i y = _mm256_sub_epi32(x, _mm256_set1_epi32(4));
__m256i z = _mm256_add_epi32(x, _mm256_set1_epi32((int)(P - 4)));
x = _mm256_min_epu32(y, z);
}
r[k] = x;
}
}
t8x8(r);
for (int j = 0; j < 8; j++)
stu256((void *)(TBUF + (cb0 + j) * TS + rb0), r[j]);
{
int cb0b = cb0 + 8;
__m256i r2[8];
if (allzero) {
for (int k = 0; k < 8; k++) r2[k] = _mm256_setzero_si256();
} else {
for (int k = 0; k < 8; k++) {
const u32 *s = src + (u64)(r0 + rb0 + k) * C + c0 + cb0b;
__m256i x = ldu256((const void *)(s));
if (ctr) {
__m256i y = _mm256_sub_epi32(x, _mm256_set1_epi32(4));
__m256i z = _mm256_add_epi32(x, _mm256_set1_epi32((int)(P - 4)));
x = _mm256_min_epu32(y, z);
}
r2[k] = x;
}
}
t8x8(r2);
for (int j = 0; j < 8; j++)
stu256((void *)(TBUF + (cb0b + j) * TS + rb0), r2[j]);
}
}
}
// emit: 256 B contiguous per destination row, from an L1-resident source
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)(c0 + cb) * R + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb += 8) stu256((void *)(d + rb), ldu256((const void *)(s + rb)));
}
}
// Same 64x64 tile loader, but the transposed tile is emitted into a block whose
// row (A) stride is `dstr` rather than the full-array stride R.
static void load_tile_avx_d(const u32 *src, u32 *S, int r0, int c0, int C, int dstr, int ctr, int allzero) {
for (int rb0 = 0; rb0 < TS; rb0 += 8) {
for (int cb0 = 0; cb0 < TS; cb0 += 8) {
__m256i r[8];
if (allzero) {
for (int k = 0; k < 8; k++) r[k] = _mm256_setzero_si256();
} else {
for (int k = 0; k < 8; k++) {
const u32 *s = src + (u64)(r0 + rb0 + k) * C + c0 + cb0;
__m256i x = ldu256((const void *)(s));
if (ctr) {
__m256i y = _mm256_sub_epi32(x, _mm256_set1_epi32(4));
__m256i z = _mm256_add_epi32(x, _mm256_set1_epi32((int)(P - 4)));
x = _mm256_min_epu32(y, z);
}
r[k] = x;
}
}
t8x8(r);
for (int j = 0; j < 8; j++)
stu256((void *)(TBUF + (cb0 + j) * TS + rb0), r[j]);
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)cb * dstr + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb += 8) stu256((void *)(d + rb), ldu256((const void *)(s + rb)));
}
}
// load one 64-column block of `src` (natural layout, C columns) transposed into
// S[u*dstr + j1], u = 0..63 local column, j1 = 0..R-1.
static void load_block(const u32 *src, u32 nsrc, u32 *S, int c0, int R, int C, int dstr, int ctr) {
for (int r0 = 0; r0 < R; r0 += TS) {
u64 firstidx = (u64)r0 * C + c0;
u64 lastidx = (u64)(r0 + TS - 1) * C + c0 + TS - 1;
if (lastidx < nsrc) { load_tile_avx_d(src, S, r0, c0, C, dstr, ctr, 0); continue; }
if (firstidx >= nsrc) { load_tile_avx_d(src, S, r0, c0, C, dstr, ctr, 1); continue; }
for (int rb = 0; rb < TS; rb++) {
u64 base = (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u64 id = base + cb;
u32 v = 0;
if (id < nsrc) { v = src[id]; if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u)); }
TBUF[cb * TS + rb] = v;
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)cb * dstr + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb++) d[rb] = s[rb];
}
}
}
// untranspose tile (plain corner pattern, no twiddles)
// straight copy of an exact u32 run (used by the phase-1 E copy)
static void copy_u32_avx(u32 *dst, const u32 *src, long n) {
long i = 0;
for (; i + 8 <= n; i += 8) stu256((void *)(dst + i), ldu256((const void *)(src + i)));
for (; i < n; i++) dst[i] = src[i];
}
// The same copy with NON-TEMPORAL stores. `c[N1_ .. N1_+M]` is written ONCE here
// and read back only much later (phase-3 region 1, `tab = c[N1_+k]`), so the
// write-allocate read it would otherwise pull is pure waste -- the line can never
// be reused out of cache (the array is 263 MB against a 6 MB L3).
static void copy_u32_nt(u32 *dst, const u32 *src, long n) {
long i = 0;
long ka = ((-(long)(((uintptr_t)dst) >> 2)) & 7);
for (; i < ka && i < n; i++) dst[i] = src[i];
for (; i + 8 <= n; i += 8) {
__m256i v = ldu256((const void *)(src + i));
_mm256_stream_si256((__m256i *)(dst + i), v);
}
for (; i < n; i++) dst[i] = src[i];
_mm_sfence();
}
static void untr_tile_avx(const u32 *in, u32 *out, int A0, int B0, int strIn, int strOut) {
for (int u0 = 0; u0 < TS; u0 += 8) {
for (int v0 = 0; v0 < TS; v0 += 16) {
__m256i r[8];
for (int k = 0; k < 8; k++)
r[k] = ldu256((const void *)(BROW((long)(A0 + u0 + k) * strIn + B0 + v0)));
t8x8(r);
for (int k = 0; k < 8; k++)
stu256((void *)(TBUF + (v0 + k) * TS + u0), r[k]);
{
int v0b = v0 + 8;
__m256i r2[8];
for (int k = 0; k < 8; k++)
r2[k] = ldu256((const void *)(BROW((long)(A0 + u0 + k) * strIn + B0 + v0b)));
t8x8(r2);
for (int k = 0; k < 8; k++)
stu256((void *)(TBUF + (v0b + k) * TS + u0), r2[k]);
}
}
}
// NT: strOut == C == 2^LGC_ is a multiple of 8 and A0/u0 are multiples of 8, so
// every 32-byte destination chunk of a row is aligned regardless of B0+v0+k.
// Emitting whole rows from TBUF makes the eight NT stores of a row consecutive,
// so they coalesce into full-line write-combining flushes.
for (int v = 0; v < TS; v++) {
u32 *d = out + (u64)(B0 + v) * strOut + A0;
const u32 *s = TBUF + v * TS;
for (int u = 0; u < TS; u += 8)
_mm256_stream_si256((__m256i *)(d + u), ldu256((const void *)(s + u)));
}
}
static void build_tab2(u32 *tw, u32 *tws, int L, u32 w) {
int off = 0;
for (int len = 2; len <= L; len <<= 1) {
u32 wl = pw(w, (u32)(L / len));
u32 c = 1;
int h = len >> 1;
for (int j = 0; j < h; j++) { tw[off + j] = c; tws[off + j] = shoup_c(c); c = mm(c, wl); }
off += h;
}
}
#pragma GCC pop_options
#endif
// ===========================================================================
// RADIX-3 COLUMN TRANSFORM. C = 3 * 2^M3L_ = NC_ .
//
// N1 = R*C = 2^LGR_ * 3 * 2^M3L_ = 3 * 2^26 = 201 326 592 >= n+m+1, so the
// cyclic convolution of length N1 IS the linear one: M = n+m-N1 < 0, the wrap
// region is empty and there is NO wrap-correction transform. Total transform
// work is 3 transforms of N1 instead of 3 x 5*2^25 + 3 x 2^26.
//
// C keeps the odd factor so that R stays a pure power of two (the corner's
// bit-reversal algebra over the R index needs that); the C index is only ever
// used as a plain exponent there, so a non-power-of-two C is free.
//
// Column DFT of length C, decimated by 3 then by 2 (index m = 3*m1 + s):
// Y_j[s] = sum_i x[j+i*M3_] w3^{i s} (3-pt DFT across the thirds)
// Z_s[j] = Y_j[s] * w^{j s} (twiddle)
// X_s[m1] = sum_j Z_s[j] (w^3)^{j m1} (length-M3_ DFT)
// with X[3*m1+s] stored at position s*M3_ + m1. The forward and the inverse use
// the SAME permutation of the C index, so the round trip is exact.
//
// The 3-pt DFT kernel (2 multiplications per 3 points):
// A=x1+x2 B=x1-x2 X0 = x0+A
// X1 = x0 + A*(w3+w3^2)/2 + B*(w3-w3^2)/2
// X2 = x0 + A*(w3+w3^2)/2 - B*(w3-w3^2)/2
// and w3+w3^2 = -1, so the two constants are cA = -1/2 and cP = (w3-w3^2)/2.
// The three inputs live in three DIFFERENT lane vectors, so every constant is a
// broadcast. For the inverse w3 -> w3^-1: cA is unchanged, cP is negated.
// ===========================================================================
static u32 R3F[2], R3FS[2];
static u32 R3I[2], R3IS[2];
static u32 T3F[2 * M3_], T3FS[2 * M3_];
static u32 T3I[2 * M3_], T3IS[2 * M3_];
static u32 tsubFw[M3_], tsubFsw[M3_], tsubIw[M3_], tsubIsw[M3_];
static void build_radix3(u32 root) { // root = the C-th root of unity (w^R)
const u32 H2 = 1006632961u; // 1/2 mod P
u32 w3 = pw(root, (u32)(NC_ / 3));
u32 w32 = mm(w3, w3);
u32 cA = (u32)(P - H2); // -1/2 = (w3 + w3^2)/2 mod P
u32 cP = mm((u32)(((u64)w3 + P - w32) % P), H2); // (w3 - w3^2)/2
R3F[0] = cA; R3F[1] = cP;
R3I[0] = cA; R3I[1] = (cP ? P - cP : 0);
for (int i = 0; i < 2; i++) { R3FS[i] = shoup_c(R3F[i]); R3IS[i] = shoup_c(R3I[i]); }
for (int s = 1; s < 3; s++) {
u32 g = pw(root, (u32)s), gi = pw(g, P - 2u), c = 1, ci = 1;
u32 *pf = T3F + (s - 1) * M3_, *pf2 = T3FS + (s - 1) * M3_;
u32 *pi = T3I + (s - 1) * M3_, *pi2 = T3IS + (s - 1) * M3_;
for (int j = 0; j < (int)M3_; j++) { pf[j] = c; pf2[j] = shoup_c(c); c = mm(c, g);
pi[j] = ci; pi2[j] = shoup_c(ci); ci = mm(ci, gi); }
}
build_tab2(tsubFw, tsubFsw, (int)M3_, pw(root, 3));
build_tab2(tsubIw, tsubIsw, (int)M3_, pw(pw(root, 3), P - 2u));
}
#pragma GCC push_options
#pragma GCC target("avx2,tune=skylake")
static inline __attribute__((always_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));
}
static inline __attribute__((always_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));
}
static inline __attribute__((always_inline)) void dft3(__m256i *v, const __m256i *c,
const __m256i *cs, __m256i Pv) {
__m256i A = vaddm(v[1], v[2], Pv);
__m256i B = vsubm(v[1], v[2], Pv);
__m256i T = shoup8(A, c[0], cs[0]);
__m256i U = shoup8(B, c[1], cs[1]);
__m256i X0 = vaddm(v[0], A, Pv);
__m256i Z = vaddm(v[0], T, Pv);
v[0] = X0; v[1] = vaddm(Z, U, Pv); v[2] = vsubm(Z, U, Pv);
}
static void col_dif3(u32 *x) {
const __m256i Pv = _mm256_set1_epi32((int)P);
__m256i c[2], cs[2];
for (int i = 0; i < 2; i++) { c[i] = _mm256_set1_epi32((int)R3F[i]); cs[i] = _mm256_set1_epi32((int)R3FS[i]); }
for (int j = 0; j < (int)M3_; j += 8) {
__m256i v[3];
for (int i = 0; i < 3; i++) v[i] = ldu256((const void *)(x + (u64)i * M3_ + j));
dft3(v, c, cs, Pv);
stu256((void *)(x + j), v[0]);
for (int s = 1; s < 3; s++) {
__m256i w = ldu256((const void *)(T3F + (s - 1) * M3_ + j));
__m256i ws = ldu256((const void *)(T3FS + (s - 1) * M3_ + j));
stu256((void *)(x + (u64)s * M3_ + j), shoup8(v[s], w, ws));
}
}
for (int s = 0; s < 3; s++) row_dif_avx(x + (u64)s * M3_, (int)M3_, tsubFw, tsubFsw);
}
static void col_dit3(u32 *x) {
for (int s = 0; s < 3; s++) row_dit_avx(x + (u64)s * M3_, (int)M3_, tsubIw, tsubIsw);
const __m256i Pv = _mm256_set1_epi32((int)P);
__m256i c[2], cs[2];
for (int i = 0; i < 2; i++) { c[i] = _mm256_set1_epi32((int)R3I[i]); cs[i] = _mm256_set1_epi32((int)R3IS[i]); }
for (int j = 0; j < (int)M3_; j += 8) {
__m256i v[3];
v[0] = ldu256((const void *)(x + j));
for (int s = 1; s < 3; s++) {
__m256i w = ldu256((const void *)(T3I + (s - 1) * M3_ + j));
__m256i ws = ldu256((const void *)(T3IS + (s - 1) * M3_ + j));
v[s] = shoup8(ldu256((const void *)(x + (u64)s * M3_ + j)), w, ws);
}
dft3(v, c, cs, Pv);
for (int i = 0; i < 3; i++) stu256((void *)(x + (u64)i * M3_ + j), v[i]);
}
}
#pragma GCC pop_options
static void row_dit(u32 *x, int L, const u32 *tw, const u32 *tws) {
#if HAVE_AVX2
row_dit_avx(x, L, tw, tws);
return;
#else
(void)tws;
#endif
int off = 0;
for (int len = 2; len <= L; len <<= 1) {
int h = len >> 1;
const u32 *t = tw + off; off += h;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j++) {
u32 u = p[j], v = mm(q[j], t[j]);
u32 s = u + v; if (s >= P) s -= P;
u32 d = u + P - v; if (d >= P) d -= P;
p[j] = s; q[j] = d;
}
}
}
}
static void pointwise(u32 *A, const u32 *B, u32 n, u32 sc, u32 sx) {
#if HAVE_AVX2
pointwise_mul_avx(A, B, n, sc, sx);
return;
#else
for (long i = 0; i < (long)n; i++) { long o = (i / (long)(sx - PADC_)) * sx + (i % (long)(sx - PADC_)); A[o] = mm(mm(A[o], B[o]), sc); }
#endif
}
static void row_dif(u32 *x, int L, const u32 *tw, const u32 *tws) {
#if HAVE_AVX2
row_dif_avx(x, L, tw, tws);
return;
#else
(void)tws;
#endif
for (int len = L; len >= 2; len >>= 1) {
int h = len >> 1;
const u32 *t = tw + h - 1;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j++) {
u32 u = p[j], v = q[j];
u32 s = u + v; if (s >= P) s -= P;
u32 d = u + P - v; if (d >= P) d -= P;
p[j] = s; q[j] = mm(d, t[j]);
}
}
}
}
static inline void bitrev_pass(u32 *x, int L, const u16 *br) {
for (int i = 0; i < L; i++) {
int k = br[i];
if (i < k) { u32 t = x[i]; x[i] = x[k]; x[k] = t; }
}
}
// ---------------------------------------------------------------------------
struct Eng {
int R, C, N, sx;
int blkr; // padded row stride of the streamed S block
int r5; // 1: column DFT length is 3*2^M3L_ (radix-3)
u32 w;
u32 *tfR; u32 *tiR; u32 *tfC; u32 *tiC;
u32 *wsR; u32 *wsC;
u16 *brR; u16 *brC;
};
// column DFT dispatch: the main engine has C = 3*2^M3L_ (radix-3); the (now
// unused) wrap-correction engine has a power-of-two C and uses the plain kernels.
static void colfwd(Eng &e, u32 *row) {
#if HAVE_AVX2
if (e.r5) { col_dif3(row); return; }
#endif
row_dif(row, e.C, e.tfC, e.wsC);
}
static void colinv(Eng &e, u32 *row) {
#if HAVE_AVX2
if (e.r5) { col_dit3(row); return; }
#endif
row_dit(row, e.C, e.tiC, e.wsC + e.C);
}
static void load_T(Eng &e, const u32 *src, u32 nsrc, u32 *S, int ctr) {
const int R = e.R, C = e.C;
#if HAVE_AVX2
// fast path: whole tiles (and whole all-zero tiles) go through registers
if (nsrc <= (u32)R * (u32)C) {
for (int r0 = 0; r0 < R; r0 += TS) {
u64 tilerow = (u64)r0 * C;
for (int c0 = 0; c0 < C; c0 += TS) {
u64 firstidx = tilerow + c0;
u64 lastidx = (u64)(r0 + TS - 1) * C + c0 + TS - 1;
if (lastidx < nsrc) { load_tile_avx(src, S, r0, c0, R, C, ctr, 0); continue; }
if (firstidx >= nsrc) { load_tile_avx(src, S, r0, c0, R, C, ctr, 1); continue; }
for (int rb = 0; rb < TS; rb++) {
u64 base = (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u64 id = base + cb;
u32 v = 0;
if (id < nsrc) { v = src[id]; if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u)); }
TBUF[cb * TS + rb] = v;
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)(c0 + cb) * R + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb++) d[rb] = s[rb];
}
}
}
return;
}
#endif
for (int r0 = 0; r0 < R; r0 += TS) {
for (int c0 = 0; c0 < C; c0 += TS) {
u64 lastidx = (u64)(r0 + TS - 1) * C + (c0 + TS - 1);
if (lastidx < nsrc) {
for (int rb = 0; rb < TS; rb++) {
const u32 *s = src + (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u32 v = s[cb];
if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u));
TBUF[cb * TS + rb] = v;
}
}
} else {
for (int rb = 0; rb < TS; rb++) {
u64 base = (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u64 id = base + cb;
u32 v = 0;
if (id < nsrc) { v = src[id]; if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u)); }
TBUF[cb * TS + rb] = v;
}
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)(c0 + cb) * R + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb++) d[rb] = s[rb];
}
}
}
}
// generic corner turn+transpose: in[A*strIn + B] -> out[B*strOut + A] * w
// (nA/nB are the dimensions of A/B; the tile is TS x TS on both)
// perm 0: w = root^{A*B}
// perm 1: w = root^{A*rev(B)} (B is a DIF output position, lg2(nB) bits)
// perm 2: w = root^{rev(A)*B} (A is a DIF output position, lg2(nA) bits)
// The permutation is folded into the twiddle table, so the load/store pattern
// is byte-for-byte the same in all three modes.
static u32 REV6[TS];
static u32 g_inv_scale = 1u;
static u32 POWT[1024], POWT64[1024];
static void build_rev6(void) {
for (int i = 0; i < TS; i++) {
int r = 0;
for (int b = 0; b < 6; b++) if (i & (1 << b)) r |= 1 << (5 - b);
REV6[i] = (u32)r;
}
}
static inline u32 revbits(u32 x, int lg) {
u32 r = 0;
for (int b = 0; b < lg; b++) if (x & (1u << b)) r |= 1u << (lg - 1 - b);
return r;
}
static void corner_generic(const u32 *in, u32 *out, int nA, int nB, int strIn, int strOut, u32 root, int perm, int split) {
u32 wv[TS];
{ u32 c = 1; for (int v = 0; v < TS; v++) { wv[v] = c; c = mm(c, root); } }
const u32 w64 = pw(root, TS);
int lgB = 0, lgA = 0;
while ((1 << lgB) < nB) lgB++;
while ((1 << lgA) < nA) lgA++;
const int S1 = nB / TS, S2 = nA / TS; // tile-index counts
const int HB = lgB - 6, HA = lgA - 6; // tile-index bit counts
u32 RV2[TS], RV2b[TS], PW2[TS], cu[TS], cu64[TS], base[TS];
if (perm == 1) {
u32 g = pw(root, (u32)S1); // root^{nB/TS}
{ u32 gM = mm(g, R32), c = R32; for (int t = 0; t < TS; t++) { RV2[REV6[t]] = c; c = mont1(c, gM); } }
{ u32 rootM = mm(root, R32), p = R32; for (int t = 0; t < S1; t++) { POWT[t] = p; p = mont1(p, rootM); } }
} else if (perm == 2) {
u32 g = pw(root, (u32)S2), g64 = pw(g, 64);
const u32 rootM2 = mm(root, R32), gM = mm(g, R32), g64M = mm(g64, R32);
{ u32 c = R32, c64 = R32;
for (int t = 0; t < TS; t++) { RV2[REV6[t]] = c; RV2b[REV6[t]] = c64; c = mont1(c, gM); c64 = mont1(c64, g64M); } }
{ u32 p = R32; for (int t = 0; t < S2; t++) { POWT[t] = p; p = mont1(p, rootM2); } }
{ u32 q = pw(root, 64), qM = mm(q, R32), p = R32; for (int t = 0; t < S2; t++) { POWT64[t] = p; p = mont1(p, qM); } }
}
u32 wr = 1; // root^{A0}
const u32 INVSCALE = g_inv_scale; // 1/N1 folded in as a uniform linear factor
for (int A0 = 0; A0 < nA; A0 += TS) {
u32 wr64 = pw(wr, TS); // root^{TS*A0}
if (perm == 1) {
u32 g = pw(wr, (u32)S1), gM2 = mm(g, R32), c = R32;
for (int t = 0; t < TS; t++) { PW2[REV6[t]] = c; c = mont1(c, gM2); }
} else if (perm == 2) {
int hi = A0 / TS;
u32 rH = POWT[revbits((u32)hi, HA)], r64H = POWT64[revbits((u32)hi, HA)];
const u32 INVSCALE_M = mm(INVSCALE, R32);
for (int u = 0; u < TS; u++) { cu[u] = mont1(RV2[u], rH); cu64[u] = mont1(RV2b[u], r64H); base[u] = INVSCALE_M; }
} else {
u32 c = 1; for (int v = 0; v < TS; v++) { PW2[v] = c; c = mm(c, wr); }
}
u32 Vloc = 1; // root^{B0}
u32 Vr = 1; // root^{A0*B0}
for (int B0 = 0; B0 < nB; B0 += TS) {
if (perm == 1) {
u32 H = revbits((u32)(B0 / TS), HB);
u32 wrH = pw(wr, H), rootH = POWT[H];
u32 RV[TS];
#if HAVE_AVX2
scale_row_mont(RV, RV2, rootH); // root^{rev(B0+v)}
scale_row_mont(TW, PW2, mm(wrH, R32)); // root^{A0*rev(B0+v)}
twiddle_rows_mont(TW, RV);
#else
for (int v = 0; v < TS; v++) RV[v] = mont1(RV2[v], rootH);
for (int v = 0; v < TS; v++) TW[v] = mont1(PW2[v], mm(wrH, R32));
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v++) cur[v] = mont1(prv[v], RV[v]);
}
#endif
} else if (perm == 2) {
if (B0 == 0) {
for (int u = 0; u < TS; u++) {
u32 *cur = TW + u * TS;
cur[0] = base[u];
for (int v = 1; v < TS; v++) cur[v] = mont1(cur[v - 1], cu[u]);
}
} else {
for (int u = 0; u < TS; u++) {
u32 *cur = TW + u * TS;
#if HAVE_AVX2
scale_row_mont(cur, cur, cu64[u]);
#else
u32 d = cu64[u];
for (int v = 0; v < TS; v++) cur[v] = mont1(cur[v], d);
#endif
}
}
for (int u = 0; u < TS; u++) base[u] = mont1(base[u], cu64[u]);
} else {
u32 RV[TS];
for (int v = 0; v < TS; v++) RV[v] = mm(Vloc, wv[v]);
for (int v = 0; v < TS; v++) TW[v] = mm(Vr, PW2[v]);
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v++) cur[v] = mm(prv[v], RV[v]);
}
}
// data path: whole tile through registers, u innermost (see NOTES: writing
// the 64 strOut-strided destination rows in v-major order is what keeps them
// from aliasing into one L1 set).
#if HAVE_AVX2
corner_tile_avx(in, out, A0, B0, strIn, strOut, split);
#else
for (int u = 0; u < TS; u++) {
const u32 *s = in + (u64)(A0 + u) * strIn + B0;
const u32 *tw = TW + u * TS;
for (int v = 0; v < TS; v++) TBUF[u * TS + v] = mont1(s[v], tw[v]);
}
for (int v = 0; v < TS; v++) {
u32 *d = out + (u64)(B0 + v) * strOut + A0;
for (int u = 0; u < TS; u++) d[u] = TBUF[u * TS + v];
}
#endif
Vr = mm(Vr, wr64);
Vloc = mm(Vloc, w64);
}
wr = mm(wr, w64);
}
}
static inline void corner_fwd(Eng &e, const u32 *S, u32 *X, u32 root) {
corner_generic(S, X, e.C, e.R, e.R, e.sx, root, 1, 0); // A = j0, B = rev'd k1 position
}
static inline void corner_inv(Eng &e, const u32 *X, u32 *S, u32 root) {
corner_generic(X, S, e.R, e.C, e.sx, e.R, root, 2, 1); // A = rev'd k1, B = j0 (S split)
}
// transpose the inverse result to natural order: S[j0*R + j1] -> dst[j1*C + j0]
static void untranspose(Eng &e, const u32 *S, u32 *dst) {
const int R = e.R, C = e.C;
#if HAVE_AVX2
for (int c0 = 0; c0 < C; c0 += TS) // j0 blocks
for (int r0 = 0; r0 < R; r0 += TS) // j1 blocks
untr_tile_avx(S, dst, c0, r0, R, e.sx); // A = j0, B = j1 (S is the split array)
return;
#endif
for (int c0 = 0; c0 < C; c0 += TS) { // j0 blocks
for (int r0 = 0; r0 < R; r0 += TS) { // j1 blocks
for (int u = 0; u < TS; u++) {
const u32 *s = S + (u64)(c0 + u) * R + r0;
u32 *d = TBUF + u * TS;
for (int v = 0; v < TS; v++) d[v] = s[v];
}
// as in corner_generic: u innermost so each destination row is written
// sequentially (C*4 is a multiple of 64KB -> one L1 set otherwise)
for (int v = 0; v < TS; v++) {
u32 *d = dst + (u64)(r0 + v) * e.sx + c0;
for (int u = 0; u < TS; u++) d[u] = TBUF[u * TS + v];
}
}
}
}
// ---------------------------------------------------------------------------
// Streamed four-step forward (perm == 1 corner): for each 64-column block A0 the
// source strip is loaded straight into an L2-resident block, the length-R row
// DFTs run there, and the corner consumes that block in place. The 671 MB
// transposed scratch S is never materialised, so the forward pays one source
// read + one spectrum write instead of load_T's write plus the corner's read.
static u32 SBLK[TS * (NRMAX_ + PADB_)] __attribute__((aligned(64)));
static void corner_fwd_stream(Eng &e, const u32 *src, u32 nsrc, u32 *X, int split, int ctr) {
const int nA = e.C, nB = e.R, strIn = e.blkr, strOut = e.sx;
const u32 root = e.w;
u32 *S = SBLK;
int lgB = 0; while ((1 << lgB) < nB) lgB++;
const int S1 = nB / TS, HB = lgB - 6;
u32 RV2[TS], PW2[TS];
{ u32 g = pw(root, (u32)S1), gM = mm(g, R32), c = R32;
for (int t = 0; t < TS; t++) { RV2[REV6[t]] = c; c = mont1(c, gM); } }
{ u32 rootM = mm(root, R32), p = R32;
for (int t = 0; t < S1; t++) { POWT[t] = p; p = mont1(p, rootM); } }
const u32 w64 = pw(root, TS);
u32 wr = 1; // root^{A0}
for (int A0 = 0; A0 < nA; A0 += TS) {
u32 wr64 = pw(wr, TS);
{ u32 g = pw(wr, (u32)S1), gM2 = mm(g, R32), c = R32;
for (int t = 0; t < TS; t++) { PW2[REV6[t]] = c; c = mont1(c, gM2); } }
{ double _t=LMD_now(); load_block(src, nsrc, S, A0, nB, nA, strIn, ctr); LMD_acc(10,_t); }
{ double _t=LMD_now(); for (int u = 0; u < TS; u++)
row_dif(S + (u64)u * strIn, nB, e.tfR, e.wsR); LMD_acc(11,_t); }
{ double _t=LMD_now(); for (int B0 = 0; B0 < nB; B0 += TS) {
u32 H = revbits((u32)(B0 / TS), HB);
u32 wrH = pw(wr, H), rootH = POWT[H];
u32 RV[TS];
scale_row_mont(RV, RV2, rootH);
scale_row_mont(TW, PW2, mm(wrH, R32));
twiddle_rows_mont(TW, RV);
corner_tile_avx_blk(S, X, A0, B0, strIn, strOut, split);
} LMD_acc(12,_t); }
wr = mm(wr, w64);
}
}
static void fwd(Eng &e, const u32 *src, u32 nsrc, u32 *S, u32 *X, int split, int ctr) {
#ifdef OLD_FWD
load_T(e, src, nsrc, S, ctr);
for (int j0 = 0; j0 < e.C; j0++) {
u32 *row = S + (u64)j0 * e.R;
row_dif(row, e.R, e.tfR, e.wsR);
}
corner_fwd(e, S, X, e.w);
#else
(void)S;
{ double _t=LMD_now(); corner_fwd_stream(e, src, nsrc, X, split, ctr); LMD_acc((e.R==4096)?4:0,_t); }
#endif
{ double _t=LMD_now(); for (int k0 = 0; k0 < e.R; k0++)
colfwd(e, split ? BROW((long)k0 * e.sx) : X + (u64)k0 * e.sx); LMD_acc(13,_t); }
}
#if 0
static void fwd_old(Eng &e, const u32 *src, u32 nsrc, u32 *S, u32 *X, int ctr) {
load_T(e, src, nsrc, S, ctr);
for (int j0 = 0; j0 < e.C; j0++) {
u32 *row = S + (u64)j0 * e.R;
row_dif(row, e.R, e.tfR, e.wsR);
}
corner_fwd(e, S, X, e.w);
for (int k0 = 0; k0 < e.R; k0++) colfwd(e, X + (u64)k0 * e.sx);
}
#endif
// Inverse: the DIT stages run on the DIF's position-indexed output, so the two
// bitrev permutations cancel and no bitrev pass is needed anywhere.
static void inv(Eng &e, u32 *X, u32 *S) { double _it=LMD_now();
for (int k0 = 0; k0 < e.R; k0++) colinv(e, X + (u64)k0 * e.sx);
corner_inv(e, X, S, pw(e.w, P - 2u));
for (int j0 = 0; j0 < e.C; j0++) {
u32 *row = BROW((long)j0 * e.R);
row_dit(row, e.R, e.tiR, e.wsR + e.R);
}
LMD_acc((e.R==4096)?5:1,_it);
}
// ---------------------------------------------------------------------------
static void eng_init(Eng &e, int lgR, int nC, u32 w, int sx,
u32 *tfR, u32 *tiR, u32 *tfC, u32 *tiC, u16 *brR, u16 *brC,
u32 *wsR, u32 *wsC) {
e.R = 1 << lgR; e.C = nC; e.N = e.R * e.C; e.w = w; e.sx = sx;
e.blkr = e.R + PADB_;
e.r5 = (nC == (int)NC_);
e.tfR = tfR; e.tiR = tiR; e.tfC = tfC; e.tiC = tiC; e.brR = brR; e.brC = brC;
e.wsR = wsR; e.wsC = wsC;
u32 wR = pw(w, (u32)e.C), wC = pw(w, (u32)e.R);
#if HAVE_AVX2
build_tab2(tfR, e.wsR, e.R, wR); build_tab2(tiR, e.wsR + e.R, e.R, pw(wR, P - 2u));
if (!e.r5) { build_tab2(tfC, e.wsC, e.C, wC); build_tab2(tiC, e.wsC + e.C, e.C, pw(wC, P - 2u)); }
#else
build_tab(tfR, e.R, wR); build_tab(tiR, e.R, pw(wR, P - 2u));
build_tab(tfC, e.C, wC); build_tab(tiC, e.C, pw(wC, P - 2u));
#endif
build_br(brR, e.R);
if (!e.r5) build_br(brC, e.C); // the radix-5 engine never bit-reverses C
build_rev6();
}
// driver for 1002e8 core. requires LGR_, LGC_, LGH_, LGR2_, LGC2_ and EXACT.
//
// N1 = 2^(LGR_+LGC_) = the cyclic length; a, b must satisfy n,m <= N1-1 and
// n+m <= 3*N1/2 so that the wrap region fits.
//
// cyc = a (*) b (cyclic, length N1). With Sa = N1-m, Sb = N1-n, M = n+m-N1:
// c_{N1+k} = sum_q a[Sa+k+q] * b[Sb+M-q] = conv(P,R)_{M+k}, q in [0,M-k]
// where P_q = a[Sa+q], R_v = b[Sb+v] (both length M+1, degree 2M <= N1-1).
// So one extra 3-transform pass gives the whole wrap correction; the cyclic
// result stored in A and the correction E kept in c[N1 .. N1+M] (read in place
// by the final combine, which walks k in increasing order).
// ===========================================================================
// VECTORISED COMBINE (phase 3).
//
// The scalar combine carried TWO i64 window sums (WA, WB) plus tt, and emitted
// c[k] = ts(k) + 4*WA(k) + 4*WB(k) - 16*tt(k)
// where WA/WB are window sums of a[]/b[]. The whole expression is only ever
// TRUNCATED to u32, so it may be evaluated in Z/2^32 -- and there the three
// running terms COLLAPSE INTO ONE ACCUMULATOR:
//
// SUM(k) = SUM(k-1) + D(k), c[k] = ts(k) + SUM(k) (mod 2^32)
// D(k) = 4*(a[k]+b[k]) - 16 (phase 0, both windows grow)
// = 4*(x[k]-x[k-1-w]) (phase 1, the longer window)
// = 16 - 4*(a[k-1-m]+b[k-1-n]) (phase 2, both windows drain)
//
// TWO carry chains become ONE, and that one chain is then carried 8 elements at
// a time by an AVX2 prefix scan whose cross-block carry is a single register.
// Everything is exact in Z/2^32, so the output is bit-identical to the scalar
// form (which only ever kept the low 32 bits anyway).
// ===========================================================================
#pragma GCC push_options
#pragma GCC target("avx2,tune=skylake")
// inclusive 8-lane prefix sum of v (mod 2^32), plus the broadcast carry-in;
// `carry` is left holding the broadcast of the new running total.
static inline __attribute__((always_inline)) __m256i scan8(__m256i v, __m256i &carry) {
v = _mm256_add_epi32(v, _mm256_slli_si256(v, 4));
v = _mm256_add_epi32(v, _mm256_slli_si256(v, 8));
__m256i t = _mm256_permute2x128_si256(_mm256_setzero_si256(), v, 0x20);
t = _mm256_shuffle_epi32(t, 0xFF); // broadcast lane 3 of the low half
v = _mm256_add_epi32(v, t);
v = _mm256_add_epi32(v, carry);
carry = _mm256_permutevar8x32_epi32(v, _mm256_set1_epi32(7));
return v;
}
#define CMBSTEP() \
{ u32 d; \
if (PH == 0) d = 4u * (a[k] + b[k]) - 16u; \
else if (PH == 1) { if (nm_lt) d = 4u * (b[k] - b[k - 1 - n]); \
else d = 4u * (a[k] - a[k - 1 - m]); } \
else d = 16u - 4u * (a[k - 1 - m] + b[k - 1 - n]); \
SUMr += d; \
u32 tab; \
if (TB == 1) { u32 x = A[AOFF(k)], y = c[N1_ + k]; tab = (x >= y) ? (x - y) : (x + (P - y)); } \
else if (TB == 2) tab = A[AOFF(k)]; \
else tab = c[k]; \
i32 ts = (i32)tab; \
if (ts > (i32)(P >> 1)) ts -= (i32)P; \
c[k] = (u32)((u32)ts + SUMr); \
}
template<int PH, int TB, int NT>
static void cmbv(long k0, long k1, const u32 *a, const u32 *b, u32 *A, u32 *c,
long n, long m, long nmax, u32 &SUMr, int nm_lt)
{
if (k0 > k1) return;
long k = k0;
{ /* align k to a multiple of 8: NC_ % 8 == 0, so k & (NC_-1) is then also a
multiple of 8 and no 8-element A-block can cross a padded row. */
long ka = k + ((-(long)k) & 7);
long kend = (ka < k1 + 1) ? ka : k1 + 1;
if (ka > k) { const __m256i Pv0 = _mm256_set1_epi32((int)P); (void)Pv0; }
for (; k < kend; k++) { CMBSTEP(); }
}
const __m256i Pv = _mm256_set1_epi32((int)P);
const __m256i Phv = _mm256_set1_epi32((int)(P >> 1));
#if HAVE_AVX2
const __m256i c16 = _mm256_set1_epi32(16);
const __m256i idx7 = _mm256_set1_epi32(7);
__m256i carry = _mm256_set1_epi32((int)SUMr);
if (NT) { /* vmovntdq faults unless 32-byte aligned; k is already 8-aligned */
carry = _mm256_set1_epi32((int)SUMr);
}
for (; k + 7 <= k1; k += 8) {
__m256i d;
if (PH == 0) {
__m256i va = ldu256(a + k), vb = ldu256(b + k);
d = _mm256_sub_epi32(_mm256_slli_epi32(_mm256_add_epi32(va, vb), 2), c16);
} else if (PH == 1) {
if (nm_lt) { __m256i v1 = ldu256(b + k), v0 = ldu256(b + k - 1 - n);
d = _mm256_slli_epi32(_mm256_sub_epi32(v1, v0), 2); }
else { __m256i v1 = ldu256(a + k), v0 = ldu256(a + k - 1 - m);
d = _mm256_slli_epi32(_mm256_sub_epi32(v1, v0), 2); }
} else {
__m256i va = ldu256(a + k - 1 - m), vb = ldu256(b + k - 1 - n);
d = _mm256_sub_epi32(c16, _mm256_slli_epi32(_mm256_add_epi32(va, vb), 2));
}
d = scan8(d, carry);
__m256i tab;
if (TB == 1) {
__m256i x = ldu256(A + AOFF(k)), y = ldu256(c + N1_ + k);
__m256i dd = _mm256_sub_epi32(x, y);
__m256i mk = _mm256_cmpgt_epi32(y, x); // x < y (all values < 2^31)
tab = _mm256_add_epi32(dd, _mm256_and_si256(mk, Pv));
} else if (TB == 2) {
tab = ldu256(A + AOFF(k));
} else {
tab = ldu256(c + k);
}
__m256i ts = _mm256_sub_epi32(tab, _mm256_and_si256(_mm256_cmpgt_epi32(tab, Phv), Pv));
__m256i out = _mm256_add_epi32(ts, d);
if (NT && (((uintptr_t)(c + k)) & 31) == 0) _mm256_stream_si256((__m256i *)(c + k), out);
else stu256(c + k, out);
}
SUMr = (u32)_mm256_extract_epi32(carry, 0);
#endif
for (; k <= k1; k++) { CMBSTEP(); }
}
#pragma GCC pop_options
static u32 Aarr[NR_ * CSX_] __attribute__((aligned(64)));
// B/S tail: the part of the shared array that is NOT c-resident (see BROW).
#if (NR_ * CSX_) > BSTAIL_SPLIT
#define BTAILN_ ((int)((NR_ * CSX_) - BSTAIL_SPLIT))
#else
#define BTAILN_ 1
#endif
static u32 BTAIL[BTAILN_] __attribute__((aligned(64)));
static Eng E1, E2;
// E1 is always the radix-5 engine, which never touches the plain C-dimension
// staged tables (colfwd/colinv use the radix-5 kernels instead). Declaring them
// at one element instead of NC_ keeps 720 KB off the 2 GiB footprint, which
// matters: the shipped Aarr+Barr already sit ~3 MB under the limit.
static u32 t1fR[NR_], t1iR[NR_], t1fC[1], t1iC[1];
static u16 b1R[NR_], b1C[1];
static u32 w1R[2 * NR_], w1C[2];
static u32 t2fR[NR2_], t2iR[NR2_], t2fC[NC2_], t2iC[NC2_];
static u16 b2R[NR2_], b2C[NC2_];
static u32 w2R[2 * NR2_], w2C[2 * NC2_];
static int inited = 0;
void poly_multiply(unsigned *a, int n, unsigned *b, int m, unsigned *c)
{
IM = ~0ull / (u64)P;
if (!inited) {
inited = 1;
{ u64 inv = 1; for (int i = 0; i < 5; i++) inv = inv * (2 - (u64)P * inv); NPI = (u32)(-(u64)inv); }
u32 w1 = pw(G, (u32)((P - 1) / N1_));
u32 w2 = pw(G, (u32)((P - 1) / N2_));
eng_init(E1, LGR_, (int)NC_, w1, (int)CSX_, t1fR, t1iR, t1fC, t1iC, b1R, b1C, w1R, w1C);
eng_init(E2, LGR2_, (1 << LGC2_), w2, (int)CSX2_, t2fR, t2iR, t2fC, t2iC, b2R, b2C, w2R, w2C);
E1.r5 = 1; E2.r5 = 0;
#if HAVE_AVX2
build_radix3(pw(w1, (u32)E1.R));
#endif
}
// The B/S array: its first BSTAIL_SPLIT elements live in `c` (dead until
// phase 3), the rest in BTAIL. `c` holds n+m+1 = 200 000 001 u32 and the
// c-resident part needs 188 805 120 of them.
g_B = (u32 *)c; g_BT = BTAIL;
// This one array is B (the b spectrum) and then S (the inverse's transposed
// scratch): every use of B ends at the pointwise, and S is written only after.
u32 *S = g_B;
u32 *A = Aarr, *B = g_B;
const int ctr = EXACT;
const u32 sc = pw((u32)(N1_ % P), P - 2u); // 1/N1 (main engine)
const u32 sc2 = pw((u32)(N2_ % P), P - 2u); // 1/N2 (wrap-correction engine) // carried by the inverse transform
const long Sa = (long)N1_ - (long)m;
const long Sb = (long)N1_ - (long)n;
const long M = (long)n + (long)m - (long)N1_;
// ---------- phase 1: wrap correction c[N1 .. N1+M] ----------
if (M >= 0) {
g_inv_scale = sc2;
fwd(E2, (const u32 *)a + Sa, (u32)(M + 1), S, A, 0, ctr); // M < 0: never runs
fwd(E2, (const u32 *)b + Sb, (u32)(M + 1), S, B, 1, ctr);
pointwise(A, B, N2_, sc2, CSX2_);
inv(E2, A, S);
untranspose(E2, S, A);
{ // A is row-padded, so walk the logical range row by row
long k = M, e = M + M + 1;
u32 *d = c + N1_;
while (k < e) {
long col = k % (long)NC2_;
long run = (long)NC2_ - col;
if (run > e - k) run = e - k;
const u32 *s = A + (k / (long)NC2_) * CSX2_ + col;
if ((d - c) % 8 == 0 && col % 8 == 0 && run % 8 == 0 && run > 0) copy_u32_nt(d, s, run);
else for (long i = 0; i < run; i++) d[i] = s[i];
d += run; k += run;
}
}
}
// ---------- phase 2: cyclic convolution ----------
g_inv_scale = sc;
fwd(E1, (const u32 *)a, (u32)n + 1, S, A, 0, ctr);
fwd(E1, (const u32 *)b, (u32)m + 1, S, B, 1, ctr);
pointwise(A, B, N1_, sc, CSX_);
inv(E1, A, S);
untranspose(E1, S, A);
// ---------- phase 3: combine (vectorised, one carry chain) ----------
long nmax = (long)n + (long)m;
#if EXACT
{
const long kA = ((long)n < (long)m) ? (long)n : (long)m; // end of "both grow"
const long kB = ((long)n < (long)m) ? (long)m : (long)n; // end of "one grows"
const int nm_lt = ((long)n < (long)m);
const long M1 = (M + 1 > 0) ? (M + 1) : 0;
u32 SUM = 0;
cmbv<0,1,1>(0, (M < kA) ? M : kA,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<0,2,1>(M1, (N1_ - 1 < kA) ? (long)(N1_ - 1) : kA,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<0,3,0>((long)N1_, kA,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<1,2,1>((kA + 1 > M1) ? kA + 1 : M1, (N1_ - 1 < kB) ? (long)(N1_ - 1) : kB,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<1,3,0>((kA + 1 > (long)N1_) ? kA + 1 : (long)N1_, kB,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<2,2,1>((kB + 1 > M1) ? kB + 1 : M1, (N1_ - 1 < nmax) ? (long)(N1_ - 1) : nmax,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
cmbv<2,3,0>((kB + 1 > (long)N1_) ? kB + 1 : (long)N1_, nmax,
a, b, Aarr, c, (long)n, (long)m, nmax, SUM, nm_lt);
}
#else
for (long k = 0; k <= nmax; k++) c[k] = c[k];
#endif
{ static int once=0; if(!once){once=1; LMD_dump();} }
_mm_sfence(); // non-temporal stores: make them globally visible
(void)E2;
}
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 4.739 s | 1581 MB + 384 KB | Accepted | Score: 100 | 显示更多 |