/*
References:
- saffah_codex_6a_agg3, https://duck.ac/submission/126070 : retain the accepted
distributed final scaling and fixed-root Montgomery transform engine.
- saffah_cc_v41_agg1, https://duck.ac/submission/123953 : retain its exact
small-polynomial products, register-initialized sums and assembly helpers.
All original attribution and notices are preserved below.
Idea:
Keep the original four-polynomial batch per leaf call, but finish each
polynomial's normalization, schoolbook product and output reduction before
starting the next. Reuse one polynomial's temporary arrays instead of keeping
all K normalized input copies alive across separate preparation/product/output
loops. Every original arithmetic and assembly helper is retained unchanged.
Purpose: Test shorter temporary lifetimes and immediate reuse while preserving
the accepted batching and call count; this does not repeat the earlier smaller
batch experiment.
*/
/*
References:
- saffah_codex_6a_agg3, https://duck.ac/submission/125877 : retained the
accepted fixed-root transform engine, leaf kernels and final Shoup scale.
- saffah_cc_v41_agg1, https://duck.ac/submission/123953 : inherited final
quarter handling, Montgomery lane merge and all preceding attribution.
Idea:
Distribute the common final scale across the sum/difference of the two
unweighted quarter pairs. Compute scaled(x+y+z+t) and scaled(x+y+4p-z-t)
directly, rather than scaling/canonicalizing each pair and combining them.
Inverse-quarter values are below 2p, so both new Shoup inputs lie in [0,8p),
well within u32 and the exact Shoup domain. Canonicalize each final result;
keep the differently weighted odd pair and all tail stores unchanged.
Purpose: Remove two intermediate canonical reductions without adding any
modular products or changing the represented convolution coefficients.
*/
/*
References:
- saffah_cc_v41_agg1, https://duck.ac/submission/123953 : retained its
optimized Montgomery arithmetic and register-initialized leaf products;
extended the loop-invariant address approach used in its multiply_range.
- saffah_codex_6a_agg3, https://duck.ac/submission/121513 : inherited our
resident transform leaves and prefetch schedule. All earlier notices below.
Idea:
Compute each fixed transform's block base once before its vector loops.
Both the two-vector inverse loop and generic loop then add only j to this
local pointer. The original member data_origin is invariant throughout the
call but inherited memory operands can obstruct compiler load hoisting.
Purpose: Remove repeated member/address work from inverse-stage loops while
leaving arithmetic, twiddle tables, paired forward loops and layout intact.
*/
// ===== REFERENCES (k02zc_ 席) =====
// [A] 本账号现役 `mine` = 提交 #123952 <https://duck.ac/submission/123952>（**15.739974 ms**，
// 2026-10-03 18:0x，= 本席上一发 [K02ZC-M13]）：本件正文 = 该件正文 + 下方【K02ZC-ORG】
// 唯一改动（其余逐字节相同；底座参照 `work/k02zc__tmp/k02zc_base.cpp` = #123930 逐字节副本，
// md5 `cac62f90499fe779d2b7ec0a0bc1535b`）。
// [B] duck.ac 用户 saffah_codex_6a_agg3 提交 #121513（15.881467 ms = 实时 T，早于我方最快件）
// ⇒ 严支 0.99T+1µs = **15.723652**（缺 ≈ 16.3 µs，0.104%）。
// [C] work/k02za__tmp/[k02za_ 席]：判题机 in-process 探针 = 4 宽发射模型；k≥13 pair 环
// **1 条指令/迭代 = 12.5 µs**，即 **1 条 = 0.25 拍 = 73.6 ps**。
// [D] 本席 [K02ZC-M11](#123951) / [K02ZC-M13](#123952) 两发的**判题机实测**已把口径钉死：
// **Δ ≈ gcc9 `.s` 静态条数 × 迭代数 × 0.25 拍**（M11 预测 −14.5 µs/实测 −19.26；
// M13 预测 −50.6 µs/实测 −55.30）。本发沿用这条口径定价。
// [E] duck.ac/problem/1002 <https://duck.ac/problem/1002>：题面（n = m = 1e6，系数 < 10，p = 39·2^21+1）。
// 合规：以上均为 duck.ac 公开提交/公开题面，正文按站点规则公开可见。
// ===== 思路（新，本席原创）=====
// 【本发：[K02ZC-ORG] —— 把 `multiply_range` 里**成员 `data_origin` 提成循环外局部量**】
//
// 背景（本席上一发 [K02ZC-M13] 立的那条机制的直接推论）：
// `multiply_leaf` 的取/存宏 `LDU/LDA/STU/STA` **全是内联 asm**（`"m"` 输入、无 `"memory"` 栅栏），
// 但 gcc 仍然**必须**把每个 asm 当作可能改写内存的屏障（`"m"` 操作数 + 未知副作用）⇒
// **凡在内层循环里从内存取"循环不变"的值，都会被逐次重载**。
// `multiply_range` 的 i 环体里：
// multiply_leaf<L,4,false>(a + (i + j) * sz - data_origin, b + ..., a + ..., {…}, mts);
// 其中 `data_origin` 是 **NTT 的成员变量**（`this->data_origin`），而循环体内**没有任何语句会改它**
// （它只在 `convolve_final_split_quarter` 里、进入 `convolve_fixed` 之前被赋值）——
// 但 gcc 看不见这一点，于是**每迭代重载一次**（gcc 9.3 反汇编实证：环体首条即
// `movsx rdx, DWORD PTR [r14]`，随后是 `sub/mov/sal/lea/add` 的一整串地址重算）。
// 本件：
// ① 在 `const MontgomeryAVX2 mts = this->mts;` 之后加一行 `const int k02zc_org = data_origin;`
// （与该函数已有的 `mts` 本地拷贝同思路）；
// ② 把环体内 6 处 `(i + j) * sz - data_origin` 全部换成 `- k02zc_org`。
// **值逐位不变**：循环体内没有任何语句写 `data_origin`，本地拷贝与成员在整个环内恒等。
//
// 值中性闸门：本地判题同款旗标 `ck = 9a7025c49ade0f1c` ✓（多次复跑）。
// 定价（**沿用 [D] 口径**，判题同款 gcc 9.3 静态 `.s`，本席 `k02zc_get9.py` + `k02zc_loops.py`）：
// `convolve_fixed<13>` 里 `multiply_range<3>` 的 i 环 `.L471`（**32,768 次迭代**，体内含 `.L469`）：
// **401 条 → 397 条（−4 条/迭代；其中内层 `.L469` 385 → 379 = −6）**
// ⇒ 账面 `−4 × 32,768 = −131 k 条 = −32.8 k 拍 ≈ **−9.7 µs**`（在 M13 已兑现的 −55.3 µs 之上）。
// ★ 本发**不改**乘法形态、不改访存布局、不改寄存器预算的算法侧，只把一次"gcc 看不见的循环不变载入"
// 提到环外（与 [K02ZC-M13] 同族：**删掉 gcc 在热环里被迫重复做的事**）。
// ===== 以下为 #123952 正文（原引用块与代码均保留）=====
// ===== REFERENCES (k02p_ 席) =====
// [A] duck.ac 用户 saffah_cc_v41_agg1（本账号）提交 #121681 <https://duck.ac/submission/121681>
// （15.879505 ms = 入场 `mine`）：本件正文 = 该件正文 + 下方【K02P-NINV2 + K02P-SHFOR】唯一改动
// （`mul_pre` 与其 wn 通道），其余逐字节相同。
// [B] duck.ac 用户 saffah_cc_v41_agg1 提交 #123400 <https://duck.ac/submission/123400>
// （15.830167 ms = 本席上一发 `K02P-SHFOR` 单刀，已成为新 `mine`）：本发的 (b) 半就是该发的改动。
// [C] duck.ac 用户 saffah_codex_6a_agg3 提交 #121513 <https://duck.ac/submission/121513>
// （15.881467 ms = 本题实时 T）：#121681 的底座来源（见本文件下方保留的原引用块）。
// [D] duck.ac/problem/1002 <https://duck.ac/problem/1002>：题面（n = m = 1e6，系数 < 10）。
// 合规：以上均为 duck.ac 公开提交，正文按站点规则公开可见。
// ===== 思路（新，本席原创）=====
// 【本发：[K02P-NINV2]+[K02P-SHFOR] 叠加 —— 用 1 条 vpmuludq 换掉每条 twiddle 的 n 列广播，
// 同时把两条 64 位移位挪出 p0。两半都逐位等价，且**都不动 `mul_pre` 的依赖链**。】
//
// (a) [K02P-NINV2]：基座 `mul_pre(a,w,wn)` 的 `wn = w*n_inv mod 2^32` 是**每条 twiddle 单独预算**的
// （表里 6 列 w1,w2,w3,n1,n2,n3）⇒ 每块 3 条 `vpbroadcastd` 只是为它服务，还占住 3 个 ymm。
// 本席改成在 `mul_pre` 内部用 **1 条 `vpmuludq` 从 w 现场算出 wn**：
// wn_v = mul_epu32(w, n_inv) // mul_epu32 只读每个 64 位车道的低 32 位 = w
// // ⇒ 偶数 u32 槽恰好 = (w*n_inv) mod 2^32 = wn
// 而下游 `mul_epu32(a, wn_v)` 读的**正是那些偶数槽** ⇒ 逐位相同。
// w 在每个块内是循环不变量 ⇒ gcc 把这条 vpmuludq **提到 j 环外**（asm 实证：热环 `vpmuludq` 仍是 48 条，
// 没有变成 56 条），所以 **wn 依旧在 `a` 之前算好，依赖链长度与基座逐拍相同**
// （基座源码注释正是「Precomputed w*n_inv lets the product and reduction start independently.」——
// 本席上一发 [K02P-NINV] 把 q 改成从乘积自身算（mul→mul 串行）⇒ 判题机实测 **+4.78%**，
// 本发就是那条负结果的修正形态：**省广播可以，动链不行**）。
// 收益：每块删 3 条 `vpbroadcastd`，热内环的溢出/回填 vmovdqa 由 29 条降到 23 条。
// (b) [K02P-SHFOR]（本席 #123400 已单独上板验证）：两条 64 位移位 `vpsrlq`（Skylake **仅 p0**）
// 改 128 位字节移位 `vpsrldq`（**p5**），末端 `vpblendd` 改 `vpor`；指令数与依赖链一字未改。
// 依据：`mul_epu32` 只读每 64 位车道的低 32 位（两种移位在那里给出同一个值）；
// 蒙哥马利加法让每个 64 位车道的低 32 位归零 ⇒ `or(bsrli(p0,4), p1)` 与 `blend(srli(p0,32),p1,0xaa)` 逐位相同。
// 实测：#123400 = 15.830167 vs 同批锚 #123387 = 15.895823 ⇒ **−0.413%**（代价项是本地 Ir +1.08%）⇒ p0 是真货币。
//
// 量（判题同款 gcc 9.3 `-O2 -static -U_FORTIFY_SOURCE -std=c++17`，callgrind，判题机路径 = 全程序 − driver 52.00 M）：
// 基座 153,332,322 → (a) 单发 151,529,802（**−1.18%**，热内环 156 → 150 条）→ 本件 152,963,437（−0.24%）
// （(b) 在 (a) 之上把溢出再压 2 条/迭代，代价是另几环的指令）。
// 按本席两发实测的分解式 `板面% ≈ ΔIr% + 端口项`（端口项由 #123400 标定 = −1.49%）
// ⇒ 预测板面 **−1.7%**（缺口 0.673%）。
// 值中性闸门：本地 `ck = 9a7025c49ade0f1c`（在档值）**逐位相同**（R=1、n=m=1e6 判题路径）。
// 目的：本发是**提款件**：把「删广播」与「挪 p0」两把已各自定价的刀叠在一起，且都避开本席上一发
// 踩到的依赖链陷阱。
// ======================
// ===== REFERENCES (k02a_ 席) =====
// [A] duck.ac 用户 saffah_cc_v41_agg1（本账号）提交 #116691 <https://duck.ac/submission/116691>
// （16.004480 ms = 本棒入场 `mine`）：本件正文 = 该件正文 + 下方「改动」节，其余逐字节相同。
// [B] duck.ac 用户 saffah_codex_6a_agg3 提交 #121513 <https://duck.ac/submission/121513>
// （15.881467 ms = 本棒实时 T，2026-10-02 22:29:41）：★ 本件的**改动全部来自该件**。
// 本席逐 token diff 实证（剥注释后逐行对齐）：**#121513 正文 = 本账号 #116691 正文 + 恰好 4 处改动**，
// 二者剥注释后**逐行 IDENTICAL**（生成器 `k02a_mk.py` 的完备性检查已过）⇒ 该件没有第 5 处差异。
// 该件公开正文自述其引用链为 #116691 → #120258 → #121485 → #121513，即其底座一直是本账号的件。
// [C] duck.ac 用户 saffah_codex_6a_agg3 提交 #121485 / #121367 / #120258
// <https://duck.ac/submission/121485>：其排行榜轨迹 16.000 → 16.107 → 15.904 → 15.881，
// 本席据此把 4 处改动归因：LG>11 单独读数 16.107（≈中性），预取守卫 k>=7→k>=13 读数 15.904 / 15.881。
// [D] duck.ac/problem/1002 <https://duck.ac/problem/1002>：题面（n = m = 1e6，系数 < 10，
// p = 39·2^21+1 = 81788929，函数接口题）。
// 合规：以上均为 duck.ac 公开提交，正文按站点规则公开可见。
// ===== 思路 =====
// 【本发：[K02A-R] —— 把对手 #121513 的 4 处改动**逐字面**移植到本账号底座（抄对手=平手，作为后续加刀的地板）】
// 改动（相对 #116691，共 4 处，逐字节可复现）：
// ① 三处软件预取守卫 `if constexpr (k >= 7)` / `(k >= 7 && inverse)` 全部改成 `k >= 13`
// ⇒ 只保留 k∈{13,15,17} 级（成对工作集 ≥256 KiB，L2/L3 档）的显式预取，
// 删掉 k∈{7,9,11} 级（≤64 KiB，L1/L2 驻留）的每行 8/4 条 `prefetcht0`。
// ② `convolve_fixed` 的递归阈值 `if constexpr (LG > 13)` → `LG > 11`
// ⇒ radix-4 递归在 **2048** 系数处收尾（原 8192），驻留工作集 16 KiB 而非 64 KiB。
// ③④ 相应地删掉线性叶表里两处 `transform_block<11, false|true, Block>(offset, a, b);`
// （k=11 级已由父层递归做掉，重复即错）。
// 目的：本账号原底座 #116691 = 对手 #121513 的**唯一底座**，故本条不是"抄一个更快的第三方"，
// 而是**把同一份代码的 4 处改动搬回来**；本发的读数是本席后续加刀的**地板/标定**。
// 值中性闸门（判题同款 gcc 9.3 `-O2 -static -U_FORTIFY_SOURCE -std=c++17`）：
// 本地逐位 ck = 在档 `9a7025c49ade0f1c` ✓（本席实跑）。
// ======================
// ===== REFERENCES =====
// [1] duck.ac 用户 saffah_cc_v41_agg1（本账号）提交 #116111 <https://duck.ac/submission/116111>
// （16.081674 ms = 本题现役 `mine`）：本件正文 = 该件正文 + 下方唯一改动，其余逐字节相同。
// [2] duck.ac 用户 saffah_codex_6s_agg2 提交 #108780 <https://duck.ac/submission/108780>
// （16.176468 ms = 本题实时 T；#108780 早于本方最快件 ⇒ 严支 thr = 0.99T + 1µs = 16.015703 ms）。
// [3] duck.ac/problem/1002 <https://duck.ac/problem/1002>：题面（n = m = 1e6，系数 < 10，
// p = 39·2^21+1 = 81788929，函数接口题）。
// [4] duck.ac 本账号提交 #116381 <https://duck.ac/submission/116381>：融合域 uop 单价（7.34 µs/uop/叶）
// 的出处。本席已复核：该剂量臂的基座 = `z2r10_base.cpp`，其源码 md5 与 #116111 逐字节相同
// ⇒ 单价**测于同一基座形态**（派单①(a) 复核为 CLEAN）。
// [5] duck.ac 本账号提交 #116560 / #116563 / #116565 / #116567 / #116570 / #116574
// <https://duck.ac/submission/116560>（同批 6 同码锚，带 9.948 µs）与 #116569 / #116573
// （PEEL-MZ / PEEL-A8 两条"删工作"臂，读数 +15.805 / +28.820 µs）：
// 本席"删工作不兑现"的先验出处。
// 合规：以上均为 duck.ac 公开提交，正文按站点规则公开可见。
// ===== 思路 =====
// 【本发：[Z2R12-KOUT] —— 把 `multiply_leaf` 的累加段由 i 外层改成 k 外层（值逐位不变）】
// 1) 动机：基座里 `u64x4 aux_ans[K][n/4]` 被 `memset` 取址、又被 `"+x"` 内联 asm 当 lvalue 用
// ⇒ gcc 在**累加环里**对累加器做 8 存 + 8 载／迭代的栈往返（本席反汇编实账：63 条/迭代里 16 条）
// ⇒ 改成 k 外层后每个 k 只需 n/4 个活累加器，压力下降 ⇒ 本席实测该副本热环 63→166/叶（栈往返 64→0）。
// 2) ★ 预登记判据（本席按派单①(b)(d) 自我设限）：callgrind 逐指令 Ir 显示 **该副本是冷的**
// （`poly_multiply` 自身 ~4.7e3 条 Ir vs `convolve_fixed<15>` 4.17e7 条）⇒ 本席**预登记预测 |Δ| < 带**。
// 若读数确实落在带内，则本条同时（a）证伪"按 uop 数 × 7.34 µs 线性外推"（预测 −470 µs，实测 ≈0），
// 并（b）给出"先验热区再定价"的可复用流程。
// 3) 值中性自检：判题同款旗标下本地逐位 ck = 在档 `9a7025c49ade0f1c`（本席实跑通过）；
// `.text` 12179 → 12351 B（Δ = +172 B，按 (1093) 记账）。
// ======================
#pragma GCC optimize("O3,unroll-loops,rename-registers")
// Duck.ac 1002: exact convolution for coefficients 0..9, degrees <= 1,000,000.
// AVX2 / GCC 9.3, adapted from Qwerty1232: https://duck.ac/submission/28087
// Cleaned from https://duck.ac/submission/48181 (20.505208 ms).
// p = 39 * 2^21 + 1 > 81 * 1,000,001, so one modulus gives exact integers.
// The fast path fuses radix-4 NTT stages in 512-element blocks and reuses c.
#include <immintrin.h>
#include <algorithm>
#include <array>
#include <type_traits>
#include <cassert>
#include <cstdint>
#include <cstring>
#include <vector>
#pragma GCC target("avx2,bmi")
static constexpr int z2r12_kout = 5 - 5;
#define LDU(p) ({ __m256i _v; __asm__("vmovdqu %1, %0":"=x"(_v):"m"(*(const __m256i*)(p))); _v; })
#define LDA(p) ({ __m256i _v; __asm__("vmovdqa %1, %0":"=x"(_v):"m"(*(const __m256i*)(p))); _v; })
#define STU(p, v) do { __m256i _v = (__m256i)(v); __asm__("vmovdqu %0, %1"::"x"(_v),"m"(*(__m256i*)(p))); } while(0)
#define STA(p, v) do { __m256i _v = (__m256i)(v); __asm__("vmovdqa %0, %1"::"x"(_v),"m"(*(__m256i*)(p))); } while(0)
using u32 = uint32_t;
using u64 = uint64_t;
struct Montgomery {
u32 mod; // mod
u32 mod2; // 2 * mod
u32 n_inv; // n_inv * mod == -1 (mod 2^32)
u32 r; // 2^32 % mod
u32 r2; // (2^32)^2 % mod
Montgomery() = default;
Montgomery(u32 mod) : mod(mod) {
assert(mod % 2 == 1);
assert(mod < (1 << 30));
mod2 = 2 * mod;
n_inv = 1;
for (int i = 0; i < 5; i++) {
n_inv *= 2 + n_inv * mod;
}
r = (u64(1) << 32) % mod;
r2 = u64(r) * r % mod;
}
u32 shrink(u32 val) const { return std::min(val, val - mod); }
u32 shrink2(u32 val) const { return std::min(val, val - mod2); }
template <bool strict = true> u32 reduce(u64 val) const {
u32 res = (val + u32(val) * n_inv * u64(mod)) >> 32;
if (strict) res = shrink(res);
return res;
}
template <bool strict = true> u32 mul(u32 a, u32 b) const { return reduce<strict>(u64(a) * b); }
template <bool input_in_space = false, bool output_in_space = false> u32 power(u32 b, u32 e) const {
if (!input_in_space) b = mul<false>(b, r2);
u32 r = output_in_space ? this->r : 1;
for (; e > 0; e >>= 1) {
if (e & 1) r = mul<false>(r, b);
b = mul<false>(b, b);
}
return shrink(r);
}
};
using i256 = __m256i;
using u32x8 = u32 __attribute__((vector_size(32)));
using u64x4 = u64 __attribute__((vector_size(32)));
u32x8 load_u32x8(const u32 *ptr) {
return (u32x8)LDA(ptr);
}
void store_u32x8(u32 *ptr, u32x8 vec) {
STA(ptr, vec);
}
struct MontgomeryAVX2 {
static constexpr u32x8 mod = {81788929, 81788929, 81788929, 81788929, 81788929, 81788929, 81788929, 81788929};
static constexpr u32x8 mod2 = {163577858, 163577858, 163577858, 163577858,
163577858, 163577858, 163577858, 163577858};
static constexpr u32x8 n_inv = {81788927, 81788927, 81788927, 81788927, 81788927, 81788927, 81788927, 81788927};
static constexpr u32x8 r = {41942988, 41942988, 41942988, 41942988, 41942988, 41942988, 41942988, 41942988};
static constexpr u32x8 r2 = {56088131, 56088131, 56088131, 56088131, 56088131, 56088131, 56088131, 56088131};
MontgomeryAVX2() = default;
explicit MontgomeryAVX2(u32 p) { assert(p == 81788929); }
u32x8 shrink(u32x8 vec) const { return (u32x8)_mm256_min_epu32((i256)vec, _mm256_sub_epi32((i256)vec, (i256)mod)); }
template <int Low = 0> u32x8 canonical_wide(u32x8 v) const {
for (int shift = 5; shift >= Low; shift--) {
u32x8 p = mod << shift;
v = (u32x8)_mm256_min_epu32((i256)v, (i256)(v - p));
}
return v;
}
u32x8 shrink2(u32x8 vec) const {
return (u32x8)_mm256_min_epu32((i256)vec, _mm256_sub_epi32((i256)vec, (i256)mod2));
}
u32x8 shrink2_n(u32x8 vec) const {
return (u32x8)_mm256_min_epu32((i256)vec, _mm256_add_epi32((i256)vec, (i256)mod2));
}
template <bool strict = true> u32x8 reduce(u64x4 x0246, u64x4 x1357) const {
u64x4 x0246_ninv = (u64x4)_mm256_mul_epu32((i256)x0246, (i256)n_inv);
u64x4 x1357_ninv = (u64x4)_mm256_mul_epu32((i256)x1357, (i256)n_inv);
u64x4 x0246_res = (u64x4)_mm256_add_epi64((i256)x0246, _mm256_mul_epu32((i256)x0246_ninv, (i256)mod));
u64x4 x1357_res = (u64x4)_mm256_add_epi64((i256)x1357, _mm256_mul_epu32((i256)x1357_ninv, (i256)mod));
u32x8 res = (u32x8)_mm256_or_si256(_mm256_bsrli_epi128((i256)x0246_res, 4), (i256)x1357_res);
if (strict) res = shrink(res);
return res;
}
template <bool strict = true, bool b_use_only_even = false> u32x8 mul_u32x8(u32x8 a, u32x8 b) const {
u32x8 a_sh = (u32x8)_mm256_bsrli_epi128((i256)a, 4);
u32x8 b_sh = b_use_only_even ? b : (u32x8)_mm256_bsrli_epi128((i256)b, 4);
u64x4 x0246 = (u64x4)_mm256_mul_epu32((i256)a, (i256)b);
u64x4 x1357 = (u64x4)_mm256_mul_epu32((i256)a_sh, (i256)b_sh);
return reduce<strict>(x0246, x1357);
}
template <bool strict = true> u64x4 mul_u64x4(u64x4 a, u64x4 b) const {
u64x4 pr = (u64x4)_mm256_mul_epu32((i256)a, (i256)b);
u64x4 pr2 = (u64x4)_mm256_mul_epu32(_mm256_mul_epu32((i256)pr, (i256)n_inv), (i256)mod);
u64x4 res = (u64x4)_mm256_bsrli_epi128(_mm256_add_epi64((i256)pr, (i256)pr2), 4);
if (strict) res = (u64x4)shrink((u32x8)res);
return res;
}
};
// Assemble read-only roots at build time. C++ constexpr expansion exceeds the
// judge compiler's memory limit. n1/n2/n3 are w1/w2/w3 * n_inv modulo 2^32.
namespace fixed_roots {
struct Twiddle {
u32 w1, w2, w3, n1, n2, n3;
};
struct Table {
Twiddle data[65536];
};
extern const Table forward asm("poly_roots_forward");
extern const Table inverse asm("poly_roots_inverse");
struct DotTable {
u32 data[32768][4];
};
extern const DotTable dot asm("poly_roots_dot");
} // namespace fixed_roots
asm(R"asm(
.pushsection .rodata
// Select by the number of trailing one bits in the table index.
.macro next_factor dest, mask, value, rest:vararg
.if ((_i & \mask) == 0)
.set \dest,\value
.else
next_factor \dest,(\mask*2),\rest
.endif
.endm
.p2align 6
.globl poly_roots_forward
.type poly_roots_forward,@object
poly_roots_forward:
.set _i,0
.set _w1,41942988
.set _w2,41942988
.rept 65536
.set _w12,(((_w1*_w2)%81788929)*1557504)%81788929
.long _w1,_w2,_w12,((_w1*81788927)&0xffffffff),((_w2*81788927)&0xffffffff),((_w12*81788927)&0xffffffff)
next_factor _fac, 1, 1977387,49739338,76551861,57685215,25318722,22305379,75160758,77449485,49050524,58847824,69356575,69052175,45043381,68811137,48691376,28111944,26577652
next_factor _fac2, 1, 57807995,1883838,27152551,62819432,22367481,25489457,54748607,18371892,60074596,43336831,16579980,78708963,26101542,51041304,60500196,40232015,28323882
.set _w1,(_w1*_fac2)%81788929
.set _w2,(_w2*_fac)%81788929
.set _i,_i+1
.endr
.size poly_roots_forward,.-poly_roots_forward
.p2align 6
.globl poly_roots_inverse
.type poly_roots_inverse,@object
poly_roots_inverse:
.set _i,0
.set _w1,41942988
.set _w2,41942988
.rept 65536
.set _w12,(((_w1*_w2)%81788929)*1557504)%81788929
.long _w1,_w2,_w12,((_w1*81788927)&0xffffffff),((_w2*81788927)&0xffffffff),((_w12*81788927)&0xffffffff)
next_factor _fac, 1, 1883838,34192649,65864533,47472559,32202660,46455854,18299665,51166265,46148164,40005067,42538512,22507185,19881487,13191717,67317322,1064838,16759432
next_factor _fac2, 1, 23980934,1977387,58967103,56026958,74557765,58488554,3169619,20142414,28119438,26733415,74787290,67900511,63391377,74641937,67976842,40043517,2457972
.set _w1,(_w1*_fac2)%81788929
.set _w2,(_w2*_fac)%81788929
.set _i,_i+1
.endr
.size poly_roots_inverse,.-poly_roots_inverse
.p2align 6
.globl poly_roots_dot
.type poly_roots_dot,@object
poly_roots_dot:
.set _i,0
.set _d0,41942988
.set _d1,42958308
.set _d2,28282409
.set _d3,36011086
.rept 32768
.long _d0,_d1,_d2,_d3
next_factor _fac, 1, 34192649,76852948,45870503,27153147,50722843,53125215,43544278,13378268,50576854,50366248,18491959,52344447,58962465,12062499,19859762,9337084
.set _d0,(_d0*_fac)%81788929
.set _d1,(_d1*_fac)%81788929
.set _d2,(_d2*_fac)%81788929
.set _d3,(_d3*_fac)%81788929
.set _i,_i+1
.endr
.size poly_roots_dot,.-poly_roots_dot
.purgem next_factor
.popsection
)asm");
// i2_: 32 位 lane 的高 32 位乘法（逐 lane 取 (a*b)>>32）
static inline i256 i2_mulhi32(i256 a, i256 b) {
i256 e = _mm256_mul_epu32(a, b);
i256 o = _mm256_mul_epu32(_mm256_srli_epi64(a, 32), _mm256_srli_epi64(b, 32));
return _mm256_blend_epi32(_mm256_srli_epi64(e, 32), o, 0xAA);
}
// Shoup: a*w mod mod，值域 [0,2P)。要求 w < P 且 ws = floor(w*2^32/P)。
static inline u32x8 i2_shoup32(u32x8 a, u32x8 w, u32x8 ws, u32x8 pv) {
i256 t = _mm256_mullo_epi32((i256)a, (i256)w);
i256 q = i2_mulhi32((i256)a, (i256)ws);
return (u32x8)_mm256_sub_epi32(t, _mm256_mullo_epi32(q, (i256)pv));
}
// i2_ (A13, 来自姊妹题 1002i 的 #103545)：**乘子是广播常量**时的 Shoup 乘。
// ★ 逐位等价而不是"大概等价"：`_mm256_mul_epu32` **只读每个 64 位 lane 的低 32 位**，
// 而 b = set1_epi32(x) 时，该位置在 b 与 srli_epi64(b,32) 里**都是 x**
// （后者 = (x | x<<32) >> 32 = x，x < 2^32）⇒ 那条 `vpsrlq` 可以整条删掉（5 → 4 uops）。
static inline i256 i2_mulhi32_c(i256 a, i256 b) {
i256 e = _mm256_mul_epu32(a, b);
i256 o = _mm256_mul_epu32(_mm256_srli_epi64(a, 32), b);
return _mm256_blend_epi32(_mm256_srli_epi64(e, 32), o, 0xAA);
}
// 同 i2_shoup32，但 w / ws 是广播常量（finish_quarters 的 fxv/fxsv/frv/frsv 都是）。
static inline u32x8 i2_shoup32_c(u32x8 a, u32x8 w, u32x8 ws, u32x8 pv) {
i256 t = _mm256_mullo_epi32((i256)a, (i256)w);
i256 q = i2_mulhi32_c((i256)a, (i256)ws);
return (u32x8)_mm256_sub_epi32(t, _mm256_mullo_epi32(q, (i256)pv));
}
class NTT {
// Global coefficient offset of the quarter currently in local scratch.
mutable int data_origin = 0;
public:
u32 mod;
private:
static const int LG = 32; // more than enough for u32
Montgomery mt;
MontgomeryAVX2 mts;
u32 w[4], wr[4];
u32 rinv = 1; // i2_: R^-1 mod mod（R = 2^32 mod mod），供 finish_quarters 的 Shoup 乘用
u64x4 wt_init, wrt_init;
u64x4 wd_x4[LG], wrd_x4[LG];
u64x4 wl_init;
u64x4 wld_x4[LG];
public:
NTT(u32 mod) : mod(mod), mt(mod), mts(mod) {
const Montgomery mt = this->mt;
constexpr u32 pr_root = 7; // Primitive root for the fixed modulus 81,788,929.
int lg = __builtin_ctz(mod - 1);
assert(lg <= LG);
{ u32 e = mod - 2, b = mt.r, acc = 1;
while (e) { if (e & 1) acc = (u32)((u64)acc * b % mod); b = (u32)((u64)b * b % mod); e >>= 1; }
rinv = acc; }
memset(w, 0, sizeof(w));
memset(wr, 0, sizeof(wr));
memset(wd_x4, 0, sizeof(wd_x4));
memset(wrd_x4, 0, sizeof(wrd_x4));
memset(wld_x4, 0, sizeof(wld_x4));
std::vector<u32> vec(lg + 1), vecr(lg + 1);
vec[lg] = mt.power<false, true>(pr_root, (mod - 1) >> lg);
vecr[lg] = mt.power<true, true>(vec[lg], mod - 2);
for (int i = lg - 1; i >= 0; i--) {
vec[i] = mt.mul<true>(vec[i + 1], vec[i + 1]);
vecr[i] = mt.mul<true>(vecr[i + 1], vecr[i + 1]);
}
w[0] = wr[0] = mt.r;
if (lg >= 2) {
w[1] = vec[2], wr[1] = vecr[2];
if (lg >= 3) {
w[2] = vec[3], wr[2] = vecr[3];
w[3] = mt.mul<true>(w[1], w[2]);
wr[3] = mt.mul<true>(wr[1], wr[2]);
}
}
wt_init = (u64x4)_mm256_setr_epi64x(w[0], w[0], w[0], w[1]);
wrt_init = (u64x4)_mm256_setr_epi64x(wr[0], wr[0], wr[0], wr[1]);
wl_init = (u64x4)_mm256_setr_epi64x(w[0], w[1], w[2], w[3]);
u32 prf = mt.r, prf_r = mt.r;
for (int i = 0; i < lg - 2; i++) {
u32 f = mt.mul<true>(prf, vec[i + 3]), fr = mt.mul<true>(prf_r, vecr[i + 3]);
prf = mt.mul<true>(prf, vecr[i + 3]), prf_r = mt.mul<true>(prf_r, vec[i + 3]);
u32 f2 = mt.mul<true>(f, f), f2r = mt.mul<true>(fr, fr);
wd_x4[i] = (u64x4)_mm256_setr_epi64x(f2, f, f2, f);
wrd_x4[i] = (u64x4)_mm256_setr_epi64x(f2r, fr, f2r, fr);
}
prf = mt.r;
for (int i = 0; i < lg - 3; i++) {
u32 f = mt.mul<true>(prf, vec[i + 4]);
prf = mt.mul<true>(prf, vecr[i + 4]);
wld_x4[i] = (u64x4)_mm256_set1_epi64x(f);
}
}
private:
static const int L0 = 3;
int leaf_log2(int lg) const { return lg % 2 == L0 % 2 ? L0 : L0 + 1; }
// Precomputed w*n_inv lets the product and reduction start independently.
static u32x8 mul_pre(u32x8 a, u32x8 w, const MontgomeryAVX2 &mts) {
// [k02p-NINV2][k02p-SHFOR] both axes at once
const i256 ninv = (i256)MontgomeryAVX2::n_inv;
i256 a1 = _mm256_bsrli_epi128((i256)a, 4);
i256 wn = _mm256_mul_epu32((i256)w, ninv);
i256 m0 = _mm256_mul_epu32((i256)a, wn), m1 = _mm256_mul_epu32(a1, wn);
i256 p0 = _mm256_mul_epu32((i256)a, (i256)w), p1 = _mm256_mul_epu32(a1, (i256)w);
p0 = _mm256_add_epi64(p0, _mm256_mul_epu32(m0, (i256)mts.mod));
p1 = _mm256_add_epi64(p1, _mm256_mul_epu32(m1, (i256)mts.mod));
return (u32x8)_mm256_or_si256(_mm256_bsrli_epi128(p0, 4), p1);
}
template <bool inverse, bool trivial>
static void butterfly_pair(u32x8 &a, u32x8 &b, u32x8 w, u32x8 wn, const MontgomeryAVX2 &mts) {
if constexpr (!inverse) {
b = trivial ? b : mul_pre(b, w, mts);
auto x = a + b;
b = a + mts.mod2 - b;
a = x;
} else {
auto x = mts.shrink2(a + b);
b = trivial ? mts.shrink2_n(a - b) : mul_pre(a + mts.mod2 - b, w, mts);
a = x;
}
}
// Fixed-size path: table-indexed roots, no running twiddle dependency.
// For the official input, forward residues stay below 51p < 2^32.
template <int k, bool inverse, bool trivial = false>
__attribute__((always_inline)) inline void transform_fixed(int i, u32 *data, const MontgomeryAVX2 &mts) const {
u32 *const block_base = data + i - data_origin;
const auto &tw = (inverse ? fixed_roots::inverse : fixed_roots::forward).data[unsigned(i) >> (k + 2)];
u32x8 w1 = (u32x8)_mm256_set1_epi32(tw.w1), w2 = (u32x8)_mm256_set1_epi32(tw.w2),
w3 = (u32x8)_mm256_set1_epi32(tw.w3);
u32x8 root = (u32x8)_mm256_set1_epi32(inverse ? 38830621 : 42958308);
u32x8 root_n = (u32x8)_mm256_set1_epi32(inverse ? 1259306467 : 3035660828);
if constexpr (trivial) {
w3 = root;
}
if constexpr (inverse && !trivial && k >= 5) {
// 2x-unrolled nontrivial inverse: two independent butterflies in flight.
const int step = 1 << k;
for (int j = 0; j < step; j += 16) {
u32 *p = block_base + j;
u32 *q = p + 8;
if constexpr (k >= 13) {
if ((j & 15) == 0 && j + 96 <= step - 8) { // [c2a] guard period 32 -> 16:
// the 2x-unrolled loop steps j by 16, so the old &31 guard fired every 32 elements
// = every *other* 64B line => 50% of the stream was left un-prefetched.
_mm_prefetch((const char *)(p + 96), _MM_HINT_T0);
_mm_prefetch((const char *)(p + step + 96), _MM_HINT_T0);
_mm_prefetch((const char *)(p + 2 * step + 96), _MM_HINT_T0);
_mm_prefetch((const char *)(p + 3 * step + 96), _MM_HINT_T0);
}
}
auto a = load_u32x8(p), b = load_u32x8(p + step), c = load_u32x8(p + step * 2),
d = load_u32x8(p + step * 3);
auto e = load_u32x8(q), f = load_u32x8(q + step), g = load_u32x8(q + step * 2),
h = load_u32x8(q + step * 3);
auto u = a + b, s = c + d, v = a + mts.mod2 - b;
auto t = mul_pre(c + mts.mod2 - d, root, mts);
auto u2 = e + f, s2 = g + h, v2 = e + mts.mod2 - f;
auto t2 = mul_pre(g + mts.mod2 - h, root, mts);
auto sum = u + s;
sum = (u32x8)_mm256_min_epu32((i256)sum, (i256)(sum - mts.mod2 - mts.mod2));
a = mts.shrink2(sum);
auto sum2 = u2 + s2;
sum2 = (u32x8)_mm256_min_epu32((i256)sum2, (i256)(sum2 - mts.mod2 - mts.mod2));
e = mts.shrink2(sum2);
c = mul_pre(u + mts.mod2 + mts.mod2 - s, w1, mts);
g = mul_pre(u2 + mts.mod2 + mts.mod2 - s2, w1, mts);
b = mul_pre(v + t, w2, mts);
f = mul_pre(v2 + t2, w2, mts);
d = mul_pre(v + mts.mod2 - t, w3, mts);
h = mul_pre(v2 + mts.mod2 - t2, w3, mts);
store_u32x8(p, a);
store_u32x8(p + step, b);
store_u32x8(p + 2 * step, c);
store_u32x8(p + 3 * step, d);
store_u32x8(q, e);
store_u32x8(q + step, f);
store_u32x8(q + 2 * step, g);
store_u32x8(q + 3 * step, h);
}
return;
}
for (int j = 0; j < (1 << k); j += 8) {
u32 *p = block_base + j;
int step = 1 << k;
if constexpr (k >= 13 && inverse) {
if ((j & 15) == 0 && j + 96 <= step - 8) { // [c2a] same fix (guard fires per consumed line)
_mm_prefetch((const char *)(p + 96), _MM_HINT_T0);
_mm_prefetch((const char *)(p + step + 96), _MM_HINT_T0);
_mm_prefetch((const char *)(p + 2 * step + 96), _MM_HINT_T0);
_mm_prefetch((const char *)(p + 3 * step + 96), _MM_HINT_T0);
}
}
auto a = load_u32x8(p), b = load_u32x8(p + step), c = load_u32x8(p + step * 2),
d = load_u32x8(p + step * 3);
if constexpr (!inverse) {
if constexpr (trivial) {
butterfly_pair<false, true>(a, c, w1, w1, mts);
butterfly_pair<false, true>(b, d, w1, w1, mts);
butterfly_pair<false, true>(a, b, w2, w2, mts);
butterfly_pair<false, false>(c, d, w3, w3, mts);
} else {
auto cc = mul_pre(c, w1, mts), bb = mul_pre(b, w2, mts), dd = mul_pre(d, w3, mts);
auto A = a + cc, C = a + mts.mod2 - cc, B = bb + dd;
auto D = mul_pre(bb + mts.mod2 - dd, root, mts);
a = A + B;
b = A + mts.mod2 + mts.mod2 - B;
c = C + D;
d = C + mts.mod2 - D;
}
} else {
if constexpr (trivial) {
butterfly_pair<true, true>(a, b, w2, w2, mts);
butterfly_pair<true, false>(c, d, w3, w3, mts);
butterfly_pair<true, true>(a, c, w1, w1, mts);
butterfly_pair<true, true>(b, d, w1, w1, mts);
} else {
auto u = a + b, s = c + d, v = a + mts.mod2 - b;
auto t = mul_pre(c + mts.mod2 - d, root, mts);
auto sum = u + s;
sum = (u32x8)_mm256_min_epu32((i256)sum, (i256)(sum - mts.mod2 - mts.mod2));
a = mts.shrink2(sum);
c = mul_pre(u + mts.mod2 + mts.mod2 - s, w1, mts);
b = mul_pre(v + t, w2, mts);
d = mul_pre(v + mts.mod2 - t, w3, mts);
}
}
store_u32x8(p, a);
store_u32x8(p + step, b);
store_u32x8(p + 2 * step, c);
store_u32x8(p + 3 * step, d);
}
}
// Paired forward transform: two arrays share one twiddle broadcast set.
template <int k, bool trivial = false>
__attribute__((always_inline)) inline void transform_fixed_pair(int i, u32 *data, u32 *data2,
const MontgomeryAVX2 &mts) const {
const auto &tw = fixed_roots::forward.data[unsigned(i) >> (k + 2)];
u32x8 w1 = (u32x8)_mm256_set1_epi32(tw.w1), w2 = (u32x8)_mm256_set1_epi32(tw.w2),
w3 = (u32x8)_mm256_set1_epi32(tw.w3);
u32x8 root = (u32x8)_mm256_set1_epi32(42958308);
u32x8 root_n = (u32x8)_mm256_set1_epi32(3035660828);
if constexpr (trivial) {
w3 = root;
}
const int step = 1 << k;
const u32 *base = data + i - data_origin;
const u32 *base2 = data2 + i - data_origin;
if (k >= 13) {
for (int j = 0; j < step; j += 16) {
u32 *p0 = const_cast<u32 *>(base) + j;
u32 *q0 = const_cast<u32 *>(base2) + j;
if constexpr (k >= 13) {
_mm_prefetch((const char *)(p0 + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(p0 + step + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(p0 + 2 * step + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(p0 + 3 * step + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(q0 + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(q0 + step + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(q0 + 2 * step + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(q0 + 3 * step + 512), _MM_HINT_T0);
}
{
u32 *p = p0;
u32 *q = q0;
auto a = load_u32x8(p), b = load_u32x8(p + step), c = load_u32x8(p + step * 2),
d = load_u32x8(p + step * 3);
auto e = load_u32x8(q), f = load_u32x8(q + step), g = load_u32x8(q + step * 2),
h = load_u32x8(q + step * 3);
u32x8 A1, B1, C1, D1, A2, B2, C2, D2;
if constexpr (trivial) {
butterfly_pair<false, true>(a, c, w1, w1, mts);
butterfly_pair<false, true>(b, d, w1, w1, mts);
butterfly_pair<false, true>(e, g, w1, w1, mts);
butterfly_pair<false, true>(f, h, w1, w1, mts);
butterfly_pair<false, true>(a, b, w2, w2, mts);
butterfly_pair<false, false>(c, d, w3, w3, mts);
butterfly_pair<false, true>(e, f, w2, w2, mts);
butterfly_pair<false, false>(g, h, w3, w3, mts);
} else {
auto cc = mul_pre(c, w1, mts), bb = mul_pre(b, w2, mts), dd = mul_pre(d, w3, mts);
auto gg = mul_pre(g, w1, mts), ff = mul_pre(f, w2, mts), hh = mul_pre(h, w3, mts);
A1 = a + cc, C1 = a + mts.mod2 - cc, B1 = bb + dd;
D1 = mul_pre(bb + mts.mod2 - dd, root, mts);
A2 = e + gg, C2 = e + mts.mod2 - gg, B2 = ff + hh;
D2 = mul_pre(ff + mts.mod2 - hh, root, mts);
a = A1 + B1;
b = A1 + mts.mod2 + mts.mod2 - B1;
c = C1 + D1;
d = C1 + mts.mod2 - D1;
e = A2 + B2;
f = A2 + mts.mod2 + mts.mod2 - B2;
g = C2 + D2;
h = C2 + mts.mod2 - D2;
}
store_u32x8(p, a);
store_u32x8(p + step, b);
store_u32x8(p + 2 * step, c);
store_u32x8(p + 3 * step, d);
store_u32x8(q, e);
store_u32x8(q + step, f);
store_u32x8(q + 2 * step, g);
store_u32x8(q + 3 * step, h);
}
{
u32 *p = p0 + 8;
u32 *q = q0 + 8;
auto a = load_u32x8(p), b = load_u32x8(p + step), c = load_u32x8(p + step * 2),
d = load_u32x8(p + step * 3);
auto e = load_u32x8(q), f = load_u32x8(q + step), g = load_u32x8(q + step * 2),
h = load_u32x8(q + step * 3);
u32x8 A1, B1, C1, D1, A2, B2, C2, D2;
if constexpr (trivial) {
butterfly_pair<false, true>(a, c, w1, w1, mts);
butterfly_pair<false, true>(b, d, w1, w1, mts);
butterfly_pair<false, true>(e, g, w1, w1, mts);
butterfly_pair<false, true>(f, h, w1, w1, mts);
butterfly_pair<false, true>(a, b, w2, w2, mts);
butterfly_pair<false, false>(c, d, w3, w3, mts);
butterfly_pair<false, true>(e, f, w2, w2, mts);
butterfly_pair<false, false>(g, h, w3, w3, mts);
} else {
auto cc = mul_pre(c, w1, mts), bb = mul_pre(b, w2, mts), dd = mul_pre(d, w3, mts);
auto gg = mul_pre(g, w1, mts), ff = mul_pre(f, w2, mts), hh = mul_pre(h, w3, mts);
A1 = a + cc, C1 = a + mts.mod2 - cc, B1 = bb + dd;
D1 = mul_pre(bb + mts.mod2 - dd, root, mts);
A2 = e + gg, C2 = e + mts.mod2 - gg, B2 = ff + hh;
D2 = mul_pre(ff + mts.mod2 - hh, root, mts);
a = A1 + B1;
b = A1 + mts.mod2 + mts.mod2 - B1;
c = C1 + D1;
d = C1 + mts.mod2 - D1;
e = A2 + B2;
f = A2 + mts.mod2 + mts.mod2 - B2;
g = C2 + D2;
h = C2 + mts.mod2 - D2;
}
store_u32x8(p, a);
store_u32x8(p + step, b);
store_u32x8(p + 2 * step, c);
store_u32x8(p + 3 * step, d);
store_u32x8(q, e);
store_u32x8(q + step, f);
store_u32x8(q + 2 * step, g);
store_u32x8(q + 3 * step, h);
}
}
} else {
for (int j = 0; j < step; j += 8) {
u32 *p = const_cast<u32 *>(base) + j;
u32 *q = const_cast<u32 *>(base2) + j;
if constexpr (k >= 13) {
if ((j & 15) == 0) { // [Z2R8D-PD] same coverage, longer lead: 512 elements (32 lines) instead of 32 (2 lines)
_mm_prefetch((const char *)(p + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(p + step + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(p + 2 * step + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(p + 3 * step + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(q + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(q + step + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(q + 2 * step + 512), _MM_HINT_T0);
_mm_prefetch((const char *)(q + 3 * step + 512), _MM_HINT_T0);
}
}
auto a = load_u32x8(p), b = load_u32x8(p + step), c = load_u32x8(p + step * 2),
d = load_u32x8(p + step * 3);
auto e = load_u32x8(q), f = load_u32x8(q + step), g = load_u32x8(q + step * 2),
h = load_u32x8(q + step * 3);
u32x8 A1, B1, C1, D1, A2, B2, C2, D2;
if constexpr (trivial) {
butterfly_pair<false, true>(a, c, w1, w1, mts);
butterfly_pair<false, true>(b, d, w1, w1, mts);
butterfly_pair<false, true>(e, g, w1, w1, mts);
butterfly_pair<false, true>(f, h, w1, w1, mts);
butterfly_pair<false, true>(a, b, w2, w2, mts);
butterfly_pair<false, false>(c, d, w3, w3, mts);
butterfly_pair<false, true>(e, f, w2, w2, mts);
butterfly_pair<false, false>(g, h, w3, w3, mts);
} else {
auto cc = mul_pre(c, w1, mts), bb = mul_pre(b, w2, mts), dd = mul_pre(d, w3, mts);
auto gg = mul_pre(g, w1, mts), ff = mul_pre(f, w2, mts), hh = mul_pre(h, w3, mts);
A1 = a + cc, C1 = a + mts.mod2 - cc, B1 = bb + dd;
D1 = mul_pre(bb + mts.mod2 - dd, root, mts);
A2 = e + gg, C2 = e + mts.mod2 - gg, B2 = ff + hh;
D2 = mul_pre(ff + mts.mod2 - hh, root, mts);
a = A1 + B1;
b = A1 + mts.mod2 + mts.mod2 - B1;
c = C1 + D1;
d = C1 + mts.mod2 - D1;
e = A2 + B2;
f = A2 + mts.mod2 + mts.mod2 - B2;
g = C2 + D2;
h = C2 + mts.mod2 - D2;
}
store_u32x8(p, a);
store_u32x8(p + step, b);
store_u32x8(p + 2 * step, c);
store_u32x8(p + 3 * step, d);
store_u32x8(q, e);
store_u32x8(q + step, f);
store_u32x8(q + 2 * step, g);
store_u32x8(q + 3 * step, h);
}
}
}
// Top-level paired forward butterfly with A in four aligned limbs.
template <int k, bool trivial = false>
__attribute__((always_inline)) inline void transform_fixed_pair_split(int i, u32 *data, u32 *limb0, u32 *limb1, u32 *limb2, u32 *limb3,
const MontgomeryAVX2 &mts) const {
const auto &tw = fixed_roots::forward.data[unsigned(i) >> (k + 2)];
u32x8 w1 = (u32x8)_mm256_set1_epi32(tw.w1), w2 = (u32x8)_mm256_set1_epi32(tw.w2),
w3 = (u32x8)_mm256_set1_epi32(tw.w3);
u32x8 root = (u32x8)_mm256_set1_epi32(42958308);
u32x8 root_n = (u32x8)_mm256_set1_epi32(3035660828);
if constexpr (trivial) {
w3 = root;
}
const int step = 1 << k;
const u32 *base = data + i - data_origin;
for (int j = 0; j < step; j += 8) {
u32 *p = const_cast<u32 *>(base) + j;
if ((j & 15) == 0 && j + 128 <= step - 8) { // [c2a] same coverage fix as transform_fixed:
// j += 8 here, so &31 fired every 32 elements = every other 64B line (50% bare).
_mm_prefetch((const char *)(p + 128), _MM_HINT_T0);
_mm_prefetch((const char *)(p + step + 128), _MM_HINT_T0);
_mm_prefetch((const char *)(p + 2 * step + 128), _MM_HINT_T0);
_mm_prefetch((const char *)(p + 3 * step + 128), _MM_HINT_T0);
_mm_prefetch((const char *)(limb0 + j + 128), _MM_HINT_T0);
_mm_prefetch((const char *)(limb1 + j + 128), _MM_HINT_T0);
_mm_prefetch((const char *)(limb2 + j + 128), _MM_HINT_T0);
_mm_prefetch((const char *)(limb3 + j + 128), _MM_HINT_T0);
}
auto a = load_u32x8(p), b = load_u32x8(p + step), c = load_u32x8(p + step * 2),
d = load_u32x8(p + step * 3);
auto e = load_u32x8(limb0 + j), f = load_u32x8(limb1 + j), g = load_u32x8(limb2 + j),
h = load_u32x8(limb3 + j);
u32x8 A1, B1, C1, D1, A2, B2, C2, D2;
if constexpr (trivial) {
butterfly_pair<false, true>(a, c, w1, w1, mts);
butterfly_pair<false, true>(b, d, w1, w1, mts);
butterfly_pair<false, true>(e, g, w1, w1, mts);
butterfly_pair<false, true>(f, h, w1, w1, mts);
butterfly_pair<false, true>(a, b, w2, w2, mts);
butterfly_pair<false, false>(c, d, w3, w3, mts);
butterfly_pair<false, true>(e, f, w2, w2, mts);
butterfly_pair<false, false>(g, h, w3, w3, mts);
} else {
auto cc = mul_pre(c, w1, mts), bb = mul_pre(b, w2, mts), dd = mul_pre(d, w3, mts);
auto gg = mul_pre(g, w1, mts), ff = mul_pre(f, w2, mts), hh = mul_pre(h, w3, mts);
A1 = a + cc, C1 = a + mts.mod2 - cc, B1 = bb + dd;
D1 = mul_pre(bb + mts.mod2 - dd, root, mts);
A2 = e + gg, C2 = e + mts.mod2 - gg, B2 = ff + hh;
D2 = mul_pre(ff + mts.mod2 - hh, root, mts);
a = A1 + B1;
b = A1 + mts.mod2 + mts.mod2 - B1;
c = C1 + D1;
d = C1 + mts.mod2 - D1;
e = A2 + B2;
f = A2 + mts.mod2 + mts.mod2 - B2;
g = C2 + D2;
h = C2 + mts.mod2 - D2;
}
store_u32x8(p, a);
store_u32x8(p + step, b);
store_u32x8(p + 2 * step, c);
store_u32x8(p + 3 * step, d);
store_u32x8(limb0 + j, e);
store_u32x8(limb1 + j, f);
store_u32x8(limb2 + j, g);
store_u32x8(limb3 + j, h);
}
}
template <bool inverse, bool trivial = false>
void transform_stage(int k, int i, u32 *data, u64x4 &wi, const MontgomeryAVX2 &mts) const {
u32x8 w1 = (u32x8)_mm256_shuffle_epi32((i256)wi, 0b00'00'00'00);
u32x8 w2 = (u32x8)_mm256_permute4x64_epi64((i256)wi, 0b01'01'01'01); // only even indices will be used
u32x8 w3 = (u32x8)_mm256_permute4x64_epi64((i256)wi, 0b11'11'11'11); // only even indices will be used
u32x8 n1 = (u32x8)_mm256_mul_epu32((i256)w1, (i256)mts.n_inv);
u32x8 n2 = (u32x8)_mm256_mul_epu32((i256)w2, (i256)mts.n_inv);
u32x8 n3 = (u32x8)_mm256_mul_epu32((i256)w3, (i256)mts.n_inv);
for (int j = 0; j < (1 << k); j += 8) {
u32 *p = data + i + j;
int step = 1 << k;
auto a = load_u32x8(p), b = load_u32x8(p + step), c = load_u32x8(p + step * 2),
d = load_u32x8(p + step * 3);
if constexpr (!inverse) {
butterfly_pair<false, trivial>(a, c, w1, w1, mts);
butterfly_pair<false, trivial>(b, d, w1, w1, mts);
butterfly_pair<false, trivial>(a, b, w2, w2, mts);
butterfly_pair<false, false>(c, d, w3, w3, mts);
} else {
if constexpr (trivial) {
butterfly_pair<true, true>(a, b, w2, w2, mts);
butterfly_pair<true, false>(c, d, w3, w3, mts);
butterfly_pair<true, true>(a, c, w1, w1, mts);
butterfly_pair<true, true>(b, d, w1, w1, mts);
} else {
auto u = a + b, v = mul_pre(a + mts.mod2 - b, w2, mts);
auto s = c + d, t = mul_pre(c + mts.mod2 - d, w3, mts);
auto sum = u + s;
sum = (u32x8)_mm256_min_epu32((i256)sum, (i256)(sum - mts.mod2 - mts.mod2));
a = mts.shrink2(sum);
b = mts.shrink2(v + t);
c = mul_pre(u + mts.mod2 + mts.mod2 - s, w1, mts);
d = mul_pre(v + mts.mod2 - t, w1, mts);
}
}
store_u32x8(p, a);
store_u32x8(p + step, b);
store_u32x8(p + 2 * step, c);
store_u32x8(p + 3 * step, d);
}
wi = mts.mul_u64x4<true>(wi, (inverse ? wrd_x4 : wd_x4)[__builtin_ctz(~i >> k + 2)]);
}
public:
// Generic forward transform; data is 32-byte aligned.
// Lazy residues grow through the stages; the product kernel normalizes them.
void transform_forward(int lg, u32 *data) const {
const MontgomeryAVX2 mts = this->mts;
const int L = leaf_log2(lg);
if (L < lg) {
const int lc = (lg - L) / 2;
u64x4 wi_data[LG / 2];
std::fill(wi_data, wi_data + lc, wt_init);
for (int k = lg - 2; k >= L; k -= 2) {
transform_stage<false, true>(k, 0, data, wi_data[k - L >> 1], mts);
}
for (int i = 1; i < (1 << lc * 2 - 2); i++) {
int s = __builtin_ctz(i) >> 1;
for (int k = s; k >= 0; k--) {
transform_stage<false>(2 * k + L, i * (1 << L + 2), data, wi_data[k], mts);
}
}
}
}
// input in [0, 2 * mod)
// output in [0, mod)
// data must be 32-byte aligned
template <bool mul_by_sc = false>
void transform_inverse(int lg, u32 *data, /* as normal number */ u32 sc = u32()) const {
const MontgomeryAVX2 mts = this->mts;
const int L = leaf_log2(lg);
if (L < lg) {
const int lc = (lg - L) / 2;
u64x4 wi_data[LG / 2];
std::fill(wi_data, wi_data + lc, wrt_init);
for (int i = 0; i < (1 << lc * 2 - 2); i++) {
int s = __builtin_ctz(~i) >> 1;
if (i + 1 == (1 << 2 * s)) {
s--;
}
for (int k = 0; k <= s; k++) {
transform_stage<true>(2 * k + L, (i + 1 - (1 << 2 * k)) * (1 << L + 2), data, wi_data[k], mts);
}
if (i + 1 == (1 << 2 * (s + 1))) {
s++;
transform_stage<true, true>(2 * s + L, (i + 1 - (1 << 2 * s)) * (1 << L + 2), data, wi_data[s],
mts);
}
}
}
const Montgomery mt = this->mt;
u32 f = mt.power<false, true>((mod + 1) >> 1, lg - L);
if (mul_by_sc) f = mt.mul<true>(f, mt.mul<false>(mt.r2, sc));
u32x8 f_x8 = (u32x8)_mm256_set1_epi32(f);
for (int i = 0; i < (1 << lg); i += 8) {
store_u32x8(data + i, mts.mul_u32x8<true, true>(load_u32x8(data + i), f_x8));
}
}
private:
// Multiply modulo x^(2^L)-w. Normalize lazy inputs; output is below 2p.
// At L=3 each sum <= 128*p*p, so Montgomery reduction gives <4p.
// O3 and the memory-operand multiply/accumulate are performance-critical.
template <int L, int K, bool remove_montgomery_reduction_factor = true>
__attribute__((optimize("O3,unroll-loops,rename-registers"))) static void
multiply_leaf(const u32 *a, const u32 *b, u32 *c, const std::array<u32x8, K> &ar_w, const MontgomeryAVX2 &mts) {
static_assert(L >= 3);
constexpr int n = 1 << L;
alignas(64) u32 aux_a[n];
alignas(64) u64 aux_b[n * 2];
#pragma GCC unroll 1
for (int k = 0; k < K; ++k) {
for (int i = 0; i < n; i += 8) {
u32x8 ai = load_u32x8(a + n * k + i);
if (remove_montgomery_reduction_factor) {
ai = mts.mul_u32x8<true, true>(ai, mts.r2);
} else {
ai = mts.canonical_wide<L == 3 ? 2 : 0>(ai);
}
store_u32x8(aux_a + i, ai);
u32x8 bi = load_u32x8(b + n * k + i);
u32x8 bi_0 = mts.canonical_wide<L == 3 ? 2 : 0>(bi);
u32x8 bi_w = mts.mul_u32x8<true, true>(bi, ar_w[k]);
store_u32x8((u32 *)(aux_b + i + 0),
(u32x8)_mm256_permutevar8x32_epi32((i256)bi_w, _mm256_setr_epi64x(0, 1, 2, 3)));
store_u32x8((u32 *)(aux_b + i + 4),
(u32x8)_mm256_permutevar8x32_epi32((i256)bi_w, _mm256_setr_epi64x(4, 5, 6, 7)));
store_u32x8((u32 *)(aux_b + n + i + 0),
(u32x8)_mm256_permutevar8x32_epi32((i256)bi_0, _mm256_setr_epi64x(0, 1, 2, 3)));
store_u32x8((u32 *)(aux_b + n + i + 4),
(u32x8)_mm256_permutevar8x32_epi32((i256)bi_0, _mm256_setr_epi64x(4, 5, 6, 7)));
}
u64x4 aux_ans[n / 4];
{
constexpr int i = 0;
u64x4 ai = (u64x4)_mm256_set1_epi32(aux_a[i]);
u64x4 ai1 = (u64x4)_mm256_set1_epi32(aux_a[i + 1]);
for (int j = 0; j < n; j += 4) {
u64x4 t0, t1;
asm("vpmuludq %2,%1,%0"
: "=&x"(t0)
: "x"(ai), "m"(*(const __m256i_u *)(aux_b + n - i + j)));
asm("vpmuludq %3,%2,%1\n\tvpaddq %1,%0,%0"
: "+x"(t0), "=&x"(t1)
: "x"(ai1), "m"(*(const __m256i_u *)(aux_b + n - i - 1 + j)));
aux_ans[j / 4] = t0;
}
}
for (int i = 2; i + 2 <= n; i += 2) {
u64x4 ai = (u64x4)_mm256_set1_epi32(aux_a[i]);
u64x4 ai1 = (u64x4)_mm256_set1_epi32(aux_a[i + 1]);
for (int j = 0; j < n; j += 4) {
u64x4 t0, t1;
asm("vpmuludq %3,%2,%1\n\tvpaddq %1,%0,%0"
: "+x"(aux_ans[j / 4]), "=&x"(t0)
: "x"(ai), "m"(*(const __m256i_u *)(aux_b + n - i + j)));
asm("vpmuludq %3,%2,%1\n\tvpaddq %1,%0,%0"
: "+x"(aux_ans[j / 4]), "=&x"(t1)
: "x"(ai1), "m"(*(const __m256i_u *)(aux_b + n - i - 1 + j)));
}
if (((i + 1) & 7) == 7 && i + 1 >= 15) {
for (int j = 0; j < n; j += 4) {
aux_ans[j / 4] = (u64x4)mts.shrink2((u32x8)aux_ans[j / 4]);
}
}
}
__asm__ __volatile__("" ::: "memory");
for (int i = 0; i < n; i += 8) {
u64x4 c0 = aux_ans[i / 4], c1 = aux_ans[i / 4 + 1];
u32x8 res = (u32x8)_mm256_permutevar8x32_epi32((i256)mts.reduce<false>(c0, c1),
_mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7));
store_u32x8(c + k * n + i, mts.shrink2(res));
}
}
}
template <int L, bool remove_montgomery_reduction_factor = true>
void multiply_leaves(int lg, const u32 *a, const u32 *b, u32 *c) const {
constexpr int sz = 1 << L;
const MontgomeryAVX2 mts = this->mts;
int cnt = 1 << lg - L;
if (cnt == 1) {
multiply_leaf<L, 1, remove_montgomery_reduction_factor>(a, b, c, {mts.r}, mts);
return;
}
if (cnt <= 8) {
for (int i = 0; i < cnt; i += 2) {
u32x8 wi = (u32x8)_mm256_set1_epi32(w[i / 2]);
multiply_leaf<L, 2, remove_montgomery_reduction_factor>(a + i * sz, b + i * sz, c + i * sz,
{wi, (mts.mod - wi)}, mts);
}
return;
}
u64x4 wi = wl_init;
for (int i = 0; i < cnt; i += 8) {
u32x8 w_ar[4] = {
(u32x8)_mm256_permute4x64_epi64((i256)wi, 0b00'00'00'00),
(u32x8)_mm256_permute4x64_epi64((i256)wi, 0b01'01'01'01),
(u32x8)_mm256_permute4x64_epi64((i256)wi, 0b10'10'10'10),
(u32x8)_mm256_permute4x64_epi64((i256)wi, 0b11'11'11'11),
};
if (L == L0) {
for (int j = 0; j < 8; j += 4) {
multiply_leaf<L, 4, remove_montgomery_reduction_factor>(
a + (i + j) * sz, b + (i + j) * sz, c + (i + j) * sz,
{w_ar[j / 2], mts.mod - w_ar[j / 2], w_ar[j / 2 + 1], mts.mod - w_ar[j / 2 + 1]}, mts);
}
} else {
for (int j = 0; j < 8; j += 2) {
multiply_leaf<L, 2, remove_montgomery_reduction_factor>(a + (i + j) * sz, b + (i + j) * sz,
c + (i + j) * sz,
{w_ar[j / 2], mts.mod - w_ar[j / 2]}, mts);
}
}
wi = mts.mul_u64x4<true>(wi, wld_x4[__builtin_ctz(~i >> 3)]);
}
}
public:
// Leaf products: normalize lazy inputs and return residues below 2p.
template <bool remove_montgomery_reduction_factor = true>
void multiply_all_leaves(int lg, const u32 *a, const u32 *b, u32 *c) const {
int L = leaf_log2(lg);
if (L == L0) {
multiply_leaves<L0, remove_montgomery_reduction_factor>(lg, a, b, c);
} else {
multiply_leaves<L0 + 1, remove_montgomery_reduction_factor>(lg, a, b, c);
}
}
template <int L> void multiply_range(int begin, int end, u32 *a, u32 *b, u64x4 &wi) const {
const MontgomeryAVX2 mts = this->mts;
// [K02ZC-ORG] hoist the (loop-invariant) member: the inline-asm store macros clobber
// memory, so GCC reloads `this->data_origin` from memory on every single iteration.
// A local copy is value-neutral here (nothing inside the loop mutates it).
const int k02zc_org = data_origin;
constexpr int sz = 1 << L;
for (int i = begin >> L; i < (end >> L); i += 8) {
u32x8 w_ar[4];
if constexpr (L == 3) {
const auto &tw = fixed_roots::dot.data[i >> 3];
for (int q = 0; q < 4; q++) w_ar[q] = (u32x8)_mm256_set1_epi32(tw[q]);
} else {
w_ar[0] = (u32x8)_mm256_permute4x64_epi64((i256)wi, 0x00);
w_ar[1] = (u32x8)_mm256_permute4x64_epi64((i256)wi, 0x55);
w_ar[2] = (u32x8)_mm256_permute4x64_epi64((i256)wi, 0xaa);
w_ar[3] = (u32x8)_mm256_permute4x64_epi64((i256)wi, 0xff);
}
if constexpr (L == L0) {
for (int j = 0; j < 8; j += 4)
multiply_leaf<L, 4, false>(
a + (i + j) * sz - k02zc_org, b + (i + j) * sz - k02zc_org, a + (i + j) * sz - k02zc_org,
{w_ar[j / 2], mts.mod - w_ar[j / 2], w_ar[j / 2 + 1], mts.mod - w_ar[j / 2 + 1]}, mts);
} else {
for (int j = 0; j < 8; j += 2)
multiply_leaf<L, 2, false>(a + (i + j) * sz - k02zc_org, b + (i + j) * sz - k02zc_org,
a + (i + j) * sz - k02zc_org, {w_ar[j / 2], mts.mod - w_ar[j / 2]},
mts);
}
if constexpr (L != 3) wi = mts.mul_u64x4<true>(wi, wld_x4[__builtin_ctz(~i >> 3)]);
}
}
template <int L> void convolve_block(int lg, int offset, u32 *a, u32 *b, u64x4 *fw, u64x4 *iw, u64x4 &dot) const {
const MontgomeryAVX2 mts = this->mts;
if (lg > 11) {
int k = lg - 2;
u64x4 w = fw[k];
if (offset == 0) {
transform_stage<false, true>(k, offset, a, w, mts);
transform_stage<false, true>(k, offset, b, fw[k], mts);
} else {
transform_stage<false>(k, offset, a, w, mts);
transform_stage<false>(k, offset, b, fw[k], mts);
}
for (int j = 0; j < 4; j++) convolve_block<L>(k, offset + (j << k), a, b, fw, iw, dot);
if (offset == 0)
transform_stage<true, true>(k, offset, a, iw[k], mts);
else
transform_stage<true>(k, offset, a, iw[k], mts);
return;
}
int end = offset + (1 << lg);
for (int k = lg - 2; k >= L; k -= 2) {
for (int i = offset; i < end; i += (1 << (k + 2))) {
u64x4 w = fw[k];
if (i == 0) {
transform_stage<false, true>(k, i, a, w, mts);
transform_stage<false, true>(k, i, b, fw[k], mts);
} else {
transform_stage<false>(k, i, a, w, mts);
transform_stage<false>(k, i, b, fw[k], mts);
}
}
}
multiply_range<L>(offset, end, a, b, dot);
for (int k = L; k <= lg - 2; k += 2)
for (int i = offset; i < end; i += (1 << (k + 2))) {
if (i == 0)
transform_stage<true, true>(k, i, a, iw[k], mts);
else
transform_stage<true>(k, i, a, iw[k], mts);
}
}
template <int K, bool Inv, int Block> __attribute__((noinline)) void transform_block(int offset, u32 *a, u32 *b) const {
const MontgomeryAVX2 mts;
for (int i = offset; i < offset + Block; i += (1 << (K + 2))) {
if constexpr (Inv) {
if (i == 0)
transform_fixed<K, true, true>(i, a, mts);
else
transform_fixed<K, true>(i, a, mts);
} else {
if (i == 0)
transform_fixed_pair<K, true>(i, a, b, mts);
else
transform_fixed_pair<K>(i, a, b, mts);
}
}
}
template <int LG, bool FWD_A = true, bool FWD_B = true, bool LEAF = true, bool INV = true>
void convolve_fixed(int offset, u32 *a, u32 *b) const {
const MontgomeryAVX2 mts;
if constexpr (LG > 11) {
constexpr int K = LG - 2;
if constexpr (FWD_A && FWD_B) {
if (offset == 0)
transform_fixed_pair<K, true>(offset, a, b, mts);
else
transform_fixed_pair<K>(offset, a, b, mts);
} else {
if constexpr (FWD_A) {
if (offset == 0)
transform_fixed<K, false, true>(offset, a, mts);
else
transform_fixed<K, false>(offset, a, mts);
}
if constexpr (FWD_B) {
if (offset == 0)
transform_fixed<K, false, true>(offset, b, mts);
else
transform_fixed<K, false>(offset, b, mts);
}
}
for (int j = 0; j < 4; j++) convolve_fixed<K, FWD_A, FWD_B, LEAF, INV>(offset + (j << K), a, b);
if constexpr (INV) {
if (offset == 0)
transform_fixed<K, true, true>(offset, a, mts);
else
transform_fixed<K, true>(offset, a, mts);
}
} else {
constexpr int Block = 1 << LG;
if constexpr (FWD_A || FWD_B) {
transform_block<9, false, Block>(offset, a, b);
transform_block<7, false, Block>(offset, a, b);
transform_block<5, false, Block>(offset, a, b);
transform_block<3, false, Block>(offset, a, b);
}
if constexpr (LEAF) {
u64x4 unused_root_an2{};
multiply_range<3>(offset, offset + Block, a, b, unused_root_an2);
}
if constexpr (INV) {
transform_block<3, true, Block>(offset, a, b);
transform_block<5, true, Block>(offset, a, b);
transform_block<7, true, Block>(offset, a, b);
transform_block<9, true, Block>(offset, a, b);
}
}
}
void convolve_final_split_quarter(int offset, u32 *a0, u32 *a1, u32 *a2, u32 *a3,
u32 *bscratch, u32 *final_dst) const {
constexpr int step = 1 << 17;
const MontgomeryAVX2 mts;
data_origin = offset;
// Both input arrays undergo the same forward butterfly; the first
// argument happens to be B so the split A limbs are second.
transform_fixed_pair_split<17>(offset, bscratch, a0, a1, a2, a3, mts);
u32 *alimb[4] = {a0, a1, a2, a3};
for (int j = 0; j < 4; j++) {
data_origin = offset + j * step;
convolve_fixed<17>(data_origin, alimb[j], bscratch + j * step);
}
data_origin = offset;
const auto &tw = fixed_roots::inverse.data[unsigned(offset) >> 19];
const u32x8 w1 = (u32x8)_mm256_set1_epi32(tw.w1), w2 = (u32x8)_mm256_set1_epi32(tw.w2),
w3 = (u32x8)_mm256_set1_epi32(tw.w3);
const u32x8 root = (u32x8)_mm256_set1_epi32(38830621);
const u32x8 root_n = (u32x8)_mm256_set1_epi32(1259306467);
for (int j = 0; j < step; j += 8) {
auto a = load_u32x8(a0 + j), b = load_u32x8(a1 + j);
auto c = load_u32x8(a2 + j), d = load_u32x8(a3 + j);
auto u = a + b, s = c + d, v = a + mts.mod2 - b;
auto t = mul_pre(c + mts.mod2 - d, root, mts);
auto sum = u + s;
sum = (u32x8)_mm256_min_epu32((i256)sum, (i256)(sum - mts.mod2 - mts.mod2));
a = mts.shrink2(sum);
c = mul_pre(u + mts.mod2 + mts.mod2 - s, w1, mts);
b = mul_pre(v + t, w2, mts);
d = mul_pre(v + mts.mod2 - t, w3, mts);
STU(final_dst + j, a);
STU(final_dst + step + j, b);
STU(final_dst + 2 * step + j, c);
STU(final_dst + 3 * step + j, d);
}
}
// Source A is read exactly once. Stores to the upper A limb begin only
// after its original low-index source has already been consumed.
void prepare_quarters_split_last(const u32 *src, int n, u32 *dst,
u32 *last0, u32 *last3, int lg) const {
auto stream = [](u32 *p, u32x8 x) { _mm256_stream_si256((i256 *)p, (i256)x); };
const int q = 1 << (lg - 2), split = 3 << 17;
alignas(32) u32 table[16];
u32 root = mt.mul(w[1], 1);
for (int i = 0; i < 16; i++) table[i] = u64(root) * i % mod;
u32x8 t0 = load_u32x8(table), t1 = load_u32x8(table + 8);
auto run = [&](int i, u32 *last, u32x8 a, u32x8 b) {
u32x8 v = (u32x8)_mm256_blendv_epi8(_mm256_permutevar8x32_epi32((i256)t0, (i256)b),
_mm256_permutevar8x32_epi32((i256)t1, (i256)b),
_mm256_cmpgt_epi32((i256)b, _mm256_set1_epi32(7)));
stream(dst + i, a + b);
stream(dst + q + i, a + mts.mod2 - b);
stream(dst + 2 * q + i, a + v);
stream(last, a + mts.mod2 - v);
};
int i = 0;
for (; i < split; i += 8)
run(i, last0 + i, (u32x8)LDU(src + i), (u32x8)LDU(src + q + i));
for (; i + 8 <= n - q; i += 8)
run(i, last3 + i - split, (u32x8)LDU(src + i), (u32x8)LDU(src + q + i));
if (i < n - q) {
alignas(32) u32 tail[8] = {};
memcpy(tail, src + q + i, (n - q - i) * 4);
run(i, last3 + i - split, (u32x8)LDU(src + i), load_u32x8(tail));
i += 8;
}
for (; i < q; i += 8)
run(i, last3 + i - split, (u32x8)LDU(src + i), u32x8{});
}
void prepare_quarters(const u32 *src, int n, u32 *dst, u32 *last, int lg) const {
auto stream = [](u32 *p, u32x8 x) { _mm256_stream_si256((i256 *)p, (i256)x); };
const int q = 1 << (lg - 2);
alignas(32) u32 table[16];
u32 root = mt.mul(w[1], 1);
for (int i = 0; i < 16; i++) table[i] = u64(root) * i % mod;
u32x8 t0 = load_u32x8(table), t1 = load_u32x8(table + 8);
auto run = [&](int i, u32x8 a, u32x8 b) {
u32x8 v = (u32x8)_mm256_blendv_epi8(_mm256_permutevar8x32_epi32((i256)t0, (i256)b),
_mm256_permutevar8x32_epi32((i256)t1, (i256)b),
_mm256_cmpgt_epi32((i256)b, _mm256_set1_epi32(7)));
stream(dst + i, a + b);
stream(dst + q + i, a + mts.mod2 - b);
stream(dst + 2 * q + i, a + v);
stream(last + i, a + mts.mod2 - v);
};
int i = 0;
for (; i + 8 <= n - q; i += 8)
run(i, (u32x8)LDU((src + i)),
(u32x8)LDU((src + q + i)));
if (i < n - q) {
alignas(32) u32 tail[8] = {};
memcpy(tail, src + q + i, (n - q - i) * 4);
run(i, (u32x8)LDU((src + i)), load_u32x8(tail));
i += 8;
}
for (; i < q; i += 8) run(i, (u32x8)LDU((src + i)), u32x8{});
}
// i2_t2w: 把 B 的 4 个 quarter 预处理所需的两个系数流打成 1 字节/下标
// (低半字节 = B[i]，高半字节 = B[q+i])，四个预处理各自只需再读 512 KiB 而不是 4 MiB。
// 系数 < 16（题面保证 < 10；原实现里 v 的查表也已经隐含 b < 16）。返回 false 表示
// 某个系数塞不下半字节，调用方回退到未打包路径。
static void store8(unsigned char *d, u32x8 z) {
i256 p = _mm256_packus_epi32((i256)z, (i256)z);
i256 q2 = _mm256_packus_epi16(p, p);
__m128i lo = _mm256_castsi256_si128(q2), hi = _mm256_extracti128_si256(q2, 1);
_mm_storel_epi64((__m128i *)d, _mm_unpacklo_epi32(lo, hi));
}
static bool pack_B(const u32 *src, int n, unsigned char *dst, u32 *first, int lg) {
const int q = 1 << (lg - 2);
int i = 0;
for (; i + 8 <= n - q; i += 8) {
u32x8 x = (u32x8)LDU(src + i), y = (u32x8)LDU(src + q + i);
store8(dst + i, x | (y << 4));
STU(first + i, x + y);
}
if (i < n - q) {
alignas(32) u32 t[8] = {};
int cnt = n - q - i;
for (int k = 0; k < cnt; k++) t[k] = src[q + i + k];
u32x8 x = (u32x8)LDU(src + i), y = load_u32x8(t);
store8(dst + i, x | (y << 4));
STU(first + i, x + y);
i += 8;
}
for (; i + 8 <= q; i += 8) {
u32x8 x = (u32x8)LDU(src + i);
store8(dst + i, x);
STU(first + i, x);
}
for (; i < q; i++) {
unsigned v = src[i];
dst[i] = (unsigned char)v;
first[i] = v;
}
return true;
}
template <int Quarter> void prepare_quarter_packed(const unsigned char *bp, u32 *dst, int lg) const {
const int q = 1 << (lg - 2);
alignas(32) u32 table[16];
u32 root = mt.mul(w[1], 1);
for (int i = 0; i < 16; i++) table[i] = u64(root) * i % mod;
u32x8 t0 = load_u32x8(table), t1 = load_u32x8(table + 8);
const i256 m15 = _mm256_set1_epi32(15);
auto run = [&](int i, u32x8 a, u32x8 b) {
u32x8 v = (u32x8)_mm256_blendv_epi8(_mm256_permutevar8x32_epi32((i256)t0, (i256)b),
_mm256_permutevar8x32_epi32((i256)t1, (i256)b),
_mm256_cmpgt_epi32((i256)b, _mm256_set1_epi32(7)));
if constexpr (Quarter == 0) store_u32x8(dst + i, a + b);
if constexpr (Quarter == 1) store_u32x8(dst + i, a + mts.mod2 - b);
if constexpr (Quarter == 2) store_u32x8(dst + i, a + v);
if constexpr (Quarter == 3) store_u32x8(dst + i, a + mts.mod2 - v);
};
for (int i = 0; i < q; i += 8) {
i256 x = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i *)(bp + i)));
run(i, (u32x8)(x & m15), (u32x8)_mm256_srli_epi32(x, 4));
}
}
template <int Quarter> void prepare_quarter(const u32 *src, int n, u32 *dst, int lg) const {
const int q = 1 << (lg - 2);
alignas(32) u32 table[16];
u32 root = mt.mul(w[1], 1);
for (int i = 0; i < 16; i++) table[i] = u64(root) * i % mod;
u32x8 t0 = load_u32x8(table), t1 = load_u32x8(table + 8);
auto run = [&](int i, u32x8 a, u32x8 b) {
u32x8 v = (u32x8)_mm256_blendv_epi8(_mm256_permutevar8x32_epi32((i256)t0, (i256)b),
_mm256_permutevar8x32_epi32((i256)t1, (i256)b),
_mm256_cmpgt_epi32((i256)b, _mm256_set1_epi32(7)));
if constexpr (Quarter == 0) store_u32x8(dst + i, a + b);
if constexpr (Quarter == 1) store_u32x8(dst + i, a + mts.mod2 - b);
if constexpr (Quarter == 2) store_u32x8(dst + i, a + v);
if constexpr (Quarter == 3) store_u32x8(dst + i, a + mts.mod2 - v);
};
int i = 0;
for (; i + 8 <= n - q; i += 8)
run(i, (u32x8)LDU((src + i)),
(u32x8)LDU((src + q + i)));
if (i < n - q) {
alignas(32) u32 tail[8] = {};
memcpy(tail, src + q + i, (n - q - i) * 4);
run(i, (u32x8)LDU((src + i)), load_u32x8(tail));
i += 8;
}
for (; i < q; i += 8) run(i, (u32x8)LDU((src + i)), u32x8{});
}
// Final inverse butterfly and scaling; the last vector may be partial.
// i2_: fx/froot 是循环不变量 ⇒ 4 个蒙哥马利乘换成 Shoup 乘（约 13~14 uops -> 10 uops）。
template<int Mode = 0> __attribute__((always_inline)) inline void finish_quarters(u32x8 x, u32x8 y, u32x8 z, u32x8 t, u32 *c, int i, int q,
int sz, u32 fxw, u32 fxs, u32 frw, u32 frs) const {
const u32x8 pv = (u32x8)_mm256_set1_epi32((int)mod);
const u32x8 fxv = (u32x8)_mm256_set1_epi32((int)fxw), fxsv = (u32x8)_mm256_set1_epi32((int)fxs);
const u32x8 frv = (u32x8)_mm256_set1_epi32((int)frw), frsv = (u32x8)_mm256_set1_epi32((int)frs);
const auto xy=x+y,zt=z+t;
auto v = mts.shrink(i2_shoup32_c(x + mts.mod2 - y, fxv, fxsv, pv));
auto r = mts.shrink(i2_shoup32_c(z + mts.mod2 - t, frv, frsv, pv));
x = mts.shrink(i2_shoup32_c(xy + zt, fxv, fxsv, pv));
z = mts.shrink(i2_shoup32_c(xy + mts.mod2 + mts.mod2 - zt, fxv, fxsv, pv));
y = mts.shrink(v + r);
if constexpr (Mode != 2) t = mts.shrink(v + mts.mod - r);
STU((c + i), x);
STU((c + q + i), y);
STU((c + 2 * q + i), z);
if constexpr (Mode == 1) STU((c + 3 * q + i), t);
else if constexpr (Mode == 0) {
if (i + 3 * q + 8 <= sz) STU((c + 3 * q + i), t);
else if (i + 3 * q < sz) memcpy(c + 3 * q + i, &t, 4 * (sz - 3 * q - i));
}
}
void convolve_inputs(const u32 *A, int n, const u32 *B, int m, u32 *c, int lg, u32 *a, u32 *b) const {
prepare_quarters(A, n, a, a + (3 << (lg - 2)), lg);
prepare_quarters(B, m, b, b + (3 << (lg - 2)), lg);
_mm_sfence();
u64x4 fw[LG], iw[LG], dot = wl_init;
std::fill(fw, fw + LG, wt_init);
std::fill(iw, iw + LG, wrt_init);
int k = lg - 2, L = leaf_log2(lg);
for (int j = 0; j < 4; j++) {
if (lg == 21)
convolve_fixed<19>(j << k, a, b);
else if (L == L0)
convolve_block<L0>(k, j << k, a, b, fw, iw, dot);
else
convolve_block<L0 + 1>(k, j << k, a, b, fw, iw, dot);
}
u32 f = mt.power<false, true>((mod + 1) >> 1, lg - L);
f = mt.mul<true>(f, mt.mul<false>(mt.r2, mt.r));
u32 fr = mt.mul(f, wr[1]);
// i2_: 折回 plain + 预算 Shoup 位移（每个进程只算一次，不在循环里）
u32 fxw = (u32)((u64)f * rinv % mod), fxs = (u32)(((u64)fxw << 32) / mod);
u32 frw = (u32)((u64)fr * rinv % mod), frs = (u32)(((u64)frw << 32) / mod);
int q = 1 << k, sz = n + m - 1;
for (int i = 0; i < q; i += 8) {
u32x8 x = load_u32x8(a + i), y = load_u32x8(a + q + i), z = load_u32x8(a + 2 * q + i),
t = load_u32x8(a + 3 * q + i);
finish_quarters(x, y, z, t, c, i, q, sz, fxw, fxs, frw, frs);
}
}
// First three A quarters live in c; one A quarter and one B quarter use
// 4 MiB of scratch. Inputs remain read-only, and c may be unaligned.
void convolve_reusing_output(const u32 *A, const u32 *B, u32 *c) const {
constexpr int lg = 21, k = 19, L = 3, n = 1000001, m = 1000001;
// i2_: 4 MiB -> 2 MiB。`last` 复用 A 的前 2^19 个元素（prepare_quarters 只读 A[i]/A[q+i]
// 再写 last[i]，写区 A[0,q) 与后续仍要读的 A[q,n) 不重叠），少触碰 512 个全新页 ≈ 125 µs。
// last 走 _mm256_stream_si256 ⇒ 必须 32B 对齐；A 不满足时回退到单独 malloc。
const bool split = (((uintptr_t)A & 31u) == 0u) && (A != B);
// q2z_ PRE: scratch 一律走 _mm_malloc(32) 的 temp（∉ A）⇒ pack_B 对 A 只读
// q2z_ BSC2: B 若 32B 对齐则直接用 B[0,2^19) 当 scratch（省掉 2 MB = 512 个新页），
// 否则回退到 _mm_malloc(32) 的 temp。判据 = 与既有 `split` 对 A 的守卫同一形式。
// 依据：`convolve_fixed<19>` 对 quarter 缓冲发 `vmovdqa`（32B 对齐取数）⇒ 缓冲必须 32B 对齐
// （本席实测：把 scratch 指向未对齐的 B ⇒ 判题机探针通道 RE，gdb 定位 `vmovdqa (%rax),%ymm4`）✓
const bool bsc = (((uintptr_t)B & 31u) == 0u) && (B != A);
u32 *temp = bsc ? nullptr : (u32 *)_mm_malloc((1 << 19) * 4, 32);
void *lastOwn = nullptr;
u32 *last = nullptr;
if (!split) {
if (((uintptr_t)A & 31u) == 0u) last = const_cast<u32 *>(A);
else { lastOwn = _mm_malloc((1 << 19) * 4, 32); last = (u32 *)lastOwn; }
}
u32 *scratch = bsc ? const_cast<u32 *>(B) : temp; // 两者都 32B 对齐 ✓
u32 *a = (u32 *)(((uintptr_t)c + 31) & ~uintptr_t(31));
u32 *last0 = a + (3 << 19);
u32 *last3 = const_cast<u32 *>(A) + (1 << 19);
if (split) prepare_quarters_split_last(A, n, a, last0, last3, lg);
else prepare_quarters(A, n, a, last, lg);
// i2_t2w: A 已读完最高区，复用 A[3q, 3q+2^18) 之后的尾部存 B 的打包副本。
unsigned char *packed = (unsigned char *)(const_cast<u32 *>(B) + (1 << 19));
const bool pkb = pack_B(B, m, packed, scratch, lg);
_mm_sfence();
for (int j = 0; j < 4; j++) {
if (pkb) {
/* quarter zero was formed while packing B */
if (j == 1) prepare_quarter_packed<1>(packed, scratch, lg);
if (j == 2) prepare_quarter_packed<2>(packed, scratch, lg);
if (j == 3) prepare_quarter_packed<3>(packed, scratch, lg);
} else {
if (j == 0) prepare_quarter<0>(B, m, scratch, lg);
if (j == 1) prepare_quarter<1>(B, m, scratch, lg);
if (j == 2) prepare_quarter<2>(B, m, scratch, lg);
if (j == 3) prepare_quarter<3>(B, m, scratch, lg);
}
data_origin = j << 19;
if (split && j == 3)
convolve_final_split_quarter(data_origin, last0, last0 + (1 << 17),
last0 + (2 << 17), last3, scratch,
scratch);
else
convolve_fixed<19>(data_origin, j == 3 ? last : a + data_origin, scratch);
}
data_origin = 0;
u32 f = mt.power<false, true>((mod + 1) >> 1, lg - L);
f = mt.mul<true>(f, mt.mul<false>(mt.r2, mt.r));
u32 fr = mt.mul(f, wr[1]);
// i2_: 折回 plain + 预算 Shoup 位移（每个进程只算一次，不在循环里）
u32 fxw = (u32)((u64)f * rinv % mod), fxs = (u32)(((u64)fxw << 32) / mod);
u32 frw = (u32)((u64)fr * rinv % mod), frs = (u32)(((u64)frw << 32) / mod);
int q = 1 << k, sz = n + m - 1;
// c may precede aligned scratch by up to seven coefficients. Writes to
// the next quarter would overwrite these tails before their final read.
alignas(32) u32x8 saved[3] = {load_u32x8(a + q - 8), load_u32x8(a + 2 * q - 8), load_u32x8(a + 3 * q - 8)};
auto run_finish = [&](int i, auto mode) {
if ((i & 15) == 0 && i + 512 < q) {
u32 *p0 = c + i + 512;
u32 *p1 = c + q + i + 512;
u32 *p2 = c + 2 * q + i + 512;
__asm__ __volatile__("prefetchw %0" :: "m"(*p0));
__asm__ __volatile__("prefetchw %0" :: "m"(*p1));
__asm__ __volatile__("prefetchw %0" :: "m"(*p2));
if ((i & 15) == 0 && i + 3 * q + 512 < sz) { // [C2A-P3] guard period 32 -> 16
u32 *p3 = c + 3 * q + i + 512;
__asm__ __volatile__("prefetchw %0" :: "m"(*p3));
}
}
u32x8 x, y, z, t = split ? load_u32x8(scratch + i) : load_u32x8(last + i);
if (i + 8 == q) {
x = saved[0];
y = saved[1];
z = saved[2];
} else {
x = load_u32x8(a + i);
y = load_u32x8(a + q + i);
z = load_u32x8(a + 2 * q + i);
}
finish_quarters<decltype(mode)::value>(x, y, z, t, c, i, q, sz, fxw, fxs, frw, frs);
};
const int fourth_count = sz - 3 * q;
const int full = (fourth_count / 8) * 8;
int i = 0;
for (; i < full; i += 8) run_finish(i, std::integral_constant<int, 1>{});
if (i < q && i < fourth_count) {
run_finish(i, std::integral_constant<int, 0>{});
i += 8;
}
for (; i < q; i += 8) run_finish(i, std::integral_constant<int, 2>{});
if (lastOwn) _mm_free(lastOwn);
if (temp) _mm_free(temp);
}
void convolve_cyclic(int lg, u32 *a, u32 *b) const {
if (lg < 7) {
transform_forward(lg, a);
transform_forward(lg, b);
multiply_all_leaves<false>(lg, a, b, a);
transform_inverse<true>(lg, a, mt.r);
return;
}
u64x4 fw[LG], iw[LG], dot = wl_init;
std::fill(fw, fw + LG, wt_init);
std::fill(iw, iw + LG, wrt_init);
int L = leaf_log2(lg);
if (L == L0)
convolve_block<L0>(lg, 0, a, b, fw, iw, dot);
else
convolve_block<L0 + 1>(lg, 0, a, b, fw, iw, dot);
u32 f = mt.power<false, true>((mod + 1) >> 1, lg - L);
f = mt.mul<true>(f, mt.mul<false>(mt.r2, mt.r));
u32x8 fx = (u32x8)_mm256_set1_epi32(f);
for (int i = 0; i < (1 << lg); i += 8) store_u32x8(a + i, mts.mul_u32x8<true, true>(load_u32x8(a + i), fx));
}
};
void poly_multiply(unsigned *A, int n, unsigned *B, int m, unsigned *c) {
n++, m++;
u32 mod = 81'788'929;
NTT ntt(mod);
int lg = 3;
while ((1 << lg) < (n + m - 1)) {
lg++;
}
auto disjoint = [](const u32 *src, const u32 *dst) {
uintptr_t s = (uintptr_t)src, d = (uintptr_t)dst;
return d + 2000001ull * 4 <= s || s + 1000001ull * 4 <= d;
};
if (n == 1000001 && m == 1000001 && disjoint(A, c) && disjoint(B, c)) {
ntt.convolve_reusing_output(A, B, c);
return;
}
u32 *a = (u32 *)_mm_malloc(4 << lg, 32);
u32 *b = (u32 *)_mm_malloc(4 << lg, 32);
if (lg >= 9 && n >= (1 << (lg - 2)) && m >= (1 << (lg - 2)) && n <= (1 << (lg - 1)) && m <= (1 << (lg - 1)) &&
n + m - 1 >= (3 << (lg - 2))) {
ntt.convolve_inputs(A, n, B, m, c, lg, a, b);
_mm_free(a);
_mm_free(b);
return;
}
std::copy(A, A + n, a);
std::copy(B, B + m, b);
std::fill(a + n, a + (1 << lg), 0);
std::fill(b + m, b + (1 << lg), 0);
ntt.convolve_cyclic(lg, a, b);
std::copy(a, a + n + m - 1, c);
_mm_free(a), _mm_free(b);
}
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 18.232 ms | 10 MB + 676 KB | Accepted | Score: 100 | 显示更多 |