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

Este artigo apresenta a portabilidade bem-sucedida do código de simulação de plasma de escala de girocinética TRIMEG para arquiteturas de GPU usando a API portátil OpenMP, demonstrando acelerações significativas e precisão física verificada através de paralelização híbrida MPI-OpenMP e simulações de modo de Gradiente de Temperatura de Íons.

O essencial

A Visão Geral: Cozinhando uma Tempestade Cósmica

Imagine tentar prever o tempo dentro de uma estrela. No mundo real, não podemos simplesmente colocar um termômetro dentro do sol ou de um reator de fusão; é quente e caótico demais. Em vez disso, os cientistas usam supercomputadores para executar "simulações virtuais" de plasma (um gás superquente e eletricamente carregado).

O código TRIMEG é uma receita específica, muito sofisticada, para simular este plasma. Ele rastreia bilhões de pequenas partículas (como grãos individuais de areia em uma tempestade) para ver como elas giram, colidem e criam turbulência. O problema? Esta receita é incrivelmente pesada. Executá-la em um computador padrão (CPU) é como tentar mover uma montanha com uma única colher. Leva tempo demais.

O Objetivo: O autor, Giorgio Daneri, queria acelerar isso usando GPUs (Unidades de Processamento Gráfico). Pense na CPU como um único chef mestre que é muito inteligente, mas só consegue picar um vegetal por vez. Uma GPU é como uma cozinha com 10.000 ajudantes de cozinha que podem picar vegetais simultaneamente. A tese trata de descobrir como fazer essa receita de um único chef mestre funcionar perfeitamente com um exército de 10.000 ajudantes, e fazendo isso de uma forma que funcione para duas marcas diferentes de cozinhas (NVIDIA e AMD).

O Desafio: O Problema do "Tradutor Universal"

O autor escolheu uma ferramenta chamada OpenMP para fazer a tradução. Pense no OpenMP como um tradutor universal que diz ao computador: "Ei, pegue esta parte da receita e entregue ao GPU".

No entanto, o autor encontrou dois grandes obstáculos:

1. O Erro do "Compilador": O software que traduz o código (o compilador) não era perfeito. Era como tentar usar um tradutor universal que às vezes esquece como dizer "sal" ou "calor". O autor teve que reescrever partes do código para se ajustar às peculiaridades do tradutor. Por exemplo, o código usava "polimorfismo" avançado (uma maneira elegante de dizer objetos que podem mudar de forma ou identidade). Os tradutores (compiladores) para as GPUs não entendiam essa mudança de forma, então o autor teve que achatar as formas em caixas rígidas para fazê-las funcionar.
2. O "Engarrafamento": Mover dados entre o computador principal (CPU) e a GPU (os ajudantes de cozinha) é lento. Se você ficar parando para entregar ingredientes de um lado para o outro, os ajudantes ficam ociosos. O autor teve que reestruturar o código para que todos os ingredientes fossem movidos para a GPU de uma só vez no início, em vez de transportá-los constantemente de um lado para o outro.

A Solução: Reestruturando a Cozinha

Para fazer o código rodar em GPUs NVIDIA e AMD, o autor teve que realizar uma "cirurgia" no código TRIMEG:

• Achatando o Mapa: O código usava um mapa complexo para encontrar onde as partículas estavam. Esse mapa era como um arquivo bagunçado. O autor o achatou em uma lista única e direta para que a GPU pudesse lê-lo instantaneamente sem se perder.
• Corrigindo a "Corrida": Às vezes, quando milhares de ajudantes tentam escrever no mesmo quadro branco ao mesmo tempo, eles escrevem uns sobre os outros (uma "condição de corrida" ou race condition). O autor encontrou pontos onde o código fazia isso e o corrigiu para que cada um escrevesse em sua própria faixa.
• O Compromisso "Um Tamanho Serve para Todos": Como as duas marcas de GPU (NVIDIA e AMD) falam linguções ligeiramente diferentes, o autor teve que criar uma versão única de código que funcione para ambas, mesmo que isso significasse usar alguns "contornos" (como usar um tipo específico de alocação de memória que funcione para ambas, mesmo que não seja o absolutamente mais rápido para uma delas).

Os Resultados: Funcionou?

O autor testou a nova versão em GPU contra a antiga versão em CPU usando dois "casos de teste" famosos (como testes de direção padrão para um carro novo):

1. O Caso do Ciclone: Uma simulação simplificada de turbulência de plasma.
2. O Caso TCV-X21: Uma simulação mais complexa e realista envolvendo a borda do plasma.

O Veredito:

• Velocidade: A versão em GPU foi significativamente mais rápida. Em alguns testes, foi quase 30 vezes mais rápida que a versão em CPU ao rodar em uma única máquina.
• Precisão: Os resultados da GPU coincidiram quase perfeitamente com os resultados da CPU. Os "padrões climáticos" (crescimento de energia e estruturas de turbulência) pareciam os mesmos.
• Portabilidade: O código rodou com sucesso tanto em hardware NVIDIA quanto AMD sem a necessidade de ser completamente reescrito para cada um.

A Ressalva (Limitações)

