OpenMP GPU Acceleration and Portability of TRIMEG-C1 for Electromagnetic Gyrokinetic Simulations in Tokamak Plasmas

Resumo Técnico: Aceleração de GPU via OpenMP e Portabilidade do TRIMEG-C1 para Simulações Girocinéticas Eletromagnéticas

Declaração do Problema
O código TRIMEG-C1 é um solver de célula de partícula-em-grade (PIC) girocinética baseado em malha triangular, projetado para simular instabilidades eletromagnéticas em plasmas de tokamak. Ele utiliza um método de elementos finitos de ordem superior $C^1$ para capturar física precisa com uma resolução de grade menor em comparação com métodos $C^0$. No entanto, à medida que os supercomputadores modernos dependem cada vez mais de arquiteturas heterogêneas, o código enfrenta o desafio de escalar para simulações multiescala envolvendo pequenas escalas eletrônicas e escalas de tempo de transporte. Embora o offloading de GPU ofereça um potencial de aceleração significativo para códigos PIC devido ao seu alto paralelismo, as abordagens existentes carecem de portabilidade. O OpenACC, utilizado em códigos como GEM e ORB5, é bem suportado em arquiteturas NVIDIA, mas carece de suporte robusto para GPUs AMD, que são cada vez mais prevalentes em clusters de Computação de Alto Desempenho (HPC) (por exemplo, o cluster Viper da instalação de Computação e Dados Max Planck). Por outro lado, CUDA ou HIP exigem uma reescrita significativa do código ou tradução manual de kernels. O problema principal abordado é alcançar uma implementação acelerada e portátil do TRIMEG-C1 tanto em arquiteturas de GPU NVIDIA quanto AMD, sem comprometer a estrutura existente do código Fortran ou exigir duplicação extensiva de código.

Metodologia
Os autores selecionaram o framework de offloading OpenMP como a estratégia de aceleração. Esta escolha foi motivada por sua capacidade de suportar GPUs NVIDIA e AMD dentro de um código Fortran unificado, minimizando mudanças estruturais e divergência de código. O processo de portabilidade focou nos componentes computacionalmente mais intensivos: o kernel do particle pusher e as operações de grade-para-partícula (G2P) (incluindo tratamento de pullback, cálculo de densidade e interpolação da função de distribuição).

As principais etapas metodológicas incluíram:

1. Adaptação de Kernel: O particle pusher, que evolui partículas de marcação (marker particles) de forma independente, foi envolvido em uma única grande região de alvo (target region) OpenMP para maximizar o paralelismo. Isso exigiu a refatoração do código para garantir que todas as funções chamadas, incluindo aquelas provenientes de uma biblioteca externa de interpolação de B-spline, fossem acessíveis pelo dispositivo (device-callable).
2. Contornos para Limitações de Compiladores: Esforços significativos foram dedicados para superar as limitações dos compiladores amdflang e nvfortran em relação ao offloading OpenMP. Desafios específicos incluíram:
   • Polimorfismo e Procedimentos Vinculados a Tipos: Nem o compilador nem o outro suportam totalmente esses recursos em código de dispositivo. Os autores modificaram o código para usar declarações de type não polimórficas em vez de class onde necessário, e contornaram as vinculações de procedimentos de tipo na biblioteca de B-spline.
   • Gerenciamento de Memória: Para lidar com arrays dinâmicos que mimetizam o std::vector do C++, os autores achataram layouts de dados de Estrutura-de-Arrays (AoS) em arrays unidimensionais para evitar problemas complexos de mapeamento de ponteiros em tempo de execução e erros de Falta de Memória (Out-Of-Memory - OOM) em APUs AMD MI300A.
   • Condições de Corrida: Uma condição de corrida no uso de arrays temporários compartilhados da biblioteca B-spline foi resolvida mudando para arrays automáticos declarados localmente.
   • Otimização de Desempenho: Em GPUs NVIDIA, o kernel sofreu inicialmente com degradação de desempenho devido à alocação de arrays temporários em tempo de execução. Isso foi mitigado ao mudar para arrays de tamanho fixo em tempo de compilação para funções específicas, uma alteração aplicada condicionalmente com base na arquitetura alvo.
3. Paralelização Híbrida: O estudo avaliou uma estratégia de paralelização híbrida MPI-OpenMP, avaliando o desempenho sob sobreassinatura de recursos (resource oversubscription, executando múltiplos processos MPI por GPU) para determinar a viabilidade para simulações de larga escala onde a alocação de nós é limitada.

Principais Contribuições

