高性能多项式计算库的设计与实现
引言
在算法竞赛和高性能计算领域,多项式运算是一个核心且高频的操作。无论是生成函数、组合数学还是信号处理,快速而高效的多项式操作都至关重要。本文将从底层优化出发,逐步构建一个支持 SIMD 加速的现代 C++ 多项式模板库,旨在提供极致性能的同时保持良好的可读性与工程结构。
本实现基于 C++20 标准,并利用 GCC 编译器对 AVX2 指令集的支持,结合编译期计算、内存池管理与融合迭代等技术,实现了包括 FFT/NTT、求逆、ln、exp、开根、插值等在内的完整"全家桶"功能。所有代码均已在洛谷平台验证并达到最优解水平。
符号约定
- 序列 $ a[l:r] $ 表示从索引 $ l $ 到 $ r-1 $ 的子序列。
- 拼接操作用 $ a + b $ 表示。
- $ \mathbb{F}_p $ 表示模质数 $ p $ 的有限域。
- $ F(x) = O(x^n) $ 表示形式幂级数中最低次项不低于 $ x^n $。
SIMD 基础与向量化编程
SIMD(Single Instruction Multiple Data)允许一条指令并行处理多个数据元素,是提升数值计算效率的关键手段。以 Intel AVX2 为例,其支持 256 位寄存器,可同时执行 8 个 32 位整数运算。
#include <immintrin.h>
#pragma GCC target("avx2")
void vector_add(const int* __restrict__ a,
const int* __restrict__ b,
int n, int* __restrict__ out) {
int i = 0;
for (; i + 8 <= n; i += 8) {
__m256i va = _mm256_loadu_si256((const __m256i*)(a + i));
__m256i vb = _mm256_loadu_si256((const __m256i*)(b + i));
__m256i vs = _mm256_add_epi32(va, vb);
_mm256_storeu_si256((__m256i*)(out + i), vs);
}
for (; i < n; ++i) out[i] = a[i] + b[i];
}
为避免未对齐访问导致的段错误,应使用 _mm256_loadu_* 和 _mm256_storeu_*。同时通过 alignas(32) 确保数据按 32 字节边界对齐,有助于提高缓存命中率。
高效模运算:Montgomery 形式
传统取模运算代价高昂。我们采用 Montgomery Reduction 技术将模乘转换为移位和加法操作。设 $ R = 2^{32} $,定义蒙域值 $ \bar{x} = xR \mod P $,则有:
预计算常量 $ P_{\text{inv}} = (-P)^{-1} \mod R $,可在无分支条件下完成规约:
struct Mont {
static constexpr u32 R = 1ULL << 32;
u32 P, P_INV;
constexpr Mont(u32 p) : P(p), P_INV(compute_inv()) {}
constexpr u32 redc(u64 x) const {
u32 m = (u32)x * P_INV;
u64 u = (u64)m * P;
return (x + u) >> 32;
}
private:
constexpr u32 compute_inv() const {
u32 x = 1;
for (int i = 0; i < 5; ++i) x *= 2 - P * x;
return -x;
}
};
配合 Lazy Reduction,允许中间结果处于 $[0, 2P)$ 范围内,进一步减少判断开销。
AVX2 加速的 NTT 实现
去除位逆序置换
标准 NTT 中的位逆序重排破坏了内存局部性。我们采用逆视角 DIF/DIT 结构,在递减步长下进行蝴蝶变换,自然输出位逆序结果,显著改善缓存行为。
Radix-4 分治
相比 Radix-2,Radix-4 将层数减少至 $ \log_4 n $,降低访存次数。每层使用如下变换矩阵:
向量化核心循环
借助 AVX2 同时处理 8 个系数。定义向量包装接口:
using vec_t = __m256i;
inline vec_t vload(const void* p) {
return _mm256_load_si256((const vec_t*)p);
}
inline vec_t vadd(vec_t a, vec_t b) {
return _mm256_add_epi32(a, b);
}
inline vec_t vmul_mont(vec_t a, vec_t b, const Mont& m) {
// 使用 _mm256_mul_epu32 实现蒙氏乘法
}
内存池设计
频繁动态分配影响性能。设计固定大小的对齐内存池,复用已释放块:
template<typename T>
class AlignedPool {
static inline std::array<std::vector<T*>, 32> pool;
public:
static T* allocate(size_t n) {
size_t cap = std::bit_ceil(std::max(n, 8uz));
int k = std::countr_zero(cap);
if (!pool[k].empty()) {
T* ptr = pool[k].back();
pool[k].pop_back();
return ptr;
}
#ifdef _WIN32
return (_aligned_malloc)(cap * sizeof(T), 32);
#else
return (T*)std::aligned_alloc(32, cap * sizeof(T));
#endif
}
static void deallocate(T* p, size_t n) {
if (!p) return;
pool[std::countr_zero(n)].push_back(p);
}
};
多项式基本操作
乘法
void polymul(Mint* f, Mint* g, size_t len) {
dif(f, len); dif(g, len);
pointwise_mul(f, g, len);
dit(f, len);
}
求逆
牛顿迭代法,结合循环卷积节省 DFT 长度:
void polyinv(const Mint* f, size_t n, Mint* out) {
out[0] = f[0].inv();
auto tmp = Pool::allocate(2 * n);
for (size_t k = 1; k < n; k <<= 1) {
copy(f, min(2*k, n), tmp, 2*k);
copy(out, k, tmp+k, k); // high part zeroed
dif(tmp, 2*k);
pointwise_mul(tmp, out, 2*k);
dit(tmp, 2*k);
clear_low(tmp, k);
dif(tmp, 2*k);
pointwise_mul(tmp, out, 2*k);
dit(tmp, 2*k);
negate_and_copy_high(tmp, out, k);
}
}
对数与指数
利用导数关系 $ \ln F(x)' = F'/F $,并通过融合迭代减少临时变量和 DFT 调用次数。对于 exp,同步更新 $ E_n $ 与其逆元 $ G_n $,避免重复求逆。
现代 C++ 封装
采用模板类封装,支持概念约束与零成本抽象:
template<u32 Mod, size_t MaxLen = SIZE_MAX>
class FPoly {
using mint = SModint<Mod>;
using pool = AlignedPool<mint, 32>;
size_t sz = 0;
pool::ptr_type data;
public:
FPoly() = default;
template<std::ranges::contiguous_range R>
requires std::convertible_to<std::ranges::range_value_t<R>, mint>
FPoly(R&& r): sz(std::ranges::size(r)), data(pool::allocate(sz)) {
bulk_convert(std::ranges::data(r), sz, data);
}
FPoly& operator*=(const FPoly& rhs);
FPoly& operator+=(const FPoly& rhs);
// ... 其他运算符
};
通过 SFINAE 或 Concepts 区分不同类型输入,实现高效初始化路径选择。
高级应用实例
线性递推 Bostan-Mori 算法
将问题转化为有理函数系数提取,利用偶奇分解递归求解:
Mint linear_recurrence(long long n, const FPoly& num, FPoly den) {
if (n == 0) return num[0] / den[0];
FPoly neg_den = den;
for (int i = 1; i < neg_den.size(); i += 2)
neg_den[i] = -neg_den[i];
FPoly new_num = num * neg_den;
FPoly new_den = den * neg_den;
extract_even_odd(new_num, n & 1);
extract_even(new_den);
return linear_recurrence(n / 2, new_num, new_den);
}
多点求值与插值
基于转置原理与分治策略,构建子树积多项式,再通过反向传播计算各点值或构造拉格朗日基底。
// 多点求值主干
void multipoint_eval(int l, int r, FPoly values, int idx = 1) {
if (l + 1 == r) { result[l] = values[0]; return; }
int mid = (l + r) / 2;
FPoly L_val = transpose_multiply(right_child[idx], values);
FPoly R_val = transpose_multiply(left_child[idx], values);
multipoint_eval(l, mid, L_val, idx * 2);
multipoint_eval(mid, r, R_val, idx * 2 + 1);
}
总结
本文介绍了一套完整的高性能多项式库构建方案,涵盖 SIMD 优化、Montgomery 模乘、Radix-4 NTT、内存池管理及 Modern C++ 接口设计。通过消除冗余操作、融合迭代逻辑与精细化控制内存访问模式,达到了当前评测系统的极限性能。该框架不仅适用于竞赛场景,也可作为教学与研究中的高效代数计算工具。