Python: JIT for known bottlenecks

Question

Is there any way to use somehow use pypy just to compile one function and not for the rest of my python program?

I have a known bottleneck where I spend 99% of my CPU time (containing mostly integer shifts and XORs) and have optimized it as much as I can in Python. I don't want to have to write and maintain C libraries unless absolutely positively necessary.

Right now I'm using Anaconda Python which is the normal python with a bunch of libraries. I would use pypy except that I don't want to have to make sure that all the rest of my program works ok w/ pypy.

Is there a way to explicitly run a JIT only on one Python function?

---

edit: the function is a modular multiplication step in GF2 (Galois field)

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

specifically:

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

It needs to work for bigints so I'm not sure how to rewrite to be Cython-compatible.

I just tried numba's @autojit and it failed because it didn't know what variable types to use and assumed small integers. I can't seem to figure out how to tell it to use standard 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

What about using Cython instead? You could convert just this one function to cython syntax, which is than compiled directly to C or so. The syntax should be close enough to python itself, probably just add a few declarations of the right types.