GPU acceleration of plane-wave density functional theory calculations in Abinit

Este artículo presenta la implementación en GPU del código Abinit para cálculos de teoría del funcional de la densidad, detallando las revisiones algorítmicas necesarias para la diagonalización iterativa y comparando el rendimiento de los nodos CPU frente a los híbridos CPU-GPU, con especial énfasis en la eficiencia de los algoritmos de gradiente conjugado bloqueado localmente óptimo y filtrado polinómico de Chebyshev.

Lo esencial

Imagina que Abinit es un chef de renombre mundial que cocina platos muy complejos: la "receta" es la estructura de los materiales (átomos, electrones) y el "plato" es la simulación de cómo se comportan. Para cocinar este plato, el chef necesita calcular la posición y el movimiento de miles de "ingredientes" (electrones) al mismo tiempo.

Hasta hace poco, este chef trabajaba con cuchillos normales (los procesadores de computadora o CPU). Funcionaba bien, pero para un banquete gigante (simulaciones a gran escala), tardaba horas o días en terminar.

Este artículo cuenta la historia de cómo el equipo de Abinit decidió equipar a su chef con brazos robóticos superpotentes (las GPUs, o tarjetas gráficas) para acelerar la cocina. Aquí te explico cómo lo hicieron y qué descubrieron, usando analogías sencillas:

1. El Reto: De Cuchillo a Brazo Robótico

Los procesadores tradicionales (CPU) son como un chef muy inteligente que hace muchas cosas a la vez, pero una por una. Las tarjetas gráficas (GPU) son como un ejército de robots que pueden hacer miles de tareas simples al mismo tiempo.

El problema es que la "receta" original de Abinit estaba escrita para el chef solitario. Tuvieron que reescribirla para que los robots pudieran entenderla. No se trata solo de poner el robot en la cocina; hay que cambiar la forma en que se cortan los ingredientes.

2. La Estrategia: El "Lote" (Batch Processing)

Imagina que tienes que cortar 1,000 patatas.

• Método antiguo (CPU): Cortas una patata, la pones en un plato, cortas la siguiente, la pones en otro... (una por una).
• Método nuevo (GPU): Pones las 1,000 patatas en una gran bandeja y usas una máquina que las corta todas de un solo golpe.

En el lenguaje de la física, esto se llama "procesamiento por lotes". En lugar de calcular la energía de un electrón a la vez, Abinit ahora agrupa a miles de electrones y les aplica las fórmulas matemáticas a todos juntos. Esto es lo que hace que la GPU sea tan rápida: no pierde tiempo en "cambiar de tarea".

3. El Viaje de los Datos: No te quedes atascado en la puerta

Una de las mayores dificultades fue el transporte de datos.

• La CPU es la oficina central (donde se planifica todo).
• La GPU es la fábrica de producción (donde ocurre la magia).

Antiguamente, la oficina enviaba un dato a la fábrica, la fábrica lo procesaba, lo devolvía a la oficina, y luego enviaba el siguiente. ¡Era como si el chef tuviera que cruzar la cocina a pie cada vez que necesitaba un cuchillo!

El equipo optimizó esto: ahora, envían todo el lote de ingredientes a la fábrica al principio, la fábrica hace todo el trabajo sin salir, y solo al final devuelve el plato terminado. Esto evita que los robots esperen a que lleguen los ingredientes.

4. Los Dos Métodos de Cocina: Dos formas de encontrar la solución

Para resolver la ecuación matemática principal (encontrar los estados de los electrones), el equipo probó dos métodos diferentes, como si fueran dos técnicas de cocina distintas:

• Método A (LOBPCG): Es como un chef que prueba un poco, ajusta, prueba de nuevo, ajusta, y así sucesivamente. Es muy preciso, pero requiere muchas "pruebas y errores" (comunicaciones entre los robots) que ralentizan el proceso. Es como si cada vez que cortaba una patata, tuviera que preguntar al jefe si va bien.
• Método B (Filtrado de Chebyshev): Es como usar un filtro mágico que separa automáticamente los ingredientes buenos de los malos en grandes cantidades. Este método hace muchas operaciones matemáticas rápidas seguidas antes de hacer una pausa para revisar.

El hallazgo clave: El Método B (Chebyshev) funcionó mucho mejor en los robots (GPUs). ¿Por qué? Porque permite hacer más trabajo "en caliente" (cálculos puros) antes de tener que detenerse a coordinarse. El Método A se quedaba atascado esperando a que los robots se hablaran entre sí.

5. El Resultado: Velocidad y Ahorro de Energía

Al probar esto en superordenadores reales (con tarjetas NVIDIA y AMD), descubrieron:

