Python：已知瓶颈的即时编译（JIT）

Question

有没有办法使用pypy只编译一个函数而非整个Python程序？

我已知的瓶颈是我花费99%的CPU时间（主要包括整数移位和异或），并已经在Python中进行了尽可能多的优化。除非绝对必要，否则我不想编写和维护C库。

现在我正在使用Anaconda Python，它是正常的Python加上一堆库。我会使用pypy，只是我不想确保所有其他部分的程序都可以与pypy正常工作。

是否有一种方法可以显式地仅在一个Python函数上运行JIT？

---

编辑：该函数是GF2（伽罗瓦域）中的模乘步骤。

https://bitbucket.org/jason_s/libgf2/src/a71a14a035468e703a7c24349be814419cdba2cb/src/libgf2/gf2.py?at=default

具体来说：

def _gf2mulmod(x,y,m):
    z = 0
    while x > 0:
        if (x & 1) != 0:
            z ^= y
        y <<= 1
        y2 = y ^ m
        if y2 < y:
            y = y2
        x >>= 1
    return z

它需要支持大整数，所以我不确定如何重写以使其兼容Cython。

我刚尝试使用numba的@autojit，但它失败了，因为它不知道要使用什么变量类型，并假设是小整数。我似乎无法弄清楚如何告诉它使用标准Python大整数。

Traceback (most recent call last):
  File "/Users/jason_s/Documents/python/libgf2/src/libgf2/gf2.py", line 440, in <module>
    dlog43 = GF2DiscreteLog(0x100000000065)
  File "/Users/jason_s/Documents/python/libgf2/src/libgf2/gf2.py", line 295, in __init__
    factors = _calculateFactors(poly)
  File "/Users/jason_s/Documents/python/libgf2/src/libgf2/gf2.py", line 264, in _calculateFactors
    if (e1 << m).value != 1:
  File "/Users/jason_s/Documents/python/libgf2/src/libgf2/gf2.py", line 379, in __lshift__
    return self._wrapraw(_gf2lshiftmod(self.value,k,self.poly))
  File "/Users/jason_s/Documents/python/libgf2/src/libgf2/gf2.py", line 195, in _gf2lshiftmod
    return _gf2mulmod(x,_gf2powmod(2,k,m),m)
  File "/Users/jason_s/Documents/python/libgf2/src/libgf2/gf2.py", line 189, in _gf2powmod
    z = _gf2mulmod(z,x,m)
  File "numbawrapper.pyx", line 193, in numba.numbawrapper._NumbaSpecializingWrapper.__call__ (numba/numbawrapper.c:3764)
OverflowError: value too large to convert to signed int

- Jason S

---

将它用Cython重写并编译成C扩展。生成的Cython代码可能与您当前的代码看起来没有太大区别。 - Blender

尝试使用 numba 可能会有帮助？http://numba.pydata.org/ - kindall

函数有多长？你能发一下代码吗？解决瓶颈问题有很多方法，最好的方法取决于你具体在做什么。 - Sven Marnach

@SvenMarnach：非常简短，请参见上文。 - Jason S

你需要支持哪个整数范围？64位足够吗？ - Sven Marnach

Answer 1

使用Cython呢？您可以将这个函数转换为Cython语法，然后直接编译成C语言。Cython的语法应该与Python非常接近，只需添加一些正确类型的声明即可。请参考： http://cython.org/

- Bas Swinckels

---

好的，一旦我安装了numba（其中包括Cython），它就可以正常工作了；函数的.pyx版本比我的纯Python版本快大约2.5倍；还不错，虽然可能没有我希望的那么好。 :-) - Jason S

Answer 2

不，你不能在PyPy中运行Python程序的一部分，然后在另一个Python中运行其他部分-这不仅仅是JIT，它具有完全不同的对象表示和许多其他内部机制。

如果你唯一关心的是不想确保其余程序与PyPy一起工作，请放心：几乎所有纯Python代码都可以与PyPy一起使用，唯一的例外是CPython实现细节。这些是晦涩的，很难意外地编写依赖于大多数细节的代码，而其他细节（例如文件不会自动及时关闭）不会破坏大多数程序。尝试使用PyPy运行整个程序即可。

如果还有其他PyPy问题，您可能希望将此函数翻译为C，并使用ctypes或cffi进行调用。麻烦的部分将是将其与Python连接起来（例如通过扩展模块），这就是ctypes和cffi为您完成的任务。您不需要完整的任意精度整数库，只需要具有几个非常简单操作的位数组：测试最低有效位、左/右移位、小于和按位异或。每个操作只需要一个简单的循环。如果原始的C实现仍然是瓶颈，您可能可以对所有这些循环进行矢量化。您还可以优化移位以避免完全复制任何内容。

