Resumen Técnico: Emulación de Componentes-Producto de BF16 para GEMM de FP32 y FP64 en Intel AMX

1. Declaración del Problema

Las CPU modernas integran cada vez más motores de matrices de alto rendimiento (como las extensiones de matriz avanzada de Intel, o AMX) optimizados para cargas de trabajo de IA de baja precisión (por ejemplo, BF16, FP8). Sin embargo, muchas aplicaciones de computación científica en campos como la dinámica de fluidos computacional y la química cuántica aún dependen de la multiplicación de matrices generales (GEMM) de FP32 y FP64 para la estabilidad numérica y la reproducibilidad. Esto crea un desajuste de hardware: los motores de baja precisión de alto rendimiento no son utilizables de forma nativa para tareas científicas de alta precisión.

La pregunta central abordada por este artículo es si el rendimiento de matrices de baja precisión diseñado para la IA puede ser reutilizado algorítmicamente para acelerar la GEMM de alta precisión en CPU. El desafío radica en emular la aritmética FP32 y FP64 utilizando componentes BF16 sin incurrir en sobrecargas prohibitivas de descomposición, movimiento de datos y reconstrucción, manteniendo al mismo tiempo una precisión numérica comparable a las bibliotecas estándar como Intel oneMKL.

2. Metodología

El artículo propone un marco de emulación de componente-producto que descompone operandos de alta precisión en múltiples matrices BF16, calcula productos de baja precisión seleccionados utilizando AMX y reconstruye el resultado en el formato objetivo.

2.1 Formatos de Punto Flotante y Dominios

• BF16: Conserva el rango de exponente de 8 bits de FP32 pero reduce la precisión de la mantisa a 8 bits.
• Emulación FP32: Opera sobre operandos estándar de FP32.
• Emulación FP64: Opera sobre un dominio restringido ($FB_{64}$) donde los valores FP64 se encuentran dentro del rango de exponente normal de BF16. Esta restricción evita la necesidad de factores de escala por panel, preservando el margen de rendimiento.

2.2 Estrategias de Descomposición

• Ruta AMX-FP32: Utiliza una descomposición de residuo de tres componentes BF16. Cada operando FP32 se divide en tres rebanadas BF16 ($A_0, A_1, A_2$) mediante redondeo de residuo directo. Esto sigue el enfoque "estilo Henry", que explota el rango de exponente compartido para asegurar la resta exacta de los residuos en FP32.
• Ruta AMX-FP64: Utiliza una descomposición Ozaki de seis rebanadas simplificada y fija. Cada operando FP64 se divide en seis rebanadas BF16. A diferencia de los esquemas Ozaki completos que utilizan escalado dinámico, este método almacena los fragmentos extraídos directamente como matrices BF16 (asumiendo que permanecen normales), eliminando la necesidad de metadatos de escala y pesos de reconstrucción.

2.3 Programación de Productos y Reconstrucción

• Reconstrucción FP32: El método evalúa seis productos de componentes seleccionados (las primeras tres diagonales de la matriz de interacción de componentes). Crucialmente, estos productos se acumulan directamente en tiles de AMX de FP32. La implementación emplea un programa de reutilización de operandos (por ejemplo, $A_1B_1 \to A_1B_0 \to \dots$) que mantiene los tiles acumuladores residentes, minimizando el tráfico de carga/almacenamiento de tiles.
• Reconstrucción FP64: El método evalúa un número variable de productos (6, 10, 15 o 21). A diferencia de la ruta FP32, cada producto BF16 seleccionado se calcula en FP32, se almacena en memoria, se ensancha a FP64 y se acumula en un bloque acumulador de FP64. Esto evita que los errores de acumulación de FP32 contaminen el resultado de alta precisión, pero incurre en una sobrecarga significativa debido a los repetidos almacenamientos y conversiones de formato.

2.4 Detalles de Implementación

La implementación utiliza una estructura de GEMM bloqueada con bloques de salida de $256 \times 256$ y micro-tiles de $32 \times 32$. Depende de buffers de componentes precomputados y contiguos en bloque para optimizar los patrones de acceso a memoria. Los operandos A se almacenan en un diseño estándar de fila principal (row-major), mientras que los operandos B se empaquetan previamente en el diseño VNNI requerido por las instrucciones de producto punto BF16 de AMX.

