Supera la multiplicación de matrices de NumPy con 150 líneas de código C

 (salykova.github.io)
1 puntos por GN⁺ 2024-07-05 | 1 comentarios | Compartir por WhatsApp
  • La multiplicación de matrices de NumPy depende de bibliotecas BLAS externas, pero esta implementación busca llevar el rendimiento de un solo hilo y multihilo al nivel de BLAS usando solo C puro y FMA3·AVX2
  • El centro del rendimiento es una estructura que divide $C$ en bloques pequeños y en la que un microkernel 16×6 repite rank-1 updates dentro de registros YMM para reducir los accesos a memoria
  • En matrices de tamaño arbitrario, el manejo de bordes puede convertirse fácilmente en un cuello de botella, por lo que se combinan almacenamientos con máscara y buffers con padding de ceros para evitar la degradación de rendimiento de las cargas con máscara
  • La reutilización de caché se asegura con blocking k_c, m_c, n_c, y el rendimiento máximo real depende en gran medida del ajuste de cantidad de hilos, tamaño del kernel y tamaño de los tiles
  • AVX-512 se excluyó para lograr soporte en una gama más amplia de CPU, así que en CPUs con AVX-512 BLAS puede ser más rápido, y la comparación con OpenBLAS también se realizó con AVX-512 desactivado

Objetivo de implementación y referencia de comparación

  • El código de la implementación está publicado en sgemm.c, y optimiza la multiplicación de matrices FP32 multihilo en procesadores modernos
  • NumPy depende de bibliotecas BLAS externas para operaciones de álgebra lineal como la multiplicación de matrices
    • Algunos ejemplos son Intel MKL, Accelerate, BLIS, GotoBLAS y OpenBLAS
    • OpenBLAS, GotoBLAS y BLIS están escritos en C/FORTRAN/Assembly e incluyen implementaciones de multiplicación de matrices optimizadas manualmente para cada microarquitectura de CPU
  • El objetivo es una implementación de multiplicación de matrices escrita en C puro, sin ensamblador de bajo nivel, que cumpla las siguientes condiciones
    • Funcione con tamaños de matriz arbitrarios
    • Se ejecute en procesadores x86-64 modernos
    • Compita con bibliotecas BLAS existentes
    • El código sea simple y fácil de extender
  • Los materiales de referencia son Fast Multidimensional Matrix Multiplication on CPU from Scratch de Simon Boehm, Matrix Multiplication de Sergey Slotin, Can you multiply a matrix? de Geohot, y artículos relacionados con GotoBLAS y BLIS

Condiciones del benchmark y cálculo de FLOPS

  • El entorno de prueba es AMD Ryzen 7 9700X, 32 GB DDR5 6000 MHz CL36, OpenBLAS 0.3.26, GCC 13.3 y Ubuntu 24.04.1 LTS
  • Se usan los flags de compilación -O3 -march=native -mno-avx512f -fopenmp
  • Para una comparación justa, al instalar OpenBLAS hay que configurar el TARGET adecuado y desactivar las instrucciones AVX-512
    • Los procesadores Zen4/5 se compilan con make TARGET=ZEN
    • De lo contrario, OpenBLAS usa instrucciones AVX-512 por defecto
  • La multiplicación de matrices FP32 de OpenBLAS se ejecuta con la API cblas_sgemm
  • El benchmark usa matrices cuadradas
    • Se evalúa desde m=n=k=200 hasta m=n=k=10000 en incrementos de 200
    • La multiplicación de matrices se repite n_iter veces, y se usa el tiempo de ejecución mediano para medir el rendimiento
  • Al multiplicar una matriz $A$ de $M \times K$ por una matriz $B$ de $K \times N$, el total de operaciones es $2MNK$ FLOP
    • El rendimiento se calcula como FLOPS=(2*m*n*k)/exec_time