• Implementação de GPU Portátil: O artigo apresenta a primeira portabilidade bem-sucedida do código TRIMEG-C1 para ambas as arquiteturas de GPU NVIDIA e AMD usando OpenMP, permitindo a execução em diversos sistemas de HPC (ex: Viper, Raven, Pitagora).
• Resolução de Desafios Específicos de Compilador: O trabalho detalha contornos específicos para as limitações de compiladores Fortran no offloading OpenMP, particularmente em relação a polimorfismo, alocação dinâmica de memória e integração de bibliotecas externas.
• Benchmarking de Desempenho: Avaliações de desempenho abrangentes foram conduzidas em múltiplos clusters (AMD MI300A, NVIDIA A100, NVIDIA H100) e tamanhos de problema, analisando tanto os ganhos de velocidade (speedups) de kernels individuais quanto a escalabilidade forte.
• Validação: A correção da implementação em GPU foi rigorosamente verificada contra resultados em CPU usando casos de teste padrão, garantindo que o processo de portabilidade não introduziu imprecisões físicas.

Resultos

• Ganho de Velocidade (Speedup): Para um tamanho de problema realista de $32 \times 10^6$ partículas de marcação em um nó completo, o kernel do particle pusher alcançou um ganho de velocidade de aproximadamente 9x tanto na AMD MI300A (Viper) quanto na NVIDIA H100 (Pitagora) em comparação com a implementação de CPU mais rápida (AMD EPYC 9754).
• Sobreassinatura: O estudo descobriu que, embora a sobreassinatura de recursos de GPU (aumentando o número de processos MPI por nó) reduza o ganho de velocidade de kernels individuais devido à serialização e troca de contexto, a abordagem híbrida permanece viável. No entanto, o ganho de velocidade global do código depende fortemente do tamanho da carga de trabalho; para cargas menores ($10^6$ partículas), a vantagem diminui significativamente quando comparada a um nó de CPU totalmente utilizado devido ao overhead de transferência de dados e intensidade aritmética limitada.
• Escalabilidade: A análise de escalabilidade forte demonstrou que o código acelerado por GPU escala bem através de múltiplos nós, embora o ganho de desempenho global seja marginal quando a fase de partícula-para-grade (P2G) domina, pois esta fase se beneficia significativamente do alto número de núcleos dos nós apenas de CPU.
• Correção: Comparações das simulações do modo de Gradiente de Temperatura de Íons (ITG) tanto para o caso base Cyclone quanto para o caso TCV-X21 mostraram boa concordância entre as versões CPU e GPU. As taxas de crescimento de energia diferiram em menos de 3% para o caso Cyclone e geralmente abaixo de 10% para o caso TCV-X21 (com maiores desvios em números de modo toroidal mais altos). As estruturas de modo bidimensional também foram preservadas.
• Profiling: O profiling em GPUs AMD revelou alta pressão de registradores limitando as Unidades de Computação (CUs) ativas e baixo uso da Local Data Share (LDS), atribuído à incapacidade da API OpenMP de gerenciar explicitamente a memória compartilhada.

Significância e Alegações
O artigo alega que a portabilidade bem-sucedida do TRIMEG-C1 para arquiteturas de GPU usando OpenMP fornece uma base computacional portátil para futuros estudos de física de larga escala. Os autores enfatizam que, embora a implementação tenha exigido uma reestruturação significativa do código para navegar pelas limitações dos compiladores, o código resultante permite a simulação de instabilidades eletromagnéticas em diversas plataformas de hardware sem a necessidade de manter bases de código separadas para diferentes fornecedores.

O trabalho estabelece que o particle pusher é o principal candidato para aceleração por GPU, rendendo ganhos de velocidade substanciais para grandes contagens de partículas. No entanto, os autores notam modestamente que o desempenho global do código é atualmente limitado por outras fases (como a resolução de equações de campo e operações P2G), que são menos propensas à aceleração por GPU ou beneficiam-se mais do paralelismo de CPU. Consequentemente, o potencial total da aceleração por GPU é realizado em simulações com grandes contagens de partículas, onde o pusher domina o tempo de execução.

O trabalho futuro, conforme delineado pelos autores, inclui a integração de solvers habilitados para GPU via PETSc, a exploração de uma porta C++ do TRIMEG para melhorar a eficiência de offloading e reduzir o esforço de portabilidade, e a implementação de offloading de GPU para operações de partícula-para-grade. O TRIMEG-C1 acelerado por GPU atual posiciona-se como um degrau para estudos de turbulência multiescala e investigações de transporte de plasma em configurações de tokamak desviadas realistas.