Faster Python calculations with Numba

It’s not really gonna be used in your database rest api is it.

The ahead of time compilation output is... Well.. let's say difficult to package _properly_ (compare it to Cython where it's well supported and documented). That makes it useless for production, unless you want to ship giant containers with compilers etc

Numba can compile in “no Python” mode only with a subset of Python. Eg classes is limited and is still experimental. Also I think string manipulation is slow but the doc has details.

If you want to specify the type (for example for aot or just because you want to make it clear) then the call signature is less flexible.

In short, pick any random Python library, you’d find there are very few places you can jit accelerate something effectively. It is for numeric.

Even for numerical code, it is more like writing C functions than say C++ (with classes etc).

But it does makes accelerating vectorized code very easy. Even if you have a function that uses Numpy, it is likely you can speed it up using Numba with a decorator only.

But when it doesn’t work, it might often be not very clear why you can’t until you get some experience.

JAX is a numerics library combined with autograd for machine learning. you string together operations in python in a functional style, which is then passed into XLA which is google's optimizing compiler that can target cpus, gpus and tpus to generate optimized machine code for those architectures.

numba lets you inline with a subset of python which is then compiled with llvm producing something very similar to what you would get if you applied a bunch of regexes to that subset of python to convert it from python to C. (with special bindings for numpy arrays, since they have special importance in these domains)

numba is specifically targeted at things like core numerical algorithms that are typically coded in C and fortran, and are typically comprised of solely for loops and basic arithmetic. JAX is more targeted at high level machine learning applications where the end user is stringing together more high level numerical algorithms.

i suspect that JAX would be a bad fit for custom computer vision or numerical algorithms that are used outside of the use-case of doing neural networks work.

JAX recompiles functions every time you call them with an array of a new shape, i.e. if called with an array of shape (7,10) and then one with shape (17, 12) the function is compiled twice. This is generally fine for deep learning and some other numerical applications where you do the same computations again and again over arrays with the same shape. But for data exploration like Pandas, in my experience, your data shapes are different with each call, so the repeated recompilations make it unattractive.

In Numba it only recompiles for dimension changes i.e. if shape changes from (7, 10) to (2, 10, 12). However numba does not integrate AD so if you need that, JAX is probably your best bet unless Enzyme matures and you want to look into integratig it with numba.

Cython does have a plugin for gdb https://cython.readthedocs.io/en/latest/src/userguide/debugg... however I have not used it. Generally I don't like writing cython code directly however and rather start with working python code and add types etc., cython -a really helps in those situations.

That's not quite true, cython supports parallel for loops either by directly calling openmp or via cython.parallel.prange etc. (which wraps openmp AFAIK) you can also get simd speed ups via directly calling simd intrinsics or compiling with simd support and relying on the compiler (which is admittedly is not as sophisticated as numbas simd support). For GPU support I would likely go and use cupy or arrayfire.

Sure :)

The ELI5 is that numba allows you designate specific functions in your application for compilation. That tells the library to run a just in time compiler on it. If you can do that successfully, you can wind up with that portion of your code base running at or near C-speed. There are a lot of gotchas and limitations, and it only works for a subset of the python language. But when it works, it can make loops over numpy arrays very fast, when they would normally be extremely slow in pure python. In a lot of cases, just a few compute bound processes are holding back the whole application. Numba allows you to improve those portions in a very targeted way, if you have a suitable use case.

I’ve got a concrete example because I just recently used numba to speed up some code. I’m working on a stereo camera object triangulation and tracking system. We wanted to make triangulation faster so I profiled the code and found a few hot spots. In some cases it was a matter of vectorizing some computations, or avoiding tearing apart numpy arrays and allocating new ones, but aside from those kinds of things, the system was spending most of its time in two functions: calculate Euclidean distance between two points and calculate distance of a point from a line (represented by two points). There wasn’t much refactoring possible because they were mostly pure numpy. I annotated them with numba’s JIT decorator and after the initial penalty of the JIT compiler, the functions are way faster.

Because my code is parallelized, I had to figure out how to trigger the compiler before forking because then each worker process had to do the compilation itself. There’s some decorator syntax to tell numba to precompile for specific numeric types, but I found it easier to just call the functions when you import them.

All told, it’s way easier than it would have been to implement this hotspot in C++ or something.

Others have good explanations, but let me put it super simply.

