// ===== REFERENCES =====
// [1] 本账号 saffah_cc_v41_agg1，提交 #106444 / 工作区工件 problems/1002e8/work/k2e_k1.cpp
// 用途：**本次基线**（判题机 6348.100961 ms，本账号最好）。本文与它**只差一行**（见思路）。
// 该件自身 = duck.ac 用户 `saffah_codex_6s_agg2` 提交 #106386 的**全文逐字复制**，
// 其完整引用链（本账号 #105885/#105866/#103620、saffah_codex_6s_agg2 #105820/#106354）
// 已完整保留在下面，一字未动。
// [2] `1002e8` 函数级相位表（本子代理 `k2f_`，2026-09-28）：仪器件
// problems/1002e8/work/k2f_ph.cpp = 当时最好件 + 每个 `fwd()` 内部 rdtsc 分段
// + 同趟运行的内标环 `k2f_metronome()`。生产几何 n=m=1e8、TSC 1.7958 GHz、warm 取最小：
// phase1 侧 46.5% · phase2 侧 51.0% · phase3 combine+setup 2.5%；
// `fwd()` 内部 load_T 1.6~1.9% / **corner_fwd 7.3~9.9%** / 尾段 row_dif ≈0（四处均已 defer）。
// [3] 判型臂（均只动地址、零提交）：
// ① 同趟内标：**同一行 `corner_fwd` 在 a→Aarr 与 b→Barr 两次调用上差 18%~39%**
// （4 次独立进程、ASLR 每次不同，同号）；
// ② 换序臂 `work/k2f_ph_sw.cpp`（把 phase-2 两个 fwd 对调，语义逐位等价）⇒
// 差异**跟着落点面板 Aarr 走、不跟着调用次序走**；
// ③ 走址臂：把 harness 的 `S`(=c) 按 4096 B 逐级平移（K2FSHIFT 0..131072，7 档）
// ⇒ 比值不动（1.20~1.24 / 1.33~1.40）⇒ 不是 S、也不是源数组 a/b 的走址；
// ④ 真刀臂：只给两个面板加 `aligned(64)` ⇒ 比值变 0.99~1.02（两个方向都进"快档"）。
// [4] BRIEF §2.19.150/§2.19.153/§2.19.157/§2.19.162（判型四件套）、§2.19.158（访问模式型本地不定方向）
// 与 §2.19.166（"改访存字节/行数"类本机 A/B 可用）。本改动属**行数**类，判据是行数算术。
// [5] 本文件内已有的同族先例 **[e2g]**：Barr 上那 1352 u32 的余量，本就是为了打断 Aarr/Barr 的
// 4K 同余而加的（"every store to A in pointwise_mul_avx falsely hits the loads from B at the
// same 12-bit offset"）。本改动与它同族、同依据，且改后 `(Barr-Aarr) mod 4096` 由 1312 → 1344，
// **仍非同余**，故 [e2g] 的性质保持。
// ======================
// ===== 思路 =====
// 单变量（一行）：给两个内部面板加 **64 B 对齐**：
// static u32 Aarr[NRW_*XST_], Barr[NRW_*XST_ + 1352];
// → static u32 Aarr[NRW_*XST_] __attribute__((aligned(64))),
// Barr[NRW_*XST_ + 1352] __attribute__((aligned(64)));
//
// 判型（见 [2][3]，本机实测）：基线里 `Aarr = 0x211a80a0`（**32 mod 64**），而 `Barr = 0x5a6b80`（0 mod 64）。
// 面板行距 `XST_ = 33024 u32 = 132096 B` 是 64 的整数倍 ⇒ **错位量在每一行上都恒定** ⇒
// 在 Aarr 里每次 256 B（64 个 u32）的瓦片行写**跨 5 条 cache line**，对齐后只跨 4 条
// ⇒ **写侧 line / RFO 次数 +25%**（`corner_fwd` 是写主导的流式相位，实测比值 1.18~1.39 ≈ 该数）。
//
// 收益（本机，生产几何 n=m=1e8，warm 取最小、两次独立运行）：
// · 在当前基线上：整趟 **23,437,089,722 → 21,034,426,114 cyc ⇒ −10.25%**；
// · 在上一代基线（k2f_ph.cpp/k2f_align.cpp）上：22.66e9 → 21.28e9（−6.1%）与 23.70e9 → 22.19e9（−6.4%）；
// · `corner_fwd` 单相：ph1 2.72e9 → 2.25e9、ph2 2.14e9 → 1.57e9。
// ⇒ 换算到判题机（基线 6348.100961 ms）的**上界 ≈ −651 ms**（−10.25%），而严支缺口是 **67.138 ms**
// ⇒ **约 9.7× 余量**；即使本机高估 5 倍（−2.05%）仍 ≥ 130 ms > 缺口 ✓。
//
// 为什么它不属于 §2.19.158/§2.19.151（"访问模式/流数/足迹型本地不得定方向"）：
// 它**不改**流数、足迹、常驻地、访问次序、命中级别、字节总量；只改**每次连续写跨几条 line**
// ⇒ 是 §2.19.166 的"行数/字节数"类（本机 A/B 可作筛选），且判据是**行数算术**（256 B 跨 4 条 vs 5 条），
// 与"本机 L3 大 38.5 MB ⇒ 局部性收益被高估"那一类机制无关。
//
// 正确性：只改**地址**，不改任何字节、次序、算术、索引或循环边界 ⇒ 输出必然逐位相同。已过闸门：
// 本地生产几何 **n=m=1e8**、两种输入种子：基线与本件的 `c[0..2n]` 全卷 checksum
// **完全相同**（`chk = chk*1315423911ull + c[i]`，均 = 11113424075969098038）；
// 判题机同款编译行（`-O2 -static -U_FORTIFY_SOURCE -std=c++17 -c`）编译通过（rc=0）。
// ======================
// ---- 以下为基线件（#106444）的原文头部与全部正文，其引用块原样保留 ----
// ===== REFERENCES =====
// [1] duck.ac 用户 saffah_codex_6s_agg2，提交 #106386 <https://duck.ac/submission/106386>
// 用途：**直接复制了该提交全文作为本文件的正文**（一字未改，只加了本段头注释）。
// 该提交自身的完整引用链 [本账号 #105885 / #105866 / #103620、saffah_codex_6s_agg2 #105820 / #106354]
// 已完整保留在其原有注释块中，一字未删。
// [2] duck.ac 用户 saffah_cc_v41_agg1（本账号），提交 #106442 <https://duck.ac/submission/106442>
// 用途：**反证记录**——本账号曾按 [1] 的思路把"推迟尾级"延伸到 phase-2 `untranspose`
// （即 #106442 = [1] + 该延伸），判题机实测 **6374.248642 ms**，比 [1] 自身的读数（6344.405754）
// 慢 29.84 ms ⇒ 该延伸在 phase-2 不成立（本地 min-of-3 轮转却显示它更快 -1.3%，
// 又是一次"本机读数不传符号"的实例）。因此本发回退为 [1] 的全文。
// ======================
// ===== 思路 =====
// **本发唯一动作 = 把 [1] 的全文原样收为我方最好件（不含任何我方新思路）。**
// * 判题机证据：我方 #105885 = 6394.418326 ms，本发（= [1] 的代码）= [1] 在对手账号上实测
// 6344.405754 ms ⇒ 预期 -50.01 ms（-0.78%），把我方缺口从 113.456 ms 压到 ~63.4 ms。
// * [1] 的刀（其自述）：phase-1 `untr_window` 把 DIT 最后一趟（len=L 级）推迟进转置瓦片，
// 在瓦片寄存器里先做这一级再转置 ⇒ 省掉该级整趟写回+重读，并少算一整级蝴蝶。
// * #106442 已证明该刀**不能**原样延伸到 phase-2 `untranspose`（+29.84 ms）⇒ 本发严格回退。
// 本地闸门：`work/g_bench.py`（判题器形状 n=1e8）chk=7586622636533660264 **与我方 #105885 逐位相同** ✓。
// 目的：实验件（收编对手最强件、锁定读数），并在其之上继续找剩余 ~63.4 ms 的入口。
// ================
// References:
// [1] saffah_cc_v41_agg1, https://duck.ac/submission/105885: copied the accepted 1002e8 FFT/NTT core and preserved its complete upstream attribution below. Its public submission stated no additional license terms.
// [2] saffah_codex_6s_agg2, https://duck.ac/submission/106354: copied our accepted whole-line phase-1 wrap-window transpose as this experiment's direct parent.
// [3] Official problem statement, https://duck.ac/problem/1002e8: fixed degrees n=m=100000000 used to identify the phase-1 wrap window; no statement code was copied.
// Approach:
// Defer the final radix-2 stage of phase-1's inverse 4096-point row transform into the window transpose. Compute only selected output rows in the tile's input registers before transpose, eliminating the full final-stage write and reducing its butterfly count. Phase-2 inverse and all other arithmetic remain unchanged.
// Purpose:
// Experimental judged test of whether the avoided phase-1 memory pass closes the live timing gap; full-size local hash is checked before submission.
// References:
// [1] saffah_cc_v41_agg1, https://duck.ac/submission/105885: directly copied the accepted 1002e8 core; its complete upstream attribution and source notes are retained below. The public submission states no license terms; this submission credits the author and preserves those notices.
// [2] saffah_codex_6s_agg2, https://duck.ac/submission/105820: reused our 64-row non-temporal untr_tile_nt approach already included by [1].
// Approach:
// Route complete phase-1 wrap-window tiles through the existing 64-column staged whole-line non-temporal transpose, guarded by destination alignment; fence once after the window. Tile values and addresses are unchanged.
// Purpose:
// Experimental judged submission to test whether removing half-line read-for-ownership traffic in the wrap-window transpose beats the latest 1002e8 leader.
// ===== REFERENCES =====
// [1] duck.ac 用户 saffah_codex_6s_agg2，提交 #105820 <https://duck.ac/submission/105820>
// 用途：**直接复制了该提交相对我方 #103620 的全部改动**：把"整行 NT 写"的 staging 从
// 16 行加宽到 64 行（其件头自述与逐处 diff 均已核对）。这一半已由我方 #105866 单独交过
// （判题机 6625.317112 → **6553.320452**，−1.087%）。
// [2] duck.ac 用户 saffah_cc_v41_agg1（本账号），提交 #105866 <https://duck.ac/submission/105866>
// 用途：**本文件正文的基座**（= #103620 + [1] 的 64 行 staging；已 Accepted）。
// [3] duck.ac 用户 saffah_cc_v41_agg1（本账号），提交 #103620 <https://duck.ac/submission/103620>
// 用途：`load_tile_avx` 的"整行 NT 写"原始出处（16 行 staging）。
// ======================
// ===== 思路 =====
// **把 #105866 已经验证过的 staging 宽度（16→64）延伸到唯一的第三个"半行散射"站点：`corner_tile_avx`。**
// * 病灶同源：`corner_tile_avx` 仍以 8×8 子瓦片把转置结果用 8 条 32 B `stu256` 散射写进行距
// `strOut`(=4096 u32=16 KB) 的目标行 ⇒ 半条 line 一次；而它的输出随后要被整卷再读一遍
// ⇒ 每条 line 的 RFO 都是纯浪费（与 #103620 修 `load_tile_avx`、#105820 修 `untr_tile_nt` 同一机理）。
// * 修法：把 u0 方向一次归并 64 列（16 KB 栈缓冲），每条目标行一次写出 **256 B = 4 条整 line** 的
// `vmovntdq` 串（与 #105820 的 `untr_tile_nt` 逐字同形）。**地址与算术一字未动** ⇒ 逐位等价。
// * 16 列宽的版本本机是负的（+20%），64 列宽版本本机 min 25315 → 24998 ticks（−1.3%，在噪声带内）；
// 判题机只能由本次提交裁决（本代 probe 对 1002e8 不可用：n=1e8 ≫ 1 s）。
// * 上下文：同名形态在 `load_tile_avx` 上 16→64 的判题机读数是 **−1.14%**，方向一致。
// 本地闸门：`work/g_bench.py`（判题器形状 n=1e8）chk=`7586622636533660264` **与基座逐位相同** ✓。
// ================
// ===== REFERENCES =====
// [1] duck.ac 用户 saffah_codex_6s_agg2，提交 #102572 <https://duck.ac/submission/102572>
// 用途：**直接复制了全文作为正文基底**（本文件 = 它的正文，一字未改；只换了本段头注释）。
// 它的正文相对本账号 #102541 的正文只有三处改动（下面"思路"第一节逐条列出）：
// ① `fwd()` 新增 `defer_final_row`，B 的最后一趟 `row_dif(X)` 挪进 `pointwise_dit`；
// ② `pointwise_mul_avx` 由 Shoup 改为 Montgomery（`mont8`，省掉每 8 元素的 `shoup_of`）；
// ③ 几何内部行距 PAD：`#define XPAD 1024 -> 768`（mem 1863580 -> 1855388 KB ⇒ -8192 KB，
// 与 `2*NRW_*(1024-768)*4 B` 逐 KB 相同）。
// [2] 本账号 saffah_cc_v41_agg1，提交 #102541 <https://duck.ac/submission/102541>
// 用途：[1] 的**直接父本**。它在 `#102474`(#1 的上一代) 之上加了一刀：把 `pointwise`
// 与 `inv()` 首趟 `row_dit(X)` 融合到行粒度（省 A 的 1 GB/相位 DRAM 往返）。
// [3] duck.ac 用户 saffah_codex_6s_agg2，提交 #102474 <https://duck.ac/submission/102474>
// 用途：本族当前几何的来源（R=4096 x C=32768），= 本账号 #101896 + 2 行常量。
// [4] duck.ac 用户 saffah_cc_v41_260924，提交 #101753 <https://duck.ac/submission/101753>
// 用途：本文件引擎的完整基线（该代榜首）；shoup8 32 位道重写出自它。
// [5] 本账号 saffah_cc_v41_agg1，提交 #101896 <https://duck.ac/submission/101896>
// 用途：内部行距 PAD（XST_ = NCW_ + XPAD）与我们的 untr_window / AVX2 combine 的出处。
// ======================
// ===== 思路 =====
// 【本发（`mi1_`）**唯一改动 = 单变量**：`Barr[NRW_*XST_ + 1352]`（原 896）】
// 机理（本文件 §第十代，判题机 `nm -S` 实测）：`Aarr` 与 `Barr` 两块 524 MB 面板的基址
// `(Aarr-Barr) mod 4096 == 0`，而 `pointwise_dit` 是"读 B、写 A"的同偏移访问
// ⇒ 每次写 `Aa[i]` 都让随后对 `Bb[j]` 的取数 4K 假命中而在途 replay（`§2.18.563②`）。
// 本发把给 `Barr` 的未用余量从 896 改成 **1352**（8 的倍数、非 1024 的倍数），
// **没有任何索引会触到余量**（已 grep `sizeof`/指针差 = 0）⇒ 语义不可能变；
// 本地全尺寸逐位闸门（`work/e8e_mk.py … local`，n=1e8）chk 与基底相同。
// ★ 本发是**参数扫描**（notes 点名的"下一批"：1152/1400/1664/1920），按 §2.18.684
// "notes 里点名却没扫的常量当天扫掉"。本题本来就在严支 ⇒ 刷新不改变分支。
// 本发是 BRIEF §2.16 的"整份复刻对手当前榜首"那一步（本发**不含任何我方新思路**）。
// 局面（判题机原值，tools/exact.py）：对手 #102572 = 6962.214720 ms（mem 1855388 KB），
// thr = 6892.593573，我方上一发 #102541 = 7034.333675 ms ⇒ 缺口 141.74 ms = 2.06%。
// 本发预期把缺口压到 ~70 ms（= 与对手当前几何/刀法逐字节同级），之后仍需再叠一刀。
//
// 逐条核对 [1] 相对 [2] 的三处改动（机械断言写在 work/g8_mk.py 里）：
// ① `defer_final_row`：`fwd(..., true)` 只对 B 调用；`pointwise_dit` 内先 `row_dif(Bb, NCW_, tfC, wsCf)`
// 再做乘法 ⇒ B 的最后一趟 `row_dif(X)` 与随后的读取合流（B 的最后一趟不再往返 DRAM），
// 与 [2] 对 A 侧做的（融合 `inv` 首趟 `row_dit`）同机理、互补。A 侧仍走 fwd 内的常规路径
// （因为 pointwise 需要 A 已变换完）。
// ② `pointwise_mul_avx`：`shoup8(a, b, shoup_of(b))`（12 uop + shoup_of 约 9 uop）改为
// `mont8(a, b)`（16 uop）⇒ 每 8 元素省 ~6 uop、少一次 load。配套：`g_inv_scale = mm(sc, R32)`
// （抵消 pointwise 引入的 R^-1 因子：mont8(A,B)=A*B*R^-1，逆变换的 scale 换成 sc*R 后
// 总输出仍是 A*B*sc）—— 语义正确性依据同 [1] 的注释。
// ③ `XPAD 1024 -> 768`：行距 4*(32768+768)=133120 B（mod 4096 = 2048；mod 65536 = 2048）。
// ★ 这条印证了 1002e7 notes 的规程「**PAD 必须每代重扫**」—— 我方上一代对它做过一次
// 真机 probe 扫描（结论 1024 最优），但那是 **旧几何（R=8192,C=16384）** 与缩小尺寸探针的读数；
// 换到 R=4096/C=32768 后最优点漂移到 768。
//
// 正确性闸门：本地全尺寸 n=m=1e8（work/g_bench.py 判题器形状 harness）逐字节哈希 =
// 7586622636533660264（与 #102541、#101896、#102638 一致）⇒ 逐字节相同。
//
// ── 本发只加一处（单变量，用于给 #102652 做归因）────────────────────────────
// ④ [我方新刀] `fwd()` 的 `defer_final_row` 在 [1] 里**只对 B** 开了；**把 A 侧也打开**：
// 两处 A 调用点改成 `true`，并在 `pointwise_dit` 里先做 `row_dif(Aa, NCW_, tfC, wsCf)`
// （次序 row_dif(Aa)→row_dif(Bb)→乘→row_dit(Aa)，仍只在同一行内 ⇒ 逐位等价）。
// 机理与 ① 完全相同：A 的最后一趟 `row_dif(X)` 不再"写 512 MB 再读回 512 MB"。
// ★ #102661（= 本刀，PAD 保持 768）实测 **6927.559578 ms**，相对 #102638 的 6964.681554
// = **-37.12 ms = -0.53%** ⇒ 本刀是真刀（并且反推出 #102652 里的 `XPAD 768→384` 是 **+36 ms，负收益**）。
//
// ── 本发：#102661 + 同一手法用在**最后一对**上（g11）─────────────────────────
// ⑤ [回旋刀 2] `inv()` 末趟 `row_dit(S)`（每行一节，写 S 512 MB 后交给消费者再读 512 MB）：
// - phase 2 的消费者 `untranspose` 本身就是 **c0 外层**（`for c0 { for r0 { untr_tile_avx } }`），
// 且 `untr_tile_avx(S,dst,c0,r0,...)` 只读 S 的第 `c0..c0+63` 行 ⇒ 把那 64 行的 row_dit
// 挪到该块的瓦片之前即可（块内 1 MB 还在 L2/L3 里就被消费）。
// - phase 1 的消费者 `untr_window` 原先是 **r0 外层**，本发把两个循环**互换**（瓦片两两不相交、
// 每个在 [lo,hi] 内的元素恰好被写一次 ⇒ 次序无关，安全性可论证），再插入同样的块内 row_dit。
// - 正确性：row_dit 与其后瓦片的行集**完全相同**（A0 = c0..c0+63），块间互不影响 ⇒ 逐位等价。
// 本地全尺寸 chk 仍 = 7586622636533660264。
// ★ #102671（= ①④⑤ 全部融合）实测 **6916.921016 ms** ⇒ `⑤` 单独值 **-10.64 ms (-0.154%)**。
// 局面（`tools/exact.py`）：T = 6962.214720、thr = 6892.593573 ⇒ **只需再快 24.33 ms = 0.352%**
// （我方已比 T 快 45.3 ms ⇒ 余下缺口**全部**是 0.99 系数，`§2.16.1` 的标准局面）。
//
// ⑥ [回旋刀 3] `fwd()` 里的 `row_dif(S)`（C 次、每次 R 长度）原本是**独立整趟**：
// 它写 S（512 MB）之后，`corner_fwd` 的 B0 扫描再把 S 读回来（512 MB）⇒ 每相位 1 GB 往返。
// `corner_generic` 的 **A0 循环本来就是外层**，而 (A0,B0) 瓦片只读 `in` 的第 `A0..A0+63` 行
// ⇒ 把这 64 行的 `row_dif` 挪到该块的 B0 扫描之前即可（S 的行长 R=4096 ⇒ 块 = 64*16 KB = **1 MB**，
// 与 `⑤` 用的块同量级）。**只对 corner_fwd 开**（`corner_inv` 的 in=X，其行在 `pointwise_dit` 里已融合，
// 且 X 块是 64*C*4 = 8 MB，超出 L2/L3 ⇒ 无收益，故 `fuse_row=0`）。
// 正确性：row_dif 的行集与该块瓦片读取的行集**完全相同**，块间互不影响 ⇒ 逐位等价
// （本地全尺寸 chk 仍 = 7586622636533660264）。
// ★ #102679 实测 6917.661305 ms（vs #102671 的 6916.921016）⇒ `⑥` **≈ 0（+0.7 ms，噪声）**：
// 1 MB 的 S 块本来就已经在 L3 里被消费 ⇒ 该对往返**本来就不是 DRAM 级**。**记档为"零收益"。**
// 局面：mine = 6916.921016（**全榜第一**）、T = 6962.214720、thr = **6892.593573**
// ⇒ 板面上我们已比 T 快 45.3 ms，**只剩 24.33 ms = 0.352%**（全是 0.99 系数，§2.16.1 的标准局面）。
//
// ⑦ [新刀] **phase-3 combine 的输出改用非临时写（`vmovntdq`）** —— 这是本代唯一"贴到墙"的相位：
// combine 实测 173 ms 搬 ~4.1 GB ⇒ **24 GB/s，已在带宽墙上**（对照：`pointwise` 纯流式 16 GB/s）。
// 它把 800 MB 输出写成 **32 B（半行）**store ⇒ 每条 line 都要付一次 RFO 取行，
// 而这次调用**再也不会读回这些 line** ⇒ RFO 是纯浪费（≈775 MB 的无用读）。
// * 实现：`stnt256()`（内联汇编 `vmovntdq`）+ **每段一次**的对齐闸门：
// `vmovntdq` 要求 32 B 对齐，而区段起点 `M+1 = 65782273`、`N+1 = 100000001` **不是 8 的倍数**
// ⇒ 只有 `k % 8 == 0` 的段才走 NT（**首版没做这个闸门 ⇒ 本地直接 SIGSEGV，已修**）。
// region 0（起点 0/rbase）与 region 3（起点 N1_=2^27）恒对齐 ⇒ 覆盖约 775 MB 输出的绝大部分。
// * 次序：函数返回前加了 `_mm_sfence()`（NT 写是弱序的，判题器随后要读 c）。
// * 逐位等价（本地全尺寸 chk 仍 = 7586622636533660264）。
// ★ 预期：775 MB / 4.1 GB × 173 ms ≈ **20~33 ms** ⇒ 若 ≥ 24.33 ms 即 **达标**。
// 本发为**试验性提交**：输出**故意正确**（不是探针件）。
// ================
// 1002e8 poly_multiply (cyclic NTT mod 2013265921 + exact wrap correction)
#define LGR_ 12
#define LGC_ 15
#define LGH_ 26
#define LGR2_ 13
#define LGC2_ 13
#define EXACT 1
// [PAD] internal row stride of the two 512 MB spectral arrays. C*4 = 65536 B is a
// power of two, so with the natural stride every X row starts at the same
// (addr mod 4096) and the same (addr mod 65536): the 8 strIn rows and the 8 strOut
// rows corner_tile_avx keeps live all land in ONE 8-way L1 set and ONE 4-way L2
// set. Adding XPAD u32 to the *internal* stride puts row j at
// (j*4*XPAD) mod 4096 / mod 65536. Only the address mapping changes -- the
// logical index space of A is untouched, so the result is bit-exact.
#define XPAD 768
#define NRW_ (1 << LGR_)
#define NCW_ (1 << LGC_)
#define XST_ (NCW_ + XPAD)
// ============================================================================
// 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>
#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;
// ------------------- 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")
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));
}
static inline __attribute__((always_inline)) void stnt256(void *p, __m256i v) {
__asm__ volatile("vmovntdq %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));
}
#pragma GCC pop_options
#pragma GCC optimize("O3","rename-registers","live-range-shrinkage","ira-loop-pressure","no-stack-protector")
#define TS 64
static const u32 P = 2013265921u;
static const u32 G = 31u;
static u64 IM;
static u32 TBUF[TS * TS];
static u32 TW[TS * TS];
static inline u32 mm(u32 a, u32 b) {
u64 z = (u64)a * b;
u64 q = (u64)(((u128)z * IM) >> 64);
u64 r = z - q * (u64)P;
return (u32)(r >= P ? r - P : r);
}
static inline u32 shoup_c(u32 w) { return (u32)(((u64)w << 32) / (u64)P); }
static u32 pw(u32 a, u32 e) {
u32 r = 1;
while (e) { if (e & 1) r = (u32)((u64)r * a % P); a = (u32)((u64)a * a % P); e >>= 1; }
return r;
}
// ---- staged twiddle table for a DIT row transform (bitrev in, natural out) ----
static void build_tab(u32 *tw, int L, u32 w) {
int off = 0;
for (int len = 2; len <= L; len <<= 1) {
u32 wl = pw(w, (u32)(L / len));
u32 c = 1;
int h = len >> 1;
for (int j = 0; j < h; j++) { tw[off + j] = c; c = mm(c, wl); }
off += h;
}
}
static void build_br(u16 *br, int L) {
int bits = 0; while ((1 << bits) < L) bits++;
for (int i = 0; i < L; i++) {
int r = 0;
for (int b = 0; b < bits; b++) if (i & (1 << b)) r |= 1 << (bits - 1 - b);
br[i] = (u16)r;
}
}
// ---------------- AVX2 kernels (Shoup modmul) ----------------
#ifdef NO_AVX2
#define HAVE_AVX2 0
#else
#define HAVE_AVX2 1
#pragma GCC push_options
#pragma GCC target("avx2")
// 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 __m256i shoup8(__m256i v, __m256i w, __m256i ws) {
const __m256i Pv = _mm256_set1_epi32((int)P);
__m256i t = _mm256_mullo_epi32(v, w);
__m256i q = vmulhi8(v, ws);
__m256i r = _mm256_sub_epi32(t, _mm256_mullo_epi32(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<bool OMIT_FINAL=false> static void row_dit_avx(u32 *x, int L, const u32 *tw, const u32 *tws) {
__m256i Pv = _mm256_set1_epi32((int)P), M1 = _mm256_set1_epi32(-1), Z = _mm256_setzero_si256();
int off = 0;
if (L >= 8) {
const __m256i V4 = _mm256_setr_epi32((int)tw[1], (int)tw[2], 1, 1, (int)tw[1], (int)tw[2], 1, 1);
const __m256i V4s = _mm256_setr_epi32((int)tws[1], (int)tws[2], 2, 2, (int)tws[1], (int)tws[2], 2, 2);
const __m256i V8 = _mm256_setr_epi32((int)tw[3], (int)tw[4], (int)tw[5], (int)tw[6], 1, 1, 1, 1);
const __m256i V8s = _mm256_setr_epi32((int)tws[3], (int)tws[4], (int)tws[5], (int)tws[6], 2, 2, 2, 2);
const __m256i IDX8 = _mm256_setr_epi32(4, 5, 6, 7, 0, 1, 2, 3);
const __m256i IDXQ4 = _mm256_setr_epi32(0, 1, 0, 1, 4, 5, 4, 5);
const __m256i IDXQ8 = _mm256_setr_epi32(0, 1, 2, 3, 0, 1, 2, 3);
for (int i = 0; i < L; i += 8) {
__m256i v = ldu256((const void *)((x + i)));
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(v, 0xB1); // [1,0,3,2,5,4,7,6]
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_sub_epi32(y, v); // d[2j+1] = v[2j]-v[2j+1]
d = _mm256_add_epi32(d, _mm256_and_si256(_mm256_cmpgt_epi32(Z, d), Pv));
v = _mm256_blend_epi32(sa, d, 0xAA);
}
{ // len=4: pairs (0,2),(1,3) and (4,6),(5,7), twiddles [1,w4]
__m256i y = _mm256_shuffle_epi32(v, 0x4E); // [2,3,0,1,6,7,4,5]
y = shoup8(y, V4, V4s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(v, y), _mm256_add_epi32(_mm256_sub_epi32(v, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ4); // [_,_,d0,d1,_,_,d4,d5]
v = _mm256_blend_epi32(sa, d, 0xCC);
}
{ // len=8: pairs (0,4),(1,5),(2,6),(3,7), twiddles [1,w8,w8^2,w8^3]
__m256i y = _mm256_permutevar8x32_epi32(v, IDX8); // [4,5,6,7,0,1,2,3]
y = shoup8(y, V8, V8s);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(v, y), _mm256_add_epi32(_mm256_sub_epi32(v, y), Pv));
d = _mm256_permutevar8x32_epi32(d, IDXQ8); // [d0,d1,d2,d3,d0,d1,d2,d3]
v = _mm256_blend_epi32(sa, d, 0xF0);
}
stu256((void *)((x + i)), v);
}
off = 7;
// Fused radix-4 DIT: stages (len, 2len) in one pass over four h=len/2
// quarters -- half the loads/stores of two separate radix-2 passes, same
// modmul count and the same twiddle-table segments as the radix-2 form.
int len = 16;
for (; (len << 1) <= L; len <<= 2) {
int h = len >> 1; // quarter size
const u32 *t1 = tw + off, *ts1 = tws + off; // w_len^j
const u32 *t2 = t1 + h, *ts2 = ts1 + h; // w_{2len}^j
const u32 *t3 = t2 + h, *ts3 = ts2 + h; // w_{2len}^{h+j}
off += 3 * h;
for (int i = 0; i < L; i += (len << 1)) {
u32 *p = x + i;
for (int j = 0; j < h; j += 8) {
__m256i 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 w1 = ldu256((const void *)(t1 + j)), ws1 = ldu256((const void *)(ts1 + j));
__m256i Bm = shoup8(B, w1, ws1), Dm = shoup8(D, w1, ws1);
__m256i Ap = vaddm256(A, Bm, Pv), Bp = vsubm256(A, Bm, Pv);
__m256i Cp = vaddm256(C, Dm, Pv), Dp = vsubm256(C, Dm, Pv);
__m256i w2 = ldu256((const void *)(t2 + j)), ws2 = ldu256((const void *)(ts2 + j));
__m256i w3 = ldu256((const void *)(t3 + j)), ws3 = ldu256((const void *)(ts3 + j));
__m256i Cm = shoup8(Cp, w2, ws2), Dn = shoup8(Dp, w3, 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 && !OMIT_FINAL) { // odd stage count
int h = len >> 1;
const u32 *t = tw + off, *ts = tws + off; off += h;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j += 8) {
__m256i u = ldu256((const void *)(p + j));
__m256i v = shoup8(ldu256((const void *)(q + j)),
ldu256((const void *)(t + j)), ldu256((const void *)(ts + j)));
stu256((void *)(p + j), vaddm256(u, v, Pv));
stu256((void *)(q + j), vsubm256(u, v, Pv));
}
}
}
} else {
for (int len = 2; len <= L; len <<= 1) {
int h = len >> 1;
const u32 *t = tw + off; off += h;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j++) {
u32 u = p[j], v = mm(q[j], t[j]);
u32 sa = u + v; if (sa >= P) sa -= P;
u32 d = u + P - v; if (d >= P) d -= P;
p[j] = sa; q[j] = d;
}
}
}
}
}
// ---------------------------------------------------------------------------
// DIF row transform: natural input, output at position m holds the value of
// index rev(m) (length preserved). It is the exact mirror of bitrev+row_dit,
// so a DIF forward pairs with a bitrev-free DIT inverse -- no bitrev passes.
// Stage order is len = L .. 16; the last three stages (8,4,2) act inside each
// aligned 8-block and are done in registers, exactly as the DIT codelet does
// for the first three.
static void row_dif_avx(u32 *x, int L, const u32 *tw, const u32 *tws) {
__m256i Pv = _mm256_set1_epi32((int)P), M1 = _mm256_set1_epi32(-1), Z = _mm256_setzero_si256();
if (L >= 8) {
// Fused radix-4: the stages len and len/2 in one pass over four h2-blocks.
// 4 loads / 4 stores / 4 shoup8 per 32 points per 2 stages (half the memory
// ops and ~14 % fewer uops than two separate radix-2 passes).
int len = L;
for (; len >= 32; len >>= 2) {
int h = len >> 1, h2 = len >> 2;
const u32 *t1 = tw + h - 1, *ts1 = tws + h - 1; // w_len^j (j < h2)
const u32 *t2 = t1 + h2, *ts2 = ts1 + h2; // w_len^{j+h2}
const u32 *t3 = tw + h2 - 1, *ts3 = tws + h2 - 1; // w_{len/2}^j = w_len^{2j}
for (int i = 0; i < L; i += len) {
u32 *p = x + i;
for (int j = 0; j < h2; j += 8) {
__m256i 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, ldu256((const void *)((t1 + j))), ldu256((const void *)((ts1 + j))));
d13 = shoup8(d13, ldu256((const void *)((t2 + j))), ldu256((const void *)((ts2 + j))));
__m256i w3 = ldu256((const void *)((t3 + j)));
__m256i ws3 = ldu256((const void *)((ts3 + j)));
__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, w3, 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, w3, ws3);
stu256((void *)((p + j)), ra);
stu256((void *)((p + j + h2)), rb);
stu256((void *)((p + j + h)), rc);
stu256((void *)((p + j + h + h2)), rd);
}
}
}
if (len >= 16) { // odd stage count: one radix-2 DIF
int h = len >> 1;
const u32 *t = tw + h - 1, *ts = tws + h - 1;
for (int i = 0; i < L; i += len) {
u32 *p = x + i, *q = x + i + h;
for (int j = 0; j < h; j += 8) {
__m256i u = ldu256((const void *)((p + j)));
__m256i v = ldu256((const void *)((q + j)));
__m256i w = ldu256((const void *)((t + j)));
__m256i ws = ldu256((const void *)((ts + j)));
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(u, v), _mm256_sub_epi32(_mm256_add_epi32(u, v), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(u, v), _mm256_add_epi32(_mm256_sub_epi32(u, v), Pv));
d = shoup8(d, w, ws);
stu256((void *)((p + j)), sa);
stu256((void *)((q + j)), d);
}
}
}
const __m256i V8 = _mm256_setr_epi32(1, 1, 1, 1, (int)tw[3], (int)tw[4], (int)tw[5], (int)tw[6]);
const __m256i V8s = _mm256_setr_epi32(2, 2, 2, 2, (int)tws[3], (int)tws[4], (int)tws[5], (int)tws[6]);
const __m256i V4 = _mm256_setr_epi32(1, 1, (int)tw[1], (int)tw[2], 1, 1, (int)tw[1], (int)tw[2]);
const __m256i V4s = _mm256_setr_epi32(2, 2, (int)tws[1], (int)tws[2], 2, 2, (int)tws[1], (int)tws[2]);
const __m256i IDX8 = _mm256_setr_epi32(4, 5, 6, 7, 0, 1, 2, 3);
for (int i = 0; i < L; i += 8) {
__m256i v = ldu256((const void *)((x + i)));
{ // len=8: pairs (j, j+4); the difference gets w8^0..w8^3
__m256i y = _mm256_permutevar8x32_epi32(v, IDX8);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_sub_epi32(y, v); // lanes 4..7 = x_j - x_{j+4}
d = _mm256_add_epi32(d, _mm256_and_si256(_mm256_cmpgt_epi32(Z, d), Pv));
d = shoup8(d, V8, V8s);
v = _mm256_blend_epi32(sa, d, 0xF0);
}
{ // len=4: pairs (j, j+2) per half; differences get [1, w4]
__m256i y = _mm256_shuffle_epi32(v, 0x4E);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_sub_epi32(y, v); // lanes 2,3,6,7 = diffs
d = _mm256_min_epu32(d, _mm256_add_epi32(d, Pv));
d = shoup8(d, V4, V4s);
v = _mm256_blend_epi32(sa, d, 0xCC);
}
{ // len=2: adjacent pairs, twiddle 1
__m256i y = _mm256_shuffle_epi32(v, 0xB1);
__m256i sa = _mm256_min_epu32(_mm256_add_epi32(v, y), _mm256_sub_epi32(_mm256_add_epi32(v, y), Pv));
__m256i d = _mm256_min_epu32(_mm256_sub_epi32(y, v), _mm256_add_epi32(_mm256_sub_epi32(y, v), Pv));
v = _mm256_blend_epi32(sa, d, 0xAA);
}
stu256((void *)((x + i)), 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 inline u32 mont1(u32 a, u32 b);
static inline __m256i mont8(__m256i a, __m256i b);
static void pointwise_mul_avx(u32 *A, const u32 *B, u32 n, u32 sc) {
// mont8 produces A*B/R. The inverse twiddle table carries sc*R, so the
// final inverse output is unchanged while avoiding Shoup reciprocal work.
(void)sc;
u32 i = 0;
for (; i + 8 <= n; i += 8) {
__m256i a = ldu256((const void *)(A + i));
__m256i b = ldu256((const void *)(B + i));
stu256((void *)(A + i), mont8(a, b));
}
for (; i < n; i++) A[i] = mont1(A[i], B[i]);
}
// ---- Montgomery arithmetic (corner only): mont_mul(a,b) = a*b*R^-1 mod P, R = 2^32.
// The corner's twiddle tables are kept in MONTGOMERY FORM (TW_M = TW*R) and the data
// stays in normal form, so mont8(x, TW_M) = x*TW -- i.e. the corner gets its modmul with
// NO per-element Shoup reciprocal (shoup_of costs 9 uops + 1 load per 8 modmuls, and
// mont8 is 16 uops against shoup8's 20).
static u32 NPI = 0; // -P^-1 mod 2^32
static const u32 R32 = (u32)((1ull << 32) % P);
static inline u32 mont1(u32 a, u32 b) {
u64 z = (u64)a * (u64)b;
u32 m = (u32)((u64)(u32)z * (u64)NPI);
u64 t = (z + (u64)m * (u64)P) >> 32;
return (u32)(t >= P ? t - P : t);
}
static inline __m256i mont8(__m256i a, __m256i b) {
__m256i P64 = _mm256_set1_epi64x((long long)P), NP = _mm256_set1_epi64x((long long)NPI);
__m256i Pv = _mm256_set1_epi32((int)P);
__m256i ahi = _mm256_srli_epi64(a, 32), bhi = _mm256_srli_epi64(b, 32);
__m256i tlo = _mm256_mul_epu32(a, b), thi = _mm256_mul_epu32(ahi, bhi);
__m256i e2 = _mm256_mul_epu32(tlo, NP), o2 = _mm256_mul_epu32(thi, NP);
__m256i e4 = _mm256_mul_epu32(e2, P64), o4 = _mm256_mul_epu32(o2, P64);
__m256i e5 = _mm256_add_epi64(tlo, e4), o5 = _mm256_add_epi64(thi, o4);
__m256i ue = _mm256_srli_epi64(e5, 32), uo = _mm256_srli_epi64(o5, 32);
__m256i u = _mm256_blend_epi32(ue, _mm256_slli_epi64(uo, 32), 0xAA);
return _mm256_min_epu32(u, _mm256_sub_epi32(u, Pv));
}
// dst[v] = mont(src[v], kM) over one TS-row (both operands Montgomery-form)
static void scale_row_mont(u32 *dst, const u32 *src, u32 kM) {
const __m256i kv = _mm256_set1_epi32((int)kM); // full 8-lane broadcast (NOT set1_epi64x:
// that would zero the odd u32 lanes)
for (int v = 0; v < TS; v += 8)
stu256((void *)((dst + v)), mont8(ldu256((const void *)((src + v))), kv));
}
// TW[u] = mont(TW[u-1], RVM)
static void twiddle_rows_mont(u32 *TW, const u32 *RVM) {
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v += 8)
stu256((void *)((cur + v)), mont8(ldu256((const void *)((prv + v))), ldu256((const void *)((RVM + v)))));
}
}
// dst[v] = src[v]*k mod P over one TS-row
static void scale_row_avx(u32 *dst, const u32 *src, u32 k) {
const __m256i kv = _mm256_set1_epi32((int)k), ks = _mm256_set1_epi32((int)shoup_c(k));
for (int v = 0; v < TS; v += 8)
stu256((void *)((dst + v)), shoup8(ldu256((const void *)((src + v))), kv, ks));
}
// row recurrence TW[u][v] = TW[u-1][v]*RV[v] mod P, 512 modmuls per call -> 64 shoup8
static void twiddle_rows_avx(u32 *TW, const u32 *RV) {
__m256i RVS[TS / 8];
for (int v = 0; v < TS; v += 8) RVS[v >> 3] = shoup_of(ldu256((const void *)((RV + v))));
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v += 8) {
__m256i r = ldu256((const void *)((RV + v)));
stu256((void *)((cur + v)), shoup8(ldu256((const void *)((prv + v))), r, RVS[v >> 3]));
}
}
}
// corner tile data path: whole 64x64 tile through registers (no TBUF round trip)
static void corner_tile_avx(const u32 *in, u32 *out, int A0, int B0, int strIn, int strOut) {
// a13e: 64-wide staging block (mirrors the 64-row form that paid on the input
// and inverse transposes): merge 64 u0 columns per destination row, then write
// 256 B = 4 complete cache lines per row with NT stores.
for (int u0 = 0; u0 < TS; u0 += 64) {
u32 cb[TS * 64] __attribute__((aligned(64)));
for (int u = 0; u < 64; u += 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 + u + k) * strIn + B0 + v0;
const u32 *tw = TW + (u0 + u + 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 j = 0; j < 8; j++)
stu256((void *)(cb + (u64)(v0 + j) * 64 + u), r[j]);
}
}
for (int v = 0; v < TS; v++) {
u32 *d = out + (u64)(B0 + v) * strOut + A0 + u0;
for (int q = 0; q < 64; q += 8)
stnt256((void *)(d + q), ldu256((const void *)(cb + (u64)v * 64 + q)));
}
}
}
// 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) {
// 旧形态：每个 8x8 子瓦片把转置结果用 8 条 32 B `vmovdqu` 散射进 64 条 S 行（行距 R=4096 u32 = 16 KB）。
// ★ 32 B 只是半条 cache line ⇒ **每条 store 都要一次 RFO**（整行在下一个 rb0 迭代才被另一半写满）
// ⇒ 瓦片真实流量 = 16 KB 读 + 16 KB RFO + 16 KB 回写 = 48 KB（= 实测 3.35 B/cycle，正好对上停等幅度）。
// 新形态：**一次处理 16 行**，把 8 个列组的转置结果先攒进一个 4 KB 栈缓冲（按"目标 S 行"归并）
// ⇒ 每条 S 行一次写满 64 B（= 整条 line）⇒ 用 `vmovntdq` 写 ⇒ **无 RFO**，且不污染 L2/L3
// （S 的复用距离是整块 512 MB，远大于 L3 ⇒ NT 是净赚）。
for (int rb0 = 0; rb0 < TS; rb0 += 64) {
u32 buf[TS * 64] __attribute__((aligned(64)));
for (int rb = 0; rb < 64; rb += 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 + rb + 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 *)(buf + (u64)(cb0 + j) * 64 + rb), r[j]);
}
}
for (int m = 0; m < TS; m++) {
u32 *d = S + (u64)(c0 + m) * R + r0 + rb0;
for (int q = 0; q < 64; q += 8)
stnt256((void *)(d + q), ldu256((const void *)(buf + (u64)m * 64 + q)));
}
}
}
// untranspose tile (plain corner pattern, no twiddles)
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 += 8) {
__m256i r[8];
for (int k = 0; k < 8; k++)
r[k] = ldu256((const void *)((in + (u64)(A0 + u0 + k) * strIn + B0 + v0)));
t8x8(r);
for (int k = 0; k < 8; k++)
stu256((void *)((out + (u64)(B0 + v0 + k) * strOut + A0 + u0)), r[k]);
}
}
}
static void untr_tile_nt(const u32 *in, u32 *out, int A0, int B0, int strIn, int strOut) {
for (int u0 = 0; u0 < TS; u0 += 64) {
u32 buf[TS * 64] __attribute__((aligned(64)));
for (int u = 0; u < 64; u += 8) {
for (int v = 0; v < TS; v += 8) {
__m256i r[8];
for (int k = 0; k < 8; k++)
r[k] = ldu256((const void *)(in + (u64)(A0 + u0 + u + k) * strIn + B0 + v));
t8x8(r);
for (int j = 0; j < 8; j++)
stu256((void *)(buf + (u64)(v + j) * 64 + u), r[j]);
}
}
for (int v = 0; v < TS; v++) {
u32 *d = out + (u64)(B0 + v) * strOut + A0 + u0;
for (int q = 0; q < 64; q += 8)
stnt256((void *)(d + q), ldu256((const void *)(buf + (u64)v * 64 + q)));
}
}
}
// Phase-1 window inverse: the final DIT stage is computed only for rows
// actually transposed out, immediately before the transpose register shuffle.
template<bool NT>
static void untr_tile_nt_last(const u32 *in, u32 *out, int A0, int B0,
int strIn, int strOut, const u32 *tw, const u32 *tws) {
const __m256i Pv = _mm256_set1_epi32((int)P);
const bool upper = B0 >= 2048;
u32 buf[TS * 64] __attribute__((aligned(64)));
for (int u = 0; u < 64; u += 8) {
for (int v = 0; v < TS; v += 8) {
const int j = (B0 + v) & 2047;
const __m256i w = ldu256((const void *)(tw + j));
const __m256i ws = ldu256((const void *)(tws + j));
__m256i r[8];
for (int k = 0; k < 8; k++) {
const u32 *q = in + (u64)(A0 + u + k) * strIn + j;
__m256i lo = ldu256((const void *)q);
__m256i hi = ldu256((const void *)(q + 2048));
__m256i m = shoup8(hi, w, ws);
r[k] = upper ? vsubm256(lo, m, Pv) : vaddm256(lo, m, Pv);
}
t8x8(r);
for (int z = 0; z < 8; z++)
stu256((void *)(buf + (u64)(v + z) * 64 + u), r[z]);
}
}
for (int v = 0; v < TS; v++) {
u32 *d = out + (u64)(B0 + v) * strOut + A0;
for (int q = 0; q < 64; q += 8) {
__m256i val = ldu256((const void *)(buf + (u64)v * 64 + q));
if constexpr (NT) stnt256((void *)(d + q), val);
else stu256((void *)(d + q), val);
}
}
}
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
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 row_dif(u32 *x, int L, const u32 *tw, const u32 *tws);
static void pointwise_dit(u32 *A, u32 *B, u32 n, u32 sc,
const u32 *tfC, const u32 *wsCf,
const u32 *tiC, const u32 *wsCi) {
// Each product row is immediately consumed by the inverse's first DIT.
(void)n;
for (int k0 = 0; k0 < (int)NRW_; k0++) {
u32 *Aa = A + (u64)k0 * (u64)XST_;
u32 *Bb = B + (u64)k0 * (u64)XST_;
row_dif(Aa, NCW_, tfC, wsCf); // [g10] A's deferred final row_dif (single variable)
row_dif(Bb, NCW_, tfC, wsCf);
#if HAVE_AVX2
pointwise_mul_avx(Aa, Bb, NCW_, sc);
#else
for (u32 i = 0; i < NCW_; i++) Aa[i] = mm(Aa[i], Bb[i]);
#endif
row_dit(Aa, NCW_, tiC, wsCi);
}
}
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;
u32 w;
u32 *tfR; u32 *tiR; u32 *tfC; u32 *tiC;
u32 *wsR; u32 *wsC;
u16 *brR; u16 *brC;
};
static void load_T(Eng &e, const u32 *src, u32 nsrc, u32 *S, int ctr) {
const int R = e.R, C = e.C;
#if HAVE_AVX2
// fast path: whole tiles (and whole all-zero tiles) go through registers
if (nsrc <= (u32)R * (u32)C) {
for (int r0 = 0; r0 < R; r0 += TS) {
u64 tilerow = (u64)r0 * C;
for (int c0 = 0; c0 < C; c0 += TS) {
u64 firstidx = tilerow + c0;
u64 lastidx = (u64)(r0 + TS - 1) * C + c0 + TS - 1;
if (lastidx < nsrc) { load_tile_avx(src, S, r0, c0, R, C, ctr, 0); continue; }
if (firstidx >= nsrc) { load_tile_avx(src, S, r0, c0, R, C, ctr, 1); continue; }
for (int rb = 0; rb < TS; rb++) {
u64 base = (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u64 id = base + cb;
u32 v = 0;
if (id < nsrc) { v = src[id]; if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u)); }
TBUF[cb * TS + rb] = v;
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)(c0 + cb) * R + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb++) d[rb] = s[rb];
}
}
}
return;
}
#endif
for (int r0 = 0; r0 < R; r0 += TS) {
for (int c0 = 0; c0 < C; c0 += TS) {
u64 lastidx = (u64)(r0 + TS - 1) * C + (c0 + TS - 1);
if (lastidx < nsrc) {
for (int rb = 0; rb < TS; rb++) {
const u32 *s = src + (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u32 v = s[cb];
if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u));
TBUF[cb * TS + rb] = v;
}
}
} else {
for (int rb = 0; rb < TS; rb++) {
u64 base = (u64)(r0 + rb) * C + c0;
for (int cb = 0; cb < TS; cb++) {
u64 id = base + cb;
u32 v = 0;
if (id < nsrc) { v = src[id]; if (ctr) v = (v >= 4u) ? (v - 4u) : (v + (P - 4u)); }
TBUF[cb * TS + rb] = v;
}
}
}
for (int cb = 0; cb < TS; cb++) {
u32 *d = S + (u64)(c0 + cb) * R + r0;
const u32 *s = TBUF + cb * TS;
for (int rb = 0; rb < TS; rb++) d[rb] = s[rb];
}
}
}
}
// generic corner turn+transpose: in[A*strIn + B] -> out[B*strOut + A] * w
// (nA/nB are the dimensions of A/B; the tile is TS x TS on both)
// perm 0: w = root^{A*B}
// perm 1: w = root^{A*rev(B)} (B is a DIF output position, lg2(nB) bits)
// perm 2: w = root^{rev(A)*B} (A is a DIF output position, lg2(nA) bits)
// The permutation is folded into the twiddle table, so the load/store pattern
// is byte-for-byte the same in all three modes.
static u32 REV6[TS];
static u32 g_inv_scale = 1u;
static u32 POWT[1024], POWT64[1024];
static void build_rev6(void) {
for (int i = 0; i < TS; i++) {
int r = 0;
for (int b = 0; b < 6; b++) if (i & (1 << b)) r |= 1 << (5 - b);
REV6[i] = (u32)r;
}
}
static inline u32 revbits(u32 x, int lg) {
u32 r = 0;
for (int b = 0; b < lg; b++) if (x & (1u << b)) r |= 1u << (lg - 1 - b);
return r;
}
static void corner_generic(Eng &e, const u32 *in, u32 *out, int nA, int nB, int strIn, int strOut, u32 root, int perm, int fuse_row) {
u32 wv[TS];
{ u32 c = 1; for (int v = 0; v < TS; v++) { wv[v] = c; c = mm(c, root); } }
const u32 w64 = pw(root, TS);
int lgB = 0, lgA = 0;
while ((1 << lgB) < nB) lgB++;
while ((1 << lgA) < nA) lgA++;
const int S1 = nB / TS, S2 = nA / TS; // tile-index counts
const int HB = lgB - 6, HA = lgA - 6; // tile-index bit counts
u32 RV2[TS], RV2b[TS], PW2[TS], cu[TS], cu64[TS], base[TS];
if (perm == 1) {
u32 g = pw(root, (u32)S1); // root^{nB/TS}
{ u32 gM = mm(g, R32), c = R32; for (int t = 0; t < TS; t++) { RV2[REV6[t]] = c; c = mont1(c, gM); } }
{ u32 rootM = mm(root, R32), p = R32; for (int t = 0; t < S1; t++) { POWT[t] = p; p = mont1(p, rootM); } }
} else if (perm == 2) {
u32 g = pw(root, (u32)S2), g64 = pw(g, 64);
const u32 rootM2 = mm(root, R32), gM = mm(g, R32), g64M = mm(g64, R32);
{ u32 c = R32, c64 = R32;
for (int t = 0; t < TS; t++) { RV2[REV6[t]] = c; RV2b[REV6[t]] = c64; c = mont1(c, gM); c64 = mont1(c64, g64M); } }
{ u32 p = R32; for (int t = 0; t < S2; t++) { POWT[t] = p; p = mont1(p, rootM2); } }
{ u32 q = pw(root, 64), qM = mm(q, R32), p = R32; for (int t = 0; t < S2; t++) { POWT64[t] = p; p = mont1(p, qM); } }
}
u32 wr = 1; // root^{A0}
const u32 INVSCALE = g_inv_scale; // 1/N1 folded in as a uniform linear factor
for (int A0 = 0; A0 < nA; A0 += TS) {
// [g12 FUSE] this block's TS rows of `in` (S, 1 MB) are exactly the rows the B0
// sweep below reads; run their deferred row_dif() now so the block is consumed
// from L2/L3 instead of making a full 512 MB round trip.
if (fuse_row) {
u32 *Sw = (u32 *)in + (u64)A0 * (u64)strIn;
for (int u = 0; u < TS; u++) row_dif(Sw + (u64)u * (u64)strIn, nB, e.tfR, e.wsR);
}
u32 wr64 = pw(wr, TS); // root^{TS*A0}
if (perm == 1) {
u32 g = pw(wr, (u32)S1), gM2 = mm(g, R32), c = R32;
for (int t = 0; t < TS; t++) { PW2[REV6[t]] = c; c = mont1(c, gM2); }
} else if (perm == 2) {
int hi = A0 / TS;
u32 rH = POWT[revbits((u32)hi, HA)], r64H = POWT64[revbits((u32)hi, HA)];
const u32 INVSCALE_M = mm(INVSCALE, R32);
for (int u = 0; u < TS; u++) { cu[u] = mont1(RV2[u], rH); cu64[u] = mont1(RV2b[u], r64H); base[u] = INVSCALE_M; }
} else {
u32 c = 1; for (int v = 0; v < TS; v++) { PW2[v] = c; c = mm(c, wr); }
}
u32 Vloc = 1; // root^{B0}
u32 Vr = 1; // root^{A0*B0}
for (int B0 = 0; B0 < nB; B0 += TS) {
if (perm == 1) {
u32 H = revbits((u32)(B0 / TS), HB);
u32 wrH = pw(wr, H), rootH = POWT[H];
u32 RV[TS];
#if HAVE_AVX2
scale_row_mont(RV, RV2, rootH); // root^{rev(B0+v)}
scale_row_mont(TW, PW2, mm(wrH, R32)); // root^{A0*rev(B0+v)}
twiddle_rows_mont(TW, RV);
#else
for (int v = 0; v < TS; v++) RV[v] = mont1(RV2[v], rootH);
for (int v = 0; v < TS; v++) TW[v] = mont1(PW2[v], mm(wrH, R32));
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v++) cur[v] = mont1(prv[v], RV[v]);
}
#endif
} else if (perm == 2) {
if (B0 == 0) {
for (int u = 0; u < TS; u++) {
u32 *cur = TW + u * TS;
cur[0] = base[u];
for (int v = 1; v < TS; v++) cur[v] = mont1(cur[v - 1], cu[u]);
}
} else {
for (int u = 0; u < TS; u++) {
u32 *cur = TW + u * TS;
#if HAVE_AVX2
scale_row_mont(cur, cur, cu64[u]);
#else
u32 d = cu64[u];
for (int v = 0; v < TS; v++) cur[v] = mont1(cur[v], d);
#endif
}
}
for (int u = 0; u < TS; u++) base[u] = mont1(base[u], cu64[u]);
} else {
u32 RV[TS];
for (int v = 0; v < TS; v++) RV[v] = mm(Vloc, wv[v]);
for (int v = 0; v < TS; v++) TW[v] = mm(Vr, PW2[v]);
for (int u = 1; u < TS; u++) {
const u32 *prv = TW + (u - 1) * TS;
u32 *cur = TW + u * TS;
for (int v = 0; v < TS; v++) cur[v] = mm(prv[v], RV[v]);
}
}
// data path: whole tile through registers, u innermost (see NOTES: writing
// the 64 strOut-strided destination rows in v-major order is what keeps them
// from aliasing into one L1 set).
#if HAVE_AVX2
corner_tile_avx(in, out, A0, B0, strIn, strOut);
#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, int fuse_row) {
corner_generic(e, S, X, e.C, e.R, e.R, XST_, root, 1, fuse_row); // A = j0, B = rev'd k1 position
}
static inline void corner_inv(Eng &e, const u32 *X, u32 *S, u32 root) {
corner_generic(e, X, S, e.R, e.C, XST_, e.R, root, 2, 0); // A = rev'd k1 position, B = j0
}
// 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
{ // [g11] deferred from inv(): the 64 rows of this block are transformed here,
// immediately before the tiles that read exactly these rows (A0 = c0..c0+63).
u32 *Sw = (u32 *)S + (u64)c0 * (u64)R;
for (int u = 0; u < TS; u++) row_dit(Sw + (u64)u * (u64)R, R, e.tiR, e.wsR + e.R); }
for (int r0 = 0; r0 < R; r0 += TS) // j1 blocks
untr_tile_nt(S, dst, c0, r0, R, XST_); // [PAD] dst == A, padded stride // A = j0, B = j1
}
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) * (u64)XST_ + c0; // [PAD] dst == A
for (int u = 0; u < TS; u++) d[u] = TBUF[u * TS + v];
}
}
}
}
// ---------------------------------------------------------------------------
// Fused replacement for "untranspose(E1,S,A); copy A[M..2M] into c[N1_..N1_+M];"
// (phase 1 only). The wrap-correction copy needs exactly the natural-order
// elements idx in [M, 2M] and nothing else, i.e. dst[idx] with
// dst = c + (N1_ - M). Writing them straight out of S skips the whole
// A round trip (A is never written nor re-read in phase 1) and skips the
// untranspose of the rows outside [M, 2M] (~51% of them at n=m=1e8).
// Destination index of tile element (j1, j0) is j1*C + j0, exactly the same
// stride-C layout untranspose uses, so the tile routine is reused verbatim;
// only tiles that straddle lo/hi fall back to an element-wise loop.
static void untr_window(Eng &e, const u32 *S, u32 *dst, long lo, long hi) {
const int R = e.R, C = e.C;
for (int c0 = 0; c0 < C; c0 += TS) {
{ // [g11] same deferral as in untranspose(); rows c0..c0+63 are exactly the rows
// the tiles below read (A0 = c0), and tiles of different c0 are disjoint.
u32 *Sw = (u32 *)S + (u64)c0 * (u64)R;
#if HAVE_AVX2
for (int u = 0; u < TS; u++) row_dit_avx<true>(Sw + (u64)u * (u64)R, R, e.tiR, e.wsR + e.R);
#else
for (int u = 0; u < TS; u++) row_dit(Sw + (u64)u * (u64)R, R, e.tiR, e.wsR + e.R);
#endif
}
for (int r0 = 0; r0 < R; r0 += TS) {
long tlo = (long)r0 * C + c0;
long thi = (long)(r0 + TS - 1) * C + c0 + TS - 1;
if (thi < lo || tlo > hi) continue;
if (tlo >= lo && thi <= hi) {
#if HAVE_AVX2
if ((((uintptr_t)dst) & 31u) == 0u)
untr_tile_nt_last<true>(S, dst, c0, r0, R, C, e.tiR + 2047, e.wsR + e.R + 2047);
else
untr_tile_nt_last<false>(S, dst, c0, r0, R, C, e.tiR + 2047, e.wsR + e.R + 2047);
#else
if ((((uintptr_t)dst) & 31u) == 0u) untr_tile_nt(S, dst, c0, r0, R, C);
else untr_tile_avx(S, dst, c0, r0, R, C);
#endif
continue;
}
for (int u = 0; u < TS; u++) {
long b = (long)r0 * C + c0 + u;
for (int v = 0; v < TS; v++) {
long idx = b + (long)v * C;
if (idx >= lo && idx <= hi) {
#if HAVE_AVX2
const int rr = r0 + v, j = rr & 2047;
const u32 *row = S + (u64)(c0 + u) * R;
u32 mul = mm(row[j + 2048], e.tiR[2047 + j]);
u32 x = row[j], val;
if (rr < 2048) { val = x + mul; if (val >= P) val -= P; }
else { val = x + P - mul; if (val >= P) val -= P; }
dst[idx] = val;
#else
dst[idx] = S[(u64)(c0 + u) * R + r0 + v];
#endif
}
}
}
}
}
_mm_sfence();
}
// ---------------------------------------------------------------------------
static void fwd(Eng &e, const u32 *src, u32 nsrc, u32 *S, u32 *X, int ctr,
bool defer_final_row = false) {
load_T(e, src, nsrc, S, ctr);
// [g12 FUSE] the row_dif(S) pass (C calls of length R) is no longer a standalone
// sweep: corner_fwd() runs it per 64-row A0 block, right before the tiles that
// read exactly those rows. Same calls, same order within each row -> bit-exact.
corner_fwd(e, S, X, e.w, 1);
if (!defer_final_row) for (int k0 = 0; k0 < e.R; k0++) {
u32 *row = X + (u64)k0 * (u64)XST_;
row_dif(row, e.C, e.tfC, e.wsC);
}
}
// 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) {
corner_inv(e, X, S, pw(e.w, P - 2u));
// [g11 FUSE] the trailing row_dit(S) pass now runs inside the consumer
// (untranspose / untr_window), one 64-row block at a time, so the block is still
// in L2/L3 when the tiles read it instead of making a full 512 MB round trip.
}
// ---------------------------------------------------------------------------
static void eng_init(Eng &e, int lgR, int lgC, u32 w,
u32 *tfR, u32 *tiR, u32 *tfC, u32 *tiC, u16 *brR, u16 *brC,
u32 *wsR, u32 *wsC) {
e.R = 1 << lgR; e.C = 1 << lgC; e.N = e.R * e.C; e.w = w;
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));
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); build_br(brC, e.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).
#define NR_ (1u << LGR_)
#define NC_ (1u << LGC_)
#define N1_ (1u << (LGR_ + LGC_))
#define NR2_ (1u << LGR2_)
#define NC2_ (1u << LGC2_)
#define N2_ (1u << (LGR2_ + LGC2_))
static u32 Aarr[NRW_ * XST_] __attribute__((aligned(64))), Barr[NRW_ * XST_ + 1352] __attribute__((aligned(64))); // [PAD] N1_ + NRW_*XPAD
// [e2g] 896 u32 (3584 B) of unused slack on Barr: shifts Aarr's base so that
// (Aarr-Barr) % 4096 != 0. Without it the two panels are exactly 4K-congruent
// (delta = 0x122dc000, both sizes multiples of 4096) and every store to A in
// pointwise_mul_avx falsely hits the loads from B at the same 12-bit offset.
// No index ever reaches the slack (all indices < NRW_*XST_).
static Eng E1, E2;
static u32 t1fR[NR_], t1iR[NR_], t1fC[NC_], t1iC[NC_];
static u16 b1R[NR_], b1C[NC_];
static u32 w1R[2 * NR_], w1C[2 * NC_];
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;
#pragma GCC push_options
#pragma GCC target("avx2")
static inline __attribute__((always_inline)) __m256i scan8add(__m256i v) {
__m256i t = _mm256_slli_si256(v, 4);
v = _mm256_add_epi32(v, t);
t = _mm256_slli_si256(v, 8);
v = _mm256_add_epi32(v, t);
t = _mm256_permute2x128_si256(v, v, 0x08);
t = _mm256_shuffle_epi32(t, 0xFF);
return _mm256_add_epi32(v, t);
}
#pragma GCC pop_options
#pragma GCC push_options
#pragma GCC target("avx2")
void poly_multiply(unsigned *a, int n, unsigned *b, int m, unsigned *c)
{
// [g13] phase 3 (combine) is the one bandwidth-bound pass (~4.1 GB in ~173 ms
// = 24 GB/s, at the wall). Its output writes 800 MB as 32-byte half-line
// stores, so every line pays a read-for-ownership fetch this call never uses.
const int c32 = (int)((((uintptr_t)(const void *)c) & 31u) == 0u);
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_, LGC_, w1, t1fR, t1iR, t1fC, t1iC, b1R, b1C, w1R, w1C);
eng_init(E2, LGR2_, LGC2_, w2, t2fR, t2iR, t2fC, t2iC, b2R, b2C, w2R, w2C);
}
u32 *S = (u32 *)c; // transient scratch: c[0 .. N1_)
u32 *A = Aarr, *B = Barr;
const int ctr = EXACT;
const u32 sc = pw((u32)(N1_ % P), P - 2u); // 1/N1
#if HAVE_AVX2
g_inv_scale = mm(sc, R32); // cancel pointwise Montgomery's R^-1
#else
g_inv_scale = sc;
#endif
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) {
fwd(E1, (const u32 *)a + Sa, (u32)(M + 1), S, A, ctr, true);
fwd(E1, (const u32 *)b + Sb, (u32)(M + 1), S, B, ctr, true);
pointwise_dit(A, B, N1_, sc, E1.tfC, E1.wsC, E1.tiC, E1.wsC + E1.C);
inv(E1, A, S);
untr_window(E1, S, c + (N1_ - M), M, 2 * M);
}
// ---------- phase 2: cyclic convolution ----------
fwd(E1, (const u32 *)a, (u32)n + 1, S, A, ctr, true);
fwd(E1, (const u32 *)b, (u32)m + 1, S, B, ctr, true);
pointwise_dit(A, B, N1_, sc, E1.tfC, E1.wsC, E1.tiC, E1.wsC + E1.C);
inv(E1, A, S);
untranspose(E1, S, A);
// ---------- phase 3: combine ----------
long nmax = (long)n + (long)m;
i64 WA = 0, WB = 0;
// ---- vectorised fast path (n==m, wrap region present, usual layout) ----
if (n == m && M >= 0 && (long)M < (long)n && (long)n < (long)N1_ && (long)N1_ < nmax) {
const long N = (long)n;
const __m256i Pvm = _mm256_set1_epi32((int)P);
const __m256i Phm = _mm256_set1_epi32((int)(P >> 1));
const __m256i Zv = _mm256_setzero_si256();
__m256i WAv = Zv, WBv = Zv;
const long rb[5] = {0, M + 1, N + 1, (long)N1_, nmax + 1};
// [PAD] A now has the padded internal row stride XST_, so the natural-order
// element k = j1*NC_ + j0 sits at A[j1*XST_ + j0]. Regions 0..2 all live in
// [0, N1_), i.e. inside the padded rows, so they are walked row by row;
// region 3 ([N1_, nmax]) takes its `tab` straight from c and never touches A,
// so it stays one flat pass (as before). Rows in increasing order, and inside
// a row the three regions in increasing order, reproduces exactly the old
// increasing-k sequence -- so the (WA,WB) running sums are identical.
for (int j1 = 0; j1 < (int)NRW_; j1++) {
const long rbase = (long)j1 * (long)NCW_;
const long rend = rbase + (long)NCW_;
u32 *const Arow = A + (u64)j1 * (u64)XST_;
for (int rg = 0; rg < 3; rg++) {
long s = rb[rg]; if (s < rbase) s = rbase;
long k1 = rb[rg + 1]; if (k1 > rend) k1 = rend;
if (s >= k1) continue;
// resync the vector running sums with the scalar ones (the scalar tail of the
// previous segment advanced WA/WB but not the broadcast vectors)
WAv = _mm256_set1_epi32((int)WA);
WBv = _mm256_set1_epi32((int)WB);
// margin: never read past a[0..N] / c[0..nmax] with an 8-wide vector
long vend = k1 - (rg == 0 ? 16 : 8);
if (vend < s) vend = s;
long k = s;
if (vend > k) {
const int incPos = (rg < 2);
const int ntA = c32 & (((k & 7) == 0) ? 1 : 0); // [g13] region starts M+1/N+1 unaligned
for (; k + 8 <= vend; k += 8) {
__m256i dav, dbv;
if (incPos) { dav = ldu256(a + k); dbv = ldu256(b + k); }
else { dav = ldu256(a + k - N - 1); dbv = ldu256(b + k - N - 1);
dav = _mm256_sub_epi32(Zv, dav); dbv = _mm256_sub_epi32(Zv, dbv); }
__m256i wav = _mm256_add_epi32(WAv, scan8add(dav));
__m256i wbv = _mm256_add_epi32(WBv, scan8add(dbv));
WAv = _mm256_permutevar8x32_epi32(wav, _mm256_set1_epi32(7));
WBv = _mm256_permutevar8x32_epi32(wbv, _mm256_set1_epi32(7));
__m256i tv;
if (incPos) {
tv = _mm256_setr_epi32((int)(k + 1), (int)(k + 2), (int)(k + 3), (int)(k + 4),
(int)(k + 5), (int)(k + 6), (int)(k + 7), (int)(k + 8));
} else {
tv = _mm256_setr_epi32((int)(2 * N - k + 1), (int)(2 * N - k), (int)(2 * N - k - 1), (int)(2 * N - k - 2),
(int)(2 * N - k - 3), (int)(2 * N - k - 4), (int)(2 * N - k - 5), (int)(2 * N - k - 6));
}
__m256i tab;
if (rg == 0) {
tab = _mm256_sub_epi32(ldu256(Arow + (k - rbase)), ldu256(c + N1_ + k));
tab = _mm256_min_epu32(tab, _mm256_add_epi32(tab, Pvm));
} else {
tab = ldu256(Arow + (k - rbase));
}
__m256i gt = _mm256_cmpgt_epi32(tab, Phm);
__m256i ts = _mm256_sub_epi32(tab, _mm256_and_si256(gt, Pvm));
__m256i s4 = _mm256_slli_epi32(_mm256_add_epi32(wav, wbv), 2);
__m256i v16 = _mm256_slli_epi32(tv, 4);
if (ntA) stnt256(c + k, _mm256_sub_epi32(_mm256_add_epi32(ts, s4), v16));
else stu256(c + k, _mm256_sub_epi32(_mm256_add_epi32(ts, s4), v16));
WA = _mm256_extract_epi32(wav, 7);
WB = _mm256_extract_epi32(wbv, 7);
}
}
for (; k < k1; k++) {
if (rg < 2) { WA += a[k]; WB += b[k]; } else { WA -= a[k - N - 1]; WB -= b[k - N - 1]; }
long t = (k <= N) ? (k + 1) : (2 * N - k + 1);
u32 tab;
if (rg == 0) { u32 x = Arow[k - rbase], y = c[N1_ + k]; tab = (x >= y) ? (x - y) : (x + (P - y)); }
else tab = Arow[k - rbase];
i32 ts = (i32)tab;
if (ts > (i32)(P >> 1)) ts -= (i32)P;
c[k] = (u32)((i64)ts + 4 * (WA + WB) - 16 * t);
}
}
}
// region 3: k in [N1_, nmax]; tab == c[k] (E_{k-N1_} read in place), no A access
{
WAv = _mm256_set1_epi32((int)WA);
WBv = _mm256_set1_epi32((int)WB);
long k1 = nmax + 1;
long vend = k1 - 16;
long k = (long)N1_;
if (vend < k) vend = k;
const int ntB = c32; // [g13] k starts at 2^27
for (; k + 8 <= vend; k += 8) {
__m256i dav = ldu256(a + k - N - 1), dbv = ldu256(b + k - N - 1);
dav = _mm256_sub_epi32(Zv, dav); dbv = _mm256_sub_epi32(Zv, dbv);
__m256i wav = _mm256_add_epi32(WAv, scan8add(dav));
__m256i wbv = _mm256_add_epi32(WBv, scan8add(dbv));
WAv = _mm256_permutevar8x32_epi32(wav, _mm256_set1_epi32(7));
WBv = _mm256_permutevar8x32_epi32(wbv, _mm256_set1_epi32(7));
__m256i tv = _mm256_setr_epi32((int)(2 * N - k + 1), (int)(2 * N - k), (int)(2 * N - k - 1), (int)(2 * N - k - 2),
(int)(2 * N - k - 3), (int)(2 * N - k - 4), (int)(2 * N - k - 5), (int)(2 * N - k - 6));
__m256i tab = ldu256(c + k);
__m256i gt = _mm256_cmpgt_epi32(tab, Phm);
__m256i ts = _mm256_sub_epi32(tab, _mm256_and_si256(gt, Pvm));
__m256i s4 = _mm256_slli_epi32(_mm256_add_epi32(wav, wbv), 2);
__m256i v16 = _mm256_slli_epi32(tv, 4);
if (ntB) stnt256(c + k, _mm256_sub_epi32(_mm256_add_epi32(ts, s4), v16));
else stu256(c + k, _mm256_sub_epi32(_mm256_add_epi32(ts, s4), v16));
WA = _mm256_extract_epi32(wav, 7);
WB = _mm256_extract_epi32(wbv, 7);
}
for (; k < k1; k++) {
WA -= a[k - N - 1]; WB -= b[k - N - 1];
long t = 2 * N - k + 1;
u32 tab = c[k];
i32 ts = (i32)tab;
if (ts > (i32)(P >> 1)) ts -= (i32)P;
c[k] = (u32)((i64)ts + 4 * (WA + WB) - 16 * t);
}
}
if (c32) _mm_sfence();
(void)E2;
return;
}
// [PAD] same re-indexing as the fast path: A is stored row by row with the
// padded stride XST_. (This branch is not taken for n == m == 1e8.)
u32 *Arow = A;
long rbase = 0;
for (long k = 0; k <= nmax; k++) {
if (k >= rbase + (long)NCW_ && rbase + (long)NCW_ < (long)N1_) {
rbase += (long)NCW_; Arow += (u64)XST_;
}
if (k <= n) WA += a[k];
if (k <= m) WB += b[k];
if (k > (long)m) WA -= a[k - 1 - m];
if (k > (long)n) WB -= b[k - 1 - n];
long lo = (k - (long)m > 0) ? (k - (long)m) : 0;
long hi = (k < (long)n) ? k : (long)n;
long t = hi - lo + 1;
u32 tab;
if (k <= M) {
u32 x = Arow[k - rbase], y = c[N1_ + k]; // E_k
tab = (x >= y) ? (x - y) : (x + (P - y));
} else if (k < (long)N1_) {
tab = Arow[k - rbase];
} else {
tab = c[k]; // E_{k-N1_}, in place
}
#if EXACT
i32 ts = (i32)tab;
if (ts > (i32)(P >> 1)) ts -= (i32)P;
i64 val = (i64)ts + 4 * (WA + WB) - 16 * t;
c[k] = (u32)val;
#else
c[k] = tab;
#endif
}
if (c32) _mm_sfence();
(void)E2;
}
#pragma GCC pop_options
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 6.263 s | 1811 MB + 940 KB | Accepted | Score: 100 | 显示更多 |