• Velocidad: Con las GPUs, las simulaciones se hicieron 10 a 17 veces más rápidas. Lo que antes tardaba un día, ahora tarda horas.
• Eficiencia: Aunque las GPUs consumen mucha energía, al terminar el trabajo tan rápido, el gasto total de energía fue mucho menor que usando solo los procesadores tradicionales durante todo el día.
• Escalabilidad: Cuando añadieron más robots (más GPUs), el sistema siguió funcionando bien, aunque hubo un límite: la parte final de "organizar el plato" (llamada procedimiento de Rayleigh-Ritz) seguía siendo un cuello de botella, como si el chef tuviera que decorar el plato manualmente al final, sin importar cuántos robots cortaran antes.

En Resumen

Este artículo es la historia de cómo un equipo de científicos tomó un software de simulación de materiales muy complejo y lo adaptó para que funcionara con la potencia de las tarjetas gráficas modernas.

No fue solo "conectar el cable"; tuvieron que reorganizar toda la cocina:

1. Agrupar tareas en lotes grandes.
2. Mantener los datos dentro de la fábrica (GPU) el mayor tiempo posible.
3. Elegir la receta matemática (Filtrado de Chebyshev) que mejor se adaptara a la velocidad de los robots.

El resultado es que ahora podemos simular materiales más complejos y grandes en menos tiempo, lo que acelera el descubrimiento de nuevos medicamentos, baterías más eficientes y materiales para la energía del futuro. ¡Es como pasar de cocinar con cuchillo a usar una fábrica de alimentos automatizada!

Resumen técnico

A continuación presento un resumen técnico detallado del artículo "GPU acceleration of plane-wave density functional theory calculations in Abinit" en español, estructurado según los puntos solicitados.

1. El Problema

Las simulaciones de gran escala en ciencia de materiales a nivel cuántico, específicamente utilizando la Teoría del Funcional de la Densidad (DFT) basada en ondas planas, requieren resolver el problema de Kohn-Sham para miles de estados electrónicos. Este proceso es computacionalmente costoso debido a:

• Diagonalización iterativa: La necesidad de diagonalizar grandes matrices de Hamiltoniano para obtener las funciones de onda electrónicas.
• Limitaciones de hardware: Las arquitecturas tradicionales de CPU (multinúcleo) están alcanzando un techo en el rendimiento para estos cálculos masivos.
• Desafío de portabilidad: Migrar estos códigos a sistemas de alto rendimiento (HPC) modernos con aceleradores GPU no es trivial. No basta con usar bibliotecas de proveedores; se requieren revisiones algorítmicas profundas para identificar operaciones matemáticas eficientes en GPU (álgebra lineal, FFTs) aplicadas a funciones de onda distribuidas en memoria.

2. Metodología

El equipo de investigación ha reescrito desde cero la versión GPU de Abinit (una herramienta de código abierto para simulación de materiales), implementando un modelo de programación híbrido MPI + OpenMP con offloading a GPU.

Estrategias Clave de Implementación:

• Procesamiento por Lotes (Batch Processing): Se agruparon los datos para operar en bloques más grandes en lugar de procesar bandas electrónicas individualmente. Esto permite ejecutar muchas FFTs (Transformadas Rápidas de Fourier) y multiplicaciones de matrices de forma asíncrona, maximizando el paralelismo SIMD y reduciendo la sobrecarga de memoria.
• Residencia de Datos en GPU: La función de onda ($\Psi$) se transfiere a la memoria de la GPU al inicio de cada iteración del campo autoconsistente (SCF) y se mantiene allí durante la diagonalización, minimizando las transferencias costosas entre anfitrión (CPU) y dispositivo (GPU). Solo se transfieren datos para calcular la densidad electrónica y para la siguiente iteración.
• Distribución de MPI: Se implementó una distribución 2D de procesadores para la función de onda. Se utilizan dos distribuciones: por filas (para solvers de álgebra lineal) y por columnas (para la aplicación del Hamiltoniano y FFTs). La transposición entre estas distribuciones se realiza mediante comunicaciones all-to-all con MPI consciente de GPU.
• Bibliotecas de Bajo Nivel: Se utilizaron bibliotecas nativas de los proveedores (cuBLAS, cuSOLVER, cuFFT para NVIDIA; rocBLAS, rocSOLVER, rocFFT para AMD) en lugar de kernels personalizados, encapsuladas en una capa de abstracción para facilitar el mantenimiento y la portabilidad.

Algoritmos Comparados:
El estudio se centra en dos algoritmos de diagonalización iterativa:

