Is collections.deque really that much faster than my jax implementations of a circular buffer?

[hericks]

I am very new to jax so please bear with me. I am currently working on a circular buffer implementation in jax that I can safely pass to jitable functions. For now, the implementation looks as follows:

import jax
import jax.numpy as jnp
import chex

from flax import struct


@struct.dataclass
class CircularBufferState:
    memory: chex.Array
    index: int
    n_elements: int


def circular_buffer_reset(
    capacity: int, dummy: chex.Array
) -> CircularBufferState:
    buffer_state = CircularBufferState(
        memory=jnp.zeros((capacity, *dummy.shape)),
        index=0,
        n_elements=0
    )
    return buffer_state

@jax.jit
def circular_buffer_push(
    state: CircularBufferState, element: chex.Array
) -> CircularBufferState:
    n_state = CircularBufferState(
        memory=state.memory.at[state.index,:].set(element),
        index=(state.index+1)%(state.memory.shape[0]),
        n_elements= jnp.maximum(state.index+1, state.n_elements)
    )
    return n_state

Next, I compared the speed of my jax buffer to collections.deque.

import timeit

capacity = 1000
n_elements = 10000
rng = jax.random.PRNGKey(42)
element = jax.random.normal(rng, (3,))


# setup deque
deque_mem = deque(maxlen=capacity)

def fill_deque_mem(mem):
    for i in range(n_elements):
        mem.append(element)

print(f"fill_deque_mem(): {timeit.timeit('fill_deque_mem(deque_mem)', globals=globals(), number=1)}")


# setup jax buffer
buffer_state = circular_buffer_reset(capacity, element)

def fill_jax_buffer(state):
    for i in range(n_elements):
        state = circular_buffer_push(state, element)

@jax.jit
def jitted_fill_jax_buffer(state):

    def body_fun(i, val):
        return circular_buffer_push(val, element)

    return jax.lax.fori_loop(0, n_elements, body_fun, state)

print(f"fill_jax_buffer(): {timeit.timeit('fill_jax_buffer(buffer_state)', globals=globals(), number=1)}")
print(f"jitted_fill_jax_buffer(): {timeit.timeit('jitted_fill_jax_buffer(buffer_state)', globals=globals(), number=1)}")

From running this script, I obtain the following output.

fill_deque_mem(): 0.0005817220080643892
fill_jax_buffer(): 0.49577246489934623
jitted_fill_jax_buffer(): 0.2540220639202744

The difference in speed between fill_jax_buffer() and jitted_fill_jax_buffer() seems reasonable. However, I am unsure about the low runtime of fill_deque_mem(), especially since deque-objects aren't restricted to hold only chex.Array-objects. What ways are there for me to accelerate my buffer implementation?

I conducted a similar test as above but first generated a random ndarray called elements of shape (n_elements, 3). In iteration i, I then appended elements.at[i,:].get() to the buffers. In this test the jitted_fill_jax_buffer() outperformed fill_deque_mem(). However, I believe that this could also be due to the fact that the outer @jit is able to further optimise the elements.at[i,:].get() calls. What do you think?

[jakevdp]

If you're interested in microbenchmarks of JAX operations, please take a look at https://jax.readthedocs.io/en/latest/faq.html#benchmarking-jax-code. In particular, I think JAX's asynchronous dispatch may be muddying your results here.

But that aside, it wouldn't surprise me if your JAX implementation is much slower than the Python implementation, because in the Python version your deque is storing pointers to an already allocated array, whereas in your JAX implementation in each step you are copying the contents of one array element-by-element into another.

[hericks]

Thank you for the reply! I'll definitely check out the benchmarking guide in more detail. However, I am not really sure about the second paragraph. The JAX buffer is also pre-allocated and I believe that no copy should be made, since the initial buffer is never used again and everything is done under jit. Am I misunderstanding something here?

[jakevdp]

The buffer is pre-allocated, but each time you add a row to the array, those values must be copied to that buffer. In the deque case, you don't have any buffer, so you're adding pointers to a linked list. Adding a pointer to a linked list is fundamentally a much cheaper operation than copying the contents of one buffer to another buffer, even if the second buffer is already allocated. Does that make sense?