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

Este artigo relata a portabilidade do código Abinit para arquiteturas multi-GPU, detalhando as revisões algorítmicas necessárias para a aceleração de cálculos de DFT em ondas planas e apresentando resultados de desempenho comparativos entre nós CPU e heterogêneos, com foco na eficiência de dois algoritmos de diagonalização.

O essencial

Imagine que você é um arquiteto tentando projetar a casa perfeita, mas em vez de tijolos e cimento, você está trabalhando com átomos e elétrons. Para prever como esses materiais se comportam, os cientistas usam um software chamado Abinit. Ele resolve equações complexas (chamadas de DFT) para descobrir como os elétrons se movem e interagem.

O problema? Essas equações são tão pesadas que, para resolver um único problema, os computadores tradicionais (que usam apenas processadores de CPU) podem levar dias ou semanas. É como tentar pintar um mural gigante usando apenas um pincel fino e lento.

Aqui entra a história deste artigo: Como acelerar esse processo usando "super-pincéis" chamados GPUs.

Aqui está a explicação simples do que os autores fizeram, usando analogias do dia a dia:

1. O Problema: A Montanha de Dados

Pense na função de onda dos elétrons (a "fotografia" de onde eles estão) como uma pilha de milhões de folhas de papel. O computador precisa ler, calcular e reorganizar essas folhas bilhões de vezes.

• CPU (Processador tradicional): É como um funcionário muito inteligente, mas que só pode segurar uma folha de cada vez. Ele é ótimo em lógica complexa, mas lento em volume.
• GPU (Placa de Vídeo): É como uma equipe de 10.000 estagiários. Cada um é menos inteligente individualmente, mas eles podem pegar milhares de folhas ao mesmo tempo e trabalhar em paralelo.

2. A Mudança de Estratégia: De "Um por Um" para "Em Lotes"

Antigamente, o código Abinit tratava os dados folha por folha (ou banda por banda).

• A Nova Abordagem (Processamento em Lotes): Os autores mudaram a lógica. Em vez de pedir para o funcionário pegar uma folha, calcular, guardar, pegar outra... eles agora dizem: "Ei, pegue mil folhas de uma vez!".
• A Analogia do Caminhão de Mudança:
  • Método Antigo: Você faz 1.000 viagens de carro pequeno para levar 1.000 caixas. Muito tempo perdido no trânsito (comunicação entre memória e processador).
  • Método Novo: Você aluga um caminhão gigante (GPU) e carrega todas as 1.000 caixas de uma vez. O trajeto é o mesmo, mas a eficiência explode.

3. Os Dois "Métodos de Corrida" (Algoritmos)

Para resolver as equações, o software precisa encontrar os "números mágicos" (autovalores) que descrevem o sistema. O artigo compara duas formas de fazer isso:

• Método A: LOBPCG (O Corredor de Obstáculos)

  • Imagine um corredor que precisa passar por vários portões, mas a cada portão ele precisa parar, verificar se está alinhado com os outros corredores (ortogonalização) e ajustar a rota.
  • No GPU: Isso é ruim. A GPU adora correr rápido, mas odeia parar para conversar ou se alinhar. Cada parada exige que os dados saiam da GPU e voltem, criando um gargalo. É como ter um carro de F1 preso no trânsito.

• Método B: Filtro de Polinômio de Chebyshev (O Trem de Alta Velocidade)

  • Imagine um trem que passa por várias estações sem parar. Ele apenas acelera e desacelera suavemente para filtrar quem entra e quem sai, sem precisar de paradas longas para "se alinhar" a cada passo.
  • No GPU: Isso é perfeito. O trem (GPU) pode aplicar o mesmo movimento em milhares de passageiros (elétrons) ao mesmo tempo, sem parar. O artigo mostra que esse método é muito mais rápido e consome menos energia em GPUs.

4. O Resultado: Velocidade e Economia

Os autores testaram isso em supercomputadores reais com placas NVIDIA e AMD.

• Velocidade: Usando GPUs, eles conseguiram fazer o trabalho de 8 computadores normais em apenas 2 máquinas com GPUs. É como ter um time de 8 pessoas trabalhando em 2 dias, mas com 2 pessoas trabalhando em 1 dia.
• Energia: Além de ser mais rápido, o método com GPU gasta menos energia elétrica para chegar ao mesmo resultado. É como trocar um carro a gasolina antigo por um elétrico moderno: mais rápido e mais barato de abastecer.