Límite teórico y base SIMD

  • Las CPUs x86-64 modernas procesan varios datos en paralelo mediante extensiones SIMD
  • Las instrucciones principales son AVX2 y FMA
    • Ambas usan registros YMM de 256 bits
    • Cada registro YMM puede contener 8 floats de 32 bits
  • La instrucción FMA VFMADD231PS realiza una operación packed single de la forma YMM1 = YMM2 * YMM3 + YMM1
  • En Ryzen 9700X, el throughput de fused multiply-add es de 0.5 ciclos/instrucción, es decir, 2 instrucciones por ciclo
  • En teoría, Ryzen 9700X puede ejecutar 32 FLOP por ciclo en un solo núcleo
    • La fórmula es 8 floats × 2(add+mul) × 2(1/TP)
    • Suponiendo una frecuencia sostenida de 4.7 GHz en 8 núcleos, el pico teórico multihilo se estima en 1203 FLOPS

Implementación básica y microkernel

  • Las matrices se almacenan en orden column-major
    • A[row][col] se accede en un puntero de C como ptr[col*M + row]
  • La implementación más simple recorre todas las filas y columnas de $C$ y calcula para cada elemento el producto punto de una fila de $A$ y una columna de $B$
  • La clave de una implementación de alto rendimiento es un microkernel que divide $C$ en submatrices de $m_R \times n_R$ y calcula cada submatriz de forma eficiente
  • El kernel inicializa $\bar{C}$ en cero dentro de registros y luego itera sobre la dimensión $K$
    • Trae a registros el vector columna de $\bar{A}$ y el vector fila de $\bar{B}$
    • Calcula el producto externo de ambos vectores y lo suma a los acumuladores de $\bar{C}$
    • Cada paso es un rank-1 update
  • Este enfoque reduce la cantidad de elementos traídos a registros a $(m_R+n_R)K$, frente a los $2K m_R n_R$ accesos a memoria del enfoque naive
  • Como las CPUs AVX tienen 16 registros YMM, el tamaño del kernel debe cumplir la siguiente restricción
    • $(m_R/8) \cdot n_R + m_R/8 + 1 \le 16$
    • $m_R$ debe ser múltiplo de 8
  • En teoría, cuanto más grandes y similares sean $m_R$ y $n_R$, mayor será la reducción de accesos a memoria, pero en un Ryzen 9700X real el kernel 16×6 muestra el mejor rendimiento
  • La implementación usa intrinsics de immintrin.h
    • __m256 es un tipo vectorial de 256 bits que representa el contenido de un registro YMM
    • _mm256_loadu_ps carga el vector columna de A
    • _mm256_broadcast_ss difunde un valor escalar de B a un vector de 8 floats
    • _mm256_fmadd_ps actualiza el acumulador
    • _mm256_storeu_ps guarda el resultado en memoria
  • El ensamblador generado incluye instrucciones SIMD FMA como vfmadd231ps y vbroadcastss

Padding para matrices de tamaño arbitrario

  • El kernel 16×6 básico funciona directamente cuando $M$ y $N$ son múltiplos de 16 y 6, respectivamente
  • En las zonas de borde, si la cantidad de columnas $n$ es menor que 6, el loop de almacenamiento se ejecuta solo hasta j < n
  • Cuando la cantidad de filas $m$ es menor que 16, como _mm256_storeu_ps guarda 8 elementos a la vez, se necesita almacenamiento con máscara
    • _mm256_maskstore_ps guarda en memoria solo los elementos cuyos bits de máscara estén activados
    • La máscara se genera según la cantidad de filas superpuestas $m$
  • Si en los bordes también se manejan las cargas con _mm256_maskload_ps, el rendimiento del kernel puede caer mucho
    • Las instrucciones adicionales para calcular la máscara agregan overhead
    • Como $n$ no es una constante en tiempo de compilación, al compilador le resulta difícil desenrollar el loop de forma eficiente
  • En su lugar, si $m \neq m_R$, se copia $\bar{A}$ a un buffer y se rellena con ceros; si $n \neq n_R$, también se copia $\bar{B}$ a un buffer y se completa con ceros
  • La implementación relacionada está en matmul_pad.h

