RSH-SpMM: A Row-Structured Hybrid Kernel for Sparse Matrix-Matrix Multiplication on GPUs

O artigo apresenta o RSH-SpMM, um framework híbrido de multiplicação de matriz esparsa por matriz densa (SpMM) para GPUs que utiliza particionamento adaptativo de linhas e uma representação RS-Tile para otimizar o uso dos Tensor Cores e manter alto desempenho em matrizes com esparsidade altamente irregular, superando as soluções existentes com acelerações de 1,27x a 6,13x.

O essencial

Imagine que você é um gerente de uma grande fábrica de montagem (o GPU) e precisa organizar a produção de milhões de peças diferentes. O problema é que as peças chegam em caixas de tamanhos e formatos totalmente aleatórios: algumas caixas têm apenas um parafuso, outras têm mil, e a maioria tem um número estranho e imprevisível.

O objetivo é montar um produto final (a Multiplicação de Matrizes Esparsas, ou SpMM) o mais rápido possível.

Aqui está a história do RSH-SpMM, a nova solução apresentada no artigo, explicada de forma simples:

1. O Problema: A Fábrica Desorganizada

Antes, as fábricas (os computadores) tentavam lidar com essa bagunça de duas formas, e nenhuma funcionava bem:

• A Abordagem "Trabalho Manual" (CUDA Cores): Usavam trabalhadores gerais que podiam pegar qualquer caixa, não importa o tamanho. Era flexível, mas lento. Eles não conseguiam usar as máquinas super-rápidas da fábrica.
• A Abordagem "Máquina Robótica" (Tensor Cores): Tiveram uma máquina super-rápida que só funcionava se as peças fossem organizadas em blocos perfeitos de 8x8. O problema? Como as caixas reais eram irregulares, eles tinham que encher os blocos vazios com "areia" (padding) para caber na máquina. Isso desperdiçava tempo e a máquina ficava parada 90% do tempo esperando algo real para processar.

O resultado: A fábrica ficava lenta, travada e ineficiente.

2. A Solução Mágica: O RSH-SpMM

Os autores criaram um novo sistema chamado RSH-SpMM. Pense nele como um gerente de logística inteligente que usa três truques principais para organizar a fábrica:

Truque 1: O "Triagem Inteligente" (Particionamento Adaptativo)

Em vez de tentar forçar todas as caixas a caberem na máquina robótica, o gerente olha para cada caixa individualmente:

• Se a caixa tem muitas peças e se parece com as vizinhas, ele a coloca na Fila Robótica (Tensor Cores).
• Se a caixa é minúscula, solitária ou muito estranha, ele a manda imediatamente para a Fila Manual (CUDA Cores).

A analogia: É como ir ao supermercado. Se você tem um carrinho cheio de compras, vai para o caixa automático (rápido). Se você tem apenas uma maçã, vai para o caixa expresso (ou o caixa manual), porque não vale a pena esperar a fila do automático só por uma coisa. Isso evita que a máquina robótica fique parada esperando.

Truque 2: A "Organização de Prateleiras" (Reordenação Local)

Antes de começar a produção, o gerente reorganiza as caixas na esteira. Ele pega caixas que são "vizinhas" (têm peças parecidas) e as coloca uma ao lado da outra.

• Por que fazer isso? Quando as caixas são parecidas, a máquina robótica consegue encaixar perfeitamente mais peças sem precisar de "areia" para preencher espaços vazios. É como organizar livros por cor na estante: fica muito mais fácil encontrar e pegar vários de uma vez.

Truque 3: O "Formato de Caixa Personalizado" (RS-Tile)

Eles criaram um novo tipo de embalagem. Em vez de usar caixas padrão que deixam muito espaço vazio, eles usam caixas que se adaptam ao conteúdo:

• Para as peças que vão para a máquina robótica, eles criam blocos compactos e perfeitos.
• Para as peças estranhas, eles usam um formato simples e leve que não gasta tempo de preparação.

3. O Resultado na Prática

Quando testaram esse sistema em computadores modernos (como o RTX 4090), os resultados foram impressionantes:

• Velocidade: O sistema ficou de 1,27 a 6 vezes mais rápido do que os melhores métodos anteriores.
• Estabilidade: Funciona bem tanto com matrizes "normais" quanto com aquelas super bagunçadas e irregulares que costumavam travar os computadores.
• Aplicação Real: Em redes neurais (a tecnologia por trás de IAs como o ChatGPT ou reconhecimento de imagens), isso significa que treinar modelos fica muito mais rápido e consome menos energia.