O autor é honesto sobre as limitações:

• O "Tradutor" ainda não é perfeito: Os compiladores (o software que transforma o código em linguagem de máquina) para essas GPUs ainda estão amadurecendo. Às vezes, eles produzem resultados matemáticos ligeiramente diferentes do CPU, o que pode causar pequenos erros ao longo do tempo.
• Descompasso de Hardware: Se você tiver um computador com muitos núcleos de CPU, mas apenas uma GPU, a GPU pode ficar sobrecarregada se você tentar alimentá-la com muitas tarefas de uma vez. O autor descobriu que, para obter os melhores resultados, é necessário equilibrar quantos "chefes" (processos MPI) você tem versus quantos "ajudantes de cozinha" (threads de GPU) estão disponíveis.
• Não existe "Solução Mágica": Embora a parte de movimentação de partículas do código tenha recebido um enorme ganho de velocidade, outras partes da simulação (como resolver as equações do campo magnético) ainda rodam na CPU porque as ferramentas para mover essas partes específicas para a GPU ainda não estão prontas.

Resumo

Em suma, esta tese é uma história de engenhosidade de engenharia. O autor pegou um código de simulação pesado, lento e complexo e conseguiu ensiná-lo a rodar em placas de vídeo modernas e poderosas. Eles navegaram por um campo minado de bugs de software e limitações de compiladores para criar uma versão que funciona em dois tipos diferentes de hardware, provando que podemos simular o plasma de fusão muito mais rápido sem perder a precisão. É um passo crucial para tornar a pesquisa de energia de fusão mais eficiente, embora a jornada para uma tradução perfeita e totalmente automatizada ainda não tenha terminado.

Resumo técnico

Resumo Técnico: Aceleração por GPU e Portabilidade do Código TRIMEG para Simulações de Plasma Girocinético usando OpenMP

Definição do Problema
Simulações de física de plasma, particularmente modelos girocinéticos usados para estudar instabilidades e turbulência em dispositivos de fusão tokamak, são computacionalmente intensivas. O código TRIMEG, um solver de partícula-em-célula (PIC) de alta precisão que utiliza um método de elementos finitos C1 em malhas triangulares não estruturadas, enfrenta desafios significativos de tempo de execução devido ao número massivo de partículas (frequentemente $10^7$ a $10^8$) necessárias para simulações realistas. Embora o código já empregue MPI para paralelismo de múltiplos nós, o empuxo de partículas (particle pushing) e as operações de grade-para-partícula (G2P) constituem o gargalo primário, consumindo até 80% do tempo total de execução. O desafio reside em acelerar esses "pontos críticos" (hotspots) específicos utilizando Unidades de Processamento Gráfico (GPUs) enquanto se mantém a portabilidade entre diferentes arquiteturas de hardware (especificamente NVIDIA e AMD) e preserva a complexa estrutura orientada a objetos do código, que inclui polimorfismo e tipos derivados.

Metodologia
O estudo foca na portabilidade do código TRIMEG para arquiteturas de GPU usando a API de offloading do OpenMP (versão 4.0 e posteriores). A metodologia envolveu:

1. Seleção de Alvos: O kernel de empuxo de partículas e as operações G2P associadas (retrocesso/pullback, cálculo de densidade e interpolação da função de distribuição) foram identificados como os alvos primários para o offloading devido à sua alta intensidade aritmética e falta de dependências entre partículas.
2. Reestruturação do Código para Portabilidade: Foi necessária uma reestruturação significativa para superar as limitações dos compiladores tanto no amdflang (AMD) quanto no nvfortran (NVIDIA). Os principais desafios incluíram:
   • Polimorfismo: Ambos os compiladores apresentaram dificuldades com tipos derivados class() e procedimentos vinculados a tipos (type-bound procedures) dentro de regiões alvo do OpenMP. A solução envolveu a refatoração do código para usar declarações type() não polimórficas sempre que possível e a implementação de um contorno para dependências circulares entre classes de partículas e de campos usando hierarquias de classe base/estendida e diretivas INCLUDE do Fortran para duplicar corpos de funções.
   • Arrays Dinâmicos: O código utilizava uma biblioteca customizada que mimetiza vetores de C++ para arrays dinâmicos. Como os kernels de GPU não conseguem lidar facilmente com alocação dinâmica ou indireções de ponteiros complexas, a estrutura de mapeamento entre caixas delimitadoras (bounding boxes) e triângulos da malha foi "achatada" de um array de estruturas para arrays 1D para facilitar transferências de memória eficientes.
   • Gerenciamento de Memória: A alocação preemptiva de memória foi implementada durante a fase de inicialização para minimizar a latência em tempo de execução. Para plataformas AMD, o Unified Shared Memory (USM) foi aproveitado onde possível, enquanto diretivas explícitas de enter data, update e exit data foram utilizadas para plataformas NVIDIA sem suporte a USM.
   • Consistência Numérica: Para garantir que os resultados da GPU coincidissem com a referência da CPU, otimizações de compilador que alteravam a aritmética de ponto flutuante (como instruções Fused-Multiply-Add) foram desativadas (-ffp-contract=off para AMD, -Mnofma para NVIDIA). Condições de corrida (race conditions) na biblioteca de B-splines externa foram resolvidas mudando de membros de objeto compartilhados para arrays automáticos declarados localmente.