Cache blocking y reutilización de datos

  • Entre los registros y la DRAM existe una jerarquía de caché de CPU, y las CPUs de escritorio modernas suelen usar cachés L1, L2 y L3
  • La caché es más rápida que la DRAM, pero su capacidad es limitada, por lo que no es posible mantener todos $A$, $B$ y $C$ completos en caché
  • La técnica de dividir las matrices en bloques pequeños, cargarlos en caché y reutilizar los mismos datos en varios rank-1 updates se llama cache blocking o tiling
  • El cache blocking de un solo hilo tiene una forma de 5 loops anidados similar a la estructura de BLIS
    • El loop más externo crea bloques $C_j$ y $B_j$ a lo largo de la dimensión $N$
    • El siguiente loop crea bloques $A_j$ y $B_p$ a lo largo de la dimensión $K$
    • $B_p$ se empaqueta para convertirse en $\tilde{B}_p$ y, si hace falta, se rellena con ceros para apuntar a la reutilización en caché L3
    • El siguiente loop crea bloques $C_i$ y $A_j$ a lo largo de la dimensión $M$, y $A_j$ se empaqueta para convertirse en $\tilde{A}_j$
    • Los últimos dos loops dividen los bloques de caché en paneles de $m_R \times k_c$ y $k_c \times n_R$, y los pasan al kernel
  • Los $\tilde{A}_j$ y $\tilde{B}_p$ empaquetados se almacenan de manera diferente
    • Los paneles internos de $\tilde{A}_j$ se almacenan en column-major
    • Los paneles internos de $\tilde{B}_p$ se almacenan en row-major
  • Los parámetros de cache blocking deben ajustarse al tamaño de caché de cada modelo de CPU
    • $k_c \times n_c$ sirve como punto de partida para llenar la caché L3
    • $m_c \times k_c$ sirve como punto de partida para llenar la caché L2
    • $k_c \times n_R$ sirve como punto de partida para llenar la caché L1
  • En la práctica, valores más grandes que los teóricos suelen dar mejor rendimiento, y como la CPU administra automáticamente la distribución en caché, hay que diseñar los loops y patrones de acceso a nivel de algoritmo
  • La implementación está en matmul_cache.h

Microoptimizaciones del kernel

  • En vez de definir los acumuladores como un arreglo tipo __m256 C_buffer[6][2], se declaran explícitamente variables de acumulación desplegadas
  • Esto ayuda a GCC a optimizar mejor el código y a evitar spilling de registros
  • El cálculo de máscaras también se modificó para usar instrucciones vectoriales
    • Se usa un arreglo estático mask[32] junto con _mm256_cvtepi8_epi32 y _mm_loadu_si64
  • Esta implementación está en matmul_micro.h

Estrategia multihilo

  • Se paralelizan tanto las operaciones aritméticas como el empaquetado
  • Los loops 5.º, 4.º y 3.º por fuera del microkernel iteran en unidades de tamaño de bloque de caché
    • Para mantener ocupados todos los hilos, la cantidad de iteraciones debe ser al menos igual a la cantidad de hilos
    • Las dimensiones de la matriz de entrada deben ser aproximadamente de al menos cantidad de hilos × tamaño del bloque de caché
  • En Ryzen 9700X, los tamaños de bloque de caché que dieron buen rendimiento en un solo hilo fueron $n_c=1535$ y $m_c=1024$
    • Para aprovechar los 8 núcleos, se necesita una dimensión mínima de tamaño $\max(m_c,n_c) \times 8 = 1535 \times 8 = 12280$
  • En cambio, los dos últimos loops iteran sobre bloques pequeños $m_R$, $n_R$, por lo que son adecuados para paralelizar
    • Generalmente $m_R$ y $n_R$ son menores que 20
    • Si $m_c$ y $n_c$ se eligen como múltiplos de la cantidad de núcleos, el trabajo puede repartirse de manera uniforme
  • En Ryzen 9700X, el mejor rendimiento se obtiene paralelizando juntos los dos loops internos con #pragma omp parallel for collapse(2) num_threads(NTHREADS)
  • En procesadores con muchos núcleos, especialmente en entornos de más de 16 núcleos, se puede considerar paralelismo anidado y paralelizar 2 o 3 loops
  • El empaquetado de $\tilde{A}$ y $\tilde{B}$ también se paraleliza con OpenMP
    • pack_blockA se paraleliza recorriendo mc en unidades de MR
    • pack_blockB se paraleliza recorriendo nc en unidades de NR
  • En la implementación multihilo, los parámetros que mostraron buen rendimiento en Ryzen 9700X son los siguientes
    • $m_c = m_R \times \text{cantidad de hilos} \times 5$
    • $n_c = n_R \times \text{cantidad de hilos} \times 50$
  • La implementación multihilo final está en matmul_parallel.h

