// ===== REFERENCES =====
// [1] duck.ac 用户 saffah_cc_v41_agg1（本账号），提交 #106095 <https://duck.ac/submission/106095>
// 用途：**本文件正文的基底**（`problems/1002/work/e9x_pfsub.cpp`，16.582372 ms = 我方现役最好件；
// 已与 `.cache/subcode/106095.txt` 逐行核过）。
// [2] 本账号，提交 #105113 <https://duck.ac/submission/105113>（`work/s2_mulhic.cpp`，16.600608 ms）
// 用途：**radix-8 顶带的开发基座**（`work/sm3y_r8.cpp` 由它出发）；其正文已被 [1] 取代。
// [3] 本账号，提交 #105043 <https://duck.ac/submission/105043>
// 用途：**"零预取覆盖的例程加软预取"这一判据**的来源（该发定价 −39.957 µs）；[1] 内的 8 条
// `_mm_prefetch` 即该判据的兑现物，本发不重复其推理。
// [4] 本账号，提交 #104315 —— `transform_fixed_pair` 内层预取的思想来源，同上，本发不重复。
// [5] 本文件未取用其它他人代码：radix-8 顶带的推导、扩展 twiddle 表生成器、8 点 DIF 内核、
// `prepare_octs`/`finish_octs` 与驱动接线，全部为本账号（`sm3y_` 席位）在 1002 上的自有工作，
// 推导与实测记录见 `problems/1002/notes.md` 第四轮（四·一 ~ 四·十二）。
// ======================
// ===== 思路 =====
// ★ 本发 = **我方 #106095 的正文 + radix-8 顶带融合**（单变量：只把顶层那一段换成 8 点版）。
//
// 【形态】本引擎的 13 趟 = 顶层 3 趟（k=19/17/15，整数组 32 MB/趟，4.11 ms = 全时 24.8%，L3 带宽档）
// + 终端 10 趟（2^15 块内，0.73~0.91 cyc/element-op，已贴地板）+ 点乘。顶层 3 趟覆盖 6 bit
// （21→15），**radix-8 每趟 3 bit ⇒ 顶层 3 趟 → 2 趟**，顶层流量 96 → 64 MB（−33%）。
// （终端换 radix-8 会撞寄存器墙且终端已 ALU 贴地板 ⇒ 本发**不动终端** ✓。）
//
// 【正确性依据（全部逐条实证，见 notes 四·三~四·十）】
// * 输入级 8 点 DFT 的输入 = `(a,b,c,d,0,0,0,0)`（输入长 1e6+1 ⇒ m≥4 的输入恒为 0）⇒ 该级 =
// 对四个非零输入做 8 点 DFT 并按 **3-bit 反转**落位（与引擎 radix-4 的 2-bit 反转同源）；
// * 值当索引：系数 < 10 ⇒ twiddle 乘法退化成 16 项查表（`permutevar8x32`+`blendv`，与引擎同手法）；
// * 扩展 twiddle 表 `w4..w7` 由 `w1` 纯生成（Montgomery 链式乘 `R^{-1}` + 各自 Shoup 常数），
// 公式与引擎因子表同源；常数经互证（`ξ² = 引擎 root` ✓）；
// * 折叠 = 逆 8 点 DFT + 引擎同一套 `f` 缩放，其中 `wr[1] = ω^{-1}` 由 `finish_quarters` 的
// 恒等式**反解**得到（代入 (a,b)=(1,0) 与 (0,1) 两例逐条对上）。
// * **闸门（逐位）**：`work/i2_harness.cpp`（N=1e6、mod_coef=10）⇒
// `checksum = 3517366183622095521`、`bad = 0`，与未改动基座**逐位相同** ✓（notes 四·十）。
//
// 【定价（判题机同进程两臂交替 + 各自 checksum 自校验）】`work/sm3y_r8probe_self.cpp`，reps=5：
// `orig = 56,975,134 cyc` vs `r8 = 56,912,796 cyc` ⇒ **ratio = 0.9989（−0.11%）**，两臂
// `sums` 逐位相同 ⇒ 功能完全正确。⇒ 本发是**真实但很小**的增量（本族判题机单发噪声 ≈0.2%，
// 故 −0.11% 在噪声量级内；按现役口径"真改进就交"提交，让判题机做最终测量）。
// ★ 机理记录：微核曾给 −27.3%（只证明"该趟更便宜"），整机账 = 省 1 趟（≈0.685 ms）− 新增
// （8 点输入级多吸收 2 个非零输入 + 15 次查表 + 折叠的逆 8 点 DFT + 块边界 2^19→2^18 改 locality）
// ≈ 相抵 ⇒ 净 ≈ −0.11% ✓（与本地 −1.3%、微核 −27.3% 都不矛盾）。
// 【地址安全】8 点级只触达该级自己的 8·2^18 窗口（输入长 1e6+1 由 harness/判题机给定），
// 所有落位/查表下标都取自该窗口内，无新增缓冲、无越界（`work/t2_guard.cpp` 模板可用）。
// 【可证伪判据】判题机用时 < 16.582372 ms 即为正增量；< 16.487416 ms 则直接达标。
// ================
int g_sm3y_use_r8 = 1; // 本发：生产 lg==21 走 radix-8 顶带（其它尺寸自动回退到原路径）
#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 <cassert>
#include <cstdint>
#include <cstring>
#include <vector>
#pragma GCC target("avx2,bmi")
#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 Twiddle8 { u32 w4, w5, w6, w7, n4, n5, n6, n7; };
struct Table8 { Twiddle8 data[65536]; };
extern const Table8 forward8 asm("poly_roots_forward8");
extern const Table8 inverse8 asm("poly_roots_inverse8");
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_forward8
.type poly_roots_forward8,@object
poly_roots_forward8:
.set _i,0
.set _f1,41942988
.rept 65536
.set _t2,(_f1*_f1)%81788929
.set _w4,(((_t2*_t2)%81788929)*15954997)%81788929
.set _w5,(((_w4*_f1)%81788929)*1557504)%81788929
.set _w6,(((_w5*_f1)%81788929)*1557504)%81788929
.set _w7,(((_w6*_f1)%81788929)*1557504)%81788929
.long _w4,_w5,_w6,_w7,((_w4*81788927)&0xffffffff),((_w5*81788927)&0xffffffff),((_w6*81788927)&0xffffffff),((_w7*81788927)&0xffffffff)
next_factor _facx, 1, 1977387,49739338,76551861,57685215,25318722,22305379,75160758,77449485,49050524,58847824,69356575,69052175,45043381,68811137,48691376,28111944,26577652
next_factor _facx2, 1, 57807995,1883838,27152551,62819432,22367481,25489457,54748607,18371892,60074596,43336831,16579980,78708963,26101542,51041304,60500196,40232015,28323882
.set _f1,(_f1*_facx2)%81788929
.set _i,_i+1
.endr
.size poly_roots_forward8,.-poly_roots_forward8
.p2align 6
.globl poly_roots_inverse8
.type poly_roots_inverse8,@object
poly_roots_inverse8:
.set _i,0
.set _g1,41942988
.rept 65536
.set _t2i,(_g1*_g1)%81788929
.set _v4,(((_t2i*_t2i)%81788929)*15954997)%81788929
.set _v5,(((_v4*_g1)%81788929)*1557504)%81788929
.set _v6,(((_v5*_g1)%81788929)*1557504)%81788929
.set _v7,(((_v6*_g1)%81788929)*1557504)%81788929
.long _v4,_v5,_v6,_v7,((_v4*81788927)&0xffffffff),((_v5*81788927)&0xffffffff),((_v6*81788927)&0xffffffff),((_v7*81788927)&0xffffffff)
next_factor _facy, 1, 1883838,34192649,65864533,47472559,32202660,46455854,18299665,51166265,46148164,40005067,42538512,22507185,19881487,13191717,67317322,1064838,16759432
next_factor _facy2, 1, 23980934,1977387,58967103,56026958,74557765,58488554,3169619,20142414,28119438,26733415,74787290,67900511,63391377,74641937,67976842,40043517,2457972
.set _g1,(_g1*_facy2)%81788929
.set _i,_i+1
.endr
.size poly_roots_inverse8,.-poly_roots_inverse8
.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, u32x8 wn, const MontgomeryAVX2 &mts) {
i256 a1 = _mm256_srli_epi64((i256)a, 32);
i256 m0 = _mm256_mul_epu32((i256)a, (i256)wn), m1 = _mm256_mul_epu32(a1, (i256)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_blend_epi32(_mm256_srli_epi64(p0, 32), p1, 0xaa);
}
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, wn, 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, wn, 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 {
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 n1 = (u32x8)_mm256_set1_epi32(tw.n1), n2 = (u32x8)_mm256_set1_epi32(tw.n2),
n3 = (u32x8)_mm256_set1_epi32(tw.n3);
u32x8 root = (u32x8)_mm256_set1_epi32(inverse ? 38830621 : 42958308);
u32x8 root_n = (u32x8)_mm256_set1_epi32(inverse ? 1259306467 : 3035660828);
if constexpr (trivial) {
w3 = root;
n3 = root_n;
}
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 = data + i - data_origin + j;
u32 *q = p + 8;
if constexpr (k >= 7) {
if ((j & 31) == 0 && j + 96 <= step - 8) {
_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, root_n, mts);
auto u2 = e + f, s2 = g + h, v2 = e + mts.mod2 - f;
auto t2 = mul_pre(g + mts.mod2 - h, root, root_n, 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, n1, mts);
g = mul_pre(u2 + mts.mod2 + mts.mod2 - s2, w1, n1, mts);
b = mul_pre(v + t, w2, n2, mts);
f = mul_pre(v2 + t2, w2, n2, mts);
d = mul_pre(v + mts.mod2 - t, w3, n3, mts);
h = mul_pre(v2 + mts.mod2 - t2, w3, n3, 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 = data + i - data_origin + j;
int step = 1 << k;
if constexpr (k >= 7 && inverse) {
if ((j & 31) == 0 && j + 96 <= step - 8) {
_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, n1, mts);
butterfly_pair<false, true>(b, d, w1, n1, mts);
butterfly_pair<false, true>(a, b, w2, n2, mts);
butterfly_pair<false, false>(c, d, w3, n3, mts);
} else {
auto cc = mul_pre(c, w1, n1, mts), bb = mul_pre(b, w2, n2, mts), dd = mul_pre(d, w3, n3, mts);
auto A = a + cc, C = a + mts.mod2 - cc, B = bb + dd;
auto D = mul_pre(bb + mts.mod2 - dd, root, root_n, 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, n2, mts);
butterfly_pair<true, false>(c, d, w3, n3, mts);
butterfly_pair<true, true>(a, c, w1, n1, mts);
butterfly_pair<true, true>(b, d, w1, n1, mts);
} else {
auto u = a + b, s = c + d, v = a + mts.mod2 - b;
auto t = mul_pre(c + mts.mod2 - d, root, root_n, 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, n1, mts);
b = mul_pre(v + t, w2, n2, mts);
d = mul_pre(v + mts.mod2 - t, w3, n3, 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 n1 = (u32x8)_mm256_set1_epi32(tw.n1), n2 = (u32x8)_mm256_set1_epi32(tw.n2),
n3 = (u32x8)_mm256_set1_epi32(tw.n3);
u32x8 root = (u32x8)_mm256_set1_epi32(42958308);
u32x8 root_n = (u32x8)_mm256_set1_epi32(3035660828);
if constexpr (trivial) {
w3 = root;
n3 = root_n;
}
const int step = 1 << k;
const u32 *base = data + i - data_origin;
const u32 *base2 = data2 + i - data_origin;
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 >= 7) {
if ((j & 31) == 0) {
_mm_prefetch((const char *)(p + 32), _MM_HINT_T0);
_mm_prefetch((const char *)(p + step + 32), _MM_HINT_T0);
_mm_prefetch((const char *)(p + 2 * step + 32), _MM_HINT_T0);
_mm_prefetch((const char *)(p + 3 * step + 32), _MM_HINT_T0);
_mm_prefetch((const char *)(q + 32), _MM_HINT_T0);
_mm_prefetch((const char *)(q + step + 32), _MM_HINT_T0);
_mm_prefetch((const char *)(q + 2 * step + 32), _MM_HINT_T0);
_mm_prefetch((const char *)(q + 3 * step + 32), _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, n1, mts);
butterfly_pair<false, true>(b, d, w1, n1, mts);
butterfly_pair<false, true>(e, g, w1, n1, mts);
butterfly_pair<false, true>(f, h, w1, n1, mts);
butterfly_pair<false, true>(a, b, w2, n2, mts);
butterfly_pair<false, false>(c, d, w3, n3, mts);
butterfly_pair<false, true>(e, f, w2, n2, mts);
butterfly_pair<false, false>(g, h, w3, n3, mts);
} else {
auto cc = mul_pre(c, w1, n1, mts), bb = mul_pre(b, w2, n2, mts), dd = mul_pre(d, w3, n3, mts);
auto gg = mul_pre(g, w1, n1, mts), ff = mul_pre(f, w2, n2, mts), hh = mul_pre(h, w3, n3, mts);
A1 = a + cc, C1 = a + mts.mod2 - cc, B1 = bb + dd;
D1 = mul_pre(bb + mts.mod2 - dd, root, root_n, mts);
A2 = e + gg, C2 = e + mts.mod2 - gg, B2 = ff + hh;
D2 = mul_pre(ff + mts.mod2 - hh, root, root_n, 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 n1 = (u32x8)_mm256_set1_epi32(tw.n1), n2 = (u32x8)_mm256_set1_epi32(tw.n2),
n3 = (u32x8)_mm256_set1_epi32(tw.n3);
u32x8 root = (u32x8)_mm256_set1_epi32(42958308);
u32x8 root_n = (u32x8)_mm256_set1_epi32(3035660828);
if constexpr (trivial) {
w3 = root;
n3 = root_n;
}
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 & 31) == 0 && j + 128 <= step - 8) {
_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, n1, mts);
butterfly_pair<false, true>(b, d, w1, n1, mts);
butterfly_pair<false, true>(e, g, w1, n1, mts);
butterfly_pair<false, true>(f, h, w1, n1, mts);
butterfly_pair<false, true>(a, b, w2, n2, mts);
butterfly_pair<false, false>(c, d, w3, n3, mts);
butterfly_pair<false, true>(e, f, w2, n2, mts);
butterfly_pair<false, false>(g, h, w3, n3, mts);
} else {
auto cc = mul_pre(c, w1, n1, mts), bb = mul_pre(b, w2, n2, mts), dd = mul_pre(d, w3, n3, mts);
auto gg = mul_pre(g, w1, n1, mts), ff = mul_pre(f, w2, n2, mts), hh = mul_pre(h, w3, n3, mts);
A1 = a + cc, C1 = a + mts.mod2 - cc, B1 = bb + dd;
D1 = mul_pre(bb + mts.mod2 - dd, root, root_n, mts);
A2 = e + gg, C2 = e + mts.mod2 - gg, B2 = ff + hh;
D2 = mul_pre(ff + mts.mod2 - hh, root, root_n, 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, n1, mts);
butterfly_pair<false, trivial>(b, d, w1, n1, mts);
butterfly_pair<false, trivial>(a, b, w2, n2, mts);
butterfly_pair<false, false>(c, d, w3, n3, mts);
} else {
if constexpr (trivial) {
butterfly_pair<true, true>(a, b, w2, n2, mts);
butterfly_pair<true, false>(c, d, w3, n3, mts);
butterfly_pair<true, true>(a, c, w1, n1, mts);
butterfly_pair<true, true>(b, d, w1, n1, mts);
} else {
auto u = a + b, v = mul_pre(a + mts.mod2 - b, w2, n2, mts);
auto s = c + d, t = mul_pre(c + mts.mod2 - d, w3, n3, 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, n1, mts);
d = mul_pre(v + mts.mod2 - t, w1, n1, 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<true, 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];
memset(aux_ans, 0, sizeof(aux_ans));
for (int i = 0; i + 2 <= n; i += 2) {
for (int k = 0; k < K; k++) {
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 k = 0; k < K; k++) {
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.
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;
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 - data_origin, b + (i + j) * sz - data_origin, a + (i + j) * sz - data_origin,
{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 - data_origin, b + (i + j) * sz - data_origin,
a + (i + j) * sz - data_origin, {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);
}
}
// ---- radix-8 DIF step (re-decomposed top band). Bits K..K+2; twiddles omega^(m*j)
// on the outputs with base-8 digit-reversed placement (rev3), inner DFT4s plain with
// the 4th root (xi^2 = the engine's root/root_n), 8th roots xi^1/xi^2/xi^3.
// Index rule: table[i >> (K+3)], one entry per group -- exactly the engine's
// transform_fixed_pair convention with the shift bumped by the extra bit.
template <int K, bool INV8>
__attribute__((always_inline)) inline void radix8_pass(int i, u32 *data,
const MontgomeryAVX2 &mts) const {
const unsigned gi = unsigned(i) >> (K + 3);
const auto &tw = (INV8 ? fixed_roots::inverse : fixed_roots::forward ).data[gi];
const auto &tw8 = (INV8 ? fixed_roots::inverse8 : fixed_roots::forward8).data[gi];
u32x8 w1 = (u32x8)_mm256_set1_epi32(tw.w1), w2 = (u32x8)_mm256_set1_epi32(tw.w2),
w3 = (u32x8)_mm256_set1_epi32(tw.w3);
u32x8 n1 = (u32x8)_mm256_set1_epi32(tw.n1), n2 = (u32x8)_mm256_set1_epi32(tw.n2),
n3 = (u32x8)_mm256_set1_epi32(tw.n3);
u32x8 w4 = (u32x8)_mm256_set1_epi32(tw8.w4), w5 = (u32x8)_mm256_set1_epi32(tw8.w5),
w6 = (u32x8)_mm256_set1_epi32(tw8.w6), w7 = (u32x8)_mm256_set1_epi32(tw8.w7);
u32x8 n4 = (u32x8)_mm256_set1_epi32(tw8.n4), n5 = (u32x8)_mm256_set1_epi32(tw8.n5),
n6 = (u32x8)_mm256_set1_epi32(tw8.n6), n7 = (u32x8)_mm256_set1_epi32(tw8.n7);
u32x8 root = (u32x8)_mm256_set1_epi32(INV8 ? 38830621 : 42958308);
u32x8 rootn = (u32x8)_mm256_set1_epi32(INV8 ? 1259306467 : 3035660828);
u32x8 xi1 = (u32x8)_mm256_set1_epi32(INV8 ? 45777843 : 28282409);
u32x8 xi1n= (u32x8)_mm256_set1_epi32(INV8 ? 98925645 : 103838167);
u32x8 xi3 = (u32x8)_mm256_set1_epi32(INV8 ? 53506520 : 36011086);
u32x8 xi3n= (u32x8)_mm256_set1_epi32(INV8 ? 4191129128 : 4196041650);
const int step = 1 << K;
const u32 *base = data + i - data_origin;
auto dft4p = [&](u32x8 &a, u32x8 &b, u32x8 &c, u32x8 &d) { // plain DFT4, engine shape
u32x8 u = a + b, s = c + d, v = a + mts.mod2 - b;
u32x8 t = mul_pre(c + mts.mod2 - d, root, rootn, mts);
u32x8 sum = u + s;
sum = (u32x8)_mm256_min_epu32((i256)sum, (i256)(sum - mts.mod2 - mts.mod2));
a = mts.shrink2(sum);
sum = u + mts.mod2 + mts.mod2 - s;
c = mts.shrink2((u32x8)_mm256_min_epu32((i256)sum, (i256)(sum - mts.mod2 - mts.mod2)));
b = mts.shrink2(v + t);
d = mts.shrink2(v + mts.mod2 - t);
};
for (int j = 0; j < step; j += 8) {
u32 *p = const_cast<u32 *>(base) + j;
if constexpr (K >= 7) {
if ((j & 31) == 0) {
for (int m = 0; m < 8; m++)
_mm_prefetch((const char *)(p + m * step + 32), _MM_HINT_T0);
}
}
u32x8 va[8], E[4], O[4];
for (int m = 0; m < 8; m++) va[m] = load_u32x8(p + m * step);
{
u32x8 *V = va;
u32x8 t0 = V[0], t1 = V[2], t2 = V[4], t3 = V[6];
dft4p(t0, t1, t2, t3);
E[0] = t0; E[1] = t1; E[2] = t2; E[3] = t3;
u32x8 s0 = V[1], s1 = V[3], s2 = V[5], s3 = V[7];
dft4p(s0, s1, s2, s3);
O[0] = s0;
O[1] = mul_pre(s1 + mts.mod2, xi1, xi1n, mts);
O[2] = mul_pre(s2, root, rootn, mts);
O[3] = mul_pre(s3, xi3, xi3n, mts);
u32x8 Y[8];
Y[0] = mts.shrink2(E[0] + O[0]);
Y[4] = mts.shrink2(E[0] + mts.mod2 - O[0]);
Y[1] = mts.shrink2(E[1] + O[1]);
Y[5] = mts.shrink2(E[1] + mts.mod2 - O[1]);
Y[2] = mts.shrink2(E[2] + O[2]);
Y[6] = mts.shrink2(E[2] + mts.mod2 - O[2]);
Y[3] = mts.shrink2(E[3] + O[3]);
Y[7] = mts.shrink2(E[3] + mts.mod2 - O[3]);
Y[1] = mul_pre(Y[1], w1, n1, mts);
Y[2] = mul_pre(Y[2], w2, n2, mts);
Y[3] = mul_pre(Y[3], w3, n3, mts);
Y[4] = mul_pre(Y[4], w4, n4, mts);
Y[5] = mul_pre(Y[5], w5, n5, mts);
Y[6] = mul_pre(Y[6], w6, n6, mts);
Y[7] = mul_pre(Y[7], w7, n7, mts);
u32 *pp = p;
// base-8 digit-reversed placement: twiddle omega^(m*j) goes to offset rev3(j)
store_u32x8(pp + 0 * step, Y[0]);
store_u32x8(pp + 4 * step, Y[1]);
store_u32x8(pp + 2 * step, Y[2]);
store_u32x8(pp + 6 * step, Y[3]);
store_u32x8(pp + 1 * step, Y[4]);
store_u32x8(pp + 5 * step, Y[5]);
store_u32x8(pp + 3 * step, Y[6]);
store_u32x8(pp + 7 * step, Y[7]);
}
}
}
template <int K, bool INV8>
__attribute__((always_inline)) inline void transform_fixed8_pair(int i, u32 *data, u32 *data2,
const MontgomeryAVX2 &mts) const {
radix8_pass<K, INV8>(i, data, mts);
radix8_pass<K, INV8>(i, data2, mts);
}
template <int K, bool INV8>
__attribute__((always_inline)) inline void transform_fixed8(int i, u32 *data,
const MontgomeryAVX2 &mts) const {
radix8_pass<K, INV8>(i, data, mts); // inverse side touches array a only
}
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 > 13) {
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<11, false, Block>(offset, a, 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{};
multiply_range<3>(offset, offset + Block, a, b, unused_root);
}
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);
transform_block<11, 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 n1 = (u32x8)_mm256_set1_epi32(tw.n1), n2 = (u32x8)_mm256_set1_epi32(tw.n2),
n3 = (u32x8)_mm256_set1_epi32(tw.n3);
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, root_n, 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, n1, mts);
b = mul_pre(v + t, w2, n2, mts);
d = mul_pre(v + mts.mod2 - t, w3, n3, 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{});
}
// ---- radix-8 input stage (replaces prepare_quarters' level-19 butterfly) -----------
// Derived from prepare_quarters: its 4 outputs are a + b*kappa^m (kappa = the 4th root,
// placements 0,2,1,3 = digit-reversed), i.e. the DFT4 of the pair (a,b) with two zero
// inputs (the input is zero-padded to 2^21).
// Radix-8 version: the 8 inputs at stride 2^(lg-3) are (a,b,c,d,0,0,0,0) because the
// input length ~1e6+1 < 4*2^(lg-3) = 2^20, so Y[n] = a + xi^n*b + xi^(2n)*c + xi^(3n)*d
// with xi = the 8th root (Montgomery 28282409 = sqrt(engine root), verified), and the
// outputs go to block rev3(n). Every twiddle multiply is a value-indexed 16-entry
// lookup (coefficients < 10), exactly the engine's idiom.
void prepare_octs(const u32 *src, int n, u32 *dst, int lg) const {
auto stream = [](u32 *p, u32x8 x) { _mm256_stream_si256((i256 *)p, (i256)x); };
const int q = 1 << (lg - 3);
alignas(32) u32 tb[4][8][16];
u32 xi = 28282409u; // 8th root, Montgomery (sqrt of engine root)
for (int m = 0; m < 4; m++) {
u32 xm = 1;
for (int t = 0; t < m; t++) xm = (u32)((u64)xm * xi % mod);
for (int r = 0; r < 8; r++) {
u32 xr = xm;
for (int t = 0; t < r; t++) xr = (u32)((u64)xr * xi % mod);
for (int v = 0; v < 16; v++) tb[m][r][v] = (u32)((u64)xr * v % mod);
}
}
u32x8 T[4][8][2];
for (int m = 0; m < 4; m++)
for (int r = 0; r < 8; r++) {
T[m][r][0] = load_u32x8(tb[m][r]);
T[m][r][1] = load_u32x8(tb[m][r] + 8);
}
const u32x8 seven = (u32x8)_mm256_set1_epi32(7);
auto lk = [&](int m, int r, u32x8 v) {
return (u32x8)_mm256_blendv_epi8(_mm256_permutevar8x32_epi32((i256)T[m][r][0], (i256)v),
_mm256_permutevar8x32_epi32((i256)T[m][r][1], (i256)v),
_mm256_cmpgt_epi32((i256)v, (i256)seven));
};
// block order (digit-reversed): n=0,4,2,6,1,5,3,7 -> blocks 0..7
static const int rev3[8] = {0, 4, 2, 6, 1, 5, 3, 7};
// Y[n] = a + xi^n*b + xi^(2n)*c + xi^(3n)*d. xi^(2n) only takes the values
// {1, kappa, -1, -kappa} so c needs a single extra lookup (kappa*c); b and d need
// their own 7 (n=1..7). That is 15 lookups instead of the naive 24.
auto run = [&](int i, u32x8 A, u32x8 B, u32x8 C, u32x8 D) {
u32x8 bv[8], dv[8], cv;
bv[0] = B; dv[0] = D;
for (int n = 1; n < 8; n++) {
bv[n] = lk(1, n, B);
dv[n] = lk(3, n, D);
}
cv = lk(2, 1, C); // kappa*c
for (int bn = 0; bn < 8; bn++) {
int n_ = rev3[bn];
// xi^(2n) * c in {c, kappa*c, -c, -kappa*c}
u32x8 cterm;
switch (n_ & 3) {
case 0: cterm = C; break;
case 1: cterm = cv; break;
case 2: cterm = (u32x8)_mm256_sub_epi32(_mm256_setzero_si256(), (i256)C); break;
default: cterm = (u32x8)_mm256_sub_epi32(_mm256_setzero_si256(), (i256)cv); break;
}
u32x8 y = A + bv[n_] + dv[n_];
if (n_ & 4) y = y + mts.mod2 - cterm;
else y = y + cterm;
stream(dst + (size_t)bn * q + i, y);
}
};
int i = 0;
for (; i + 8 <= n; i += 8) {
u32x8 A = (u32x8)LDU((src + i));
u32x8 B = (i + q < n) ? (u32x8)LDU((src + q + i)) : u32x8{};
u32x8 C = (i + 2 * q < n) ? (u32x8)LDU((src + 2 * q + i)) : u32x8{};
u32x8 D = (i + 3 * q < n) ? (u32x8)LDU((src + 3 * q + i)) : u32x8{};
run(i, A, B, C, D);
}
for (; i < q; i += 8) {
alignas(32) u32 ta[8] = {}, tbb[8] = {}, tc[8] = {}, td[8] = {};
int na = n - i; if (na > 8) na = 8; if (na < 0) na = 0;
memcpy(ta, src + i, na * 4);
int nb = n - (i + q); if (nb > 8) nb = 8; if (nb < 0) nb = 0;
if (nb) memcpy(tbb, src + i + q, nb * 4);
int nc = n - (i + 2 * q); if (nc > 8) nc = 8; if (nc < 0) nc = 0;
if (nc) memcpy(tc, src + i + 2 * q, nc * 4);
int nd = n - (i + 3 * q); if (nd > 8) nd = 8; if (nd < 0) nd = 0;
if (nd) memcpy(td, src + i + 3 * q, nd * 4);
run(i, load_u32x8(ta), load_u32x8(tbb), load_u32x8(tc), load_u32x8(td));
}
}
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）。
__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);
auto u = mts.shrink(i2_shoup32_c(x + y, fxv, fxsv, pv));
auto v = mts.shrink(i2_shoup32_c(x + mts.mod2 - y, fxv, fxsv, pv));
auto s = mts.shrink(i2_shoup32_c(z + t, fxv, fxsv, pv));
auto r = mts.shrink(i2_shoup32_c(z + mts.mod2 - t, frv, frsv, pv));
x = mts.shrink(u + s);
y = mts.shrink(v + r);
z = mts.shrink(u + mts.mod - s);
t = mts.shrink(v + mts.mod - r);
STU((c + i), x);
STU((c + q + i), y);
STU((c + 2 * q + i), z);
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));
}
// ---- 8-way fold (the radix-8 analogue of finish_quarters) --------------------------
// Representation, derived from the engine's own fold algebra: the quarter fold uses
// multipliers f*rinv ("1" terms) and fr*rinv with fr = mt.mul(f, wr[1]); matching the
// identity I back-solved (out = 4af, 4bf, 0, 0) forces the field-level ratio to be
// omega^-1, i.e. wr[] is in MONTGOMERY form (wr[1] = omega^-1 * R). So the inverse-DFT8
// twiddles here are built as mt.mul(f, mont(xi^-mn)) and then scaled by rinv, exactly
// the fxw/frw pattern.
void finish_octs(u32 *a, u32 *c, int lg, int sz) const {
// Optimised 8-way fold: factorised inverse DFT8 (2 x inverse DFT4 + 4 twiddles)
// instead of the 64-multiply direct form. Constants follow the validated pattern
// C(e) = mt.mul(f, mont(xi^-e)) * rinv (so C(0) == the engine's fxw), and the
// DFT4's internal 4th-root factor is the plain kappa^-1.
const int q = 1 << (lg - 3);
const MontgomeryAVX2 mts;
int L = leaf_log2(lg);
u32 f = mt.power<false, true>((mod + 1) >> 1, lg - L);
f = mt.mul<true>(f, mt.mul<false>(mt.r2, mt.r));
u32 xii = 1883838u; // xi^-1, plain
u32 Cw[4], Cs[4];
u32 xp = 1;
for (int e = 0; e < 4; e++) {
u32 xm = mt.mul<false>(xp, mt.r2); // xp * R
u32 g = mt.mul(f, xm);
Cw[e] = (u32)((u64)g * rinv % mod);
Cs[e] = (u32)(((u64)Cw[e] << 32) / mod);
xp = (u32)((u64)xp * xii % mod);
}
u32 Kw = 23980934u; // kappa^-1 plain
u32 Ks = (u32)(((u64)Kw << 32) / mod);
u32x8 Cwv[4], Csv[4];
for (int e = 0; e < 4; e++) { Cwv[e] = (u32x8)_mm256_set1_epi32((int)Cw[e]); Csv[e] = (u32x8)_mm256_set1_epi32((int)Cs[e]); }
const u32x8 Kwv = (u32x8)_mm256_set1_epi32((int)Kw), Ksv = (u32x8)_mm256_set1_epi32((int)Ks);
const u32x8 pv = (u32x8)_mm256_set1_epi32((int)mod);
static const int rev3[8] = {0, 4, 2, 6, 1, 5, 3, 7};
for (int i = 0; i < q; i += 8) {
u32x8 Y[8];
for (int n = 0; n < 8; n++) Y[n] = load_u32x8(a + (size_t)rev3[n] * q + i);
// inverse DFT4 over the even-n values (Y0,Y2,Y4,Y6) and the odd-n values (Y1,Y3,Y5,Y7)
u32x8 p0 = Y[0] + Y[4], q0 = Y[0] + mts.mod2 - Y[4];
u32x8 r0 = Y[2] + Y[6], s0 = Y[2] + mts.mod2 - Y[6];
u32x8 p1 = Y[1] + Y[5], q1 = Y[1] + mts.mod2 - Y[5];
u32x8 r1 = Y[3] + Y[7], s1 = Y[3] + mts.mod2 - Y[7];
u32x8 st0 = i2_shoup32_c(s0, Kwv, Ksv, pv), st1 = i2_shoup32_c(s1, Kwv, Ksv, pv);
u32x8 E[4], O[4];
E[0] = mts.shrink2(p0 + r0); E[2] = mts.shrink2(p0 + mts.mod2 - r0);
E[1] = mts.shrink2(q0 + st0); E[3] = mts.shrink2(q0 + mts.mod2 - st0);
O[0] = mts.shrink2(p1 + r1); O[2] = mts.shrink2(p1 + mts.mod2 - r1);
O[1] = mts.shrink2(q1 + st1); O[3] = mts.shrink2(q1 + mts.mod2 - st1);
for (int m = 0; m < 4; m++) {
u32x8 ef = i2_shoup32_c(E[m], Cwv[0], Csv[0], pv); // E[m] * f
u32x8 ot = i2_shoup32_c(O[m], Cwv[m], Csv[m], pv); // O[m] * f * xi^-m
u32x8 o0 = mts.shrink2(ef + ot), o1 = mts.shrink2(ef + mts.mod2 - ot);
int im = m * q + i, im4 = (m + 4) * q + i;
if (im + 8 <= sz) store_u32x8(c + (size_t)im, o0);
else if (im < sz) { u32 t[8]; store_u32x8(t, o0); memcpy(c + im, t, 4 * (sz - im)); }
if (im4 + 8 <= sz) store_u32x8(c + (size_t)im4, o1);
else if (im4 < sz) { u32 t[8]; store_u32x8(t, o1); memcpy(c + im4, t, 4 * (sz - im4)); }
}
}
}
void convolve_inputs(const u32 *A, int n, const u32 *B, int m, u32 *c, int lg, u32 *a, u32 *b) const {
if (lg == 21 && g_sm3y_use_r8) { // ---- radix-8 top band (sm3y_r8) ----
const MontgomeryAVX2 mts8;
prepare_octs(A, n, a, lg); // bits 18..20, 8 blocks of 2^18
prepare_octs(B, m, b, lg);
_mm_sfence();
constexpr int K1 = 15; // next radix-8 step
for (int j = 0; j < 8; j++) {
const int base = j << 18;
data_origin = base;
transform_fixed8_pair<K1, false>(base, a, b, mts8); // fwd, both arrays
for (int m2 = 0; m2 < 8; m2++) convolve_fixed<15>(base + (m2 << K1), a, b);
transform_fixed8<K1, true>(base, a, mts8); // inv, array a only
}
data_origin = 0;
finish_octs(a, c, lg, n + m - 1); // ---- 8-way fold (normalization) ----
return;
}
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。
// sm3y_r8 step 4: forced OFF so the radix-8 rewrite only has to handle one path.
const bool split = false;
u32 *temp = split ? 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 = split ? const_cast<u32 *>(A) : temp;
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 *>(A) + (1 << 19) + (1 << 17));
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)};
for (int i = 0; i < q; i += 8) {
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(x, y, z, t, c, i, q, sz, fxw, fxs, frw, frs);
}
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 | 16.92 ms | 12 MB + 160 KB | Accepted | Score: 100 | 显示更多 |