mimalloc: Un Deep Dive en la Asignación de Memoria Concurrente a Escala

La gestión eficiente de la memoria en sistemas distribuidos y concurrentes a gran escala es un problema fundamental de la computación. Con la proliferación de aplicaciones multihilo y el crecimiento exponencial de los requisitos de memoria, especialmente en cargas de trabajo de Machine Learning (LLMs), los asignadores de memoria tradicionales se convierten en cuellos de botella críticos debido a la contención por recursos compartidos. mimalloc aborda este desafío mediante un diseño que prioriza la localidad de caché y minimiza la sincronización explícita, permitiendo un alto rendimiento en entornos con cientos de hilos y cientos de gigabytes de memoria.

Históricamente, los asignadores de memoria han evolucionado desde simples listas de libres hasta complejos sistemas que intentan equilibrar la fragmentación, el rendimiento y la sobrecarga. Sin embargo, la escala de concurrencia actual exige un replanteamiento. mimalloc se posiciona como una solución que, inspirada en principios de algoritmos aleatorizados y estructuras de datos distribuidas, logra una escalabilidad casi lineal al reducir la probabilidad de contención en las rutas de código más frecuentes.

Arquitectura del Sistema

La arquitectura de mimalloc se basa en el principio de 'heaps' thread-local, denominados 'theaps'. Cada hilo mantiene su propio 'theap', que a su vez gestiona un conjunto de 'páginas' de mimalloc, típicamente de 64 KiB. Cada página contiene bloques de un tamaño fijo, organizados en 'size classes' para mitigar la fragmentación interna. Esta segregación thread-local permite que la mayoría de las operaciones de asignación y liberación se realicen sin sincronización explícita, mejorando la localidad de caché y reduciendo la latencia.

Para cada página de mimalloc, existen tres listas de libres: una 'free list' principal para asignaciones, una 'local_free' list para bloques liberados por el mismo hilo, y una 'thread_free' (atómica) list para bloques liberados por otros hilos. Las operaciones de liberación entre hilos utilizan un 'compare-and-swap' atómico para añadir bloques a la 'thread_free' list de la página. La baja probabilidad de contención en estas listas atómicas, debido a la gran cantidad de páginas y, por ende, de listas de libres, es clave para la escalabilidad. Cuando una página se vacía, se puede liberar al sistema operativo o se puede 'robar' por otro hilo que necesite memoria, un mecanismo similar al 'work-stealing' en los 'thread pools' modernos, que optimiza el uso de la memoria sin sacrificar la escalabilidad. La función 'mi_ptr_page' en mimalloc v3 utiliza un mapa de memoria bajo demanda para recuperar metadatos de página, mejorando la robustez ante punteros inválidos.

void* mi_malloc( size_t size )
{
  mi_theap_t* const theap = mi_get_thread_local_theap();
  if (size > MI_MAX_SMALL_SIZE) return mi_malloc_generic(theap,size); // slow generic path
  const size_t index = (size + sizeof(void*))/sizeof(void*); // round size
  mi_page_t* const page = theap->small_pages[index];
  mi_block_t* const block = page->free; // head of free list
  if (block == NULL) return mi_malloc_generic(theap,size); // slow generic path
  page->free = block->next; // pop free list
  page->used++;
  return block;
}

void mi_free(void* p)
{
  mi_page_t* const page = mi_ptr_page(p); // get the page meta-data that contains p
  if (page==NULL) return;
  if (mi_thread_id() == page->thread_id) { // do we own this page?
    mi_block_t* const block = (mi_block_t*)p;
    block->next = page->local_free; // push on the `local_free` list
    page->local_free = block;
    if (--page->used == 0) mi_page_free(page); // is the entire page free?
  }
  else {
    mi_free_cross_thread(page, p); // free in a page owned by another thread
  }
}

void mi_free_cross_thread(mi_page_t* page, mi_block_t* block)
{
  mi_block_t* tfree = mi_atomic_load(&page->thread_free); // head of the thread free list
  do {
    block->next = tfree; // push our block in front
  } while (!mi_atomic_compare_and_swap(&page->thread_free, &tfree /*expect*/, block /*new*/))
}

Fundamentos Teóricos

El diseño de mimalloc, particularmente su uso de múltiples listas de libres por página y la dependencia de operaciones atómicas poco contendidas, tiene paralelismos con los principios de algoritmos aleatorizados. Así como la aleatorización puede simplificar el balanceo de estructuras de datos complejas como árboles binarios (por ejemplo, árboles de búsqueda binarios aleatorizados), mimalloc utiliza la distribución de miles de listas de libres para reducir la probabilidad de contención. Este enfoque contrasta con asignadores que dependen de estructuras de datos concurrentes más complejas y costosas.

La tensión entre escalabilidad y eficiencia en el uso de la memoria, resuelta en mimalloc mediante el 'page stealing', es un problema clásico en sistemas distribuidos y concurrentes. Se relaciona con el concepto de 'load balancing' y 'resource sharing' en entornos multihilo, donde la asignación de recursos exclusivos a cada trabajador (hilo) maximiza el rendimiento individual pero puede llevar a la subutilización de recursos, mientras que la compartición global maximiza la utilización pero introduce contención. mimalloc busca un punto óptimo, permitiendo la compartición de páginas solo cuando es necesario y de forma eficiente.