3. Contribuciones Clave

1. Puente Algorítmico: Demuestra un método práctico para emular GEMM de FP32 y FP64 utilizando instrucciones BF16 de Intel AMX, cerrando la brecha entre el hardware optimizado para IA y los requisitos de la computación científica.
2. Emulación FP32 con Precisión de Nivel FP32: Muestra que un programa de seis productos BF16 con acumulación en FP32 puede lograr una precisión comparable a oneMKL SGEMM (medida por el error relativo normativo y el error por componente escalado por GEMM), sin pretender identidad de bits.
3. Emulación FP64 Controlada: Introduce un compromiso ajustable entre precisión y rendimiento para la emulación FP64 variando el número de productos de componentes retenidos (de 6 a 21) dentro de una descomposición fija de seis rebanadas.
4. Movimiento de Datos Optimizado: Desarrolla estrategias de programación específicas (reutilización de operandos para FP32) y diseños de datos (empaquetado VNNI, buffers precomputados) para mitigar la sobrecarga de la descomposición y la reconstrucción en los motores de matrices de CPU.

4. Resultados Experimentales

Los experimentos se realizaron en un servidor Intel Xeon Platinum 8462Y+ de doble socket, comparando los kernels propuestos contra Intel oneMKL SGEMM y DGEMM.

• Rendimiento FP32: La implementación AMX-FP32 supera a oneMKL SGEMM en todos los tamaños de matrices cuadradas probados ($N=256$ a $N=32768$). Las aceleraciones oscilan entre 1.14$\times$ y 2.56$\times$, observándose las mayores ganancias en matrices de mayor tamaño, donde la ventaja de rendimiento de BF16 amortiza la sobrecarga de la descomposición.
• Rendimiento FP64: Los resultados están más limitados.
  • AMX-FP64-6 (6 productos) logra aceleraciones de hasta 1.72$\times$ sobre oneMKL DGEMM para matrices grandes ($N \ge 16384$).
  • AMX-FP64-10 proporciona ganancias marginales solo para las matrices más grandes.
  • AMX-FP64-15 y AMX-FP64-21 no superan a DGEMM en las configuraciones probadas, ya que los costos de almacenar, convertir y acumular productos adicionales consumen el margen de rendimiento disponible.
• Precisión:
  • AMX-FP32 logra errores en el rango de $10^{-7}$, consistente con una precisión de nivel FP32.
  • Las variantes AMX-FP64 muestran una disminución del error a medida que aumenta el número de productos, siendo la variante de 21 productos la que se acerca a la precisión de nivel DGEMM, pero con un costo de rendimiento que anula la aceleración.

5. Significado y Reivindicaciones

El artículo posiciona este trabajo como un estudio algorítmico temprano más que como un reemplazo completo de las bibliotecas nativas de alta precisión. Su importancia radica en demostrar que:

• Los motores de matrices de baja precisión en CPU pueden utilizarse eficazmente para la computación científica de alta precisión mediante una cuidadosa descomposición algorítmica y la optimización del movimiento de datos.
• Existe un techo de rendimiento para la emulación: mientras que la emulación FP32 es altamente efectiva debido a la acumulación eficiente en FP32, la emulación FP64 está limitada por la sobrecarga de ensanchar y acumular los resultados fuera de los registros de tile de AMX.
• El método ofrece un compromiso ajustable para FP64, permitiendo a las aplicaciones seleccionar un número de productos que equilibre sus requisitos específicos de precisión frente a las necesidades de rendimiento.

Los autores señalan explícitamente las limitaciones: el prototipo actual solo soporta matrices cuadradas, requiere entradas dentro del rango de exponente de BF16 (para FP64) y no maneja desbordamientos/subdesbordamientos (overflow/underflow) ni valores excepcionales. Se identifica como trabajo futuro la extensión del marco a matrices rectangulares, dominios de entrada más amplios mediante descomposición escalada y la adaptación a otros motores de matrices de CPU (por ejemplo, Arm SME).