GPU Acceleration and Portability of the TRIMEG Code for Gyrokinetic Plasma Simulations using OpenMP

Este artículo presenta la exitosa portabilidad del código de simulación de plasma de la función de girocinética TRIMEG a arquitecturas de GPU utilizando la API portátil OpenMP, demostrando aceleraciones significativas y una precisión física verificada mediante la paralelización híbrida MPI-OpenMP y simulaciones del modo de gradiente de temperatura de iones.

Lo esencial

La visión general: Cocinando una tormenta cósmica

Imagina intentar predecir el clima dentro de una estrella. En el mundo real, no podemos simplemente meter un termómetro dentro del sol o de un reactor de fusión; es demasiado caliente y caótico. En su lugar, los científicos utilizan supercomputadoras para ejecutar "simulaciones virtuales" de plasma (gas supercaliente y eléctricamente cargado).

El código TRIMEG es una receta específica, muy sofisticada, para simular este plasma. Rastrea miles de millones de partículas diminutas (como granos de arena individuales en una tormenta) para ver cómo se arremolinan, chocan y crean turbulencia. ¿El problema? Esta receta es increíblemente pesada. Ejecutarla en una computadora estándar (CPU) es como intentar mover una montaña con una sola cuchara. Toma demasiado tiempo.

El Objetivo: El autor, Giorgio Daneri, quería acelerar esto utilizando GPUs (Unidades de Procesamiento Gráfico). Piensa en una CPU como un único chef maestro que es muy inteligente, pero que solo puede picar un vegetal a la vez. Una GPU es como una cocina con 10,000 sous-chefs que pueden picar vegetales simultáneamente. La tesis trata sobre cómo lograr que la receta de ese único chef maestro funcione perfectamente con un ejército de 10,000 sous-chefs, y hacerlo de una manera que funcione para dos marcas de cocinas diferentes (NVIDIA y AMD).

El Desafío: El problema del "Traductor Universal"

El autor eligió una herramienta llamada OpenMP para realizar la traducción. Piensa en OpenMP como un traductor universal que le dice a la computadora: "Oye, toma esta parte de la receta y dásela a la GPU".

Sin embargo, el autor se topó con dos grandes obstáculos:

1. El error del "Compilador": El software que traduce el código (el compilador) no era perfecto. Era como intentar usar un traductor universal que a veces olvida cómo decir "sal" o "calor". El autor tuvo que reescribir partes del código para adaptarse a las peculiaridades del traductor. Por ejemplo, el código utilizaba "polimorfismo" avanzado (una forma elegante de decir objetos que pueden cambiar su forma o identidad). Los traductores (compiladores) para las GPUs no entendían este cambio de forma, por lo que el autor tuvo que aplanar las formas en cajas rígidas para que funcionaran.
2. El "Atasco de Tráfico": Mover datos entre la computadora principal (CPU) y la GPU (los sous-chefs) es lento. Si te detienes constantemente para entregar ingredientes de un lado a otro, los sous-chefs se quedan inactivos. El autor tuvo que reestructurar el código para que todos los ingredientes se movieran a la GPU una sola vez al principio, en lugar de transportarlos constantemente de un lado a otro.

La Solución: Reestructurando la Cocina

Para que el código funcione tanto en GPUs NVIDIA como AMD, el autor tuvo que realizar cierta "cirugía" al código TRIMEG:

• Aplanar el Mapa: El código utilizaba un mapa complejo para encontrar dónde estaban las partículas. Este mapa era como un archivador desordenado. El autor lo aplanó en una lista única y recta para que la GPU pudiera leerla instantáneamente sin perderse.
• Corregir la "Carrera": A veces, cuando miles de sous-chefs intentan escribir en la misma pizarra al mismo tiempo, escriben unos sobre otros (una "condición de carrera" o race condition). El autor encontró puntos donde el código hacía esto y lo arregló para que cada uno escribiera en su propio carril.
• El Compromiso de "Talla Única": Debido a que las dos marcas de GPU (NVIDIA y AMD) hablan lenguajes ligeramente distintos, el autor tuvo que crear una única versión de código que funcione para ambas, incluso si eso significaba usar algunos "trucos" (como usar un tipo específico de asignación de memoria que funcione para ambas, aunque no sea la más rápida para una de ellas).

