加速 CPython 中的模运算

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) 来避免除法的方法：

A prime modulus requires the computation of a double-width product and an explicit reduction step. If a modulus just less than a power of 2 is used (the Mersenne primes 2³¹−1 and 2⁶¹−1 are popular, as are 2³²−5 and 2⁶⁴−59), reduction modulo m = 2^e − d can be implemented more cheaply than a general double-width division using the identity 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):degraded:

$ 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（是我对 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-整数。

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

只是为了给人一种感觉，如果不使用 Python-generators - 生成所有数字的函数（但不要将它们转换为 Python-objects 因此无开销):

%%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   271, %edi, %edi
        movslq  %edi, %rdx
        movq    %rdx, %rax
        salq    , %rax
        addq    %rdx, %rax
        movl    %edi, %edx
        sarl    , %edx
        sarq    , %rax
        subl    %edx, %eax
        movl    %eax, %edx
        sall    , %edx
        subl    %eax, %edx
        movl    %edi, %eax
        subl    %edx, %eax

正如你所见，没有div。

这里是第二个版本的汇编程序（操作更少）：

        imull   271, %edi, %eax
        movl    %eax, %edx
        sarl    , %eax
        andl    47483647, %edx
        addl    %edx, %eax
        movl    %eax, %edx
        sarl    , %eax
        andl    47483647, %edx
        addl    %edx, %eax

显然，更少的操作并不总是意味着更快的代码，但在本例中似乎是这样。

加速 CPython 中的模运算

Speeding up modulo operations in CPython

python

optimization

ctypes

numpy

lcg