加速CPython中的模操作

Question

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

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

783只是一个任意的种子。48271是由Park和Miller在原论文中推荐的系数(PDF: Park, Stephen K.;基思·w·米勒(1988)。"随机数生成器:好的很难找到")

我想提高这个LCG的性能。文献描述了一种使用按位技巧来避免除法的方法(来源):

素模需要计算双宽积和显式约简步骤。如果使用一个刚好小于2的幂的模(梅森素数2³¹−1和2⁶¹−1是常用的，还有2³²−5和2⁶⁴−59)，化简模m = 2^e−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中的模操作还要慢?
• 你能在CPython下以其他方式提高这个PRNG的性能吗，也许使用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 -int。

然而，如果你真的想获得最好的性能-使用生成器是一个杀手，因为它意味着大量的开销(参见这个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)为原始版本生成的汇编程序(在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

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