Resumo em uma frase

O RSH-SpMM é como um gerente de fábrica que não tenta forçar todas as peças a se encaixarem em um molde rígido; em vez disso, ele separa o que é "fácil e regular" para as máquinas rápidas e o que é "difícil e irregular" para os trabalhadores manuais, organizando tudo de forma que ninguém fique esperando, resultando em uma produção muito mais rápida e eficiente.

Resumo técnico

Resumo Técnico: RSH-SpMM

1. O Problema

A Multiplicação de Matriz Esparsa por Matriz Densa (SpMM) é uma primitiva computacional fundamental em redes neurais gráficas (GNNs), análise de grafos, simulação científica e inferência de Grandes Modelos de Linguagem (LLMs). No entanto, atingir alto desempenho em GPUs modernas é extremamente desafiador devido à irregularidade extrema das matrizes esparsas do mundo real.

As principais dificuldades identificadas são:

• Heterogeneidade Estrutural: Distribuições de comprimento de linha "cauda pesada" (algumas linhas têm muitos não-zeros, outras muito poucos), densidades locais que mudam rapidamente e padrões de não-zeros fragmentados.
• Desalinhamento com Tensor Cores: Os Tensor Cores (TC) das GPUs modernas exigem operandos densos e alinhados em blocos (tiles) para alta eficiência. A irregularidade das matrizes esparsas impede o preenchimento eficiente desses blocos, levando a baixa utilização dos TCs e desperdício de ciclos de computação.
• Limitações das Abordagens Atuais:
  • Métodos baseados apenas em CUDA Cores: Robustos para irregularidade, mas limitados pela taxa de transferência aritmética escalar, não aproveitando o potencial dos TCs.
  • Métodos baseados apenas em Tensor Cores: Frequentemente exigem preenchimento (padding) excessivo ou reordenação rígida, resultando em baixa densidade de tiles e desempenho instável quando a esparsidade varia localmente.
  • Métodos Híbridos Existentes: Geralmente usam particionamento em granularidade grosseira (nível de matriz ou grandes blocos), falhando em capturar a coerência estrutural fina dentro de grupos de linhas, deixando oportunidades de aceleração não exploradas.

2. Metodologia: RSH-SpMM

O RSH-SpMM é um framework híbrido de execução em GPU, estruturado por linhas, projetado para alinhar a esparsidade irregular com os pipelines de execução modernos. A solução integra três componentes principais:

A. Representação RS-Tile (Row-Structured Tile)

• O formato de armazenamento compresso divide a matriz esparsa em duas partes disjuntas:
  1. Parte TC (Tensor Core): Agrupa linhas estruturalmente coerentes em "janelas de linha" que são compactadas em blocos fixos de 8x8 (alinhados com a geometria de MMA dos TCs). Utiliza mapas de bits (bitmaps) para codificar a localização dos não-zeros, eliminando a necessidade de ponteiros para cada elemento.
  2. Parte Residual CUDA: Isola linhas curtas, isoladas ou estruturalmente incompatíveis que degradariam a densidade dos tiles TC. Essas linhas são processadas por um caminho de execução leve baseado em CUDA Cores, sem a sobrecarga de reconstrução de tiles.

B. Estratégia de Execução Híbrida de Granularidade Fina

• Particionamento Adaptativo de Linhas: Antes da execução, o sistema analisa cada linha. Linhas com muitos não-zeros ou alta cobertura de colunas são direcionadas para os TCs. Linhas com poucos não-zeros ou que aumentariam significativamente o custo de formação de tiles (sem contribuir proporcionalmente) são roteadas para o caminho CUDA.
• Kernel de Pipeline de Tensor Core: O kernel para a parte TC utiliza um pipeline de dois estágios com double-buffering. Ele sobrepõe a transferência de dados da memória global para a memória compartilhada, a decodificação baseada em bitmap (transformando dados esparsos em fragmentos densos nos registradores) e a execução da instrução MMA (Multiply-Accumulate). Isso esconde a latência de memória e mantém os TCs ocupados.
• Kernel CUDA de Baixa Sobrecarga: Para as linhas residuais, utiliza-se um kernel simples de "carregar-computar" (fused-load-compute), evitando a sobrecarga de sincronização e pressão de registradores associada à preparação de tiles TC.

C. Reordenação Consciente de Localidade (Locality-aware Reordering)