Los Resultados: ¿Funcionó?

El autor probó la nueva versión de la GPU contra la antigua versión de la CPU utilizando dos "casos de prueba" famosos (como pruebas de conducción estándar para un auto nuevo):

1. El Caso del Ciclón: Una simulación simplificada de la turbulencia del plasma.
2. El Caso TCV-X21: Una simulación más compleja y realista que involucra el borde del plasma.

El Veredicto:

• Velocidad: La versión de la GPU fue significativamente más rápida. En algunas pruebas, fue casi 30 veces más rápida que la versión de la CPU al ejecutarse en una sola máquina.
• Precisión: Los resultados de la GPU coincidieron casi perfectamente con los resultados de la CPU. Los "patrones climáticos" (crecimiento de energía y estructuras de turbulencia) se veían iguales.
• Portabilidad: El código se ejecutó con éxito tanto en hardware NVIDIA como AMD sin necesidad de ser reescrito completamente para cada uno.

La Captura (Limitaciones)

El autor es honesto sobre las limitaciones:

• El "Traductor" aún no es perfecto: Los compiladores (el software que convierte el código en lenguaje de máquina) para estas GPUs aún están madurando. A veces producen resultados matemáticos ligeramente diferentes a los de la CPU, lo que puede causar errores minúsculos con el tiempo.
• Desajuste de Hardware: Si tienes una computadora con muchos núcleos de CPU pero solo una GPU, la GPU podría verse abrumada si intentas alimentarla con demasiadas tareas a la vez. El autor descubrió que, para obtener los mejores resultados, es necesario equilibrar cuántos "chefs" (procesos MPI) tienes frente a cuántos "sous-chefs" (hilos de GPU) hay disponibles.
• No hay una "Solución Mágica": Aunque la parte de movimiento de partículas del código recibió un enorme aumento de velocidad, otras partes de la simulación (como resolver las ecuaciones del campo magnético) todavía se ejecutan en la CPU porque las herramientas para mover esas partes específicas a la GPU aún no están listas.

Resumen

En resumen, esta tesis es una historia de ingenio de ingeniería. El autor tomó un código de simulación pesado, lento y complejo y logró enseñarle a ejecutarse en tarjetas gráficas modernas y potentes. Navegó por un campo de minas de errores de software y limitaciones del compilador para crear una versión que funciona en dos tipos diferentes de hardware, demostando que podemos simular el plasma de fusión mucho más rápido sin perder precisión. Es un paso crucial hacia la realización de una investigación de energía de fusión más eficiente, aunque el viaje hacia una traducción perfecta y totalmente automatizada aún no ha terminado.

Resumen técnico

Resumen Técnico: Aceleración por GPU y Portabilidad del Código TRIMEG para Simulaciones de Plasma Girocinéticos mediante OpenMP

Planteamiento del Problema
Las simulaciones de física de plasmas, particularmente los modelos girocinéticos utilizados para estudiar inestabilidades y turbulencia en dispositivos de fusión tokamak, son computacionalmente intensivas. El código TRIMEG, un resolvedor de partículas en celda (PIC) de alta precisión que utiliza un método de elementos finitos C1 sobre mallas triangulares no estructuradas, enfrenta desafíos significativos de tiempo de ejecución debido al número masivo de partículas (a menudo de $10^7$ a $10^8$) requeridas para simulaciones realistas. Aunque el código ya emplea MPI para el paralelismo multinodo, el empuje de partículas (particle pushing) y las operaciones de la rejilla a la partícula (G2P) constituyen el cuello de botella principal, consumiendo hasta el 80% del tiempo total de ejecución. El desafío radica en acelerar estos "puntos calientes" (hotspots) específicos utilizando Unidades de Procesamiento Gráfico (GPU), manteniendo al mismo tiempo la portabilidad entre diferentes arquitecturas de hardware (específicamente NVIDIA y AMD) y preservando la compleja estructura orientada a objetos del código, que incluye polimorfismo y tipos derivados.

