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

 (salykova.github.io)
1 pontos por GN⁺ 2024-07-05 | 1 comentários | Compartilhar no WhatsApp
  • A multiplicação de matrizes do NumPy depende de bibliotecas BLAS externas, mas esta implementação busca elevar o desempenho single-thread e multithread ao nível de BLAS usando apenas C puro e FMA3·AVX2
  • O centro do desempenho é dividir $C$ em pequenos blocos e usar um microkernel 16×6 que repete rank-1 updates dentro dos registradores YMM, reduzindo acessos à memória
  • Em matrizes de tamanho arbitrário, o tratamento das bordas tende a virar gargalo, então a implementação combina armazenamento com máscara e buffers com preenchimento por zero para evitar a perda de desempenho dos carregamentos com máscara
  • A reutilização de cache é garantida com blocking em k_c, m_c, n_c, e o pico real de desempenho depende fortemente do ajuste fino de número de threads, tamanho do kernel e tamanho dos tiles
  • O AVX-512 foi excluído para dar suporte a uma gama maior de CPUs, então em CPUs com AVX-512 o BLAS pode ser mais rápido, e a comparação com OpenBLAS também foi feita com o AVX-512 desativado

Objetivo da implementação e base de comparação

  • O código da implementação está disponível em sgemm.c e otimiza a multiplicação de matrizes FP32 multithread em processadores recentes
  • O NumPy depende de bibliotecas BLAS externas para operações de álgebra linear como multiplicação de matrizes
    • Exemplos incluem Intel MKL, Accelerate, BLIS, GotoBLAS e OpenBLAS
    • OpenBLAS, GotoBLAS e BLIS são escritos em C/FORTRAN/Assembly e incluem implementações de multiplicação de matrizes ajustadas manualmente para cada microarquitetura de CPU
  • O objetivo é uma implementação de multiplicação de matrizes escrita em C puro, sem assembly de baixo nível, que ainda satisfaça as seguintes condições
    • Funciona com matrizes de tamanho arbitrário
    • Roda em processadores x86-64 modernos
    • Compete com bibliotecas BLAS existentes
    • O código é simples e fácil de expandir
  • As referências incluem Fast Multidimensional Matrix Multiplication on CPU from Scratch, de Simon Boehm, Matrix Multiplication, de Sergey Slotin, Can you multiply a matrix?, de Geohot, e artigos sobre GotoBLAS e BLIS

Condições de benchmark e cálculo de FLOPS

  • O ambiente de teste é AMD Ryzen 7 9700X, 32GB DDR5 6000 MHz CL36, OpenBLAS 0.3.26, GCC 13.3 e Ubuntu 24.04.1 LTS
  • As flags de compilação usadas são -O3 -march=native -mno-avx512f -fopenmp
  • Para uma comparação justa, é preciso definir o TARGET apropriado ao instalar o OpenBLAS e desativar instruções AVX-512
    • Processadores Zen4/5 são compilados com make TARGET=ZEN
    • Caso contrário, o OpenBLAS usa instruções AVX-512 por padrão
  • A multiplicação de matrizes FP32 do OpenBLAS é executada pela API cblas_sgemm
  • O benchmark usa matrizes quadradas
    • A avaliação vai de m=n=k=200 até m=n=k=10000, em passos de 200
    • A multiplicação de matrizes é repetida n_iter vezes, e o tempo mediano de execução é usado para medir desempenho
  • Ao multiplicar uma matriz $M \times K$ $A$ por uma matriz $K \times N$ $B$, o total de operações é $2MNK$ FLOP
    • O desempenho é calculado como FLOPS=(2*m*n*k)/exec_time