5. O Grande Desafio: O "Gargalo" Final

Mesmo com a GPU sendo incrível, existe uma etapa final (chamada Rayleigh-Ritz) que ainda é um pouco lenta, como se fosse o momento em que o trem precisa entrar na garagem e se organizar.

• A Solução Futura: Os autores sugerem que, no futuro, eles vão dividir esse problema final em pedaços menores para que a GPU possa resolver tudo em paralelo, eliminando o último gargalo.

Resumo Final

Este artigo é sobre como os cientistas pegaram um software de física muito pesado e o "traduziram" para a linguagem das placas de vídeo (GPUs).

• Eles mudaram a forma de carregar os dados (de um por um para em lotes gigantes).
• Eles escolheram o método de corrida (algoritmo) que menos exige paradas e conversas (Filtro de Chebyshev).
• O resultado é que simulações que antes levavam dias agora levam horas, economizando tempo e dinheiro, permitindo que cientistas descubram novos materiais mais rápido.

É basicamente a diferença entre tentar mover uma montanha de areia com uma colher (CPU antigo) e usar uma escavadeira gigante (GPU otimizada).

Resumo técnico

Título: Aceleração por GPU de cálculos de Teoria do Funcional da Densidade (DFT) de ondas planas no Abinit

1. O Problema

Cálculos de estrutura eletrônica em larga escala, baseados na Teoria do Funcional da Densidade (DFT) de ondas planas, exigem a resolução do problema de Kohn-Sham para milhares de estados eletrônicos. Isso envolve operações algébricas lineares intensivas, especificamente a diagonalização de grandes matrizes Hamiltonianas.

• Desafio: Aumento do tamanho dos problemas e a necessidade de reduzir o tempo de simulação em arquiteturas de Computação de Alto Desempenho (HPC).
• Limitação Atual: A transição de arquiteturas puramente multi-core (CPU) para arquiteturas híbridas (CPU-GPU) exige não apenas o uso de bibliotecas de fornecedores, mas também revisões algorítmicas profundas. Operações tradicionais, como a diagonalização iterativa, muitas vezes são limitadas por largura de banda de memória ou comunicação, não aproveitando plenamente o poder de processamento massivo das GPUs.

2. Metodologia

Os autores realizaram uma portabilidade completa do código Abinit (versão 10.6+) para arquiteturas multi-GPU, utilizando o modelo de programação MPI + OpenMP com offloading para GPU.

• Estratégia de Portabilidade:

  • Processamento em Lotes (Batch Processing): O código foi reestruturado para agrupar dados (funções de onda) em lotes maiores. Em vez de processar bandas eletrônicas individualmente, as operações (como FFTs e multiplicações de matrizes) são aplicadas a blocos de vetores simultaneamente. Isso maximiza a paralelização de dados e reduz a sobrecarga de chamadas de kernel.
  • Memória Residente na GPU: A função de onda ($\Psi$) é transferida da CPU para a GPU apenas no início de cada iteração de campo autoconsistente (SCF) e mantida na memória da GPU durante todo o processo de diagonalização, minimizando transferências host-dispositivo.
  • Bibliotecas de Baixo Nível: Utilização de bibliotecas nativas de fornecedores (cuBLAS, cuSOLVER, cuFFT para NVIDIA; rocBLAS, rocSOLVER, rocFFT para AMD) através de uma camada de abstração em Fortran, evitando kernels personalizados e garantindo portabilidade.
  • Algoritmos Iterativos: Comparação e implementação de dois algoritmos principais para a diagonalização:
    1. LOBPCG (Locally Optimal Block Preconditioned Conjugate Gradient): Baseado em iteração de subespaço com ortogonalização de blocos.
    2. Filtragem Polinomial de Chebyshev: Baseado em filtragem espectral que aplica o Hamiltoniano repetidamente sem ortogonalização inter-bloco frequente.

• Modelo de Comunicação: Uso de MPI GPU-aware para transferências diretas entre GPUs, com transposições de matriz (distribuição de linhas vs. colunas) otimizadas para operações coletivas all-to-all.

3. Principais Contribuições