Metodología
El estudio se centra en el portado del código TRIMEG a arquitecturas de GPU utilizando la API de offloading de OpenMP (versión 4.0 y posteriores). La metodología consistió en:

1. Selección de Objetivos: Se identificaron el núcleo del empuje de partículas (particle pusher kernel) y las operaciones G2P asociadas (retroceso o pullback, cálculo de densidad e interpolación de la función de distribución) como los objetivos primarios para el offloading debido a su alta intensidad aritmética y la falta de dependencias entre partículas.
2. Reestructuración del Código para la Portabilidad: Se requirió una reestructuración significativa para superar las limitaciones de los compiladores tanto en amdflang (AMD) como en nvfortran (NVIDIA). Los desafíos clave incluyeron:
   • Polimorfismo: Ambos compiladores tuvieron dificultades con los tipos derivados class() y los procedimientos vinculados a tipos (type-bound procedures) dentro de las regiones objetivo de OpenMP. La solución consistió en refactorizar el código para utilizar declaraciones de tipo type() no polimórficas siempre que fuera posible e implementar un método de trabajo para las dependencias circulares entre las clases de partículas y de campos utilizando jerarquías de clases base/extendidas y directivas INCLUDE de Fortran para duplicar los cuerpos de las funciones.
   • Arreglos Dinámicos: El código utilizaba una biblioteca personalizada que emula los vectores de C++ para arreglos dinámicos. Dado que los núcleos de GPU no pueden manejar fácilmente la asignación dinámica o las indirecciones de punteros complejas, la estructura de mapeo entre las cajas delimitadoras (bounding boxes) y los triángulos de la malla se "aplanó" de un arreglo de estructuras a arreglos unidimensionales para facilitar transferencias de memoria eficientes.
   • Gestión de Memoria: Se implementó la asignación de memoria preventiva durante la fase de inicialización para minimizar la latencia en tiempo de ejecución. Para plataformas AMD, se aprovechó la Memoria Compartida Unificada (USM) donde fue posible, mientras que para plataformas NVIDIA se utilizaron las directivas explícitas enter data, update y exit data ante la falta de soporte para USM.
   • Consistencia Numérica: Para asegurar que los resultados de la GPU coincidieran con la referencia de la CPU, se desactivaron las optimizaciones del compilador que alteran la aritmética de punto flotante (como las instrucciones Fused-Multiply-Add), utilizando -ffp-contract=off para AMD y -Mnofma para NVIDIA. Las condiciones de carrera en la biblioteca externa de B-splines se resolvieron cambiando los miembros de objetos compartidos por arreglos automáticos declarados localmente.
3. Evaluación del Rendimiento: La implementación se probó en los clústeres Viper (AMD MI300A), Raven (NVIDIA A100) y Pitagora (NVIDIA H100). El rendimiento se evaluó mediante:
   • Perfilado de Núcleos (Kernel Profiling): Uso de rocprof-compute y nsys para analizar la ocupación de recursos, el ancho de banda de memoria y la mezcla de instrucciones.
   • Estudios de Escalabilidad: Se realizaron pruebas de escalabilidad fuerte para evaluar la eficiencia del offloading híbrido MPI-OpenMP, examinando específicamente el impacto de la sobre suscripción (oversubscribing) de GPUs con múltiples procesos MPI.
   • Exploración del Tamaño de la Rejilla: Ajuste del número de equipos (teams) de OpenMP y de hilos por equipo para maximizar la utilización del hardware.

Contribuciones Clave