Suppose we have two very large vectors and we want to take the dot product of them. If we store these vectors in a standard Python list, write pure Python, and try to run it in CPython, it is very, very slow, at least relative to how fast almost any other non-dynamic scripting language could do it. (I have to add all those caveats because this is easy mode for a JIT and PyPy might blast through this. But CPython will be very slow.)

Suppose instead we have a data type that is a C array under the hood, and all we do is this:

    array1 = load_from_file("something")
    array2 = load_from_file("other_file")
    dot_product = array1.dot_product(array2)

Suppose "load_from_file" loads into an optimized data structure that provides Python access, but under the hood is stored efficiently, and the "dot_product" method leaves Python and directly performs the dot product in some more efficient language.

In this situation, you're running a bare handful of slow Python opcodes, a few dozen at the most, and doing vast quantities of numerical computation at whatever the speed of the underlying language is, which in this case can easily use SIMD or whatever other acceleration technologies are available, even though Python has no idea what those are.

That is the core technique used to make Python go fast. There are several ways that things get from Python to that fast code and underlying structure. The simplest and the oldest is to simply do what I laid out above directly, with external, non-Python implementations of all the really fast stuff. I believe this is the core of NumPy, though I'm not sure, but it's certainly been the core of a lot of other things in Python for decades.

There are also approaches involving writing in a Python-like language, but one that can be compiled down (Cython), approaches around trying to compile Python directly with a JIT into something that under the hood is no longer using CPython's interpretation structure (PyPy), and other interesting hybrid approaches.

CPython is generally a slow language, but it isn't always a big deal in practice because if you write Python code where a reasonable amount of most lines you write go out to faster code, the slowness of Python itself is no big deal. For example, you can write decent, performant image manipulation code in Python as long as you're stringing together image manipulation commands to some underlying library, because in that case, for each "expensive" Python opcode, you're doing a lot of "real work" at max speed. I scare-quote "expensive" because in that case, the Python opcodes may be a vanishing percentage of the cost of the work. The problem with numerical manipulation in pure CPython is that it swings the other way; each Python opcode requires many assembly-level instructions to implement, but it's inefficient to run all those instructions just to add two integers together, a single assembly opcode.

Arguably the dominant factor in sitting down with any programming language and writing reasonably perfomant code on your first pass is understanding the relative costs of things and just being careful not to write code that unnecessarily amplifies the costs of doing the "real work" with too much bookkeeping. The same 1000 cycle's worth of "bookkeeping" can be anything from a total irrelevancy if your "payload" for that bookkeeping is many millions of cycles worth of work, to a crippling amplified slowdown if your payload is only one cycle's worth of work. Even a "fast" language like C++ can be brought to a crawl if you manage to use enough abstractions and bounce through enough non-inlinable method calls and through the requisite creation and teardown of function frames, etc. just to add two numbers together, and then do it all again a few million times.

I never understood the use case for embedding python in rust. In my experience the actual high performance code is a small portion of the overall code base, and writing all the glue code file processing, plotting is typically much faster in python due to the excellent excellent ecosystem. I further forgo the ability to use the python repl if I do it this way around.

Don't get me wrong I like rust, but I would always write the high performance code to be a module that can be imported in python, write everything in rust and the use python for some specific functions only.

Oh, if you check my comment history, you'll see me complaining a bit about Julia. I'm not saying it's perfect. But numba is so clunky and fickle that I think, just in terms of comparing the two, Julia is pretty clearly superior if you don't have a strong reason to use Python.

It has excellent packages for non linear programming though. Both JuMP and JuliaSmoothOptimizers are an excellent eco system for optimization. I find it odd though that Non linear programming is a subset AI and not the reverse.

There's a package for using R libraries (https://juliainterop.github.io/RCall.jl/stable/) and one for Python libraries (https://github.com/JuliaPy/PyCall.jl). It's pretty seamless.

If you write quite a bit of numerical code you can often achieve impressive speed ups using pythran with very little changes, it's like numexpr + cython on steroids. I highly recommend checking it out. The community is quite small, so there are still some bugs, Serge the main author is really responsive.

I would not be surprised. I did post it to HN at the time, and received some feedback on how I should better optimise (which I mention at the bottom of the post).

The exercise for me was to see how I fast I could get the julia code with my background, I am quite honest about the fact that I have less experience in julia, however for me the whole reason to look for something else then cython is that it takes so much background knowledge to get the ultimate speed, if I have to do the same for julia I don't really gain anything. I would argue the amount of optimisations in the julia code are already more than what I had to do in pythran.