SIMD in Pure Python | Blog

SIMD in Pure Python

By David Buchanan, 4th January 2024

First of all, this article is an exercise in recreational "because I can" programming. If you just want to make your Python code go fast, this is perhaps not the article for you. And perhaps Python is not the language you want, either!

By the end, I'll explain how I implemented Game of Life in pure Python (plus pysdl2 for graphics output) running in 4K resolution at 180fps, which represents a ~3800x speedup over a naive implementation.

If you're already familiar with SIMD and vectorization, you might want to skip this next section.

A Brief Introduction to SIMD

If you want to get the most performance out of a modern CPU, you're probably going to be using SIMD instructions - Single Instruction, Multiple Data.

For example, if you have two arrays of length 4, containing 32-bit integers, most modern CPUs will let you add the pairwise elements together in a single machine instruction (assuming you've already loaded the data into registers). Hopefully it's obvious why this is more efficient than looping over the individual array elements. Intel CPUs have had this specific capability since SSE2 was introduced in the year 2000, but SIMD as a concept predates it by a lot.

Newer CPUs cores have been expanding on these capabilities ever since, meaning SIMD instructions are more relevant than ever for maximising CPU throughput.

If you're programming in a language like C, an optimising compiler will recognise code that can be accelerated using SIMD instructions, and automatically emit appropriate machine code. Compilers can't optimise everything perfectly, though, so anyone who wants to squeeze out the maximum performance might end up using Intrinsics to explicitly tell the compiler which instructions to use. And if that still isn't enough, you might end up programming directly in assembly.

There are many ways to express SIMD programs, but the rest of them are out of scope for this article!

"Vectorization" is the process of transforming a typical program into one that operates over whole arrays of data (i.e. vectors) at once (for example, using SIMD). The work done by the optimising compiler described above, or the human writing Intrinsic operations, is vectorization.

A Brief Introduction to CPython

CPython is the reference implementation of the Python language, written mostly in C, hence the name. Other implementations exist, but when people say "Python" they're often implicitly referring to CPython. I'll try to only say Python when I'm referring to the language as a whole, and CPython when I'm talking about a CPython implementation detail (which we'll be getting into later).

The TL;DR is that CPython compiles your code into a bytecode format, and then interprets that bytecode at run-time. I'll be referring to that bytecode interpreter as the "VM".

SIMD in Python

Python does not natively have a concept of SIMD. However, libraries like NumPy exist, allowing for relatively efficient vectorized code. NumPy lets you define vectors, or even n-dimensional arrays, and perform operations on them in a single API call.

import numpy as np

a = np.array([1, 2, 3, 4])<br>b = np.array([2, 4, 6, 8])

print(a + b) # [ 3 6 9 12]

Without NumPy, the above example would require a loop over array elements (or a list comprehension, which is fancy syntax for the same thing, more or less).

Internally, NumPy is implemented using native C extensions, which in turn use Intrinsics to express SIMD operations. I'm not an expert on NumPy implementation details, but you can peruse their SIMD code here. Note that the code has been customised for various CPU architectures.

CPython itself, being an interpreted Python implementation, is slow. But if you can structure your program so that all the "real work" gets done inside a library like NumPy, it can be surprisingly efficient overall.

NumPy is excellent and widely used for getting real work done. However, NumPy is not "pure" Python!

SIMD in Pure Python

By "pure," I mean using only functionality built into the Python language itself, or the Python standard library.

This is an entirely arbitrary and self-imposed constraint, but I think it's a fun one to work within. It's also vaguely useful, since libraries like NumPy aren't available in certain environments.

Earlier, I said Python doesn't natively have a concept of SIMD. This isn't entirely true; otherwise the article would end here. Python supports bitwise operations over pairs of integers: AND (&), OR (|), XOR (^). If you think about these as operations over vectors of booleans, each bit being one bool, it is SIMD!

Unlike many other programming languages, Python integers have unlimited precision. That is, they can accurately represent integers containing arbitrarily many digits—at least, until you run out of memory. This means we can evaluate an unlimited number of conceptually-parallel boolean operations with a single python operator.

SIMD...