/*
References:
- saffah_codex_6a_agg3, https://duck.ac/submission/126070 : retain the accepted
distributed-final-scale NTT, leaf multiplication and complete prior notices.
- saffah_cc_v41_agg1, https://duck.ac/submission/123953 : inherited the wide
Montgomery multiply and lazy degree-7 leaf accumulator implementation.
Idea:
For degree-7 leaves only, retain the weighted B product in its lazy
Montgomery range instead of applying the final subtraction of P. The
forward operand is below 51P and the multiplier below P, so the product
is below 2P. This is already smaller than the existing unweighted B bound
of 4P; A remains below 4P. Thus every eight-term accumulator stays below
128P^2, its Montgomery reduction below 3.44P, and the existing subtraction
of 2P still returns a value below 2P. Other leaf lengths keep strict products.
Purpose: Remove a redundant canonicalization per vector without changing
convolution scheduling, product representation or inverse-stage bounds.
*/
/*
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[K][n];
alignas(64) u64 aux_b[K][n * 2];
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[k] + 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<(L != 3), true>(bi, ar_w[k]);
store_u32x8((u32 *)(aux_b[k] + i + 0),
(u32x8)_mm256_permutevar8x32_epi32((i256)bi_w, _mm256_setr_epi64x(0, 1, 2, 3)));
store_u32x8((u32 *)(aux_b[k] + i + 4),
(u32x8)_mm256_permutevar8x32_epi32((i256)bi_w, _mm256_setr_epi64x(4, 5, 6, 7)));
store_u32x8((u32 *)(aux_b[k] + n + i + 0),
(u32x8)_mm256_permutevar8x32_epi32((i256)bi_0, _mm256_setr_epi64x(0, 1, 2, 3)));
store_u32x8((u32 *)(aux_b[k] + n + i + 4),
(u32x8)_mm256_permutevar8x32_epi32((i256)bi_0, _mm256_setr_epi64x(4, 5, 6, 7)));
}
}
u64x4 aux_ans[K][n / 4];
// [K02ZC-M13] no zero-fill at all: the i==0 pass *writes* the accumulator
// (0+p0+p1 == p0+p1 mod 2^64). The fence below replaces the one the memset
// call used to provide for GCC (the STU/STA store macros are inline asm).
for (int k = 0; k < K; k++) {
{
constexpr int i = 0;
u64x4 ai = (u64x4)_mm256_set1_epi32(aux_a[k][i]);
u64x4 ai1 = (u64x4)_mm256_set1_epi32(aux_a[k][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[k] + 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[k] + n - i - 1 + j)));
aux_ans[k][j / 4] = t0;
}
}
for (int i = 2; i + 2 <= n; i += 2) {
u64x4 ai = (u64x4)_mm256_set1_epi32(aux_a[k][i]);
u64x4 ai1 = (u64x4)_mm256_set1_epi32(aux_a[k][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[k][j / 4]), "=&x"(t0)
: "x"(ai), "m"(*(const __m256i_u *)(aux_b[k] + n - i + j)));
asm("vpmuludq %3,%2,%1\n\tvpaddq %1,%0,%0"
: "+x"(aux_ans[k][j / 4]), "=&x"(t1)
: "x"(ai1), "m"(*(const __m256i_u *)(aux_b[k] + n - i - 1 + j)));
}
if (((i + 1) & 7) == 7 && i + 1 >= 15) {
for (int j = 0; j < n; j += 4) {
aux_ans[k][j / 4] = (u64x4)mts.shrink2((u32x8)aux_ans[k][j / 4]);
}
}
}
}
// n is even (L >= 3): the unrolled loop above consumed rows in pairs
// and advanced i past the final pair; nothing remains.
__asm__ __volatile__("" ::: "memory"); // [K02ZC-M13] the fence the memset call used to give
for (int k = 0; k < K; k++) {
for (int i = 0; i < n; i += 8) {
u64x4 c0 = aux_ans[k][i / 4], c1 = aux_ans[k][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 | 15.592 ms | 10 MB + 676 KB | Accepted | Score: 100 | 显示更多 |