Accelerating finite-element-based projector augmented-wave density functional theory calculations with scalable GPU-centric computational methods

Este artigo apresenta um método escalável de ondas aumentadas projetadas por elementos finitos centrado em GPU (PAW-FE) que aproveita inovações algorítmicas, como aritmética de precisão mista e iteração de subespaço filtrada por Chebyshev, para alcançar acelerações significativas e desempenho pronto para exascala em simulações de teoria do funcional da densidade quimicamente precisas e de grande escala.

O essencial

Imagine que você está tentando prever como uma máquina complexa, como um motor de carro ou um novo tipo de bateria, se comportará. Para fazer isso com precisão, você precisa entender o comportamento de cada elétron individual dentro dos materiais que compõem a máquina. Este é o trabalho de um campo chamado Teoria do Funcional da Densidade (DFT). É como tentar simular uma pista de dança massiva e intrincada, onde bilhões de elétrons se movem em sincronia.

Por muito tempo, os cientistas tiveram um problema: simular essas danças para pequenos grupos de átomos é fácil, mas assim que você tenta simular um sistema grande e complexo (como uma nanopartícula metálica minúscula ou uma folha de material torcida), o computador fica sobrecarregado. É como tentar dirigir uma dança para 100.000 pessoas usando um método projetado para 100; as instruções se emaranham, a memória se enche e a simulação leva uma eternidade para terminar.

Este artigo apresenta uma nova maneira super-rápida de executar essas simulações, especificamente projetada para computadores modernos e poderosos que usam GPUs (os mesmos chips que alimentam videogames de ponta e IA). Aqui está como eles fizeram isso, dividido em conceitos simples:

1. O Jeito Antigo vs. O Novo Mapa

• O Jeito Antigo (Ondas Planas): Imagine tentar mapear uma cidade usando uma grade gigante e uniforme onde cada polegada quadrada tem o mesmo tamanho. Se você quiser ver um detalhe minúsculo (como um único tijolo em um prédio), você precisa fazer a inteira grade incrivelmente fina, mesmo para o céu vazio acima da cidade. Isso desperdiça uma quantidade massiva de poder de computador. É assim que a maioria dos softwares atuais funciona.
• O Jeito Novo (Elementos Finitos): Os autores usam uma abordagem de "mapa inteligente". Imagine um mapa que dá zoom apenas onde é necessário (como o centro da cidade movimentado) e dá zoom para fora onde está vazio (como o céu). Isso é chamado de discretização por Elementos Finitos (EF). Isso permite que eles concentrem seu poder de computação exatamente onde os elétrons estão fazendo coisas interessantes, economizando enormes quantidades de tempo e memória.

2. O Truque "PAW": O Disfarce Mágico

Para tornar a matemática ainda mais fácil, eles usam um método chamado Onda Aumentada por Projetor (PAW).

• O Problema: Elétrons perto do centro de um átomo (o núcleo) tremem e vibram selvagemente, tornando-os difíceis de calcular.
• A Solução: PAW é como colocar um "disfarce suave" nos elétrons. Ele finge que os elétrons são suaves e fáceis de manusear para a maior parte do cálculo, mas mantém um "truque mágico" secreto que permite revelar instantaneamente o comportamento verdadeiro e selvagem dos elétrons exatamente quando precisa verificar os detalhes perto do núcleo. Isso permite que eles usem um mapa muito mais grosso (simplificado) sem perder precisão.

3. O Impulso de Velocidade da GPU: A Linha de Montagem

Os autores não apenas mudaram o mapa; eles mudaram como o computador faz a matemática para se adequar às GPUs modernas.

• O Gargalo: Geralmente, os computadores passam muito tempo esperando que os dados se movam entre a memória e o processador.
• A Correção: Eles redesenharam a matemática para que o computador possa fazer muitos cálculos ao mesmo tempo (como uma linha de montagem) em vez de um por um. Eles também usaram uma técnica inteligente chamada Filtragem de Chebyshev, que é como uma peneira que separa rapidamente os elétrons "importantes" dos "não importantes", para que o computador não desperdice tempo com aqueles que não precisa.

4. Os Atalhos "Bons o Suficiente" (Precisão Mista)

Esta é talvez a parte mais criativa.

• A Analogia: Imagine que você está pintando um mural gigante. Para o céu de fundo, você não precisa misturar a tinta com precisão microscópica; uma mistura "boa o suficiente" funciona bem e é muito mais rápida. Você só precisa de precisão extrema para os detalhes minúsculos de um rosto.
• A Aplicação: Os autores perceberam que, para as partes do cálculo que apenas precisam acertar a forma geral, eles podem usar matemática de precisão mais baixa (como usar uma régua com menos marcações). Isso é muito mais rápido em chips modernos. Eles só mudam para matemática "ultraprecisa" para as etapas finais e críticas.
• O Resultado: Ao misturar matemática de alta precisão e baixa precisão, e ao sobrepor transferências de dados com cálculos (fazendo duas coisas ao mesmo tempo), eles fizeram a simulação rodar 8 a 20 vezes mais rápido do que antes.