3. Avaliação de Desempenho: A implementação foi testada nos clusters Viper (AMD MI300A), Raven (NVIDIA A100) e Pitagora (NVIDIA H100). O desempenho foi avaliado através de:
   • Perfilamento de Kernel: Uso de rocprof-compute e nsys para analisar ocupação de recursos, largura de banda de memória e misturas de instruções.
   • Estudos de Escalabilidade: Testes de escalabilidade forte foram conduzidos para avaliar a eficiência do offloading híbrido MPI-OpenMP, examinando especificamente o impacto da sobrecarga (oversubscribing) de GPUs com múltiplos processos MPI.
   • Exploração de Tamanho de Grade: Ajuste do número de equipes (teams) OpenMP e threads por equipe para maximizar a utilização do hardware.

Principais Contribuições

• Primeira Portabilidade entre Fornecedores: Este trabalho apresenta um esforço pioneiro na portabilidade de um código girocinético Fortran complexo e orientado a objetos para GPUs NVIDIA e AMD usando um único código base via offloading OpenMP.
• Contornos de Compilador: A tese documenta contornos específicos e não triviais para as limitações de compiladores em relação a polimorfismo, arrays dinâmicos e ponteiros de procedimento em regiões alvo do OpenMP. Destaca a falta de documentação abrangente para nvfortran e amdflang a respeito dessas funcionalidades.
• Análise de Paralelismo Híbrido: O estudo fornece uma análise detalhada das trocas (trade-offs) no offloading híbrido MPI-OpenMP, demonstrando que, embora a aceleração por GPU seja eficaz para o empuxo de partículas, a falta de multithreading OpenMP no código original exige a sobrecarga de processos, o que pode limitar a escalabilidade em nós com alta contagem de núcleos, mas recursos de GPU limitados.
• Verificação Numérica: A implementação inclui um processo rigoroso de verificação comparando taxas de crescimento de energia e estruturas de modo 2D contra resultados de CPU, confirmando que a versão GPU reproduz a física com alta fidelidade, apesar de pequenas variações numéricas causadas pelo tratamento de ponto flutuante específico de cada compilador.

Resultados

• Aceleração (Speedup): Para uma carga de trabalho realista de $32 \times 10^6$ elétrons, a implementação em GPU alcançou um speedup de aproximadamente 14.8x no nó AMD Viper e 29.6x no nó NVIDIA Pitagora em comparação com a versão de CPU compilada com GCC no cluster TOK.
• Eficiência de Kernel: O kernel de empuxo de partículas representou a maior parte do tempo de execução. O perfilamento no AMD MI300A mostrou alta intensidade aritmética com taxas de acerto de cache L1/L2 superiores a 80%, embora apenas 18% dos acessos à memória tenham sido coalescentes.
• Limitações de Escalabilidade: Testes de escalabilidade forte revelaram que, embora a parte acelerada por GPU escale bem, o aumento de velocidade global da aplicação é limitado pelas partes não aceleradas (ex: solvers de campo usando PETSc) e pelo overhead de sobrecarga de GPUs. No cluster NVIDIA Pitagora, o suporte a múltiplas GPUs via OpenMP foi considerado não funcional na versão testada do compilador (nvfortran 24.9), limitando a capacidade de utilizar todos os GPUs de um nó simultaneamente.
• Corretude: Simulações do caso Cyclone (modo ITG) e do caso TCV-X21 (instabilidade ITG não linear) confirmaram que a versão GPU reproduz corretamente as taxas de crescimento de energia e as estruturas de modo observadas na versão de CPU, com diferenças atribuídas à inicialização do gerador de números aleatórios e variações de ponto flutuante específicas do compilador, em vez de erros algorítmicos.

Significância e Alegações
O artigo afirma que, embora o offloading OpenMP ofereça um caminho promissor para a portabilidade entre diferentes arquiteturas de HPC, ele não é uma solução "transparente" para códigos legados complexos. O trabalho demonstra que alcançar uma versão de alto desempenho e funcional para GPU exige uma exploração extensiva de compiladores e uma reestruturação significativa do código para contornar as limitações atuais no suporte a recursos avançados de Fortran.

Os autores enfatizam que o sucesso dessa portabilidade depende fortemente da versão específica do compilador, e não apenas do paradigma de programação. Eles concluem que a implementação de GPU do TRIMEG é uma ferramenta funcional e precisa para simulações girocinéticas, capaz de entregar acelerações substanciais para as partes computacionalmente mais caras. No entanto, eles notam modestamente que o potencial total do hardware (particularmente nós com múltiplas GPUs) é atualmente dificultado pelo suporte imaturo de compiladores para offloading de múltiplos dispositivos e pela falta de multithreading OpenMP na estrutura de código da CPU subjacente. O trabalho serve como um guia prático e uma "documentação substituta" para outros que tentam portar códigos Fortran complexos para arquiteturas heterogêneas.