TEC Techniques Technique

Vectorization

Also known as: vectorization · vectorised code · array programming · numpy · SIMD · vector instructions · ufunc

A Python loop adds up a million numbers in 68 milliseconds, numpy in under half a millisecond, though the number of additions is the same. Where the interpreter’s time goes, which operations on a real photograph vectorise hundreds of times over, which do not vectorise at all, and why adding fractions cannot simply be reordered.

Intermediate 20 minutes array measurement images parallel numerical precision

The photograph “Water and mountains” is 600 by 305 pixels, 183 thousand points. Turning it grey takes a few multiplications and additions per pixel. A Python loop does it in 92 milliseconds, numpy in under a millisecond, though the arithmetic is the same. Below are four operations on this picture, and for each a race of methods on real measurements: a Python loop, numpy, a Go loop and your browser. Choose an operation and a size: one photograph, a phone snapshot, or an album of a hundred pictures.

Grey, histogram and blur in numpy catch up with the Go loop, but dithering does not: numpy has no ready-made operation for it, and a Python loop over a numpy array turns out two and a half times slower than a loop over an ordinary list. Let us work out why, step by step.

The idea

In a Python loop each number is a separate object somewhere in memory, and each iteration is dozens of interpreter steps: read the bytecode, find out the type, look up addition for that type, create a result object. The arithmetic itself is a small share of this work. Vectorization moves the loop inside a compiled library: the program says “add these two arrays”, and the loop over a million elements spins in machine code over a dense array of numbers of one type. There the processor can also apply vector instructions, which handle 16–64 bytes at a time.

The word “vectorization” means two different things, and it helps to tell them apart. For a programmer in Python, R or MATLAB it is a way of writing: whole arrays instead of loops. For a compiler and a processor it is a way of executing: one instruction over several numbers. The first almost always leads to the second, but the main gain comes from the first alone: escaping the interpreter.

Where a loop’s time goes

The simplest task is to count the spaces in War and Peace. That is five million bytes, one comparison each.

A Python loop spends about 40 nanoseconds per byte, on the measuring machine more than a hundred clock cycles per comparison. A call to bytes.count is the same loop, but in C, and it is twenty times faster. numpy compares the whole array with a space using vector instructions and counts the ones, seven times faster still: 0.3 nanoseconds per byte, 130 times faster than the Python loop. In Go an ordinary loop is already machine code, 0.9 nanoseconds per byte, while bytes.Count is written in assembly and compares 16 bytes per instruction on arm64 and 32 on amd64: 15 times faster again.

All the lines on the chart are parallel: the time of every method grows as n. Vectorization does not change the asymptotics; it changes the constant factor, sometimes by hundreds of times.

Code: what vectorises and what does not

The four operations from the scene are four different cases. The listing does each one twice, with a Python loop and with whole numpy arrays, and checks that the answers agree to the pixel.

photo 600 x 305 = 183000 pixels, 549000 bytes of colour
brightness: mean 122.78, from 0 to 228
histogram: the most common brightness 211, 2384 pixels
3 x 3 blur of the inner pixels: sum 22272633
Floyd-Steinberg: 96337 black and 86663 white pixels
  • An elementwise operation: grey. Each pixel is independent: brightness = (299r + 587g + 114b + 500) div 1000. In numpy it is the same formula, only over three columns of the array at once. Almost 120 times faster than the loop.
  • Counting: the histogram. The loop “counts[v] += 1” writes to scattered cells and cannot be written as a formula over an array. But numpy has a ready-made function for this operation, np.bincount: the same loop, only in C. 21 times faster.
  • Neighbours: a 3 × 3 blur. A new pixel is the average of nine neighbours. If you shift the whole picture by a pixel left, right, up and down, nine slices of one array without copying, then the sum of the nine slices is the sum of neighbours for all pixels at once. Almost 500 times faster than the loop, and this is the best case: there is a lot of work in the loop, and in numpy it turns into nine additions of whole arrays.
  • A dependency: Floyd–Steinberg dithering. A pixel is rounded to black or white, and the rounding error is passed to the neighbours to the right and below. So the next pixel depends on how the previous one was rounded, and all pixels cannot be computed at once. numpy has no “dither” operation, and nothing is gained.

The last case is the most instructive. If you take a numpy array and walk it with the same Python loop, it gets not faster but slower, two and a half times, for both dithering and the histogram. Each access v[i] to a numpy array creates a new numpy number object, while a Python list holds ready-made objects, and for numbers from −5 to 256 they are even the very same pre-created objects. numpy is fast only when the loop is inside it, not outside. Such operations need a compiled language, Go runs the dithering 40 times faster than the Python loop, or a compiler for Python such as Numba or Cython.

Sometimes the dependency can be sidestepped by rewriting the order of traversal. A dithering pixel waits for four neighbours: the one on the left and three in the row above, directly above, to the left and to the right. If each next row lags behind the previous one by two pixels, then pixels with the same x + 2y no longer depend on each other and can be computed simultaneously. This is how dithering is parallelised when it really has to be.

Vector instructions and a sum of fractions

Adding a million fractions with vector instructions means spreading them over lanes: eight accumulators, the first adding numbers 0, 8, 16, …, the second 1, 9, 17, …, and at the end the eight partial sums are added up. It would seem a C or Go compiler could do this by itself. It does not, and here is why.

