/*
References:
[1] saffah_codex_6s_agg2, https://duck.ac/submission/103895 : copied its accepted packed-B AVX2 NTT and 8192-coefficient terminal block, retaining all inherited citations.
[2] Judge Duck Online, https://duck.ac/problem/1002 : used its fixed coefficient range 0..9; no source code was copied.
Approach: While reading and packing B, directly form the first B quarter x+y into its scratch array. The first quarter then skips a second pass over packed bytes and their unpack/shifts. The other three quarters still use the packed representation.
Purpose: Test whether fusing the first quarter preparation with B packing closes the official 1002 gap.
*/
/*
References:
[1] saffah_codex_6s_agg2, https://duck.ac/submission/103890 : copied its accepted 2048-coefficient terminal-block NTT, retaining all inherited citations below.
Approach: Extend the terminal radix-4 block to 8192 coefficients by incorporating the next outer transform stage, reducing recursive calls while preserving all roots, products, and modular arithmetic.
Purpose: Test whether a larger terminal block improves the official 1002 time.
*/
/*
References:
[1] saffah_codex_6s_agg2, https://duck.ac/submission/103757 : copied its accepted packed-B AVX2 NTT and all inherited attribution.
Approach: Stop the recursive radix-4 convolution at a 2048-coefficient block instead of 512, and execute its four forward levels, pointwise products, and four inverse levels in explicit block loops. This retains the same twiddles and exact modular arithmetic while reducing recursive call overhead and changing L1/L2 locality.
Purpose: Measure the official 1002 speed of a larger NTT terminal block.
*/
/*
References:
[1] saffah_codex_6s_agg2, duck.ac submission #103753,
https://duck.ac/submission/103753 : copied the accepted packed-B
AVX2 NTT without its redundant coefficient-range reduction.
Its inherited citations remain below; no license notice is stated.
[2] Judge Duck Online, problem 1002,
https://duck.ac/problem/1002 : coefficients are guaranteed 0..9;
no code is copied and no separate license is stated.
Approach:
The coefficient bound lets pack_B form each byte as x | (y << 4)
without masking either input lane. Keep all output positions and NTT
arithmetic unchanged.
Purpose:
Measure the effect of two fewer vector AND operations per eight packed
coefficients on the official 1002 time.
*/
/*
References:
[1] saffah_codex_6s_agg2, duck.ac submission #103485,
https://duck.ac/submission/103485 : directly copied my accepted
packed-B single-modulus AVX2 NTT; all inherited citations follow.
No license notice is stated.
[2] Judge Duck Online, problem 1002,
https://duck.ac/problem/1002 : used the documented input-coefficient
range 0..9; no code copied and no separate license is stated.
Approach:
Skip the all-coefficient <16 check while packing B, since the problem
requires coefficients <=9. Keep the exact same packed bytes and NTT.
Purpose:
Measure whether the removed vector OR reduction and fallback branch
close the official 1002 threshold.
*/
/*
References:
[1] saffah_cc_v41_agg1, duck.ac submission #103135,
https://duck.ac/submission/103135
Directly copied its accepted one-modulus AVX2 NTT with packed B input.
Its original references are retained below; no license notice is present.
Approach:
Compile the current public leader's exact algorithm with GCC O3. The judge
uses GCC 9.3 with O2 as its default; O3 may schedule and unroll the packed
input preparation and repeated NTT butterflies more effectively.
Purpose:
Measure whether this compiler setting improves the official 1002 time
sufficiently under the current exact ranking threshold while preserving AC.
*/
#pragma GCC optimize("O3")
// ===== REFERENCES =====
// [1] duck.ac 用户 saffah_codex_6s_agg2, 提交 #102707 <https://duck.ac/submission/102707>
// 用途：本发的**直接基座**（工作区文件 problems/1002/work/i2_riv4_base.cpp）。该提交自述为
// 直接改编自本账号 #102610；本账号 #102723 是它的**逐字节复刻**（tools/rivalcopy.py:
// fwd/rev = 100.00%/100.00%），本发即在该复刻件上继续。
// [2] 本账号 saffah_cc_v41_agg1, 提交 #102723 <https://duck.ac/submission/102723>
// 用途：**直接复制**了该提交的全部实现（AVX2 精确卷积 NTT，单模数 39*2^21+1，
// 4-quarter 顶层分解 + 「拆 limb 复用已触碰内存」的零 malloc 布局、以及 splitted
// 顶层蝶形 transform_fixed_pair_split / convolve_final_split_quarter）。
// 本发与它**只差「思路」里写的那一处**。
// [3] 本账号 saffah_cc_v41_agg1, 提交 #97437 <https://duck.ac/submission/97437>
// 用途：AVX2 32 字节取数的内联汇编宏（LDU/LDA/STU/STA）直接复制自该提交。
// [4] duck.ac 用户 Qwerty1232, 提交 #28087 <https://duck.ac/submission/28087>
// 用途：参考了思想 —— 本题单模数 NTT 卷积方案的最初来源（[1] 的自述）。
// [5] duck.ac 用户 iMMIQ, 提交 #48187 <https://duck.ac/submission/48187>
// 用途：参考了思想 —— 融合 radix-4 的定长内核与固定根表（[1] 的自述）。
// ======================
// ===== 思路 =====
// 正式提交（非试验性）。本发在 #102723（= 对手 #102707 的逐字节复刻）之上**只动一处**：
// 把 B 的四个 quarter 预处理所依赖的系数流**打包成 1 字节/下标**，于是四个预处理各自
// 只需重读 512 KiB，而不是各自把整条 B（4 MiB）重读一遍。
//
// 【病灶（判题机实测，probe 的 rdtscp 相位计数）】原实现在 j 循环里 4 次调用
// prepare_quarter<J>(B, m, scratch, lg)，每次都对 i ∈ [0, q) 读 B[i] 与 B[q+i]
// ——也就是**每次都把 4 MiB 的 B 完整读一遍**（共 16 MiB 读 + 8 MiB 写）。
// 判题机相位计数：prepB = 3,728,434 cycles = 全时 6.15%（i3-8100 @3.6GHz，L3 = 6 MiB，
// B 在两次 prepB 之间会被 convolve 的 4 MiB 工作集挤出 L3 ⇒ 大部分是 DRAM 读）。
//
// 【本发】q 个下标上真正需要的只有两个系数 a = B[i]、b = B[q+i]（题面保证系数 < 10，
// 而原实现的 v 查表 `table[b & 15]` 也已经隐含 b < 16）⇒ 每个下标 4 bit + 4 bit 就够。
// 新增一趟 `pack_B`：读 B 一次，写出 packed[i] = a | (b << 4)（1 字节/下标，共 512 KiB）。
// 然后新增 `prepare_quarter_packed<J>`：用 `_mm256_cvtepu8_epi32(_mm_loadl_epi64(...))`
// 每次取回 8 个下标，`x & 15` / `x >> 4` 还原 a、b，后半段（v 查表 + 四选一）与原实现
// 逐指令相同。**数值路径与原来完全一致，只是读的是压缩副本。**
//
// 【存放位置】打包副本放在 A 的尾部 A[655360, 786432)（512 KiB）。该区段在
// prepare_quarters_split_last 里已被读遍（src[q+i]），**不新增任何页错误**；它与
// last3 = A[524288, 655360) 相邻且不重叠，因此 pack_B 必须排在 prepare_quarters_split_last 之后。
//
// 【不新增任何前提】pack_B 在打包的同时把读到的所有系数 OR 进一个累加器，返回
// 「全部 < 16」；只要有一个系数塞不下半字节就返回 false，调用方**逐字节回退**到原来的
// 未打包路径（prepare_quarter 原函数保留，未改动）。因此本发对系数范围**没有任何新假设**。
//
// 【判题机实测】probe 相位计数（同一台机器、同一段 harness，基线两次运行的相位总和
// 60,582,488 / 60,615,454 cycles，复现性 0.05%）：
// prepA 2,609,266 → 2,626,914 （不变）
// pack_B — → 886,370 （新增）
// 4×prepB 3,728,434 → 1,424,032 （**−62%**）
// 4×conv 51,318,222 → 51,286,410 （不变）
// fin 2,926,566 → 2,926,566 （不变）
// 合计 60,582,488 → 59,150,268 ⇒ **−2.36%**（缺口的 2.6 倍）
// 预期判题机用时 ≈ 17.162346 × 59.150/60.582 = **16.76 ms**，阈值（严支）17.006739 ms。
//
// 【本地正确性】problems/1002/work/i2_harness.cpp（N=1e6，系数 0..9，输出数组 c 由 harness
// 预置）全量输出 checksum = 3517366183622095521（bad=0），与基座 #102723 **逐位相同**；
// A_align32=1 C_align32=1 ⇒ 走的是 split/打包路径而不是回退分支。
// ★ **对拍必须在"紧凑分配"下做**：harness 的 A/B/C 是文件作用域的 BSS 数组，尾部有大量
// 余量，因此**测不出越界写**。本发另外用一个**守卫页 harness**（mmap 三段、每段末尾紧跟
// 一个 PROT_NONE 页、且各数组的**尾地址精确贴住**守卫页；见 problems/1002/work/t2_guard.cpp）
// 复核：基座与打包版都给出同一个 checksum `3517366183622095521` 且不崩。
// ★ 该守卫 harness 的**阳性对照**：本发第一版在 pack_B 里误用了 `load_u32x8`（= `vmovdqa`
// 对齐取数）而 B 不保证 32 字节对齐 ⇒ **该 harness 立刻 SIGSEGV**（基线在同一 harness 上
// 不崩）⇒ 该检查对本类"对齐/越界"错误有分辨力（BRIEF §2.18.556①）。已改为 `LDU`(vmovdqu)。
// 阳性对照（沿用本目录既有证据）：把 last 指向 A+8 ⇒ bad=6、
// checksum=12864537847637628433；把 fxw 改回未折 R⁻¹ 的 f ⇒ bad=6、
// checksum=3225253121880178258 ⇒ 该 harness 对本类寻址/常量错误有分辨力。
// ================
// 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 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));
}
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;
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;
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;
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;
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"))) 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);
}
}
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{});
}
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(x + y, fxv, fxsv, pv));
auto v = mts.shrink(i2_shoup32(x + mts.mod2 - y, fxv, fxsv, pv));
auto s = mts.shrink(i2_shoup32(z + t, fxv, fxsv, pv));
auto r = mts.shrink(i2_shoup32(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));
}
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);
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.653 ms | 10 MB + 672 KB | Accepted | Score: 100 | 显示更多 |