5. O Que Eles Realmente Conquistaram

O artigo afirma que, com esses novos métodos:

• Velocidade: Eles agora podem simular sistemas com 10.000 a 130.000 elétrons em um tempo prático (minutos a horas) em supercomputadores.
• Comparação: Seu método é cerca de 8 vezes mais rápido do que o software padrão líder (Quantum ESPRESSO) para sistemas desse tamanho.
• Escala: Eles executaram com sucesso uma simulação de um material de "bicamada torcida" (duas folhas de átomos torcidas juntas) contendo 130.000 elétrons. Este é um tamanho que era anteriormente impossível de simular com esse nível de precisão usando métodos padrão.

Resumo

Em resumo, os autores construíram um novo motor altamente eficiente para simular materiais. Eles combinaram um "mapa inteligente" que dá zoom apenas onde é necessário, um truque de "disfarce mágico" para simplificar a matemática e um modo de "avanço rápido" que usa precisão mais baixa para etapas não críticas. O resultado é uma ferramenta que pode modelar materiais massivos e complexos em supercomputadores modernos em uma fração do tempo que levava antes, abrindo a porta para projetar novos materiais para baterias, eletrônicos e catalisadores muito mais rápido.

Resumo técnico

1. Declaração do Problema

A modelagem precisa de sistemas materiais complexos (por exemplo, interfaces, defeitos, nanoclusters, heteroestruturas torcidas) requer simulações de Teoria do Funcional da Densidade (DFT) envolvendo $10^4$ a $10^5$ elétrons. No entanto, implementações existentes enfrentam gargalos significativos:

• Limitações de Onda Plana (PW): Códigos PW-PAW amplamente utilizados (por exemplo, VASP, Quantum ESPRESSO) dependem de Transformadas Rápidas de Fourier (FFTs), que incurram em comunicação massiva do tipo "todos-para-todos". Isso os torna limitados pela largura de banda e com baixa escalabilidade em arquiteturas modernas de GPU de exaescala.
• Limitações do Espaço Real: Métodos de espaço real existentes frequentemente utilizam Pseudopotenciais de Conservação de Norma (ONCV), que exigem um grande número de funções de base (altos graus de liberdade) para alcançar precisão química, especialmente para metais de transição, levando a grandes pegadas de memória e custos computacionais elevados.
• Incompatibilidade de Hardware: Solucionadores de autovalores atuais são frequentemente inadequados para a alta intensidade aritmética e requisitos de baixa latência de clusters de GPU modernos, falhando em explorar capacidades de precisão mista ou sobrepor comunicação com computação de forma eficaz.

2. Metodologia

Os autores apresentam o PAW-FE, uma formulação de Onda Projetora Aumentada (PAW) discretizada por Elementos Finitos (FE) otimizada para arquiteturas de GPU multi-nó.

A. Formulação Matemática

• PAW-GHEP: O método resolve o Problema de Autovalores Hermitiano Generalizado (GHEP): $H\tilde{\Psi} = S\tilde{\Psi}\Lambda$, onde $H$ é o Hamiltoniano e $S$ é a matriz de sobreposição PAW.
• Formalismo de Spin Colinear: As equações são derivadas dentro de um framework de spin colinear para lidar com sistemas magnéticos.
• Condições de Contorno: A discretização FE acomoda naturalmente condições de contorno periódicas, semi-periódicas (lâminas) e não periódicas (nanoclusters) sem artefatos de periodicidade artificial.

B. Inovações Computacionais

Para resolver o PAW-GHEP eficientemente em GPUs, os autores desenvolveram várias estratégias algorítmicas-chave:

1. Iteração de Subespaço Filtrado por Chebyshev Baseada em Resíduo (R-ChFSI):

   • Em vez do ChFSI padrão, eles utilizam uma formulação baseada em resíduo ($R = H\tilde{\Psi} - S\tilde{\Psi}\Lambda$).
   • Vantagem Chave: Esta formulação é tolerante a produtos matriz-vetor inexatos, permitindo o uso de inversos aproximados para a matriz de sobreposição PAW ($S^{-1}$) e aritmética de precisão reduzida sem sacrificar a convergência.

2. Matriz de Sobreposição Aproximada:

   • Em vez de inverter explicitamente a matriz esparsa $S$, eles utilizam uma aproximação computacionalmente barata ($\tilde{S}^{-1}$) derivada de uma aproximação diagonal da matriz de massa e correções localizadas. Isso evita inversões globais caras.

