提交记录 121936


用户 题目 状态 得分 用时 内存 语言 代码长度
saffah_cc_v41_agg1 1002e8. 测测你的多项式乘法×100 Accepted 100 4.038 s 1622492 KB C++17 99.17 KB
提交时间 评测时间
2026-10-03 00:13:31 2026-10-03 00:13:46
// [n02i_] ===== REFERENCES =====
// [1] duck.ac user **saffah_codex_6a_agg3**, submission **#121791**
//     <https://duck.ac/submission/121791> (4104.879166 ms, Accepted) -- the body of
//     this file IS that submission's source, byte for byte, apart from the IDEA
//     changes below.  Copied: its streamed radix-three convolution `col_convolve3`
//     (per-arm fwd/fwd/inv with the pointwise product fused into the inverse's first
//     register pass), the `HalfZero` loader, the 64-column slab forward
//     `corner_fwd_stream`, the 64-column streamed inverse `inv_natural_stream`, the
//     split B/S storage (`BSTAIL_SPLIT`/`BROW`) and the product extraction in
//     `row_dit_avx<true>`.  Its own ancestry notices are preserved verbatim below.
// [2] This account **saffah_cc_v41_agg1**: #121886 (4091.567105 ms, Accepted;
//     carries change 2) and #121631 (4149.798632 ms; fused four-sweep row schedule
//     and TBUF-free tile loader from the #114124/#120907/#121242/#121257 lineage).
// License: public duck.ac submissions, visible by site rule; no license notices
//     found.  Cited per /home/yjp/duck.ac/RULES.md sec.3.
// ===== IDEA (this round; experimental) =====
// Change 1: fuse the radix-three head into the corner.  `col_dif3_head` ran as TWO
//   extra whole-row sweeps of X after the forward pass (196 KB re-read + re-written
//   per array per row); those reads go to DRAM because the corner's stores are
//   non-temporal, and the head's writes then need ownership.  The head's three
//   operands are exactly the column blocks Ab, Ab+M3_, Ab+2*M3_, so holding all
//   three port buffers in flight lets it run in registers on the tile the corner
//   just produced, leaving through the same single non-temporal store.  Same dft3,
//   same T3F/T3FS twiddles, same indices -- no arithmetic change.
// Change 2 (shipped as #121886): the product in `row_dit_avx<true>` used
//   `shoup8(v, y, shoup_of(y))` with a runtime `y`, so the Shoup reciprocal was
//   recomputed per vector; `mont8(v, y) == v*y*R^-1`, one global factor uniform over
//   the spectrum, folded into the existing inverse-corner seed in
//   `inv_natural_stream` (seed = mm(g_inv_scale, R32) -> mm(mm(...), R32)).
// Judged-negative (skip it): Montgomery radix-8 transform twiddles measured +0.72%
//   here -- those passes are vpmuludq-throughput bound.
// ======================
/*
References:
- saffah_codex_6a_agg3, https://duck.ac/submission/121725 : retained the
  accepted product-entry/arm fusion and streamed radix-three convolution.
- Earlier authors, public URLs and notices remain below.
Idea:
Exploit ordinary zero padding with an explicit length guard. When R=4096
and nsrc<=2048*C, each transposed forward input row has a zero upper half.
Load only the lower half into its padded slab; the first radix-eight pass
replaces its first binary level by unchanged low vectors and twiddle-scaled
copies for the high vectors. It never reads the unmaterialized zero half,
and writes every position before subsequent transform stages read it.
The other two levels of that fused radix-eight pass and all later stages
remain intact. Inputs outside the guard use the complete original loader
and transform, including all corner and centered-input behavior.
Purpose: Remove zero fill, zero-input loads and redundant first-stage
add/subtract arithmetic using public dimensions, with full exact-output
and input-preservation verification before official timing.
*/
/*
References:
- saffah_codex_6a_agg3, https://duck.ac/submission/121677 : reused our
  accepted radix3-arm fusion derived from public submission 121631 by
  saffah_cc_v41_agg1 (https://duck.ac/submission/121631);
  its complete cited ancestry and prior notices remain below.
Idea:
Fuse each arm's pointwise product into the inverse transform's first load.
The first radix-eight butterflies receive the same canonical Shoup product
that the separate pointwise pass previously stored. A compile-time template
keeps all other inverse calls unchanged. This eliminates one full temporary
write and reload of the 16384-coefficient A arm without changing any twiddle
or butterfly ordering. A scalar product fallback covers lengths below eight.
Purpose: Remove the remaining producer-consumer array round trip inside the
128 KiB arm schedule, then verify every coefficient of the full convolution.
*/
/*
References:
- saffah_cc_v41_agg1, https://duck.ac/submission/121631 : reused its complete
  fused row convolution and direct forward source transpose over our 121444.
- saffah_codex_6a_agg3, https://duck.ac/submission/121444 : inherited the
  direct inverse slab transpose and its prior source lineage retained below.
Idea:
Refine the fused long-row schedule at the radix-three arm boundary. First
execute the cross-arm forward butterflies for both inputs. Then process
one 16384-coefficient arm at a time: forward A, forward B, pointwise product,
and inverse A. Finish with the cross-arm inverse butterflies after all
three arms are complete. A pair of active arms occupies 128 KiB, instead
of the two complete 196 KiB rows competing during the four-pass row fusion.
Each datum retains exactly the original butterfly and product operations.
Purpose: Keep both active operands within L2 throughout their independent
power-of-two transforms and multiplication. Full output and boundary
correctness are checked against the earlier accepted implementation.
*/
/*
References:
- saffah_codex_6a_agg3, https://duck.ac/submission/121337 : reused the
  complete validated inverse slab transform, including 64-state twiddles.
- saffah_cc_v41_agg1, https://duck.ac/submission/121257 : inherited
  mixed-radix engine, storage geometry and source notices below.
Idea:
Store each 8x8 transposed corner directly in the inverse slab, eliminating
its intermediate 64x64 TBUF write and read. Unlike the full-size unpadded
scratch used by the older corner, this slab has R+16 elements per row, so
successive destination rows land in distinct L1 cache sets. Keep the final
output transpose's TBUF so its non-temporal stores remain contiguous.
Purpose: Remove one tile-buffer round trip per coefficient without changing
twiddle generation, transform order or the in-place overwrite proof.
*/
/*
References:
- saffah_cc_v41_agg1, https://duck.ac/submission/121257 : copied the public
radix-3 engine, split B/c storage, and non-temporal corner stores.
- saffah_codex_6a_agg3, https://duck.ac/submission/121223 : retained prepared
factors, specialized corners and radix-8 row kernels (also present in 121257).
- Further source ancestry and any original notices are preserved below.
Idea: Experimental purpose and implementation:
Fuse inverse corner, the short row transforms, and final transpose in a
64-column slab. After all original input rows in a slab have been consumed,
write its natural-order result back to those same columns. Future slabs use
disjoint columns, so this removes the full-size inverse scratch write/read
without overwriting live input. Keep the per-row Montgomery twiddle state in
a compact table across slabs; use ordinary stores for the immediately read
slab and retain non-temporal stores for its final output.
*/
// [n02b_] —— 取回对手**再下一代**最新件 #121223 + 补回对手丢掉的我方独占刀(NT store 消 RFO)
// ===== REFERENCES =====
// [1] duck.ac 用户 **saffah_codex_6a_agg3**,提交 **#121223** <https://duck.ac/submission/121223>
//     (4480.158163 ms,Accepted)—— **本文件正文 = 该提交的完整源码**(逐字节副本,仅在**最前面**插入本注释块;
//     另有**唯一一处**正文改动,见"思路")。**抄了什么**:其本代整代重写 ——
//     ① **列几何由 radix-5 改 radix-3**(`N3_=3, M3L_=14 ⇒ NC_=3*16384=49152`)使 `N1_ = NR_*NC_ = 201326592 > n+m = 2e8`
//     ⇒ **`M = n+m-N1_ < 0`,整个 wrap-correction(E2)相位被消掉**(其一相位 671 MB 级变换 + 一个 2^26 点变换全部不再执行);
//     ② `BSTAIL_SPLIT` 的 B/S 数组切分(前 188805120 个元素落在 `c` 里、其余落 `BTAIL`,`BROW(el)` 做地址映射);
//     ③ `corner_tile_avx`/`corner_tile_avx_blk` 增加 `split` 形参走 `BROW(...)`;④ `dif_pass8<2048>` 覆盖新的 L=16384 行。
//     ★ 其自述世系(其原注释块随正文保留):其本件的前身 = **本账号 #120907**、**#114124** 与本文件"世系"块所列各件。
// [2] 本账号 **saffah_cc_v41_agg1**,提交 **#121242** <https://duck.ac/submission/121242>(4888.186934 ms)
//     —— 我方上一发;本文件重新补回的 NT 出口趟即该发对该世代(对手 #120978 <https://duck.ac/submission/120978>,
//     5125.777078 ms)的唯一改动,在那一代判题机定价为 **−6.674%**(5237.493157 → 4888.186934)。
// [3] 本账号既有"整代复刻 + 署名 + 自加一刀"的合规先例:**#107081**、**#114077**、**#120907**、**#121242** ⇒ 本发沿用同一形式。
// **许可合规**:以上均为 duck.ac 公开提交,其正文按站点规则公开可见、可直接取用;两份源件均未见许可证声明;
//   本文件按 `/home/yjp/duck.ac/RULES.md` 第 3 节逐项署名引用作者账号与原提交地址。
// ======================
/*
References:
- saffah_cc_v41_260924, https://duck.ac/submission/121149: copied the radix-three
length-201326592 convolution and c-plus-tail split workspace. This removes
the separate wrap-correction transform required by the previous geometry.
- saffah_codex_6a_agg3, https://duck.ac/submission/120978: reused forward/inverse
radix-eight row kernels, prepared Shoup factors and specialized corners.
- saffah_cc_v41_agg1, https://duck.ac/submission/120907: inherited forward
radix-eight DIF kernel; its ancestry includes https://duck.ac/submission/114124.
- saffah_codex_6a_agg3, https://duck.ac/submission/120819 and
https://duck.ac/submission/120776: prepared-factor and corner-specialization
implementation inherited through 120978. These public submissions supplied
no separate license notices; source authors and copied components are cited.
Idea:
- Combine the new radix-three geometry with both validated radix-eight row
schedules and predecoded Shoup recurrences. Extend three-stage fusion to
16384-element rows, leaving the final two binary levels in the existing
radix-four pass. Split-workspace addresses remain exactly as in 121149.
- Remove inherited timing instrumentation and dead corner-base updates.
Purpose:
- Validate the complete 200000001-coefficient output, split-storage boundaries
and preserved inputs locally, then measure this combined ordinary algorithm
on the official judge. No new assembly or system calls are introduced.
*/
// [n02d_] ===== REFERENCES =====
// [1] duck.ac user **saffah_codex_6a_agg3**, submission **#121444**
//     <https://duck.ac/submission/121444> (4228.716313 ms, Accepted) -- the body
//     of this file IS that submission's source byte-for-byte (only its leading
//     comment block was replaced by this one).  Copied from it: everything it
//     carries, notably the inverse slab transform `inv_natural_stream` with the
//     TBUF staging removed from its gather (its own change over #121337).
// [2] same user, submission **#121337**
//     <https://duck.ac/submission/121337> (4254.810910 ms) -- the generation
//     #121444 inherits from (64-column slab streaming of the inverse).
// [3] this account **saffah_cc_v41_agg1**, submissions **#121601 / #121257 /
//     #121242 / #120907 / #114124 / #114077** -- #121601 is our previous shot
//     (this same fusion on top of #121337, 4210.020775 ms); the rest supplied
//     the NT corner stores, the fused radix-8 row kernels and the radix-3 split
//     geometry that all later generations inherit from us.
// License: all of the above are public duck.ac submissions, visible by site
//     rule; neither source carries a license notice.  Cited per
//     /home/yjp/duck.ac/RULES.md sec.3.
// ===== 思路 =====
// [n02d_] Two knives, both pure traffic/staging cuts, no arithmetic change:
// (A) FUSE4 -- phase 2's four full-array passes
//       fwd(A){colfwd(A)} ; fwd(B){colfwd(B)} ; pointwise(A,B) ; inv{colinv(A)}
//     become ONE pass over the row index r:
//       for r: col_dif3(A_r) ; col_dif3(B_r) ; A_r *= B_r ; col_dit3(A_r)
//     Rows are independent, so the result is byte-identical; the row (196 KB)
//     stays in L2 across all four steps, so the DRAM traffic of that phase
//     drops from 8 x 805 MB to 3 x 805 MB.  This is the knife of our #121601
//     (measured on the judge: 4254.810910 -> 4210.020775 ms, -1.05%).
// (B) the same TBUF-staging removal #121444 applied to the inverse gather,
//     applied to `load_tile_avx_d` (the forward source transpose): the 8x8
//     register tile goes straight into its S row instead of TBUF + a second
//     64x64 copy loop.
// ======================
// 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>
#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];
// [n02d] fuse4: while 1, fwd() leaves the colfwd sweep and
// inv_natural_stream() leaves the colinv sweep to the single fused row pass.
static int g_fuse4 = 0;
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.
template<int H> static void dit_pass8(u32 *x,int L,const u32 *tw,const u32 *tws){
const __m256i Pv=_mm256_set1_epi32((int)P);
const u32 *t1=tw+H-1,*s1=tws+H-1;
const u32 *t2=tw+2*H-1,*s2=tws+2*H-1;
const u32 *t3=tw+4*H-1,*s3=tws+4*H-1;
for(int i=0;i<L;i+=8*H){
u32 *p=x+i;
for(int j=0;j<H;j+=8){
__m256i v0=ldu256(p+j),v1=ldu256(p+j+H),
v2=ldu256(p+j+2*H),v3=ldu256(p+j+3*H),
v4=ldu256(p+j+4*H),v5=ldu256(p+j+5*H),
v6=ldu256(p+j+6*H),v7=ldu256(p+j+7*H);
auto dit2=[&](__m256i &a,__m256i &b,__m256i w,__m256i ws){
__m256i bm=shoup8(b,w,ws),old=a;
a=vaddm256(old,bm,Pv);b=vsubm256(old,bm,Pv);
};
__m256i w1=ldu256(t1+j),ws1=ldu256(s1+j);
dit2(v0,v1,w1,ws1);dit2(v2,v3,w1,ws1);
dit2(v4,v5,w1,ws1);dit2(v6,v7,w1,ws1);
__m256i w20=ldu256(t2+j),ws20=ldu256(s2+j);
__m256i w21=ldu256(t2+j+H),ws21=ldu256(s2+j+H);
dit2(v0,v2,w20,ws20);dit2(v4,v6,w20,ws20);
dit2(v1,v3,w21,ws21);dit2(v5,v7,w21,ws21);
dit2(v0,v4,ldu256(t3+j),ldu256(s3+j));
dit2(v1,v5,ldu256(t3+j+H),ldu256(s3+j+H));
dit2(v2,v6,ldu256(t3+j+2*H),ldu256(s3+j+2*H));
dit2(v3,v7,ldu256(t3+j+3*H),ldu256(s3+j+3*H));
stu256(p+j,v0);stu256(p+j+H,v1);
stu256(p+j+2*H,v2);stu256(p+j+3*H,v3);
stu256(p+j+4*H,v4);stu256(p+j+5*H,v5);
stu256(p+j+6*H,v6);stu256(p+j+7*H,v7);
}
}
}
static inline __m256i shoup_of(__m256i w);
static inline __m256i mont8(__m256i a, __m256i b);
template<bool Product=false>
static void row_dit_avx(u32 *x, int L, const u32 *tw, const u32 *tws, const u32 *other=nullptr) {
auto initial_load=[&](int offset) {
 __m256i v=ldu256(x+offset);
 if constexpr(Product){__m256i y=ldu256(other+offset);v=mont8(v,y);}
 return v;
};
if constexpr(Product){if(L<8)for(int i=0;i<L;++i)x[i]=mm(x[i],other[i]);}
__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 = initial_load(i4 + 0);
__m256i vr1 = initial_load(i4 + 8);
__m256i vr2 = initial_load(i4 + 16);
__m256i vr3 = initial_load(i4 + 24);
__m256i vr4 = initial_load(i4 + 32);
__m256i vr5 = initial_load(i4 + 40);
__m256i vr6 = initial_load(i4 + 48);
__m256i vr7 = initial_load(i4 + 56);
{
{
__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 d0 = _mm256_sub_epi32(y, vr0);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr0 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__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 d0 = _mm256_sub_epi32(y, vr1);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr1 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__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 d0 = _mm256_sub_epi32(y, vr2);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr2 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__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 d0 = _mm256_sub_epi32(y, vr3);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr3 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__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 d0 = _mm256_sub_epi32(y, vr4);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr4 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__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 d0 = _mm256_sub_epi32(y, vr5);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr5 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__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 d0 = _mm256_sub_epi32(y, vr6);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr6 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__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 d0 = _mm256_sub_epi32(y, vr7);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
vr7 = _mm256_blend_epi32(sa, d, 0xAA);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr0, 0x4E);
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);
vr0 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr1, 0x4E);
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);
vr1 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr2, 0x4E);
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);
vr2 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr3, 0x4E);
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);
vr3 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr4, 0x4E);
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);
vr4 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr5, 0x4E);
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);
vr5 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr6, 0x4E);
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);
vr6 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_shuffle_epi32(vr7, 0x4E);
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);
vr7 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr0, IDX8);
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);
vr0 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr1, IDX8);
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);
vr1 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr2, IDX8);
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);
vr2 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr3, IDX8);
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);
vr3 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr4, IDX8);
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);
vr4 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr5, IDX8);
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);
vr5 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr6, IDX8);
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);
vr6 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__m256i y = _mm256_permutevar8x32_epi32(vr7, IDX8);
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);
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 = initial_load(i4);
{
__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 d0 = _mm256_sub_epi32(y, v);
__m256i d = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
v = _mm256_blend_epi32(sa, d, 0xAA);
}
{
__m256i y = _mm256_shuffle_epi32(v, 0x4E);
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);
v = _mm256_blend_epi32(sa, d, 0xCC);
}
{
__m256i y = _mm256_permutevar8x32_epi32(v, IDX8);
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);
v = _mm256_blend_epi32(sa, d, 0xF0);
}
stu256((void *)((x + i4)), v);
}
off = 7;
int len = 16;
if(L==4096 || L==8192 || L==16384){
dit_pass8<8>(x,L,tw,tws);
dit_pass8<64>(x,L,tw,tws);
dit_pass8<512>(x,L,tw,tws);
len=8192;off=4095;
}
for (; (len << 1) <= L; len <<= 2) {
int h = len >> 1;
const u32 *t1 = tw + off,  *ts1 = tws + off;
const u32 *t2 = t1 + h,    *ts2 = ts1 + h;
const u32 *t3 = t2 + h,    *ts3 = ts2 + h;
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) {
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.
template<int H3, bool HalfZero=false> static void dif_pass8(u32 *x, int L, const u32 *tw, const u32 *tws) {
const __m256i Pv = _mm256_set1_epi32((int)P);
const int len = H3 << 3;
const u32 *t1 = tw + (len >> 1) - 1, *s1 = tws + (len >> 1) - 1;
const u32 *t2 = tw + (len >> 2) - 1, *s2 = tws + (len >> 2) - 1;
const u32 *t3 = tw + H3 - 1,        *s3 = tws + H3 - 1;
for (int i = 0; i < L; i += len) {
u32 *p = x + i;
for (int j = 0; j < H3; j += 8) {
__m256i v0=ldu256(p+j),v1=ldu256(p+j+H3),
v2=ldu256(p+j+2*H3),v3=ldu256(p+j+3*H3);
__m256i v4,v5,v6,v7;
auto dif2=[&](__m256i &a,__m256i &b,__m256i w,__m256i ws) {
__m256i sm=vaddm256(a,b,Pv),df=vsubm256(a,b,Pv);
a=sm;b=shoup8(df,w,ws);
};
if constexpr(HalfZero) {
 v4=shoup8(v0,ldu256(t1+j),ldu256(s1+j));
 v5=shoup8(v1,ldu256(t1+j+H3),ldu256(s1+j+H3));
 v6=shoup8(v2,ldu256(t1+j+2*H3),ldu256(s1+j+2*H3));
 v7=shoup8(v3,ldu256(t1+j+3*H3),ldu256(s1+j+3*H3));
} else {
 v4=ldu256(p+j+4*H3);v5=ldu256(p+j+5*H3);
 v6=ldu256(p+j+6*H3);v7=ldu256(p+j+7*H3);
 dif2(v0,v4,ldu256(t1+j),ldu256(s1+j));
 dif2(v1,v5,ldu256(t1+j+H3),ldu256(s1+j+H3));
 dif2(v2,v6,ldu256(t1+j+2*H3),ldu256(s1+j+2*H3));
 dif2(v3,v7,ldu256(t1+j+3*H3),ldu256(s1+j+3*H3));
}
__m256i w20=ldu256(t2+j),s20=ldu256(s2+j);
__m256i w21=ldu256(t2+j+H3),s21=ldu256(s2+j+H3);
dif2(v0,v2,w20,s20);dif2(v4,v6,w20,s20);
dif2(v1,v3,w21,s21);dif2(v5,v7,w21,s21);
__m256i w3=ldu256(t3+j),ws3=ldu256(s3+j);
dif2(v0,v1,w3,ws3);dif2(v2,v3,w3,ws3);
dif2(v4,v5,w3,ws3);dif2(v6,v7,w3,ws3);
stu256(p+j,v0);stu256(p+j+H3,v1);
stu256(p+j+2*H3,v2);stu256(p+j+3*H3,v3);
stu256(p+j+4*H3,v4);stu256(p+j+5*H3,v5);
stu256(p+j+6*H3,v6);stu256(p+j+7*H3,v7);
}
}
}
template<bool HalfZero=false>
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) {
int len = L;
if (L == 4096) {                       // [rt8b] 2^12: 3 radix-8 passes -> len 8
dif_pass8<512,HalfZero>(x, L, tw, tws); dif_pass8<64>(x, L, tw, tws); dif_pass8<8>(x, L, tw, tws);
len = 8;
} else if (L == 8192) {                // [rt8b] 2^13: 3 radix-8 passes -> len 16
dif_pass8<1024>(x, L, tw, tws); dif_pass8<128>(x, L, tw, tws); dif_pass8<16>(x, L, tw, tws);
len = 16;
}
if (L == 16384) {
dif_pass8<2048>(x,L,tw,tws); dif_pass8<256>(x,L,tw,tws); dif_pass8<32>(x,L,tw,tws);
len=32; // the existing radix-four pass completes 32 -> 8
}
if (L != 4096 && L != 8192)
for (; len >= 32; len >>= 2) {
int h = len >> 1, h2 = len >> 2;
const u32 *t1 = tw + h - 1,  *ts1 = tws + h - 1;
const u32 *t2 = t1 + h2,    *ts2 = ts1 + h2;
const u32 *t3 = tw + h2 - 1, *ts3 = tws + h2 - 1;
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) {
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)));
{
{
__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);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr0 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__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);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr1 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__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);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr2 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__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);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr3 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__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);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr4 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__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);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr5 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__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);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr6 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__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);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
vr7 = _mm256_blend_epi32(sa, d, 0xF0);
}
}
{
{
__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);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr0 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__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);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr1 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__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);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr2 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__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);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr3 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__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);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr4 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__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);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr5 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__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);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr6 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__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);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
vr7 = _mm256_blend_epi32(sa, d, 0xCC);
}
}
{
{
__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);
}
}
{
{
__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);
}
}
{
{
__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);
}
}
{
{
__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);
}
}
{
{
__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);
}
}
{
{
__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);
}
}
{
{
__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);
}
}
{
{
__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)));
{
__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);
d0 = _mm256_min_epu32(d0, _mm256_add_epi32(d0, Pv));
__m256i d = shoup8(d0, V8, V8s);
v = _mm256_blend_epi32(sa, d, 0xF0);
}
{
__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);
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
v = _mm256_blend_epi32(sa, d, 0xCC);
}
{
__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_prepared(u32 *dst, const u32 *src, u32 k, u32 ks) {
const __m256i kv = _mm256_set1_epi32((int)k);
const __m256i ksv = _mm256_set1_epi32((int)ks);
for (int v = 0; v < TS; v += 8)
stu256(dst + v, shoup8(ldu256(src + v), kv, ksv));
}
static void scale_row_mont(u32 *dst, const u32 *src, u32 kM) {
const u32 k = mont1(kM, 1u);
const __m256i kv = _mm256_set1_epi32((int)k);
const __m256i ks = _mm256_set1_epi32((int)shoup_c(k));
for (int v = 0; v < TS; v += 8)
stu256(dst + v, shoup8(ldu256(src + v), kv, ks));
}
// TW[u] = mont(TW[u-1], RVM)
static void twiddle_rows_mont(u32 *TW, const u32 *RVM) {
// RVM contains Montgomery residues. Decode each reusable factor once,
// retaining every TW value in the original Montgomery representation.
__m256i factor[TS / 8], factor_shoup[TS / 8];
const __m256i one = _mm256_set1_epi32(1);
for (int v = 0; v < TS; v += 8) {
factor[v / 8] = mont8(ldu256(RVM + v), one);
factor_shoup[v / 8] = shoup_of(factor[v / 8]);
}
for (int u = 1; u < TS; ++u) {
const u32 *prev = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v += 8)
stu256(cur + v, shoup8(ldu256(prev + v),
factor[v / 8], factor_shoup[v / 8]));
}
}
// 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;
// [n02b] NT store: written once here, re-read only a whole phase later
// (colfwd / row_dit / pointwise), so the RFO traffic is pure waste.
// shape = this file's own untr_tile_avx NT tail.
if ((((uintptr_t)d) & 31u) == 0u) {
for (int u = 0; u < TS; u += 8)
_mm256_stream_si256((__m256i *)(d + u), ldu256((const void *)(s + u)));
} else {
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;
// [n02b] NT store: written once here, re-read only a whole phase later
// (colfwd / row_dit / pointwise), so the RFO traffic is pure waste.
// shape = this file's own untr_tile_avx NT tail.
if ((((uintptr_t)d) & 31u) == 0u) {
for (int u = 0; u < TS; u += 8)
_mm256_stream_si256((__m256i *)(d + u), ldu256((const void *)(s + u)));
} else {
for (int u = 0; u < TS; u += 8) stu256((void *)(d + u), ldu256((const void *)(s + u)));
}
}
}
// [n02i] same 8x8-block corner multiply+transpose as corner_tile_avx_blk, but the
// finished 64x64 tile lands in a caller-supplied buffer (tile[v*TS+u]) instead of
// being NT-stored into X.  The fused head in corner_fwd_stream consumes it from
// L1 and performs the single NT store for all three radix-three ports at once.
static void corner_tile_avx_to(const u32 *in, u32 *tile, int B0, int strIn) {
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++)
r[k] = mont8(ldu256((const void *)(in + (u64)(u0 + k) * strIn + B0 + v0)),
ldu256((const void *)(TW + (u0 + k) * TS + v0)));
t8x8(r);
for (int k = 0; k < 8; k++)
stu256((void *)(tile + (v0 + k) * TS + u0), r[k]);
}
}
}
// 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);
// [n02d] same staging removal the rival applied in the inverse gather: the
// transposed 8x8 tile goes straight to its S row (u32 *)(S + (cb0+j)*dstr +
// r0 + rb0) instead of TBUF -> 64x64 copy loop.  Layout is identical.
for (int j = 0; j < 8; j++)
stu256((void *)(S + (u64)(cb0 + j) * dstr + r0 + rb0), r[j]);
}
}
}
// 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_head(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));
}
}
}
static void col_dif3(u32 *x) {
col_dif3_head(x);
for (int s = 0; s < 3; s++) row_dif_avx(x + (u64)s * M3_, (int)M3_, tsubFw, tsubFsw);
}
static void col_dit3_tail(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)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]);
}
}
static void col_dit3(u32 *x) {
for (int s = 0; s < 3; s++) row_dit_avx(x + (u64)s * M3_, (int)M3_, tsubIw, tsubIsw);
col_dit3_tail(x);
}
// [n02i] the radix-three head (identical arithmetic to col_dif3_head) applied in
// registers to one already-cornered 64x64 tile per port, storing all three ports'
// rows straight into X.  Defined here because it must sit inside the AVX2 target
// region (corner_fwd_stream itself is outside it and cannot hold intrinsics).
static void corner_head_store_avx(const u32 *t0, const u32 *t1, const u32 *t2,
u32 *d0, u32 *d1, u32 *d2, int Ab) {
const __m256i Pvh = _mm256_set1_epi32((int)P);
__m256i hc[2], hcs[2];
for (int i = 0; i < 2; i++) { hc[i] = _mm256_set1_epi32((int)R3F[i]); hcs[i] = _mm256_set1_epi32((int)R3FS[i]); }
for (int u = 0; u < TS; u += 8) {
__m256i vv[3];
vv[0] = ldu256((const void *)(t0 + u));
vv[1] = ldu256((const void *)(t1 + u));
vv[2] = ldu256((const void *)(t2 + u));
dft3(vv, hc, hcs, Pvh);
__m256i w1 = ldu256((const void *)(T3F + Ab + u));
__m256i s1 = ldu256((const void *)(T3FS + Ab + u));
__m256i w2 = ldu256((const void *)(T3F + M3_ + Ab + u));
__m256i s2 = ldu256((const void *)(T3FS + M3_ + Ab + u));
_mm256_stream_si256((__m256i *)(d0 + u), vv[0]);
_mm256_stream_si256((__m256i *)(d1 + u), shoup8(vv[1], w1, s1));
_mm256_stream_si256((__m256i *)(d2 + u), shoup8(vv[2], w2, s2));
}
}
// [n02i] the two radix-three head sweeps are no longer here: corner_fwd_stream
// applies the identical head in registers as it produces each tile group.
static void col_convolve3(u32 *a,u32 *b) {
for(int arm=0;arm<3;++arm) {
 u32 *pa=a+(u64)arm*M3_,*pb=b+(u64)arm*M3_;
 row_dif_avx(pa,(int)M3_,tsubFw,tsubFsw);
 row_dif_avx(pb,(int)M3_,tsubFw,tsubFsw);
 row_dit_avx<true>(pa,(int)M3_,tsubIw,tsubIsw,pb);
}
col_dit3_tail(a);
}
#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 col_dif3_arms(u32 *x) {
for (int s = 0; s < 3; s++) row_dif_avx(x + (u64)s * M3_, (int)M3_, tsubFw, tsubFsw);
}
static void colfwd(Eng &e, u32 *row) {
#if HAVE_AVX2
if (e.r5) { col_dif3_arms(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;
}
template<int perm>
static void corner_generic(const u32 *in, u32 *out, int nA, int nB, int strIn, int strOut, u32 root, 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];
u32 cu64_plain[TS],cu64_shoup[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;
cu64_plain[u]=mont1(cu64[u],1u); cu64_shoup[u]=shoup_c(cu64_plain[u]);
}
} 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_prepared(cur, cur, cu64_plain[u], cu64_shoup[u]);
#else
u32 d = cu64[u];
for (int v = 0; v < TS; v++) cur[v] = mont1(cur[v], d);
#endif
}
}
// base is used only to initialize the first tile. Later tiles update TW directly.
} 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<1>(S, X, e.C, e.R, e.R, e.sx, root, 0);   // A = j0, B = rev'd k1 position
}
static inline void corner_inv(Eng &e, const u32 *X, u32 *S, u32 root) {
corner_generic<2>(X, S, e.R, e.C, e.sx, e.R, root, 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.
// [n02i] three port buffers so the radix-three head can be fused into the corner
// (see corner_fwd_stream).  N02I_SB is the stride between ports.
#define N02I_SB (TS * (NRMAX_ + PADB_))
static u32 SBLK[3 * N02I_SB] __attribute__((aligned(64)));
static u32 TILE3[3][TS * TS] __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;
int lgB = 0; while ((1 << lgB) < nB) lgB++;
const int S1 = nB / TS, HB = lgB - 6;
u32 RV2[TS], PW2[3][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 bool halfzero=(nB==4096 && (u64)nsrc<=(u64)(nB/2)*nA);
const int AM = nA / 3;                        // 2^14 = M3_: columns per radix-three port
// [n02i] The radix-three head used to run as two extra whole-row sweeps of X after
// this pass, re-reading (from DRAM -- the writes below are non-temporal) and
// re-writing 196 KB per array per row.  The three ports of a head group are the
// three column blocks Ab, Ab+M3_, Ab+2*M3_, so holding all three in flight lets the
// head run in registers on the tile it just produced, and the output leaves through
// the same single NT store the corner used to issue.
for (int Ab = 0; Ab < AM; Ab += TS) {
for (int p = 0; p < 3; p++) {
u32 *S = SBLK + (u64)p * N02I_SB;
const int A0 = Ab + p * AM;
const u32 wr = pw(root, (u32)A0);             // root^{A0}
{ u32 g = pw(wr, (u32)S1), gM2 = mm(g, R32), c = R32;
for (int t = 0; t < TS; t++) { PW2[p][REV6[t]] = c; c = mont1(c, gM2); } }
load_block(src, nsrc, S, A0, halfzero?nB/2:nB, nA, strIn, ctr);
for (int u=0;u<TS;++u){
 if(halfzero)row_dif_avx<true>(S+(u64)u*strIn,nB,e.tfR,e.wsR);
 else row_dif(S+(u64)u*strIn,nB,e.tfR,e.wsR);
}
}
for (int B0 = 0; B0 < nB; B0 += TS) {
u32 H = revbits((u32)(B0 / TS), HB);
u32 rootH = POWT[H];
for (int p = 0; p < 3; p++) {
const u32 wr = pw(root, (u32)(Ab + p * AM));
u32 wrH = pw(wr, H);
u32 RV[TS];
scale_row_mont(RV, RV2, rootH);
scale_row_mont(TW, PW2[p], mm(wrH, R32));
twiddle_rows_mont(TW, RV);
corner_tile_avx_to(SBLK + (u64)p * N02I_SB, TILE3[p], B0, strIn);
}
for (int v = 0; v < TS; v++) {
const u32 *t0 = TILE3[0] + v * TS, *t1 = TILE3[1] + v * TS, *t2 = TILE3[2] + v * TS;
const u64 rowbase = (u64)(B0 + v) * strOut + Ab;
u32 *d0 = split ? BROW((long)rowbase) : X + rowbase;
u32 *d1 = split ? BROW((long)(rowbase + AM)) : X + rowbase + AM;
u32 *d2 = split ? BROW((long)(rowbase + 2 * (long)AM)) : X + rowbase + 2 * AM;
corner_head_store_avx(t0, t1, t2, d0, d1, d2, Ab);
}
}
}
}
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;
{ corner_fwd_stream(e, src, nsrc, X, split, ctr);  }
#endif
if (!g_fuse4) { for (int k0 = 0; k0 < e.R; k0++)
colfwd(e, split ? BROW((long)k0 * e.sx) : X + (u64)k0 * e.sx);  }
}
#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) { 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);
}
}
// Each slab owns the same column interval before and after the inverse corner.
// It is safe to overwrite X only after all R input rows in that interval were read.
static u32 INV_TW[NRMAX_ * TS] __attribute__((aligned(64)));
static u32 INV_STEP[NRMAX_], INV_STEP_SHOUP[NRMAX_];
#pragma GCC push_options
#pragma GCC target("avx2,tune=skylake")
static void inv_natural_stream(Eng &e, u32 *X) {
if (!g_fuse4)
for (int k = 0; k < e.R; ++k) colinv(e, X + (u64)k * e.sx);
const u32 root = pw(e.w, P - 2u);
const u32 seed = mm(mm(g_inv_scale, R32), R32);
int lg = 0; while ((1 << lg) < e.R) ++lg;
for (int k = 0; k < e.R; ++k) {
const u32 factor = pw(root, revbits(k, lg));
const u32 factorM = mm(factor, R32);
INV_STEP[k] = pw(factor, TS);
INV_STEP_SHOUP[k] = shoup_c(INV_STEP[k]);
u32 *tw = INV_TW + k * TS;
tw[0] = seed;
for (int v = 1; v < TS; ++v) tw[v] = mont1(tw[v - 1], factorM);
}
for (int c0 = 0; c0 < e.C; c0 += TS) {
for (int r0 = 0; r0 < e.R; r0 += TS) {
if (c0) for (int u = 0; u < TS; ++u) {
const int k = r0 + u;
u32 *tw = INV_TW + k * TS;
scale_row_prepared(tw, tw, INV_STEP[k], INV_STEP_SHOUP[k]);
}
for (int u0 = 0; u0 < TS; u0 += 8) {
for (int v0 = 0; v0 < TS; v0 += 8) {
__m256i rows[8];
for (int k = 0; k < 8; ++k) {
const u32 *src = X + (u64)(r0 + u0 + k) * e.sx + c0 + v0;
const u32 *tw = INV_TW + (r0 + u0 + k) * TS + v0;
rows[k] = mont8(ldu256(src), ldu256(tw));
}
t8x8(rows);
for (int k = 0; k < 8; ++k)
stu256(SBLK + (u64)(v0 + k) * e.blkr + r0 + u0, rows[k]);
}
}
}
for (int v = 0; v < TS; ++v)
row_dit(SBLK + (u64)v * e.blkr, e.R, e.tiR, e.wsR + e.R);
// Every X row in c0..c0+TS has now been consumed, so this is in-place safe.
for (int r0 = 0; r0 < e.R; r0 += TS) {
for (int u0 = 0; u0 < TS; u0 += 8) {
for (int v0 = 0; v0 < TS; v0 += 8) {
__m256i rows[8];
for (int k = 0; k < 8; ++k)
rows[k] = ldu256(SBLK + (u64)(u0 + k) * e.blkr + r0 + v0);
t8x8(rows);
for (int k = 0; k < 8; ++k)
stu256(TBUF + (v0 + k) * TS + u0, rows[k]);
}
}
for (int v = 0; v < TS; ++v) {
u32 *dst = X + (u64)(r0 + v) * e.sx + c0;
for (int u = 0; u < TS; u += 8)
_mm256_stream_si256((__m256i *)(dst + u), ldu256(TBUF + v * TS + u));
}
}
}
_mm_sfence();
}
// [n02d] one row of the fused phase-2 sweep: pa *= pb, byte-identical to
// pointwise_mul_avx's per-row body (which is what the flat pointwise ran).
static void fuse4_row_avx(u32 *pa, const u32 *pb, int n) {
int i = 0;
for (; i + 8 <= n; i += 8) {
__m256i x = ldu256((const void *)(pa + i));
__m256i y = ldu256((const void *)(pb + i));
stu256((void *)(pa + i), shoup8(x, y, shoup_of(y)));
}
for (; i < n; i++) pa[i] = mm(pa[i], pb[i]);
}
#pragma GCC pop_options
// ---------------------------------------------------------------------------
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;
g_fuse4 = 1;
fwd(E1, (const u32 *)a, (u32)n + 1, S, A, 0, ctr);
fwd(E1, (const u32 *)b, (u32)m + 1, S, B, 1, ctr);
// [n02d] fused sweep: colfwd(A_r) | colfwd(B_r) | pointwise(A_r,B_r) | colinv(A_r)
for (int r = 0; r < (int)NR_; r++) {
u32 *pa = A + (u64)r * CSX_;
u32 *pb = BROW((long)r * CSX_);
if(E1.r5) col_convolve3(pa,pb);
else {
colfwd(E1, pa);
colfwd(E1, pb);
fuse4_row_avx(pa, pb, (int)NC_);
colinv(E1, pa);
}
}
inv_natural_stream(E1, A);
g_fuse4 = 0;
// ---------- 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
_mm_sfence();   // non-temporal stores: make them globally visible
(void)E2;
}

CompilationN/AN/ACompile OKScore: N/A

Testcase #14.038 s1584 MB + 476 KBAcceptedScore: 100


Judge Duck Online | 评测鸭在线
Server Time: 2026-11-02 14:43:27 | Loaded in 2 ms | Server Status
个人娱乐项目,仅供学习交流使用 | 捐赠