Resumo Técnico: DCC – Compilação Centrada em Dados de Kernels de Aprendizado de Máquina para Arquiteturas de Processamento em Memória

1. Declaração do Problema

As arquiteturas de Processamento em Memória (PIM) oferecem uma solução promissora para o gargalo de memória no Aprendizado de Máquina (ML) ao colocar unidades de processamento próximas aos bancos de DRAM. No entanto, existe uma incompatibilidade fundamental entre os requisitos de layout de dados de processadores Host de alto desempenho (por exemplo, GPUs) e os núcleos PIM.

• Requisito do Host: Para explorar o paralelismo em nível de banco e a alta largura de banda, os processadores Host exigem que elementos de dados consecutivos sejam distribuídos entre múltiplos bancos de DRAM.
• Requisito do PIM: Os núcleos PIM só podem acessar dados de seus bancos locais. Para maximizar a largura de banda local e permitir acesso eficiente a múltiplos elementos, elementos consecutivos devem ser colocados dentro do mesmo banco.

Isso necessita de rearranjos de dados (entrada e saída) realizados via o barramento de memória do Host, o que incorre em custos significativos de movimentação de dados que podem dominar o tempo de execução de ponta a ponta. As abordagens atuais de compilação para sistemas PIM sofrem de duas limitações principais:

1. Falta de Otimização Sistemática: Ferramentas existentes frequentemente visam backends específicos (por exemplo, UPMEM) ou kernels específicos, carecendo de suporte para cargas de trabalho de ML diversas e múltiplos dispositivos PIM (por exemplo, Samsung HBM-PIM, SK Hynix GDDR6-AiM).
2. Ajuste Centrado em Computação: Compiladores existentes (por exemplo, ATiM) otimizam cronogramas de computação independentemente dos custos de rearranjo de dados. Os autores demonstram que transformações de computação e rearranjos de dados são interdependentes; otimizá-los isoladamente produz desempenho subótimo porque um cronograma eficiente em computação pode exigir rearranjos de dados caros, enquanto um cronograma de computação menos eficiente pode permitir movimentação de dados mais barata.

2. Metodologia: DCC

Os autores propõem o DCC, o primeiro compilador de ML centrado em dados para sistemas PIM. O DCC co-otimiza conjuntamente estratégias de rearranjo de dados e cronogramas de computação em nível de laço dentro de um processo de ajuste unificado. O sistema compreende quatro componentes-chave:

2.1 Abstração Multicamada

O DCC introduz uma abstração genérica que mapeia a hierarquia de memória PIM para uma hierarquia de computação, desacoplando o compilador de semânticas específicas do backend. Esta abstração opera em três níveis:

• Nível de Sistema: O xPU Host (por exemplo, GPU), espaço de memória Host e dispositivos PIM.
• Nível de Grupo PIM: Grupos de núcleos PIM (e seus bancos locais) pertencentes ao mesmo canal ou rank de memória. Isso permite otimizações de granularidade grossa e transmissão de instruções em nível de grupo.
• Nível de Núcleo PIM: Unidades de processamento individuais e seus bancos locais, permitindo otimizações de granularidade fina para acesso à memória e layout de dados local.

2.2 Gerador de Cronograma Centrado em Dados

Diferentemente de compiladores tradicionais que começam com transformações de laço, o DCC adota uma estratégia primeiro-dados:

1. Construção do Conjunto de Dimensões: Dimensões de tensores são correlacionadas com variáveis de laço para formar "conjuntos de dimensões complexas".
2. Construção de Blocos de Dados: O sistema gera exaustivamente partições de dados candidatas (blocos) entre grupos e núcleos PIM.
3. Poda de Rascunhos: Três regras eliminam rascunhos pouco promissores:
   • Redundância: Removendo rascunhos com partições de dados e mapeamentos idênticos.
   • Alinhamento SIMD: Podando rascunhos onde os tamanhos de partição não se alinham com a largura SIMD do backend (por exemplo, 16 vias).
   • Equivalência de Desempenho: Mantendo apenas um rascunho se múltiplos rascunhos produzirem o mesmo desempenho por núcleo no pior caso.
4. Mapeamento: Os rascunhos restantes são mapeados para blocos de computação, determinando intervalos de laço e índices de tensores para cada núcleo.

2.3 Otimizador de Código Específico para PIM

Este componente aplica otimizações adaptadas às restrições PIM:

