This article is the 18th day article of MicroAd Advent Calendar 2020.

Introduction

numba is a library for speeding up python code.

Simply annotate the function you want to speed up as follows, and you will be able to execute it at the speed of native machine code.

import numba

@numba.jit
def f(x, y):
    return x + y

You can easily speed up the code by using numba, but it will compile unlike normal python code. If you use it without understanding the mechanism, the execution speed of the entire process may slow down.

In this article, I'll show you some surprising numba pitfalls and how to avoid them.

no use python mode

There are two modes in numba: __no python __ where optimization is performed and __object __ which is the same compilation method as usual.

The default setting is to switch between these two modes automatically. no python mode takes precedence, but if this fails it will automatically switch to object mode.

In addition to not being optimized in object mode The execution time may take longer due to the compilation overhead.

So, in order to utilize numba, understand the processing that does not support optimization, It is important to devise a function such as cutting out a function so that such processing is not included.

no python mode non-compliant processing
・ Dict
・ Set
・ Pandas
・ Yield from
・ Set ・ dict ・ Generator comprehension

Also, in numba, you can set an error if you can not compile in no python mode, so Basically, it is recommended to use the following writing style.

@jit(nopython=True)
def f(x, y):

@njit
def f(x, y):

Avoid compilation overhead

In numba, the timing to compile the function is automatically switched according to the implementation.

If you use it without knowing it, compilation will run in situations where processing speed is required, and The numba used for speeding up may return and become a bottleneck.

Since you are consciously manipulating the compile timing, let's understand its variations and how to write it.

Compile when the function is called for the first time

This is the default behavior.

@njit
def f(x, y):
  return x + y

Compile when the function is loaded

By specifying the type, the function will be compiled at the defined timing. This writing method is the most recommended because the type specification also has the effect of speeding up.

@njit(int32(int32, int32))
def f(x, y):
    return x + y

Use cache

Also, by setting to use cache, you can do without compiling except for the first startup. This method can be used with other type specifications, so there is no loss even if you add options for the time being.

@njit(cache=True)
def f(x, y):
  return x + y

Ahead of compile

You can also prevent compilation from happening at run time by automatically creating a compiled module. To be honest, this method is a bit cumbersome and makes the code verbose. It's a last resort when you don't want to incur compilation overhead at any time.

from numba.pycc import CC

cc = CC('my_module')

@cc.export('multi', 'i4(i4, i4)')
def f(a, b):
    return a + b

if __name__ == "__main__":
    cc.compile()
>>> from my_module import f
>>> f(3, 4)

Summary

How was that. This time, you can easily achieve high speed, but I have summarized the points to note about numba, which has a slight quirk in operation. We hope that this article will be helpful to everyone who has read it.

Tips for speeding up python code correctly with numba