Lecturas relacionadas

1 comentarios

 
GN⁺ 2024-07-05
Opiniones de Hacker News

  • Si el punto de este artículo es que, por lo general, queda margen de rendimiento, más bien se queda corto al estimar cuánto se puede mejorar. Y eso aun cuando el esfuerzo que se invierte en las bibliotecas de multiplicación de matrices es mucho mayor que en la mayoría del software.
    Si no se trata de código ya muy optimizado, es común lograr mejoras de 10 a 1000 veces o más sobre el código existente sin demasiado esfuerzo. En orden aproximado de importancia, lo más importante es si la elección del algoritmo es adecuada y si se puede eliminar el trabajo en sí; también pesa mucho si se pueden reducir operaciones costosas como idas y vueltas al kernel o malloc.
    La vectorización con intrínsecos vectoriales explícitos está bien, pero muchas veces basta con reorganizar los datos de arreglos de estructuras a estructuras de arreglos para obtener el mismo código de máquina. La eficiencia de caché también es importante y, en código paralelo, se vuelve más complejo cuando no hay aislamiento de datos por hilo, como ocurre con el false sharing. Por último, también son posibles optimizaciones específicas del hardware, como usar intrínsecos o ensamblador escrito a mano.

    • Tampoco hay que dejar fuera el impacto de la red. Una vez descubrimos que una consulta distribuida traía por la red alrededor de un millón de filas y luego hacía un join que dejaba solo 5 a 10 filas; al corregirlo logramos una mejora de rendimiento de cientos de veces.
      Cambiamos la consulta para que el join ocurriera en el servidor remoto y por la red solo se enviaran 5 a 10 filas, y se aceleró de inmediato. Siempre hay overhead fijo y latencia, pero si mandas por una conexión de red muchos más datos de los necesarios, al final el rendimiento se arruina. También vale la pena leer “It's the latency, stupid”, sobre el impacto de la latencia: http://www.stuartcheshire.org/rants/latency.html
      En general, estoy de acuerdo con las consideraciones anteriores y con ese orden aproximado.
    • “Si la elección del algoritmo es adecuada” en la práctica se ha convertido en una especie de cargo cult. Muchas veces el algoritmo “más rápido” tiene constantes reales horribles, y la opción que hace más trabajo termina rindiendo mejor.
      Muchas entrevistas, en vez de evaluar cómo razonar por qué una implementación es lenta, medirla con benchmarks y corregirla, se han vuelto trivias de memorización de algoritmos oscuros al estilo “porque Google lo hace así”.

  • Los patrones de programación comunes no se especializan lo suficiente para el hardware y dejan mucho rendimiento sobre la mesa. Este artículo es un ejemplo interesante, y otra demostración clásica es “There's plenty of room at the top”.
    https://www.science.org/doi/10.1126/science.aam9744

  • Para entender esto, los artículos del repositorio de BLIS son casi la referencia canónica. No entiendo por qué alguien pensaría que un BLAS optimizado no rinde, y con matrices lo suficientemente grandes deberías esperar más del 90% del pico de la CPU.
    La última vez que lo revisé, OpenBLAS en serie era en general parecido a MKL, y BLAS implementa GEMM como bloque básico de álgebra lineal, no matmul. Tampoco entiendo por qué usar numpy en vez de un framework de benchmarks, y en Zen creo que habría que comparar contra el BLAS de AMD, es decir, la implementación basada en BLIS. Antes BLIS tenía una historia mejor que OpenBLAS en cuanto a paralelización, y AMD BLIS también tiene un cambio de implementación para dimensiones “pequeñas”; no sé si eso existe actualmente en OpenBLAS.
    Para vectorizar el microkernel no son estrictamente necesarios los intrínsecos SIMD; un compilador de C decente lo vectoriza por completo y también desenrolla los bucles. El microkernel en C puro de BLIS alcanza más del 80% del rendimiento de una implementación optimizada a mano para Haswell con tamaños de bloque adecuados. La diferencia probablemente se deba al prefetch, aunque no lo entiendo con precisión.

    • Los intrínsecos SIMD y el desenrollado manual de bucles definitivamente sí son necesarios. Esa es la razón por la que todas las bibliotecas BLAS vectorizan y desenrollan manualmente los bucles.
      Incluso los compiladores modernos no logran hacer bien la autovectorización y el desenrollado de bucles con una tasa de éxito del 100%.

  • El artículo y la implementación se ven bien, pero me intriga cuál es el “secreto”. OpenBLAS lleva décadas optimizado en ensamblador+C para este problema exacto; ¿cómo puede ganarle?
    Entra en detalle sobre caché y demás, pero me pregunto si BLAS no aprovecha estas cosas, o si está mejor ajustado para un procesador específico.

    • OpenBLAS no está tan optimizado para arquitecturas modernas específicas. La matriz tampoco era tan grande, y numpy tiene sobrecarga de cffi.
      La diferencia de rendimiento fue mucho más marcada en el throughput pico que en el throughput promedio, y casi no hay aplicaciones donde el pico sea lo importante. El código de benchmark mostrado parece pasar por el asignador de Python del lado de numpy, mientras que la implementación en C no pasa por el asignador, así que ahí es donde primero revisaría si hay un error o una inconsistencia en el microbenchmark. Muchas rutinas de numpy admiten operaciones in-place, así que habría que ver explícitamente un benchmark de versiones in-place en ambos lados.
      numpy también tiene verificaciones de límites y manejo de errores que se ejecutan independientemente de la implementación subyacente, lo que explica por qué con matrices pequeñas puede ser muy lento incluso frente a listas de Python puras. Si se agregan unos miles de ciclos de sobrecarga pura, es difícil hacerlo rápido.
      Esta implementación adopta un enfoque bastante metódico para saturar las cachés relevantes y, aunque en cierto sentido sea obvio, las mejoras de ingeniería claras merecen destacarse en discusiones como esta. OpenBLAS también ha tenido mucha gente trabajando en él, pero es poco probable que hayan pensado en todo. Para explicarlo bien haría falta un análisis profundo del código de ambos lados.
    • Ganarle a OpenBLAS no es sorprendente ni tampoco algo sin precedentes. Por ejemplo, la biblioteca de álgebra lineal Mir del lenguaje D ya lo hizo hace algunos años [1].
      Para implementaciones en C++ y C, ver los enfoques de metaprogramación [2], [3]. Lo realmente sorprendente es que muchos lenguajes modernos como Matlab, Julia y Mojo todavía dependan de OpenBLAS, aunque claro, cada uno tendrá sus razones.
      [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 puede darle ventaja porque compila para el modelo exacto de CPU. Lo más probable es que numpy haya sido compilado para un objetivo x86-64 más genérico y antiguo.
      En CPUs Ryzen, -march=native probablemente use v4, mientras que numpy apuntaría a v1 o v2.
      https://en.wikipedia.org/wiki/X86-64#Microarchitecture_level...
    • numpy 2.0 integra Google Highway para usar mejor SIMD en varias microarquitecturas, así que la comparación del lado de numpy debería mejorar.

  • El artículo es bueno y también es excelente que hayan hecho que el benchmark sea fácil de reproducir. En mi Xeon W-2245 de 16 núcleos a 3.90 GHz, matmul.c hizo la multiplicación de matrices 8192x8192 en 1.41 segundos con gcc -O3 y en 1.47 segundos con clang -O2, mientras que NumPy tardó 1.07 segundos.
    Creo que con un kernel AVX-512 sería mucho más rápido. Otra razón por la que el rendimiento puede no ser ideal es OpenMP; por experiencia, administrar explícitamente un pool de hilos con pthreads puede reducir la sobrecarga. También sería mejor usar sysconf(_SC_NPROCESSORS_ONLN) en vez de hardcodear la cantidad de CPUs.

  • No hay razón para darle una carga distinta a un lado usando Python y al otro C. Se podría haber hecho una comparación justa escribiendo ambos en C: uno llamando a una biblioteca BLAS y el otro llamando a esta implementación.

    • Aquí sí tiene sentido compararlo con Python. Hoy en día, la forma más popular de hacer este tipo de cálculo es usar Python con numpy.
      La sobrecarga no es enorme, pero como se dijo en otras partes de este hilo, es importante llamarlo correctamente. Enfrentar código numpy ingenuo contra código C ajustado claramente no es una comparación justa.

  • Aunque no es una ruta crítica, me molesta la ineficiencia en la generación de máscaras, es decir, el uso de bit_mask. Formas más eficientes serían crear un arreglo constante global con la forma {-1,-1,...,0,0,...} y cargar desde los offsets de elemento 16-m, 8-m, o comparar un vector constante {0,1,2,3,4,...} contra m y m-8 difundidos
    Pero esto aplica solo a una columna de la matriz, y el bucle posterior de maskload/maskstore tarda mucho más, así que es una observación muy menor. En particular, el store sigue siendo lento incluso en Zen 4[1], y las instrucciones AVX-512 son 6 veces más rápidas aunque la única diferencia sea que reciben la máscara desde registros de máscara. clang de todos modos autovectoriza el shift, así que probablemente sea apenas 2 o 3 veces más lento que mi sugerencia
    [1]: https://uops.info/table.html?search=vmaskmovps&cb_lat=on&cb_...

    • Soy el autor. Es realmente mi primera vez optimizando código C y usando intrínsecos, así que no soy experto en esta área, pero quiero aprender más
      Aprecio mucho el feedback que aporta nuevas perspectivas. Si mal no recuerdo, cuando probé “crear un arreglo constante global y cargarlo” era un poco más lento que hacer shifts de la máscara de bits, pero lo voy a probar de nuevo para estar seguro. La idea de “comparar un vector constante {0, 1, 2, 3, 4, ...} contra m y m-8 difundidos” es buena; voy a intentarla
    • Al crear el arreglo constante global, puedes poner los elementos como int8_t y, al cargar, hacer extensión de signo de bytes a int32_t. La combinación _mm_loadu_si64 / _mm256_cvtepi8_epi32 debería compilarse como una sola instrucción vpmovsxbd con operando de memoria
      Así, si lo alineas correctamente con alignas(32), todo el arreglo constante cabe en una sola línea de caché. En el caso de uso del texto original se necesitan dos máscaras, así que la segunda instrucción vpmovsxbd será un acierto seguro en caché L1D, lo cual encaja bien

  • ¿Qué tal tinyBLAS de jart?
    https://hacks.mozilla.org/2024/04/llamafiles-progress-four-m...
    Y https://justine.lol/matmul/

    • Ayer hablé bastante con Justine, y en esa workstation esta implementación parece ser al menos 2 veces más rápida que tinyBLAS. La discusión completa está en el Discord de Mozilla AI: https://discord.com/invite/NSnjHmT5xY

  • Más allá de los benchmarks, ¿cuál es el motivo para hacer multihilo la multiplicación de matrices en sí? En la práctica, ¿no sería más conveniente usar múltiples hilos en el algoritmo que usa la multiplicación?

    • En HPC, de hecho, por lo general se hace así. Sin embargo, simplemente reemplazarlo por un BLAS paralelo puede ayudar fácilmente a ciertos tipos de código en R
      Pero en el código HPC, por lo general GEMM no es el cuello de botella

  • Todavía solo lo hojeé, pero este artículo tiene muchos detalles y explicaciones. Parece un texto bastante bueno sobre cómo se implementa la multiplicación rápida de matrices teniendo en cuenta consideraciones de arquitectura, así que lo agregué a mi lista de lectura