加速 Python 中的模运算

Question

这是一个 Park-Miller 伪随机数生成器：

def gen1(a=783):
    while True:
        a = (a * 48271) % 0x7fffffff
        yield a

783 只是一个任意种子。 48271 是 Park 和 Miller 在原论文中推荐的系数（PDF：Park, Stephen K.; Miller, Keith W. (1988). "Random Number Generators: Good Ones Are Hard To Find"）

我想改进这款 LCG 的性能。文献描述了一种使用按位技巧来避免除法的方法 (source)：

素数模数需要计算双倍宽度乘积和显式缩减步骤。如果使用的模数刚好小于 2 的幂（梅森素数 2³¹-1 和 2⁶¹-1 很流行，2^{32-5 和 264-59)，减少模 m = 2e - d 可以比使用恒等式 2 的一般双宽度除法更便宜地实现e ≡ d (mod m)。}

注意模数 0x7fffffff 实际上是梅森素数 2**32 - 1，这是用 Python 实现的想法：

def gen2(a=783):
    while True:
        a *= 48271
        a = (a & 0x7fffffff) + (a >> 31)
        a = (a & 0x7fffffff) + (a >> 31)
        yield a

基本基准测试脚本：

import time, sys

g1 = gen1()
g2 = gen2()

for g in g1, g2:
    t0 = time.perf_counter()
    for i in range(int(sys.argv[1])): next(g)
    print(g.__name__, time.perf_counter() - t0)

在 pypy (7.3.0 @ 3.6.9) 中性能有所提升，例如生成 100 M 词条：

$ pypy lcg.py 100000000
gen1 0.4366550260456279
gen2 0.3180829349439591

不幸的是，在 CPython (3.9.0 / Linux) 中性能实际上降级：

$ python3 lcg.py 100000000
gen1 20.650125587941147
gen2 26.844335232977755

我的问题：

• 为什么通常被吹捧为优化的按位算术实际上比 CPython 中的模运算还要慢？
• 您能否通过其他方式提高此 PRNG 在 CPython 下的性能，例如使用 numpy 或 ctypes？

请注意，此处不一定需要任意精度整数，因为此生成器生成的数字永远不会超过：

>>> 0x7fffffff.bit_length()
31

Answer 1

我的猜测是，在 CPython 版本中，大部分时间用于开销（解释器、动态调度）而不是实际的算术运算。因此，添加更多步骤（即更多开销）并没有多大帮助。

PyPy 的运行时间看起来更像是使用 C 整数进行 10^8 模运算所需的时间，因此它可能能够使用 JIT，它没有太多开销，因此我们可以看到加速算术运算。

减少开销的一种可能方法是使用 Cython（here 是我对 Cython 如何帮助减少解释器和调度开销的一项调查），并且对生成器开箱即用：

%%cython
def gen_cy1(int a=783):
    while True:
        a = (a * 48271) % 0x7fffffff
        yield a
        
def gen_cy2(int a=783):
    while True:
        a *= 48271
        a = (a & 0x7fffffff) + (a >> 31)
        a = (a & 0x7fffffff) + (a >> 31)
        yield a

我使用以下函数进行测试：

def run(gen,N):
    for i in range(N): next(gen)

测试显示：

N=10**6
%timeit run(gen1(),N)   #  246 ms
%timeit run(gen2(),N)   #  387 ms
%timeit run(gen_cy1(),N)   # 114 ms
%timeit run(gen_cy2(),N)   # 107 ms

两个 Cython 版本都同样快（并且比原来的版本快一些），因为具有更多的操作，并不会真正花费更多的开销，因为算术运算是使用 C-int 完成的，而不再使用 Python-ints。

但是，如果一个人真的很想获得最佳性能 - 使用生成器是一个杀手，因为这意味着很多开销（例如，参见 SO-post）。

只是为了给人一种感觉，如果不使用 Python 生成器可能会发生什么 - 生成所有数字的函数（但不将它们转换为 Python 对象，因此没有开销）：

%%cython
def gen_last_cy1(int n, int a=783):
    cdef int i
    for i in range(n):
        a = (a * 48271) % 0x7fffffff
    return a

def gen_last_cy2(int n, int a=783):
    cdef int i
    for i in range(n):
        a *= 48271
        a = (a & 0x7fffffff) + (a >> 31)
        a = (a & 0x7fffffff) + (a >> 31)
    return a

导致以下时机：

N=10**6
%timeit gen_last_cy1(N)  # 7.21 ms
%timeit gen_last_cy2(N)  # 2.59 ms

这意味着如果不使用生成器，可以节省超过 90% 的运行时间！

---

我有点惊讶，调整后的第二个版本优于原来的第一个版本。通常，如果可能，C 编译器不会直接执行模运算，而是自己使用位技巧。但在这里，至少在我的机器上，C 编译器技巧不如。

由 gcc (-O2) 为原始版本生成的汇编器 (live on gotbold.org)：

        imull   $48271, %edi, %edi
        movslq  %edi, %rdx
        movq    %rdx, %rax
        salq    $30, %rax
        addq    %rdx, %rax
        movl    %edi, %edx
        sarl    $31, %edx
        sarq    $61, %rax
        subl    %edx, %eax
        movl    %eax, %edx
        sall    $31, %edx
        subl    %eax, %edx
        movl    %edi, %eax
        subl    %edx, %eax

可以看到，没有div。

这里是第二个版本的汇编器（操作少得多）：

        imull   $48271, %edi, %eax
        movl    %eax, %edx
        sarl    $31, %eax
        andl    $2147483647, %edx
        addl    %edx, %eax
        movl    %eax, %edx
        sarl    $31, %eax
        andl    $2147483647, %edx
        addl    %edx, %eax

显然，更少的操作并不总是意味着更快的代码，但在这种情况下似乎确实如此。