`all_bitstrings` golf

[Zacharias030]

Hi Freddi,

I was very intrigued by the simplicity of the problem statement and wanted to learn a little bit more about high performance numpy code :)

First of all, while the iterator solution is very fast, the returned data type is different and I didn't find a way to preserve its speed if I cast to np.ndarray in the end. Maybe there are smarter ways to create an array from the iterator though. np.fromiter doesn't directly apply because the iterator returns tuples, so I will consider the iterator solution as the winner of the brevity / elegance section until I learn how to generate a numpy array with it efficiently :)

Secondly, I am contributing my own solution that is closely based on your fast solution.

# reproduction of Frederik's fast solution
def all_bitstrings_old(size):
    bitstrings = np.ndarray((2**size, size), dtype=int)
    for i in range(size):
        bitstrings[:, i] = np.tile(np.repeat(np.array([0, 1]), 2 ** (size-i-1)), 2**i)
    return bitstrings

# Zach's solution
def all_bitstrings(size):
    bitstrings = np.ndarray((size,2**size), dtype=int)
    a = np.array([0,1], dtype=int)
    for i in range(size):
        bitstrings[i] = a.repeat(2**(size-i-1)).reshape(-1, 2**(size-i)).repeat(2**i, 0).reshape(-1)
    return bitstrings.T

# Zach's solution based on cartesian products
def all_bitstrings_cartesian(size):
    a = np.array((0,1), dtype=np.int64)
    bitstrings = np.empty((2,) * size + (size,), dtype=np.int64)
    for i in range(size):
        ai = a.reshape((1,)*i + (2,) + (1,)*(size-i-1))
        bitstrings[..., i] = ai
    return bitstrings.reshape(-1, size)

# Performance measurements
In [192]: %timeit all_bitstrings_old(10)
105 µs ± 803 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [304]: %timeit all_bitstrings(10)
45 µs ± 1.66 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit all_bitstrings_cartesian(10)
34.6 µs ± 1.98 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)


In [194]: %timeit all_bitstrings_old(20)
304 ms ± 9.15 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [305]: %timeit all_bitstrings(20)
134 ms ± 1.23 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit all_bitstrings_cartesian(20)
228 ms ± 6.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)


In [334]: %timeit all_bitstrings_old(24)  # 3.2GB array for int64
5.73 s ± 79.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [333]: %timeit all_bitstrings(24)
2.63 s ± 18.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit all_bitstrings_cartesian(24)
4.37 s ± 63.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

My picture of python and C kernel performance is extremely incomplete, so I will just venture wild guesses in the following.

A few ideas I had below, narrated.

Numpy

• First and most important idea: The access pattern over the columns of the numpy array looks suspicious. I change it to working primarily on rows and amortising a single big transposition in the end. This gave a performance improvement of 10% over all_bitstrings_old(10) and 100%+ on larger instances.
  A closer look at the circumstances under which numpy returns a view or copy of an array reveals that the bitstrings.T call is cheating, because the base array underneath is still transposed w.r.t. your fast solution. Adding .copy() after it ate up most of the benefits that were gained by writing the function transposed:

In [162]: b = all_bitstrings_old(10)

In [163]: %timeit b.T.copy()
7.99 µs ± 124 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [164]: b = all_bitstrings_old(20)

In [165]: %timeit b.T.copy()
250 ms ± 16.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

• Using method calls on arrays instead of function calls (np.repeat(a, ...) -> a.repeat(...)). I was surprised to see this shave off ~13µs from the all_bitstrings(10) solution (~78µs -> ~65µs).

• Not instantiating an np.array from a python list size times, but passing it by reference helps quite a bit.

In [114]: a
Out[114]: array([0, 1])

In [115]: %timeit np.repeat(a, 2 ** (10-3-1))
1.66 µs ± 18.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [116]: %timeit np.repeat(np.array([0,1]), 2 ** (10-3-1))
2.45 µs ± 221 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

• Later it turned out that even reducing the complexity of the slicing syntax improves speed in a measurable way:

In [111]: %timeit bitstrings[0]
119 ns ± 1.3 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

In [112]: %timeit bitstrings[0,:]
172 ns ± 1.54 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

• Finally, I was working on understanding the implementation of np.tile (because numba.jit doesn't support np.tile apparently). np.tile is actually implemented as a call to repeat and reshape and I felt that there is something to be gained maybe.
  On very small arrays, I can do:

In [292]: b
Out[292]: array([0, 1, 1, 1])

In [291]: np.all(np.tile(b, 3) == b.reshape(-1, b.size).repeat(3, 0).reshape(-1))
Out[291]: True

In [285]: %timeit np.tile(b, 3)
3.85 µs ± 146 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [286]: %timeit b.reshape(-1, b.size).repeat(3, 0).reshape(-1)
1.53 µs ± 37.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [287]: %timeit b.reshape(-1, b.size).repeat(3, 0).reshape(-1).copy()
1.84 µs ± 9.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [288]: %timeit np.tile(b, 3).copy()
4.24 µs ± 189 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Replacing np.tile in the solution gave another ~30% speedup (~65µs -> 45µs on all_bitstrings(10)), but not so much for bigger instances. However, I really like the symmetry that the final form of the algorithm with reshape and repeat calls exposes :)

I wonder if there are even simpler formulations that we are missing.

---

Numba

The point of the whole exercise for me was to also have a look at numba. The support of numpy is more limited than I expected, because concretely it doesn't support np.tile and also doesn't support the rework with a.repeat(2**i, axis=0) because the dispatch on axes is apparently hard to get right for numba.

