速度是编译器实现两倍的取模算法

platelett · 2023-01-17 18:33:05 · 个人记录

CF blog

不是直接取模，是解决模乘的问题：

给定 k,m 和 n 个数 a_i，算每个 a_i\times k\bmod m

笔者想到了一个计算方法，比编译器的实现（当 m 是常量）快了几乎一倍，这可以有效加速 NTT 和一些 DP。

首先

其中 \{x\}=x-\lfloor x\rfloor 是小数部分函数。上面式子的原理是 a_i\times k\bmod m 只和 \frac{a_i\times k}m 的小数部分有关。

设 \frac km\approx\frac p{2^{64}}，那么

这里有两个选择题。

• 最后一个式子算出来不一定是整数，要下取整还是上取整。

我们选择 \frac p{2^{64}}\ge\frac km，最后式子下取值。前者会使结果偏大，后者会使结果偏小，感性理解一下这样结果会比较准确。

定理：令 p=\lceil\frac km\times2^{64}\rceil, 当 a_i\le \frac{2^{64}}m 时，下取整的计算结果是准确的。

证明：写外面了

代码实现：

const int P = 998244353;

void calc(int n, int k, int a[]) {
    unsigned long long p = (((unsigned __int128)k << 64) + P - 1) / P;
    for(int i = 0; i < n; i++)
        a[i] = (unsigned)a[i] * p * (unsigned __int128)P >> 64;
}

几点说明：

• 代码中用了 unsigned __int128 所以必须要 64 位架构才能用。
• (unsigned)a[i] * p 是 unsigned long long 自然溢出，底层是 64 位乘法只保留低 64 位。
• * (unsigned __int128)P >> 64 底层是 64 位乘法只保留高 64 位（rdx 寄存器），和自然溢出速度一样。
• 把 a[i] 转成 unsigned 是因为 int 先转 unsigned 再转 unsigned long long 比直接转 unsigned long long 更快，后者需要 符号扩展。

速度测试。

测试方式和代码，包含两个部分：

• 第一个部分是吞吐量 (Throughput)，CPU 高度并行的用时（现代 CPU 单核也能指令级并行），共包含 50000\times50000 次模乘。
• 第二个部分是延迟 (Latency)，每次模乘依次进行而不并行的用时，共包含 50000\times25000 次模乘。

可能的输出：

Throughput test(50000 * 50000):
Compiler's signed modulo:   1954.83 ms
Compiler's unsigned modulo: 1746.73 ms
My modulo:                  1160.47 ms

Latency test(50000 * 25000):
Compiler's signed modulo:   4329.33 ms
Compiler's unsigned modulo: 3945.29 ms
My modulo:                  2397.97 ms

顺便普及一下几个常识：

• int 先转 unsigned 再转 long long 比直接转 long long 更快，但负数会出错。
• unsigned long long 对 常量 取模比 long long 快。
• 常量模乘的并行速度几乎是串行的 5 倍（笔者发明的模乘也是）。

扩展

还可以计算 (a_1b_1+a_2b_2+\cdots+a_nb_n)\bmod m，但 \sum a_i 不能超过 \frac{2^{64}}m。

令 p_i=\lceil\frac{b_i}m\times2^{64}\rceil。