Python:已知瓶颈的JIT

时间:2013-10-13 21:17:20

标签: python jit pypy numba

有没有办法以某种方式使用pypy来编译一个函数而不是我的python程序的其余部分?

我有一个已知的瓶颈,我花了99%的CPU时间(主要包含整数移位和XOR)并且在Python中尽可能地优化它。除非绝对必要,否则我不想编写和维护C库。

现在我正在使用Anaconda Python,它是带有一堆库的普通python。我会使用pypy,除非我不想确保我的所有其他程序都能正常工作。

有没有办法只在一个Python函数上显式运行JIT?


编辑:该函数是GF2(Galois字段)中的模块化乘法步骤

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

它需要为bigints工作,所以我不确定如何重写为Cython兼容。

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

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

4 个答案:

答案 0 :(得分:6)

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

如果您唯一关注的是不想确保程序的其余部分与PyPy一起使用,请放心:几乎所有纯Python代码都与PyPy一起使用,唯一的例外是CPython实现细节。这些是模糊的,很难意外地编写依赖于大多数代码的代码,而其他代码(例如文件未被及时自动关闭)将不会破坏大多数程序。试试用PyPy运行整个程序。

如果PyPy存在其他问题,您可能只想将此功能转换为C并使用ctypescffi进行调用。令人讨厌的部分是使用Python(例如通过扩展模块)将其连接起来,这是ctypescffi为你做的事情。您不需要完整的仲裁精度整数库,只需要一个具有一些非常简单操作的位数组:测试最低有效位,左/右移位,小于和按位异或。每个都只是一个微不足道的循环。如果天真的C实现仍然是瓶颈,你可以对所有这些循环进行矢量化。您也可以优化班次以避免复制任何内容。

答案 1 :(得分:6)

使用Cython怎么样?您可以将这一个函数转换为cython语法,而不是直接编译为C左右。语法应该足够接近python本身,可能只是添加一些正确类型的声明。

答案 2 :(得分:4)

Cython已经被提及了几次,但我想为那些在搜索中发现这个问题的人指出另一种选择,他们可能没有任意精度整数要求:{{3 }}。它通过类型推断将Python转换为C ++。换句话说,与Cython不同,您不会向Python添加类型声明;你只需编写Python。但是您必须确保您的变量不会在程序中更改类型,否则无法推断类型。 Python的某些其他更动态的功能也不受支持,但这通常不是计算密集型代码的问题。

对于任意精度整数要求:对于适合“原生”(固定大小)整数的整数,您可能会获得比注释中记录的2.5倍大得多的速度改进。这种改进是如此激烈,以至于如果计算的很大一部分可以用原生整数完成,那么就可以为你的函数设置一个(极快)本机整数版本的函数,并且仅将函数的常规版本用于不适合的值。 (并非所有问题都可以分解为这样的单独案例,但对于那些可能的问题,至少可以检查这种方法。)

答案 3 :(得分:1)

我实际上创建了一个 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)