< Previous post | Next post >

[2026-09-30][Recap]

Using Arena allocators to reduce synchronization barriers.

Context

When developing my machine learning framework, I quickly realized I would have to implement a memory management policy. My framework (DAHL) is based on a task-based parallel computing library (StarPU). Globally, a main thread submit tasks synchronously, and worker threads execute them asynchronously. The tasks execute operations on data structures (vectors, matrices, ...) in the main memory (RAM). Obviously, we have to synchronize at some point, particularly when freeing memory. Before freeing a block of memory, we have to ensure that every task working with this block have finished.

To this extent, we thought that using Arenas looked fitting. Particularly because we can group related memory together (Datasets, Weights and Biases, Temporary computation, layers results...), as they have similar lifetimes. However, as stated before, we need to synchronize all tasks together. In our implementation, we simply keep track of the data handles inside of each arenas, then we call blocking functions to ensure no task is working on this part of memory.

Here is a schema of the principle of arenas, in this implementation I use a linked list to arbitrary grow the size of the buffer:

Arena principle: 1-2. Allocating only advances the offset pointer; 3. new allocation doesn't fit the buffer, we create a new block; 4. reset only consists of moving the offset pointer; 5. previous memory gets overwritten.

It has a few advantages: the code is way simpler to write and memory in managed efficiently (Allocate one time, and reuse the same buffers). However, we still have to deal with synchronization barriers which slows down the training in the transition between each batch. The next gantt diagram shows the parallelization of a CNN tasks during 5 batches. Green is tasks being executed, red is sleeping time when waiting for other tasks completion.

Single/Naive: Parallelization of my CNN without any synchronization optimization.

On the figure above, we can clearly identify the transition between each batch, as we need to synchronize tasks to clear Arenas containing batch data. A way to fix this is to use two arenas for the batch computations:

Using a single arena, we need to wait for the completion of each tasks.
Using two arenas, the we block the previous batch, which will be already finished.

While gantt bellow does not indicate clear advantages, we will see that this trick shows more effectiveness paired with another technique.

Double/Naive: Parallelization of my CNN with double batch arena.

Another obvious problem concerns loading images data for each batch. Here we have similar constraints: we want to load images at the beginning of a batch and free them at the end. Here, we can use the same trick as before by using two arenas. But what could be useful would be to prefetch image batches so that they are ready in advance. For that we can use a third arena to hold prefetch'd batches. We name each version of the implementation as naive, prefetch, advanced. The advanced implementation actually uses the previous trick by keeping the previous arena to reset it after.

Depiction of each image loader types implemented: naive (load when needed), prefetch (but still blocking), advanced (prefetch + non-blocking).

The first gantt used the naive loader, so let's look at the prefetch and advanced loaders. We look at results for double batch arenas.

Using the prefetch implementation we don't notice that much improvements. However when combining double batch arenas + the advanced image loader we get near perfect scheduling:

Some workers might still have some sleep time in some areas, but we don't notice synchronization barriers as we did before. The prefetch implementation had no apparent effect on the task dependencies as, anyways we are blocked by the previous batch. Prefetching just allows to directly start computing after the synchronization.

For later experiments, we should rework the prefetch loader into a "double arena" implementation instead, just to evaluate the trick separately.

Experiments

I ran a bunch of experiments to compare the combination of batch arena/loaders. On Rammus we don't notice anything in particular. It seems that these techniques do not improve the performances on this server. This is probably linked to the fact that my tasks are too tiny for this server, as noticed in my previous paper. Thus StarPU already struggles anyways to schedule correctly and isn't bound by synchronizations.

Results are way more interesting on my tinier servers:

Some observations:

However on poppy, we notice small differences:

< Previous post | Next post >