• Primer Portado entre Proveedores: Este trabajo presenta un esfuerzo pionero en el portado de un complejo código girocinético de Fortran, orientado a objetos, tanto a GPUs NVIDIA como AMD utilizando un único código base mediante offloading de OpenMP.
• Soluciones de Trabajo del Compilador (Compiler Workarounds): La tesis documenta soluciones específicas y no triviales para las limitaciones de los compiladores respecto al polimorfismo, arreglos dinámicos y punteros de procedimiento en las regiones objetivo de OpenMP. Destaca la falta de documentación exhaustiva de nvfortran y amdflang respecto a estas características.
• Análisis de Paralelismo Híbrido: El estudio proporciona un análisis detallado de las compensaciones en el offloading híbrido MPI-OpenMP, demostrando que, si bien la aceleración por GPU es efectiva para el empuje de partículas, la falta de multihilo de OpenMP en el código original requiere la sobre suscripción, lo que puede limitar la escalabilidad en nodos con un alto número de núcleos pero recursos de GPU limitados.
• Verificación Numérica: La implementación incluye un proceso de verificación riguroso comparando las tasas de crecimiento de energía y las estructuras de modos 2D contra los resultados de la CPU, confirmando que la versión de GPU reproduce la física con alta fidelidad a pesar de las menores desviaciones numéricas causadas por el manejo de punto flotante específico de cada compilador.

Resultados

• Aceleración (Speedup): Para una carga de trabajo realista de $32 \times 10^6$ electrones, la implementación de GPU logró una aceleración de aproximadamente 14.8x en el nodo AMD Viper y 29.6x en el nodo NVIDIA Pitagora en comparación con la versión de CPU compilada con GCC en el clúster TOK.
• Eficiencia del Núcleo: El núcleo del empuje de partículas representó la mayor parte del tiempo de ejecución. El perfilado en el AMD MI300A mostró una alta intensidad aritmética con tasas de acierto en caché L1/L2 superiores al 80%, aunque solo el 18% de los accesos a memoria fueron coalescentes.
• Limitaciones de Escalabilidad: Las pruebas de escalabilidad fuerte revelaron que, si bien la parte acelerada por GPU escala bien, la velocidad del aplicativo en su totalidad está limitada por las partes no aceleradas (por ejemplo, los resolvedores de campo usando PETSc) y la sobrecarga de la sobre suscripción de GPUs. En el clúster NVIDIA Pitagora, se encontró que el soporte para múltiples GPUs vía OpenMP no era funcional en la versión del compilador probada (nvfortran 24.9), lo que limita la capacidad de utilizar todos los GPUs disponibles en un nodo simultáneamente.
• Corrección: Las simulaciones del caso Cyclone (modo ITG) y del caso TCV-X21 (inestabilidad ITG no lineal) confirmaron que la versión de GPU reproduce correctamente las tasas de crecimiento de energía y las estructuras de modo observadas en la versión de CPU, atribuyendo las diferencias a la inicialización del generador de números aleatorios y a las variaciones de punto flotante específicas del compilador, en lugar de errores algorítmicos.

Significancia y Reivindicaciones
El artículo sostiene que, si bien el offloading de OpenMP ofrece un camino prometedor para la portabilidad entre diferentes arquitecturas de HPC, no es una solución "transparente" para códigos heredados complejos. El trabajo demuestra que lograr una versión de GPU funcional y de alto rendimiento requiere una exploración extensiva de los compiladores y una reestructuración significativa del código para sortear las limitaciones actuales en el soporte de los compiladores para funciones avanzadas de Fortran.

Los autores enfatizan que el éxito de esta portabilidad depende fuertemente de la versión específica del compilador más que del paradigma de programación en sí. Concluyen que la implementación de la GPU de TRIMEG es una herramienta funcional y precisa para simulaciones girocinéticas, capaz de entregar aceleraciones sustanciales para las partes computacionalmente más costosas del código. Sin embargo, señalan modestamente que el pleno potencial del hardware (particularmente los nodos con múltiples GPUs) se ve actualmente obstaculizado por el soporte de compiladores inmaduro para el offloading de múltiples dispositivos y la falta de multihilo de OpenMP en la estructura de código de la CPU subyacente. El trabajo sirve como una guía práctica y una "documentación sustituta" para otros que intenten realizar portados similares de códigos Fortran complejos a arquitecturas heterogéneas.