La evolución de las cargas de trabajo, particularmente en Machine Learning, ha impulsado la necesidad de acelerar operaciones fundamentales de álgebra lineal directamente en la CPU. Tradicionalmente, las arquitecturas de CPU se han centrado en productos internos (dot products) debido a la optimización del uso de registros, un recurso escaso. Sin embargo, la computación moderna, especialmente con matrices grandes y datos cuantificados, se beneficia enormemente de la eficiencia de los productos exteriores.
ACE (Advanced Matrix Extensions) aborda este problema fundamental al introducir un conjunto de instrucciones y un modelo de programación optimizado para productos exteriores dentro del ecosistema x86. Esto permite un procesamiento más eficiente de matrices, reduciendo la latencia y el consumo de energía asociado con el movimiento de datos, y ofreciendo una alternativa de baja latencia a los aceleradores externos como las GPUs. La relevancia actual radica en la omnipresencia de modelos de IA que dependen intensamente de estas operaciones, y la necesidad de ejecutar inferencia de manera eficiente en el borde o en servidores con restricciones de energía y latencia.
Arquitectura del Sistema
ACE se construye sobre la extensión AMX (Advanced Matrix Extensions) de Intel, utilizando sus registros de 'tile' de 8 KB para almacenar valores de matriz. A diferencia de AMX TMUL, que se centra en productos internos y permite configuraciones flexibles de tiles, ACE estandariza los tiles a 64 bytes por 16 filas y se especializa en instrucciones de producto exterior. Esto simplifica la lógica de control y optimiza el flujo de datos para este tipo de operación.
Una característica clave de ACE es su enfoque en la conversión y des-cuantificación de tipos de datos. Utiliza instrucciones AVX-512/AVX10 como VPERMB y VPERMI2B, junto con la nueva VUNPACKB, para manejar eficientemente la conversión de datos cuantificados (2 a 7 bits) a tipos nativos (8 bits). Esto proporciona una flexibilidad significativa para soportar diversos formatos de cuantificación definidos por software. Además, ACE introduce un registro Block Scale Register (BSR0) de 1024 bits para aplicar factores de escala a bloques de valores, mitigando el bajo rango dinámico de tipos de datos de baja precisión como FP8. Las instrucciones de ACE toman entradas de registros vectoriales AVX-512/AVX10 y acumulan en los registros de tile, lo que permite un mayor tamaño de tile de salida (32x64 elementos) en comparación con AMX, reduciendo la presión sobre el ancho de banda de la memoria L1D al minimizar las cargas de datos de entrada en los registros de tile.
Flujo de Procesamiento de Datos Cuantificados con ACE
- 1 Carga de Datos Cuantificados Carga de vectores de entrada con pesos de modelo cuantificados (2-7 bits) en ...
- 2 Des-cuantificación Uso de VUNPACKB para extraer elementos, seguido de VPERMB/VPERMI2B para conve...
- 3 Aplicación de Escala Carga de factores de escala en BSR0 y aplicación a los datos de entrada.
- 4 Operación de Producto Exterior Instrucciones ACE realizan el producto exterior entre vectores de entrada y a...
- 5 Almacenamiento de Resultado El resultado acumulado en los registros de tile se escribe de vuelta a memoria.
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | Intel AMX (Advanced Matrix Extensions) | Proporciona los registros de 'tile' de 8 KB que ACE utiliza para la acumulación de resultados de matrices. |
| compute | Intel AVX-512/AVX10 | Utilizado para cargar los vectores de entrada en las operaciones de producto exterior y para las instrucciones de des-cuantificación (VPERMB, VPERMI2B, VUNPACKB). |
| compute | OCP Microscaling Formats | Especificación para formatos de escalado de baja precisión, implementada por ACE con el Block Scale Register (BSR0). |
Trade-offs
Ganancias
- ▲ Reducción del tráfico de caché
- ▲ Flexibilidad en tipos de datos cuantificados
- △ Mayor tamaño de tile de salida
- △ Potencial para mayores frecuencias de reloj en escenarios compute-bound
Costes
- △ Complejidad de la ISA
- △ Latencia de context switch (si los registros de tile se usan como scratchpad)
- △ Menor flexibilidad en la configuración de tiles (vs. AMX TMUL)
Fundamentos Teóricos
El concepto de productos exteriores como bloque fundamental para la multiplicación de matrices y otras operaciones de álgebra lineal tiene raíces profundas en la teoría de tensores y el álgebra multilineal. La descomposición en valores singulares (SVD), por ejemplo, se expresa naturalmente como una suma de productos exteriores, lo que demuestra la versatilidad de esta operación como primitiva computacional. Algoritmos como la Transformada Rápida de Fourier (FFT) también han sido reformulados en términos de productos exteriores para aprovechar arquitecturas que los aceleran.
La optimización del acceso a la memoria y el uso de la jerarquía de caché, fundamental para el rendimiento de ACE, se relaciona con trabajos clásicos en la optimización de algoritmos de multiplicación de matrices, como la familia de algoritmos MOMMS (Memory-Optimized Matrix Multiplication Algorithms). Estos estudios, que datan de las décadas de 1980 y 1990, exploran cómo el 'tiling' (división de matrices en submatrices) y la reutilización de datos pueden reducir drásticamente los requisitos de ancho de banda de memoria, un principio que ACE y SME aplican directamente a nivel de ISA para maximizar la eficiencia de los registros y las cachés.