• Para maximizar a formação de tiles de alta qualidade, o sistema aplica uma reordenação prévia das linhas.
• Utiliza uma similaridade Jaccard ponderada (w-Jaccard), onde colunas com menor frequência global (mais informativas) recebem maior peso.
• O algoritmo constrói um grafo de vizinhos mais próximos (kNN), extrai uma Árvore Geradora Mínima (MST) e realiza uma travessia DFS para gerar uma ordem global, seguida de refinamento local (trocas 2-opt). Isso agrupa linhas estruturalmente similares, reduzindo a fragmentação dos tiles.

D. Balanceamento de Carga Adaptativo

• Diferente de métodos rígidos, o RSH-SpMM ajusta o particionamento dinamicamente. Linhas "super-longas" podem ser divididas se necessário, enquanto linhas curtas são isoladas cedo para evitar desbalanceamento dentro dos blocos de threads (warps), garantindo que os Streaming Multiprocessors (SMs) permaneçam bem ocupados.

3. Contribuições Principais

1. RS-Tile: Uma representação compacta que expõe fragmentos alinhados aos Tensor Cores enquanto roteia linhas irregulares para um caminho CUDA de baixa sobrecarga, reduzindo a sobrecarga de metadados e aumentando a densidade de tiles.
2. Estratégia Híbrida de Granularidade Fina: Integra particionamento adaptativo de linhas, um kernel de pipeline de Tensor Core balanceado e um kernel CUDA mínimo, eliminando a necessidade de manter estruturas de indexação esparsa duplas.
3. Técnica de Reordenação: Um método consciente de localidade que agrupa linhas estruturalmente similares para reduzir a fragmentação e melhorar a utilização de MMA.
4. Desempenho Robusto: Avaliação abrangente demonstrando ganhos consistentes em diversas cargas de trabalho esparsas, superando os métodos mais avançados (SOTA).

4. Resultados Experimentais

Os autores avaliaram o RSH-SpMM em GPUs NVIDIA RTX 4090 (Ada Lovelace) e RTX 3090 (Ampere), utilizando matrizes do conjunto SuiteSparse e cargas de trabalho reais de GNNs.

• Aceleração Geral: O RSH-SpMM superou consistentemente as abordagens SOTA, alcançando acelerações de 1,27x a 6,13x em comparação com os melhores concorrentes.
• Comparação com Baselines:
  • Vs. cuSPARSE (CUDA Core): Aceleração média de 2,35x na RTX 4090 e 2,86x na RTX 3090.
  • Vs. Métodos Híbridos (HC-SpMM): Aceleração média de 2,10x.
  • Vs. Métodos Puros de Tensor Core (DTC-SpMM, Acc-SpMM): O RSH-SpMM manteve ganhos estáveis (ex: 1,91x sobre DTC-SpMM), demonstrando superioridade em matrizes com alta variabilidade de densidade onde métodos fixos falham.
• Eficiência de Hardware:
  • Aumento na utilização dos Tensor Cores de 5,6% (Acc-SpMM) para 8,8% (RSH-SpMM).
  • Aumento no throughput dos SMs de 28% para 33%.
• Cenário de Ponta a Ponta (GNN): Em um treinamento de GCN de 6 camadas, o RSH-SpMM reduziu o tempo total de treinamento em 1,49x comparado ao cuSPARSE e 1,09x comparado ao TC-GNN, validando o impacto prático na aplicação final.
• Eficiência de Armazenamento: O formato RS-Tile reduziu o tamanho dos metadados em aproximadamente 15% comparado a formatos híbridos balanceados existentes, eliminando a necessidade de tabelas de remapeamento complexas.

5. Significado e Impacto

O RSH-SpMM preenche uma lacuna crítica na computação de alto desempenho para GPUs: a capacidade de lidar com a heterogeneidade estrutural fina das matrizes esparsas do mundo real sem sacrificar a eficiência dos aceleradores modernos (Tensor Cores).

Ao invés de forçar uma estrutura rígida sobre dados irregulares ou depender exclusivamente de núcleos escalares mais lentos, o RSH-SpMM adota uma abordagem híbrida e adaptativa. Ele reconhece que a esparsidade varia em nível de linha e ajusta a estratégia de execução dinamicamente. Isso permite que as GPUs mantenham alta utilização dos recursos computacionais mais potentes (TCs) mesmo em cenários de dados desordenados, tornando-o uma solução robusta e essencial para o futuro de redes neurais gráficas e sistemas de aprendizado de máquina esparsos.