Limites teóricos e base em SIMD

  • CPUs x86-64 modernas usam extensões SIMD para processar vários dados em paralelo
  • As principais instruções são AVX2 e FMA
    • Ambas usam registradores YMM de 256 bits
    • Cada registrador YMM pode armazenar 8 floats de 32 bits
  • A instrução FMA VFMADD231PS executa uma operação packed single no formato YMM1 = YMM2 * YMM3 + YMM1
  • No Ryzen 9700X, a vazão de fused multiply-add é de 0,5 ciclos por instrução, ou seja, 2 instruções por ciclo
  • Em teoria, o Ryzen 9700X pode executar 32 FLOP por ciclo em um único núcleo
    • O cálculo é 8 floats × 2(add+mul) × 2(1/TP)
    • Assumindo clock sustentado de 4,7GHz em 8 núcleos, o pico teórico multithread é estimado em 1203 FLOPS

Implementação básica e microkernel

  • As matrizes são armazenadas em ordem column-major
    • A[row][col] é acessado no ponteiro C como ptr[col*M + row]
  • A implementação mais simples percorre todas as linhas e colunas de $C$ e calcula o produto interno da linha de $A$ com a coluna de $B$ para cada elemento
  • O núcleo da implementação de alto desempenho é o microkernel, que divide $C$ em submatrizes de tamanho $m_R \times n_R$ e calcula cada uma delas com eficiência
  • O kernel inicializa $\bar{C}$ com zero nos registradores e depois itera ao longo da dimensão $K$
    • Traz para os registradores o vetor coluna de $\bar{A}$ e o vetor linha de $\bar{B}$
    • Calcula o produto externo dos dois vetores e soma aos acumuladores de $\bar{C}$
    • Cada etapa é um rank-1 update
  • Em comparação com a abordagem ingênua, que faz $2K m_R n_R$ acessos à memória, esse método reduz o número de elementos carregados para os registradores a $(m_R+n_R)K$
  • Como CPUs AVX têm 16 registradores YMM, o tamanho do kernel precisa obedecer à seguinte restrição
    • $(m_R/8) \cdot n_R + m_R/8 + 1 \le 16$
    • $m_R$ deve ser múltiplo de 8
  • Em teoria, quanto maiores e mais próximos entre si forem $m_R$ e $n_R$, maior a redução de acessos à memória, mas no Ryzen 9700X real o kernel 16×6 apresentou o melhor desempenho
  • A implementação usa intrinsics de immintrin.h
    • __m256 é um tipo vetorial de 256 bits que representa o conteúdo de registradores YMM
    • _mm256_loadu_ps carrega o vetor coluna de A
    • _mm256_broadcast_ss faz broadcast de um valor escalar de B para um vetor com 8 floats
    • _mm256_fmadd_ps atualiza os acumuladores
    • _mm256_storeu_ps grava o resultado na memória
  • O assembly gerado inclui instruções SIMD FMA como vfmadd231ps e vbroadcastss