In the end, I rewrote the function again using very primitive operations and being not afraid of for loops. The speedup was still nice, but only for small instances. For larger ones numpy is apparently much smarter than my for loop :)

import numpy as np
import numba

@numba.jit(nopython=True)
def all_bitstrings_jit(size):
    bitstrings = np.empty((size,2**size), dtype=np.int64)
    a = np.array([0,1], dtype=np.int64)
    for i in range(size):
        # The problem is that `.repeat()` is not supported by numba with an axis argument.
        #    bitstrings[i] = a.repeat(2**(size-i-1)).reshape(-1, 2**(size-i)).repeat(2**i, 0).reshape(-1)
        b = a.repeat(2**(size-i-1))
        
        # therefore manually copy `b` 2**i times into the rows of `c`. This is where compilation should shine.
        c = np.empty((2**i, 2**(size-i)), dtype=np.int64)
        for j in range(2**i):
            c[j] = b
        bitstrings[i] = c.reshape(-1)
    return bitstrings.T

# run once to compile, takes ~0.3 seconds
l = all_bitstrings_jit(4)

Performance on all_bitstrings(10) around 3x, but not much on the larger instances:

%timeit all_bitstrings_jit(10)
15.6 µs ± 318 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit all_bitstrings_jit(20)
130 ms ± 3.17 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit all_bitstrings_jit(24)
3.46 s ± 143 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

---

Solution via cartesian product

Adapting a (stackoverflow answer)[https://stackoverflow.com/questions/11144513/cartesian-product-of-x-and-y-array-points-into-single-array-of-2d-points] for cartesian products, here is another pretty efficient and elegant solution that uses a different idea:

Interesting to think about why this is faster for small sizes, but slower for big ones...

size = 4

def all_bitstrings_cartesian(size):
    a = np.array((0,1), dtype=np.int64)
    bitstrings = np.empty((2,) * size + (size,), dtype=np.int64)
    for i in range(size):
        ai = a.reshape((1,)*i + (2,) + (1,)*(size-i-1))
        bitstrings[..., i] = ai
    return bitstrings.reshape(-1, size)

%timeit all_bitstrings_cartesian(10)
34.6 µs ± 1.98 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit all_bitstrings_cartesian(20)
228 ms ± 6.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit all_bitstrings_cartesian(24)
4.37 s ± 63.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

I didn't manage to numba.njit this function though.
Interestingly, the memory access pattern again looks suspicious, and using the transposition trick made large instances faster, but small instances slower:

size = 4

def all_bitstrings_cartesian(size):
    a = np.array((0,1), dtype=np.int64)
    bitstrings = np.empty((size,)+ (2,) * size, dtype=np.int64)   # note the inverted (size,) dimension!
    for i in range(size):
        bitstrings[i] = a.reshape((1,)*i + (2,) + (1,)*(size-i-1))
    return bitstrings.reshape(size, -1)

%timeit all_bitstrings_cartesian(10)
60.3 µs ± 3.12 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit all_bitstrings_cartesian(20)
143 ms ± 5.74 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit all_bitstrings_cartesian(24)
2.72 s ± 10.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[frederikwilde]

Hi Zac,

Thanks for this wonderful input. Firstly, allocating the small array np.array([1, 0]) first is definitely something one should always do. Thanks for pointing this out.

I really like your version of the solution. Looping over "columns" of the final array makes a lot more sense. Here is a version of this that I find slightly easier to read. It does the same operations and takes the same amount of time:

def all_bitstrings_zac(size):
    bitstrings = np.ndarray((size,2**size), dtype=int)
    a = np.array([0,1], dtype=int)
    for i in range(size):
        bitstrings[i] = a.repeat(2**(size-i-1))[np.newaxis, :].repeat(2**i, 0).flatten()
    return bitstrings.T

And I think returning a view is totally legitimate. I tested only looking up a single row of that view and not surprisingly, this preserves the advantage.

In [7]: %timeit all_bitstrings(20)[512]                                                                    
300 ms ± 7.33 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [8]: %timeit all_bitstrings_zac(20)[512]                                                                
138 ms ± 1.31 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

So this solution is always at least as good as the previous one.

The Numba example is great as well! I will try my luck with JAX and see how far I can push it with JIT and vectorization there. It's remarkable that for large instances the Numpy code is so much faster.

Lastly, the cartesian solution is really cool! It took me a bit to understand what is happening there. My guess is that working with very high dimensional arrays and creating them via reshaping is probably inefficient.

Would you want to open a pull request to add your functions to this little collection? Or I can do it for you, if you don't mind. I would add the new all_bitstrings function, the cartesian one and the Numba jitted one. I think returning a transposed view instead of a copy is the best solution, as one might only want to access a few rows.

[frederikwilde]

I just realized that only accessing a few bit strings is not a reasonable use case. In that case you'd just generate those bit strings individually.

[frederikwilde]

I added an implementation with JAX. For small sizes it's very slow, but it seems to be the fastest for 24 bits:

In [2]: %timeit all_bitstrings_jax(10)                                                                     
1.21 ms ± 39.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [3]: %timeit all_bitstrings_jax(20)                                                                     
60.2 ms ± 1.39 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [4]: %timeit all_bitstrings_jax(24)                                                                     
1.06 s ± 47.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

I suppose the asynchronous dispatch provided by jax.vmap only really gives an advantage when the underlying function runs sufficiently long.

Additionally, one obviously saves a lot of time by working with 8-bit integers, which I didn't do here for comparability.

[Zacharias030]

Ha, this is getting really good :)
Thanks for showing such a cool JAX implementation. Really really feel that I need to learn JAX and understand precisely what is going on!

I'm happy to open a PR to add my things too.