1. LOBPCG (Locally Optimal Block Preconditioned Conjugate Gradient): Basado en vectores, requiere ortogonalización entre bloques y múltiples pasos de Rayleigh-Ritz.
2. Filtrado Polinómico de Chebyshev: Basado en espectro, filtra los autovalores indeseados mediante polinomios, evitando dependencias entre bloques y reduciendo la comunicación MPI.

3. Contribuciones Clave

• Portado Completo de Abinit: Se logró una implementación moderna y eficiente en arquitecturas multi-GPU, superando intentos anteriores que usaban métodos de recursión o bibliotecas específicas de NVIDIA.
• Análisis Teórico de Costos: Se desarrollaron modelos teóricos para estimar el costo computacional (FLOPS), el volumen de comunicación MPI y la intensidad aritmética de los algoritmos LOBPCG y Chebyshev.
• Capacidad de Comparación entre Proveedores: Se establecieron métricas para comparar el rendimiento y la eficiencia energética entre nodos con GPUs NVIDIA (A100, H100) y AMD (MI250X, MI300).
• Optimización de la Intensidad Aritmética: Se demostró que el filtrado de Chebyshev maximiza la intensidad aritmética al aplicar el Hamiltoniano $k$ veces por transposición de datos, mientras que LOBPCG requiere más comunicaciones y operaciones de ortogonalización que son limitadas por la memoria.

4. Resultados

Los resultados se obtuvieron en supercomputadores como Jean Zay (NVIDIA), Adastra (AMD) y Topaze.

• Aceleración y Escalabilidad:
  • Las GPUs NVIDIA mostraron factores de aceleración superiores a los de AMD en general.
  • En el caso de prueba de un sistema de Ti (255 átomos, 4096 bandas), 2 nodos con GPU NVIDIA superaron a 8 nodos solo CPU.
  • El paso de Filtrado se beneficia enormemente de la aceleración GPU, mientras que el paso de Rayleigh-Ritz (diagonalización en el subespacio) es un cuello de botella, especialmente en arquitecturas AMD debido a un rendimiento inferior de la rutina hegvd de LAPACK.
• Eficiencia Energética:
  • Los nodos híbridos CPU+GPU en NVIDIA mostraron factores de ahorro energético significativos (hasta 15x) en comparación con nodos solo CPU.
  • En AMD, el ahorro fue menor (3x-5x) debido a la menor aceleración del paso de Rayleigh-Ritz.
• Comparación de Algoritmos (LOBPCG vs. Chebyshev):
  • Chebyshev es superior en GPU: Aumentar el grado del polinomio en Chebyshev mejora la precisión de los autovectores sin degradar significativamente el rendimiento, permitiendo una convergencia SCF más rápida con menos iteraciones totales.
  • Limitaciones de LOBPCG: Aumentar las líneas de minimización en LOBPCG tiene un efecto modesto en la aceleración y añade carga computacional innecesaria sin reducir las iteraciones SCF. Además, LOBPCG sufre más por la sobrecarga de comunicación MPI.
  • Modelo Roofline: Los análisis mostraron que la aplicación del Hamiltoniano es compute-bound (limitada por cómputo) y se beneficia de los núcleos Tensor, mientras que la ortogonalización y Rayleigh-Ritz son memory-bound (limitadas por memoria) y de baja intensidad aritmética.

5. Significado e Impacto

Este trabajo representa un avance significativo en la simulación de materiales a gran escala:

• Viabilidad de Sistemas Grandes: Permite realizar cálculos DFT de ondas planas en sistemas mucho más grandes y complejos en tiempos razonables, aprovechando la potencia de los superordenadores modernos.
• Guía de Algoritmos: Establece que para arquitecturas GPU, los algoritmos basados en filtrado espectral (como Chebyshev) son preferibles a los métodos de gradiente conjugado por bloques (LOBPCG) debido a su menor dependencia de la comunicación y mayor intensidad aritmética.
• Eficiencia Energética: Demuestra que la computación en GPU no solo acelera los cálculos, sino que reduce drásticamente el consumo energético por simulación, un factor crítico para la sostenibilidad del HPC.
• Disponibilidad: El código está disponible en la versión 10.6 o posterior de Abinit, permitiendo a la comunidad científica aprovechar estas optimizaciones en hardware heterogéneo (NVIDIA y AMD).

En conclusión, el éxito del portado de Abinit no se debió solo a la transferencia de código, sino a una reingeniería algorítmica que priorizó operaciones intensivas en cómputo y minimizó las operaciones limitadas por memoria y comunicación, adaptando la lógica de la DFT a la arquitectura de los aceleradores modernos.