Padding para matrizes de tamanho arbitrário

  • O kernel 16×6 básico funciona diretamente quando $M$ e $N$ são múltiplos de 16 e 6, respectivamente
  • Na região de borda, quando o número de colunas $n$ é menor que 6, o loop de armazenamento é executado apenas até j < n
  • Quando o número de linhas $m$ é menor que 16, _mm256_storeu_ps grava 8 elementos por vez, então é necessário usar armazenamento com máscara
    • _mm256_maskstore_ps grava na memória apenas os elementos cujos bits de máscara estão ativados
    • A máscara é gerada de acordo com o número de linhas sobrepostas $m`
  • Se até os carregamentos na borda forem tratados com _mm256_maskload_ps, o desempenho do kernel pode cair bastante
    • Instruções extras para calcular a máscara geram overhead
    • Como $n$ não é uma constante em tempo de compilação, o compilador tem mais dificuldade para desenrolar o loop com eficiência
  • Em vez disso, quando $m \neq m_R$, $\bar{A}$ é copiada para um buffer e preenchida com zeros; quando $n \neq n_R$, $\bar{B}$ também é copiada para um buffer e preenchida com zeros
  • A implementação relacionada está em matmul_pad.h

Cache blocking e reutilização de dados

  • Entre os registradores e a DRAM existe a hierarquia de cache da CPU, e CPUs desktop modernas normalmente usam caches L1, L2 e L3
  • O cache é mais rápido que a DRAM, mas tem capacidade limitada, então não é possível manter toda a $A$, $B$ e $C$ no cache ao mesmo tempo
  • Dividir as matrizes em pequenos blocos, carregá-los no cache e reutilizar os mesmos dados em vários rank-1 updates é o que se chama cache blocking ou tiling
  • O cache blocking single-thread segue uma estrutura de 5 loops semelhante à do BLIS
    • O loop mais externo cria blocos $C_j$ e $B_j$ ao longo da dimensão $N$
    • O loop seguinte cria blocos $A_j$ e $B_p$ ao longo da dimensão $K$
    • $B_p$ é empacotado em $\tilde{B}_p$ e, quando necessário, preenchido com zeros para favorecer a reutilização no cache L3
    • O loop seguinte cria blocos $C_i$ e $A_j$ ao longo da dimensão $M$, e $A_j$ é empacotado em $\tilde{A}_j$
    • Os dois loops finais dividem os blocos de cache em painéis $m_R \times k_c$ e $k_c \times n_R$ para enviá-los ao kernel
  • Os buffers empacotados $\tilde{A}_j$ e $\tilde{B}_p$ são armazenados de formas diferentes
    • Os painéis internos de $\tilde{A}_j$ são armazenados em column-major
    • Os painéis internos de $\tilde{B}_p$ são armazenados em row-major
  • Os parâmetros de cache blocking precisam ser ajustados ao tamanho de cache de cada modelo de CPU
    • $k_c \times n_c$ é o ponto de partida para preencher o cache L3
    • $m_c \times k_c$ é o ponto de partida para preencher o cache L2
    • $k_c \times n_R$ é o ponto de partida para preencher o cache L1
  • Na prática, valores maiores que os teóricos muitas vezes trazem desempenho melhor, e como a CPU gerencia a disposição no cache automaticamente, o nível algorítmico deve focar no desenho dos loops e no padrão de acesso
  • A implementação está em matmul_cache.h

Micro-otimizações do kernel

  • Em vez de definir acumuladores como um array, por exemplo __m256 C_buffer[6][2], as variáveis acumuladoras são explicitamente abertas uma a uma
  • Essa abordagem ajuda o GCC a otimizar melhor o código e evitar register spilling
  • O cálculo de máscara também foi alterado para usar instruções vetoriais
    • Um array estático mask[32] é usado junto com _mm256_cvtepi8_epi32 e _mm_loadu_si64
  • A implementação correspondente está em matmul_micro.h

Estratégia de multithreading

  • Tanto as operações aritméticas quanto o empacotamento são paralelizados
  • O 5º, 4º e 3º loops fora do microkernel iteram em unidades do tamanho dos blocos de cache
    • Para manter todas as threads ocupadas, o número de iterações precisa ser pelo menos igual ao número de threads
    • As dimensões da matriz de entrada precisam ser, aproximadamente, maiores ou iguais a número de threads × tamanho do bloco de cache
  • No Ryzen 9700X, os tamanhos de bloco de cache com bom desempenho em thread única foram $n_c=1535$, $m_c=1024$
    • Para usar todos os 8 núcleos, é necessária uma dimensão mínima de $\max(m_c,n_c) \times 8 = 1535 \times 8 = 12280$
  • Em contrapartida, os dois últimos loops repetem blocos pequenos de $m_R$, $n_R$, o que os torna adequados para paralelização
    • Em geral, $m_R$, $n_R$ são menores que 20
    • Escolher $m_c$, $n_c$ como múltiplos do número de núcleos ajuda a distribuir o trabalho de forma uniforme
  • No Ryzen 9700X, o melhor desempenho veio da paralelização conjunta de dois loops internos com #pragma omp parallel for collapse(2) num_threads(NTHREADS)
  • Em processadores com muitos núcleos, especialmente acima de 16, pode valer a pena considerar paralelismo aninhado e paralelização de 2 a 3 loops
  • O empacotamento de $\tilde{A}$ e $\tilde{B}$ também é paralelizado com OpenMP
    • pack_blockA é paralelizado ao percorrer mc em unidades de MR
    • pack_blockB é paralelizado ao percorrer nc em unidades de NR
  • Na implementação multithread, os parâmetros que mostraram bom desempenho no Ryzen 9700X foram os seguintes
    • $m_c = m_R \times \text{number of threads} \times 5$
    • $n_c = n_R \times \text{number of threads} \times 50$
  • A implementação multithread final está em matmul_parallel.h

Leituras relacionadas

1 comentários

 
GN⁺ 2024-07-05
Opiniões do Hacker News

  • Se o ponto deste texto é que, em geral, ainda há margem de desempenho, então ele até subestima o tamanho da melhoria possível. Isso apesar de o esforço investido em bibliotecas de multiplicação de matrizes ser muito maior do que na maior parte dos softwares.
    Quando o código ainda não é fortemente otimizado, é comum conseguir melhorias de 10 a mais de 1000 vezes sobre o código existente sem um esforço enorme. Em ordem aproximada de importância, o mais importante é saber se a escolha do algoritmo é adequada e se é possível eliminar a própria tarefa; também pesa bastante reduzir operações caras como idas e voltas ao kernel ou malloc.
    A vetorização pode se beneficiar de intrínsecos vetoriais explícitos, mas muitas vezes apenas reorganizar os dados de um array de structs para um struct de arrays/arrays já gera o mesmo código de máquina. A eficiência de cache também é importante e, em código paralelo, isso fica mais complexo quando não há isolamento de dados por thread, como em falso compartilhamento. Por fim, também são possíveis otimizações específicas de hardware, como intrínsecos ou assembly escrito à mão

    • O impacto da rede também não pode ser ignorado. Certa vez encontrei uma consulta distribuída que trazia cerca de 1 milhão de linhas pela rede e depois fazia um join, sobrando apenas 5 a 10 linhas; ao corrigir isso, obtivemos uma melhoria de desempenho de centenas de vezes.
      Ao mudar a consulta para que o join acontecesse no servidor remoto e apenas 5 a 10 linhas fossem enviadas pela rede, ela ficou rápida imediatamente. Sempre há overhead fixo e latência, mas, se você manda por uma conexão de rede muito mais dados do que o necessário, o desempenho acaba desandando. Também vale ler “It's the latency, stupid”, sobre o impacto da latência: http://www.stuartcheshire.org/rants/latency.html
      No geral, concordo com as considerações acima e com a ordem aproximada delas
    • “A escolha do algoritmo é adequada?” acabou, na prática, virando uma espécie de culto cargo. Muitas vezes um algoritmo “mais rápido” tem constantes reais horríveis, e a opção que faz mais trabalho acaba tendo desempenho melhor.
      Muitas entrevistas, em vez de avaliarem como raciocinar sobre por que uma implementação é lenta, fazer benchmarks e corrigi-la, viraram quizzes de memorização de algoritmos obscuros no estilo “porque o Google faz assim”

  • Padrões comuns de programação não são suficientemente especializados para o hardware e deixam muito desempenho na mesa. Este texto é um exemplo interessante, e outra demonstração clássica é “There's plenty of room at the top”.
    https://www.science.org/doi/10.1126/science.aam9744

  • Para entender isso, os artigos no repositório do BLIS são praticamente a referência canônica. Não sei por que alguém acharia que um BLAS otimizado não entrega desempenho; para matrizes suficientemente grandes, eu esperaria mais de 90% do pico da CPU.
    Da última vez que olhei, o OpenBLAS serial era, em geral, parecido com o MKL, e BLAS implementa GEMM como bloco básico de álgebra linear, não matmul. Também não entendo muito o uso de numpy em vez de um framework de benchmark, e, no Zen, acho que a comparação deveria ser com o BLAS da AMD, ou seja, a implementação baseada em BLIS. Antigamente, o BLIS tinha uma história melhor do que o OpenBLAS em paralelização, e o AMD BLIS também tem alternância de implementação para dimensões “pequenas”; não sei se isso existe hoje no OpenBLAS.
    Intrínsecos SIMD não são estritamente necessários para vetorizar o microkernel, e um bom compilador C vetoriza completamente e também desenrola os loops. O microkernel em C puro do BLIS chega a mais de 80% do desempenho da implementação otimizada à mão para Haswell com tamanhos de bloco adequados. A diferença provavelmente se deve a prefetch, mas não entendo exatamente

    • Intrínsecos SIMD e desenrolamento manual de loops são claramente necessários. É por isso que todas as bibliotecas BLAS vetorizam e desenrolam loops manualmente.
      Mesmo compiladores modernos não conseguem acertar autovetorização e desenrolamento de loops corretamente com 100% de taxa de sucesso

  • O artigo e a implementação parecem bons, mas fico curioso sobre qual é o “segredo”. O OpenBLAS vem sendo otimizado em assembly+C há décadas para exatamente esse problema; como seria possível vencê-lo?
    O texto trata de caching etc. em detalhes, e fico me perguntando se o BLAS não aproveita essas coisas ou se isso foi ajustado melhor para um processador específico

    • O OpenBLAS não é tão otimizado assim para certas arquiteturas modernas específicas. As matrizes também não eram tão grandes, e o numpy tem overhead de cffi
      A diferença de desempenho ficou muito mais evidente na vazão de pico do que na vazão média, e quase nenhuma aplicação se importa com pico. O código de benchmark mostrado parece passar pelo alocador do Python no lado do numpy, enquanto a implementação em C não passa por um alocador; então esse é o primeiro lugar para verificar erro ou inconsistência de microbenchmark. Muitas rotinas do numpy dão suporte a operações in-place, então acho que seria preciso ver explicitamente benchmarks das versões in-place dos dois lados
      O numpy também tem verificações de limites e tratamento de erros que rodam independentemente da implementação subjacente, o que explica por que ele pode ser muito mais lento até do que listas Python puras em matrizes pequenas. Se você adiciona alguns milhares de ciclos de overhead puro, fica difícil torná-lo rápido
      Esta implementação adota uma abordagem bastante criteriosa para saturar os caches relevantes e, em certo sentido, é óbvia, mas melhorias de engenharia claras merecem ser destacadas em discussões como esta. O OpenBLAS também teve muita gente trabalhando nele, mas é pouco provável que tenham pensado em tudo. Para explicar direito, seria necessária uma análise profunda do código dos dois lados
    • Vencer o OpenBLAS não é surpreendente nem algo sem precedentes. Por exemplo, a biblioteca de álgebra linear Mir, da linguagem D, já fez isso alguns anos atrás [1]
      Para implementações em C++ e C, veja a abordagem de metaprogramação [2], [3]. O que realmente surpreende é que muitas linguagens modernas como Matlab, Julia e Mojo ainda dependam do OpenBLAS, embora certamente cada uma tenha seus motivos
      [1] Numeric age for D: Mir GLAS is faster than OpenBLAS and Eigen (2016):
      http://blog.mir.dlang.io/glas/benchmark/openblas/2016/09/23/...
      [2] Vastly outperforming LAPACK with C++ metaprogramming (2018):
      https://wordsandbuttons.online/vastly_outperforming_lapack_w...
      [3] Outperforming LAPACK with C metaprogramming (2018):
      https://wordsandbuttons.online/outperforming_lapack_with_c_m...
    • -march=native compila para o modelo exato da CPU, então pode haver uma vantagem. É bem provável que o numpy tenha sido compilado para um alvo x86-64 mais genérico e antigo
      Em CPUs Ryzen, -march=native provavelmente usa v4, enquanto o numpy deve mirar v1 ou v2
      https://en.wikipedia.org/wiki/X86-64#Microarchitecture_level...
    • O numpy 2.0 integra o Google Highway para usar melhor SIMD em várias microarquiteturas, então a comparação do lado do numpy deve melhorar

  • O artigo é bom, e também é ótimo que tenham tornado os benchmarks fáceis de reproduzir. No meu Xeon W-2245 de 16 núcleos a 3,90 GHz, o matmul.c fez a multiplicação de matrizes 8192x8192 em 1,41 s com gcc -O3 e em 1,47 s com clang -O2, enquanto o NumPy levou 1,07 s
    Acho que um kernel AVX-512 seria muito mais rápido. Outro motivo para o desempenho deixar a desejar pode ser o OpenMP; pela minha experiência, gerenciar explicitamente um pool de threads com pthreads pode reduzir o overhead. Também seria melhor usar sysconf(_SC_NPROCESSORS_ONLN) em vez de codificar o número de CPUs diretamente

  • Não há motivo para colocar cargas diferentes, com um lado em Python e o outro em C. Daria para fazer uma comparação justa escrevendo ambos em C: um chamando uma biblioteca BLAS e o outro chamando esta implementação

    • Aqui faz sentido comparar com Python. Hoje, a forma mais popular de fazer esse tipo de cálculo é usando Python com numpy
      O overhead não é enorme, mas, como já foi dito em outros pontos desta thread, é importante chamar corretamente. Colocar um código numpy ingênuo contra um código C ajustado claramente não é uma comparação justa

  • Não é um caminho quente, mas a ineficiência na geração da máscara — ou seja, o uso de bit_mask — incomoda. Uma forma mais eficiente seria criar um array constante global no formato {-1,-1,...,0,0,...} e carregar a partir dos offsets de elemento 16-m, 8-m, ou então comparar um vetor constante {0,1,2,3,4,...} com m e m-8 transmitidos por broadcast
    Ainda assim, isso se aplica só a uma coluna da matriz, e o loop subsequente de maskload/maskstore demora muito mais, então é uma implicância bem pequena. Em especial, o armazenamento continua lento até no Zen 4[1], e as instruções AVX-512 são 6 vezes mais rápidas mesmo com a única diferença de receberem a máscara de um registrador de máscara. De qualquer forma, o clang autovetoriza os shifts, então deve ficar só umas 2 a 3 vezes mais lento que a minha sugestão
    [1]: https://uops.info/table.html?search=vmaskmovps&cb_lat=on&cb_...

    • Sou o autor. É realmente a primeira vez que otimizo código C e uso intrinsics, então não sou especialista nessa área, mas quero aprender mais
      Agradeço muito pelo feedback que traz uma nova perspectiva. Pelo que lembro, “criar um array global constante e carregá-lo” parecia um pouco mais lento que deslocar a máscara de bits quando testei, mas vou testar de novo para ter certeza. A abordagem de “comparar o vetor constante {0, 1, 2, 3, 4, ...} com m e m-8 transmitidos por broadcast” é uma boa ideia, vou tentar
    • Ao criar o array global constante, você pode usar int8_t para os elementos e, ao carregar, fazer extensão de sinal dos bytes para int32_t. A combinação _mm_loadu_si64 / _mm256_cvtepi8_epi32 deve compilar para uma única instrução vpmovsxbd com operando de memória
      Assim, quando alinhado corretamente com alignas(32), o array constante inteiro cabe em uma única linha de cache. Como no caso de uso do texto original são necessárias duas máscaras, a segunda instrução vpmovsxbd certamente será um acerto no cache L1D, o que funciona bem

  • E o tinyBLAS do jart?
    https://hacks.mozilla.org/2024/04/llamafiles-progress-four-m...
    E https://justine.lol/matmul/

    • Conversei bastante com a Justine ontem, e nessa workstation esta implementação parece ser pelo menos 2 vezes mais rápida que o tinyBLAS. A discussão completa está no Discord da Mozilla AI: https://discord.com/invite/NSnjHmT5xY

  • Tirando benchmarks, qual é o motivo para tornar a própria multiplicação de matrizes multithread? Na prática, não seria mais vantajoso usar múltiplas threads no algoritmo que usa a multiplicação?

    • Em HPC, de fato, isso geralmente é feito assim. Dito isso, apenas trocar por um BLAS paralelo pode ajudar facilmente certos tipos de código em R
      Mas, em geral, em código HPC, GEMM não é o gargalo

  • Ainda só dei uma passada de olhos, mas o artigo tem muitos detalhes e explicações. Parece um texto bastante bom explicando como uma multiplicação de matrizes rápida é implementada levando em conta considerações de arquitetura, então coloquei na minha lista de leitura