Superando a multiplicação de matrizes do NumPy com 150 linhas de código C

Resumo

Introdução

• A multiplicação de matrizes é um elemento essencial nas redes neurais modernas
• O NumPy alcança alto desempenho usando bibliotecas BLAS externas
• Este artigo explica como implementar uma multiplicação de matrizes de alto desempenho que seja simples, portátil e escalável

Desempenho do NumPy

• O NumPy usa OpenBLAS em CPUs AMD
• A medição de desempenho é calculada em FLOP/s
• O desempenho de thread única e multithread do NumPy foi medido em uma CPU Ryzen 7 7700

Limites teóricos

• Explica a hierarquia de memória da CPU e as extensões SIMD
• Teoricamente, é possível atingir 163 GFLOPS em thread única e 1203 GFLOPS em multithread

Implementação simples

• Explica o algoritmo básico de multiplicação de matrizes e mede o desempenho de uma implementação simples
• A implementação simples atinge 2,7 GFLOPS

Kernel

• Explica como resolver a multiplicação de matrizes dividindo o problema em subproblemas menores
• Otimiza o kernel usando instruções SIMD
• Com um kernel 16x6, atinge 147 GFLOPS

Mascaramento e empacotamento

• Explica como lidar com casos de borda para suportar tamanhos arbitrários de matrizes
• Otimiza o desempenho com mascaramento e empacotamento
• A nova implementação atinge 56 GFLOPS

Cache

• Explica o sistema de memória do cache da CPU
• Otimiza a reutilização de dados e o gerenciamento de cache aproveitando o cache

Opinião do GN⁺

• Este artigo é muito didático ao explicar passo a passo como implementar uma multiplicação de matrizes de alto desempenho
• É possível aprender métodos de otimização que aproveitam instruções SIMD e o cache da CPU
• Ajuda a entender o funcionamento interno de bibliotecas como o NumPy
• Outros projetos com funcionalidade semelhante incluem Intel MKL, OpenBLAS etc.
• Ao adotar novas tecnologias ou open source, é preciso considerar desempenho e portabilidade

[GN⁺]

Comentários do Hacker News

• A maior parte do software não é otimizada, então há bastante margem para melhorar o desempenho

  • A escolha do algoritmo é o mais importante
  • É preciso verificar se dá para reduzir operações pesadas, como chamadas de kernel
  • É possível melhorar o desempenho por meio de vetorização
  • É preciso verificar se dá para otimizar a eficiência de cache
  • É preciso analisar a possibilidade de otimizações específicas para o hardware

• Os artigos citados no repositório do BLIS são materiais de referência para entender este tema

  • Não entende por que alguém acharia que um BLAS otimizado tem desempenho ruim
  • Em vez de NumPy, deveria usar o BLAS da AMD
  • O BLIS tem paralelização melhor que o OpenBLAS

• Instruções SIMD não são necessárias para a vetorização do microkernel

  • Com o tamanho de bloco adequado, o microkernel em C puro do BLIS alcança mais de 80% do desempenho de uma implementação otimizada manualmente

• A maioria dos padrões de código não é totalmente específica ao hardware e por isso perde muito desempenho

  • Vale consultar o artigo clássico de ciência da computação "There's plenty of room at the top"

• É elogiável que tenham facilitado a repetição dos benchmarks

  • Em uma CPU Xeon de 16 núcleos, matmul.c leva 1,41 s quando compilado com gcc -O3, 1,47 s com clang -O2, e o NumPy leva 1,07 s
  • Acredita que um kernel avx512 seria mais rápido
  • Usar pthreads em vez de omp para gerenciar explicitamente um pool de threads pode reduzir a sobrecarga

• Há dúvida se a implementação do NumPy realmente usa multithreading

• Há curiosidade sobre o motivo de ter desempenho melhor que o OpenBLAS

  • Está lidando com detalhes como cache
  • Fica a dúvida se é mais otimizado para um processador específico

• Comparar Python de um lado e C do outro não é justo

  • Seria melhor escrever ambos em C para comparar

• A ineficiência na geração de máscaras incomoda

  • Há formas mais eficientes de criar um array global de constantes ou de comparar com um vetor constante
  • Mas isso é um problema pequeno e, na prática, provavelmente não fará grande diferença

• Há dúvida sobre a utilidade prática de fazer multithreading da própria multiplicação de matrizes

  • O multithreading provavelmente seria mais útil no algoritmo que usa a multiplicação de matrizes

• Menção ao tinyBLAS do jart

  • Link relacionado fornecido