Wednesday, November 10, 2021

Speed up your Python code with Pyjion

There are many ways of running Python code. The most commonly used is CPython, the official Python interpreter. CPython is the command you get when you type python at the command line, the one you download from python.org.

 

CPython is an interpreter. It compiles your Python code into an intermediary language called byte code and evaluates this code whenever your code is executed. This evaluation process allows for Python's endless flexibility, but it does have a downside. Simple calculations that CPUs are very efficient at, like arithmetic of integers and floating point numbers are significantly slower in CPython than other interpreters like JavaScript's node.js. I presented a case for this problem using a famous benchmark "the n-body problem" at PyCon US 2020:

 

 

In practice, Python developers adopt tools like Cython to make numerical operations more efficient by compiling code ahead of time, using C extensions. This is an effective solution for performance, but isn't ideal for rapid iteration or experimentation, especially in Jupyter notebooks.

 

This is the problem that Pyjion aims to solve. 

 

Pyjion is a Python 3.10 package that analyses Python code and optimizes it on the fly, using a JIT, or "Just-in-Time" compiler. Pyjion can be pip installed into a CPython 3.10 installation on Linux, Mac OS X, or Windows. Pyjion can make your Python code execute faster without any code changes. For the n-body benchmark showcased in the 2019 talk, Pyjion solves the problem 3x faster than standard CPython 3.10.

 

Pyjion's main features are:

  • Profile Guided JIT Compiler
  • Native 64-bit float and integer support
  • Small, fast compiler
  • Windows, macOS and Linux
  • Intel and ARM CPU support
  • Builtin IL and ASM disassembler
  • Support for native debugging and profiling tools

Installing Pyjion

Pyjion requires:
After following the installation steps, pyjion is a python module that you can import a Python 3.10 environment:
python -m pip install pyjion​

After importing pyjion, enable the JIT using the enable function:
import pyjion; pyjion.enable()
 
Any Python code you define or import after enabling pyjion will be JIT compiled. You don't need to execute functions in any special API, it's completely transparent:
>>> def half(x):
...    return x/2
>>> half(2)
1.0
You can also execute Pyjion against any script or module:
pyjion my_script.py​

Or, for an existing Python module:
pyjion -m calendar

That's it! For the full API, checkout the documentation.

Performance improvements

In these simple functions, Pyjion will detect that the variables y, x, and z are numbers that are scoped to the function test_floats and test_ints:

 

def test_floats(n=100000):
    """ Test float/integer arithmetic """
    for y in range(n):
        x = 0.1
        z = y * y + x - y
        x *= z


def test_ints(n=100000):
    """ Test integer arithmetic """
    for y in range(n):
        x = 2
        z = y * y + x - y
        x *= z

 

These variables are never referenced outside the function, therefore it optimizes these operations by compiling them directly into the assembly instructions to leverage native CPU instructions for adding, multiplying 64-bit integers and floating point numbers. This is accomplished by using .NET 6's cross-platform JIT compiler, "ryuJIT".

 

The resulting performance is that Pyjion will execute both functions 90-95% faster, a 10-20x performance gain over Python 3.10.

Posted at https://sl.advdat.com/3mZvBna
Labels: ,