• Otimização de Rearranjo de Dados: Explora memória on-chip controlável (por exemplo, memória compartilhada da GPU) para reorganizar dados. Realiza leituras seriadas dos canais de memória do Host seguidas por escritas paralelas para os canais PIM, maximizando o paralelismo de canal para ambas as operações de leitura e escrita.
• Otimização de Código de Computação: Utiliza geração hierárquica de comandos. Prioriza comandos no estilo SIMD em nível de grupo (transmitindo uma única instrução para todos os núcleos de um grupo) para reduzir a sobrecarga do barramento de comandos, recorrendo a comandos em nível de banco apenas quando necessário.

2.4 Preditor de Desempenho Acoplado

Para evitar o custo proibitivo de perfilar milhares de rascunhos candidatos, o DCC emprega um preditor baseado em aprendizado (XGBoost).

• Função: Estima conjuntamente custos de rearranjo de dados e tempos de execução de computação para selecionar a configuração de melhor desempenho de ponta a ponta.
• Treinamento: O modelo é treinado offline em um subconjunto diversificado de rascunhos.
• Tempo de Execução: Para tamanhos de tensores estáticos, o melhor cronograma é pré-selecionado. Para tamanhos dinâmicos (por exemplo, em LLMs), o preditor estima o desempenho em tempo real.

2.5 Interface de Programação

O DCC fornece uma interface Python compatível com PyTorch. Usuários anotam kernels com @DCC_kernel. O compilador analisa essas definições de alto nível, gera a Representação Intermediária (IR) envolvendo blocos de dados e computação, e emite instruções de backend de baixo nível.

3. Resultados Chave

Os autores avaliaram o DCC em dois backends PIM de última geração: Samsung HBM-PIM e AttAcc (um protótipo de pesquisa com unidades especializadas GEMV/Softmax), comparando com execução apenas em GPU e implementações PIM de base.

• Kernels de ML Individuais:
  • HBM-PIM: O DCC alcança até 7,68× de aceleração (média de 2,21×) sobre a execução em GPU.
  • AttAcc: O DCC alcança até 13,17× de aceleração (média de 3,92×) sobre a execução em GPU.
  • Comparação com PIM de Base: O DCC melhora a implementação AttAcc de base em até 2,74× (média de 1,75×) e o HBM-PIM em até 2,09× (média de 1,53×) para kernels específicos como GEMV e RED.
• Inferência de LLM de Ponta a Ponta:
  • No AttAcc, o DCC acelera a inferência de GPT-3 e LLaMA-2 em uma média de 4,52× (até 7,71× para LLaMA-2) sobre a execução apenas em GPU.
  • O DCC supera significativamente a implementação AttAcc de base (que usa tiling fixo) ao permitir execução eficiente das camadas de geração e projeção QKV, que anteriormente eram gargalos na GPU.
• Eficiência Energética: O DCC reduz o consumo de energia em uma média de 1,24× comparado ao sistema AttAcc de base devido à redução da movimentação de dados.
• Sobrecarga de Compilação: O tempo de compilação offline é baixo (segundos), e o esquema de poda reduz o espaço de busca em ~9×. O preditor acoplado alcança alta precisão (97,37% em LLMs) e reduz o tempo de compilação em ~3,8× comparado à busca exaustiva.

4. Significado e Alegações

O artigo afirma que o DCC é o primeiro compilador de ML centrado em dados para sistemas PIM que:

1. Otimiza Conjuntamente rearranjos de dados e cronogramas de computação, demonstrando que esses dois fatores são interdependentes e devem ser ajustados juntos para evitar desempenho subótimo.
2. Suporta Múltiplos Backends através de uma abstração multicamada, permitindo portabilidade entre diversas arquiteturas PIM (por exemplo, HBM-PIM, AttAcc) sem exigir ajuste manual específico do backend.
3. Habilita Alto Desempenho para kernels de ML intensivos em memória e inferência de LLM em grande escala, superando significativamente tanto a execução apenas em GPU quanto abordagens existentes de compilação PIM.
4. Melhora a Programabilidade fornecendo uma interface Python de alto nível que abstrai a complexidade do gerenciamento de layout de dados de baixo nível e do agendamento de instruções para dispositivos PIM.

Os autores concluem que o DCC fornece um caminho prático e eficaz para aceleração heterogênea de ML, equilibrando alto desempenho com programabilidade em ecossistemas PIM emergentes. O código é de código aberto para facilitar pesquisas futuras.