3. Quadratura Multi-Resolução:

   • Para lidar com integrais centradas no átomo (densidades pseudo-núcleo) em grades FE grosseiras, emprega-se um esquema de quadratura multi-resolução. Ele utiliza regras de quadratura refinadas apenas dentro das esferas de aumento dos átomos, enquanto usa regras mais grosseiras em outros lugares, garantindo precisão sem refinamento de malha.

4. Precisão Mista e Comunicação de Baixa Precisão:

   • Cálculo: Utiliza aritmética FP32/TF32 para etapas de filtragem Chebyshev (multiplicações matriz-matriz) e BF16 para comunicação entre vizinhos mais próximos.
   • Robustez: A dependência do algoritmo R-ChFSI em resíduos garante que essas reduções de precisão não degradem a precisão final de dupla precisão do estado fundamental.

5. Sobreposição de Computação-Comunicação:

   • A filtragem Chebyshev é realizada em blocos. Enquanto um bloco de funções de onda sofre comunicação entre processadores (MPI), outro bloco realiza computação (operações GEMM) na GPU. Isso oculta a latência de comunicação, um fator crítico para a escalabilidade em exaescala.

6. Operações Densas ao Nível da Célula:

   • Em vez de construir matrizes esparsas globais, o método reformula as operações como multiplicações densas matriz-matriz ao nível da célula. Isso maximiza a intensidade aritmética e a localidade de cache, tornando-o ideal para execução em GPU.

3. Contribuições Principais

• Formulação PAW-FE: Uma formulação PAW discretizada por FE inovadora dentro de um formalismo de spin colinear que suporta condições de contorno genéricas.
• Solucionador de Autovalores R-ChFSI: Uma extensão da Iteração de Subespaço Filtrado por Chebyshev baseada em Resíduo para resolver o PAW-GHEP, permitindo o uso de inversos aproximados e precisão mista.
• Implementação Escalável em GPU: Uma estratégia de implementação abrangente que apresenta aritmética de precisão mista (FP32/TF32/BF16), sobreposição de computação-comunicação e álgebra linear densa ao nível da célula.
• Integração Multi-Resolução: Um esquema de quadratura que permite a avaliação precisa de integrais PAW centradas no átomo em malhas grosseiras.

4. Resultados e Desempenho

O método foi testado em supercomputadores de classe líder (OLCF Frontier, ALCF Aurora, ALCF Polaris) utilizando GPUs AMD, Intel e NVIDIA.

• Precisão: Validado contra códigos de onda plana (Abinit, Quantum ESPRESSO) para moléculas ($O_2$, $NO_2$) e cristais (Cr CCB). Erros em energia e forças são ordens de magnitude abaixo dos requisitos de precisão química ($O(10^{-12})$ Ha/átomo para energia, $O(10^{-6})$ Ha/bohr para forças).
• Aceleração CPU-GPU: Alcançou aceleração de 8× a 20× em GPUs em comparação com execuções apenas em CPU (arquiteturas Intel e AMD).
• Comparação com Onda Plana (QE):
  • Para sistemas com ~10.000 elétrons, o PAW-FE alcança uma redução de 8× no tempo mínimo de parede em comparação com o Quantum ESPRESSO.
  • Para sistemas maiores (>10.000 elétrons), a aceleração aumenta ainda mais devido à localidade das funções de base FE versus a comunicação global dos métodos PW.
• Comparação com ONCV-FE:
  • O PAW-FE requer ~6× menos recursos computacionais e alcança ~5× mais rápido tempo de solução em comparação com métodos FE usando Pseudopotenciais de Conservação de Norma (ONCV), principalmente devido aos graus de liberdade reduzidos habilitados pelo método PAW.
• Escalabilidade em Exaescala:
  • Demonstrou com sucesso um cálculo de estado fundamental para um sistema de WTe2 bicamada torcida compreendendo 130.000 elétrons (11.000 átomos).
  • Alcançou um tempo de solução de ~2 minutos por iteração SCF em 400 nós, provando a viabilidade do método para simulações realistas em grande escala.

5. Significado

Este trabalho estabelece o PAW-FE como um método pronto para exaescala para simulações de primeiros princípios. Ao preencher a lacuna entre a alta precisão do método PAW e a eficiência dos elementos finitos de espaço real, ele supera os gargalos de comunicação dos métodos de onda plana. A implementação bem-sucedida de estratégias de precisão mista e sobreposição em diversas arquiteturas de GPU demonstra um caminho para simular rotineiramente sistemas materiais complexos com $10^5$ elétrons, permitindo novas descobertas em catálise, interfaces de baterias e materiais quânticos que anteriormente eram computacionalmente intratáveis.