• Implementação Completa no Abinit: Apresentação da primeira versão robusta e de alto desempenho do Abinit para GPUs, suportando tanto NVIDIA quanto AMD.
• Análise Teórica de Algoritmos: Desenvolvimento de modelos teóricos para estimar o custo computacional (FLOPS) e de comunicação (latência e largura de banda) para LOBPCG e Chebyshev.
  • Demonstraram que o LOBPCG sofre de sobrecarga de comunicação (escala linearmente com o número de blocos e iterações), enquanto o Chebyshev minimiza a comunicação, realizando-a apenas no início e no fim do processo de filtragem.
• Otimização de Intensidade Aritmética: Provar que a filtragem de Chebyshev atinge uma intensidade aritmética $k$ vezes maior (onde $k$ é o grau do polinômio) do que o LOBPCG por rank MPI, tornando-a mais adequada para GPUs.
• Métricas de Desempenho e Eficiência Energética: Introdução de métricas para avaliar não apenas o tempo de execução, mas também o consumo de energia e a eficiência de portabilidade entre diferentes fornecedores de hardware.

4. Resultados

Os testes foram realizados em supercomputadores (Jean Zay, Adastra, Topaze) com sistemas de titânio (255 átomos) e óxido de gálio (320 átomos).

• Aceleração (Speedup):
  • As GPUs NVIDIA (A100, H100) mostraram speedups consistentemente superiores às GPUs AMD (MI250X).
  • O passo de filtragem (aplicação do Hamiltoniano) foi altamente acelerado, enquanto o passo de Rayleigh-Ritz (diagonalização do subespaço) mostrou aceleração modesta, especialmente em GPUs AMD, devido a limitações nas bibliotecas LAPACK (routine hegvd).
  • Em um cenário de 2 nós GPU (NVIDIA), o desempenho superou o de 8 nós CPU, permitindo economias significativas de recursos.
• Eficiência Energética:
  • O uso de GPUs resultou em reduções drásticas no consumo de energia por cálculo. Os nós NVIDIA no Jean Zay apresentaram fatores de economia de energia superiores aos nós AMD no Adastra.
  • O consumo de energia em nós NVIDIA permaneceu quase constante ao aumentar o número de nós (devido à baixa escalabilidade do passo de Rayleigh-Ritz), enquanto em AMD cresceu quase linearmente.
• Comparação de Algoritmos:
  • Chebyshev vs. LOBPCG: O algoritmo de filtragem de Chebyshev demonstrou ser superior em GPUs. Aumentar o grau do polinômio em Chebyshev melhora a precisão dos autovetores sem degradar o desempenho, permitindo convergência mais rápida do SCF.
  • Em contraste, aumentar as linhas de minimização no LOBPCG adicionou carga computacional desnecessária sem melhorar significativamente a convergência do SCF, tornando-o menos eficiente em GPUs.
• Modelo Roofline: A análise mostrou que as operações de aplicação do Hamiltoniano são compute-bound (limitadas por computação) e atingem o pico de desempenho, enquanto a ortogonalização e o Rayleigh-Ritz são memory-bound (limitados por memória) e de baixa intensidade aritmética.

5. Significado e Perspectivas

• Viabilidade de Cálculos em Larga Escala: O trabalho demonstra que cálculos DFT de ondas planas em larga escala são viáveis e eficientes em arquiteturas heterogêneas modernas, reduzindo drasticamente o tempo de simulação e o custo energético.
• Mudança de Paradigma Algorítmico: O estudo conclui que, para GPUs, algoritmos que priorizam operações limitadas por computação (como a aplicação repetida do Hamiltoniano em grandes lotes) são superiores aos que dependem fortemente de ortogonalização de blocos e comunicação frequente.
• Futuro: Os autores sugerem que a implementação de métodos de "Spectrum Slicing" (fatia de espectro) poderia mitigar o gargalo do passo de Rayleigh-Ritz, permitindo simulações de sistemas ainda maiores ao distribuir a diagonalização entre GPUs de forma independente.

Em resumo, o artigo estabelece um novo estado da arte para o código Abinit, provando que a combinação de uma arquitetura de software bem projetada (processamento em lotes, memória residente) com a escolha correta de algoritmos (filtragem de Chebyshev) é essencial para explorar plenamente o potencial das GPUs na física computacional.