- user395760

---

我喜欢这个答案，因为值得指出的是，PyPy非常接近于Python 2.7，因此几乎任何纯Python 2程序都可以使用PyPy工作，而无需进行任何更改。但OP没有明确表明他依赖各种模块或软件包的程度；如果其中任何一个不是纯Python（包括一些标准库模块），则很可能它们不会与PyPy一起工作。（实际上，大多数这样的包甚至不适用于不同版本的CPython。）这就是C、ctypes和cffi进入画面的时候，你也提到了这一点。 - John Y

根据我的原始帖子：“我正在使用Anaconda Python，这是普通的Python并带有一堆库”（我强调）。我不能依赖pypy支持所有这些软件包，即使我可以，我也与不擅长安装的人合作，因此我使用Anaconda，因为它使该过程变得简单。 - Jason S

Answer 3

Cython已被提到几次，但对于那些在搜索中找到此问题的人，可能没有任意精度整数要求的人们，我想指出另一种选择：ShedSkin。它通过类型推断将Python翻译为C ++。换句话说，与Cython不同，您不需要向Python添加类型声明；您只需编写Python代码。但是，您必须确保您的变量不会在程序内更改类型，否则类型无法推断。某些Python的较动态特性也不受支持，但这通常不是计算密集型代码的问题。

至于任意精度整数的要求：对于适合于“本地”（固定大小）整数的整数，您可能会获得比您在评论中提到的2.5倍更大的速度提升。速度提升如此巨大，以至于如果您的大部分计算可以使用本地整数完成，那么仅针对该情况拥有一个（极快的）本地整数版本的函数几乎是值得的，并且仅在值不匹配时使用函数的一般版本。（并非所有问题都可以分解成单独的情况，但对于可以这样做的问题，至少可以尝试这种方法。）

- John Y

Answer 4

我实际上创建了一个Python包，可以实现你想要完成的功能。它在Galois域上实现了numpy数组。我使用numba进行JIT编译来优化它。它还支持像你提到的任意大的整数。https://github.com/mhostetter/galois

In [1]: import numpy as np                                                                     

In [2]: import galois                                                                          

In [3]: GF = galois.GF(2**100)                                                                 

In [4]: print(GF.properties)                                                                   
GF(2^100):
  characteristic: 2
  degree: 100
  order: 1267650600228229401496703205376
  irreducible_poly: Poly(x^100 + x^57 + x^56 + x^55 + x^52 + x^48 + x^47 + x^46 + x^45 + x^44 + x^43 + x^41 + x^37 + x^36 + x^35 + x^34 + x^31 + x^30 + x^27 + x^25 + x^24 + x^22 + x^20 + x^19 + x^16 + x^15 + x^11 + x^9 + x^8 + x^6 + x^5 + x^3 + 1, GF(2))
  is_primitive_poly: True
  primitive_element: GF(2, order=2^100)

In [5]: A = GF.Random((2,2)); A                                                                
Out[5]: 
GF([[853109427014456778157610146134, 957579797461316189986596054496],
    [619061399056446662243502059294, 762244553049699515226100523552]],
   order=2^100)

In [6]: B = GF.Random((2,2)); B                                                                
Out[6]: 
GF([[511709585928908961018234228239, 206374347029181859532039074035],
    [795530021671674913470994904012, 918203712488921499667394325749]],
   order=2^100)

In [7]: A + B                                                                                  
Out[7]: 
GF([[1005796869832339943233379227481, 1152773228746217240881872950547],
    [1097497412292991510222532386002, 160953539027946640002225213141]],
   order=2^100)

In [8]: A * B                                                                                  
Out[8]: 
GF([[296314095771552265037299061152, 688536035673482273067277820628],
    [1177970297984569800118703939222, 537328370564266356643331706738]],
   order=2^100)

In [9]: A / A                                                                                  
Out[9]: 
GF([[1, 1],
    [1, 1]], order=2^100)

# Fermat's Little Theorem
In [10]: A ** (GF.order - 1)                                                                   
Out[10]: 
GF([[1, 1],
    [1, 1]], order=2^100)

In [11]: A @ B                                                                                 
Out[11]: 
GF([[1027776659503691614378629238339, 470187664463292378234435322369],
    [86178777179053209582733631256, 172677144553521647820627674227]],
   order=2^100)

In [12]: np.linalg.inv(A) @ A                                                                  
Out[12]: 
GF([[1, 0],
    [0, 1]], order=2^100)

- Matt Hostetter