Python：已知瓶颈的 JIT

Question

有什么方法可以以某种方式使用 pypy 来编译一个函数而不是用于我的 python 程序的其余部分？

我有一个已知的瓶颈，我花费了 99% 的 CPU 时间（主要包含整数移位和 XOR），并在 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

它需要为 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

Answer 1

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

Answer 2

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

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

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

Answer 3

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

至于任意精度整数要求：对于适合“本机”（固定大小）整数的整数，您可能会获得比您在评论中提到的 2.5 倍更大的速度提升。改进是如此之大，以至于如果您的大部分计算可以使用本机整数完成，那么可能值得为这种情况使用（极快）您的函数的本机整数版本，并且仅将函数的通用版本用于不适合的值。 （并非所有问题都可以像这样分解成不同的情况，但对于那些可以分解的情况，至少值得一试这种方法。）

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)