exact:            499774542.9491545
1 accumulator:    499774542.94915307  off by -1.43e-06
4 accumulators:   499774542.949157  off by +2.5e-06
8 accumulators:   499774542.94915396  off by -5.36e-07

The answers differ. Adding fractions is not associative: (1016 + 1) − 1016 in a double gives 0, because the one is lost when the sum is rounded, while (1016 − 1016) + 1 gives 1. The order of additions changes the result, and a compiler, which must execute the program as written, has no right to reorder them. C compilers do it only with permission, with flags such as -ffast-math, while integers are added in lanes freely: there the order changes nothing. numpy, on the contrary, adds fractions its own way on purpose, pairwise, in chunks, and this is not only faster but usually more accurate than a plain loop.

The Go compiler does not generate vector instructions for such loops, but lanes help even without them.

With one accumulator Go spends about a nanosecond per number: each addition waits for the result of the previous one, and adding fractions takes the processor several cycles. With four accumulators, 0.24; with eight, about 0.2: the processor runs four or eight independent chains simultaneously. These are not vector instructions but parallelism within one core, yet the reason for the gain is the same: independence. And so is the trap: eight accumulators put into an array [8]float64 give 0.53, twice as bad as eight separate variables. The Go compiler keeps such an array in memory rather than in registers, and each addition becomes a read and a write. Only a measurement shows this.

Python loses to everyone in this comparison: the loop takes 65 nanoseconds per number, the built-in sum 11, numpy 0.5.

What it costs

The asymptotics do not change: a vector operation on n numbers is the same O(n) work as a loop. What changes is the cost of one step, and it has its own pitfalls.

  • A call costs. Each numpy call, parsing arguments, checking types, allocating memory for the result, takes a few microseconds. On the sums chart this shows on the left: on a thousand numbers numpy spends 6.5 nanoseconds per number, only half as much as sum, and on a million 0.5. Vectorising a hundred operations on arrays of ten numbers is pointless; combining them into one operation on a 100 × 10 array is useful.
  • Temporary arrays. The expression (a * b + c) // d on numpy arrays creates two intermediate arrays the size of the original and one more for the result. On a million numbers this is unnoticeable, but on hundreds of millions memory runs out before time does. In the article’s second problem the obvious vector formulation asks for almost a gigabyte.
  • Memory, not arithmetic. When an operation is simple, such as adding two arrays, numpy is limited not by the processor but by how fast data arrives from memory. Then two passes over an array instead of one cost twice as much, and one operation that does everything in a single pass is better. More on this in the articles on the cache and on latencies.

When it will not help

  • Each step depends on the previous one. Dithering, generations of a cellular automaton, time steps of a simulation, dynamic programming where the next cell waits for the previous. You can vectorise only along the direction with no dependency: in the Life of the first problem the generations go strictly one after another, but within a generation all cells are computed at once.
  • Branches on data. “If the element is above a threshold, one thing, otherwise another” turns into a mask: compute both branches for all elements and select by the mask. If the branches are expensive and one of them is almost never needed, that is wasted work.
  • The code is already compiled. In Go, C, Rust or Java a loop is already machine code. There “writing with arrays” is not needed for speed, and the compiler inserts vector instructions itself where it can, for integers but not for a sum of fractions, or they are taken from libraries written in assembly.

Where it works

  • Scientific computing in Python. numpy, pandas, SciPy and PyTorch are operations on whole arrays written in C, C++ and Fortran. Good code on them contains almost no loops over elements.
  • Images and sound. Filters are sums of shifted copies of a signal, like the blur above; brightness, contrast, blending are elementwise operations.
  • Standard libraries. Finding a byte, comparing and copying memory, counting: memchr and memcpy in C, bytes.Count and bytes.IndexByte in Go are written with vector instructions separately for each processor.
  • Machine learning. Matrix multiplication, the basis of neural networks, runs on the processor’s vector instructions, and more often on a graphics card, where thousands of simple cores work on the same principle: one instruction, much data.

History

The idea of writing computations over whole arrays is older than vector processors. In 1962 Kenneth Iverson described in the book “A Programming Language” the APL language, where adding two arrays, summing along an axis and selecting by a mask are ordinary operations of the language, without loops. APL gave rise to J, K and q, while MATLAB, R and numpy made this style familiar to scientists and analysts.

Vector instructions appeared in supercomputers. Seymour Cray’s Cray-1 (1976) held 64 numbers in each vector register and processed them in a pipeline, one result per cycle. They came to mass-market processors for the sake of video and sound: MMX in 1997, SSE in 1999, NEON on ARM in the mid-2000s, AVX in 2011.

numpy grew out of two libraries. Numeric was written by Jim Hugunin in 1995, later came Numarray, and in 2005–2006 Travis Oliphant merged them into numpy. In 2020 its authors described in Nature how the library became the foundation of scientific Python, from the image of a black hole’s shadow to the detection of gravitational waves.

Problems

In the first problem the generations of Life cannot be computed simultaneously, but the cells within a generation can: vectorise over the field, loop over generations. In the second a photograph is converted to a palette, and the main thing is not only to remove the loop but also not to create an array that will not fit in memory.

Life on a large field A photo in sixteen colours

Data: Photographs by the site's author War and Peace by Leo Tolstoy, public domain