Extensiones AVX, SSE,…: ¿Cuáles son las equivalentes en Apple? ¿Qué es Apple AMX?

27 de febrero de 2023

Como sabes, en la ISA AM64, EM64T o x86-64, como la quieras llamar, además del set básico de instrucciones también te encuentras con extensiones como AVX, MMX, SSE, FMA4, XOP, etc. Sin embargo, esto también ha sido así en otras ISAs más allá de la x86, como puede ser SPARC, PA-RISC, PowerPC, Arm, etc. Y por eso en este artículo te mostraremos cuáles son las extensiones equivalentes de Apple Silicon, concretamente nos centraremos en Apple AMX, uno de los grandes desconocidos.

Quizás también te pueda interesar:

Contenido

Apple Silicon (Arm Inside)

Como sabes, Apple decidió abandonar Intel y pasar a crear sus propios chips basados en la ISA Arm, y por tanto, también ha heredado de esta algunas de sus extensiones. Las máquinas con Apple Silicon, ya sean los A-Series para dispositivos móviles o las M-Series para ordenadores, usan el set de instrucciones completo de la ISA, además de algunas de sus extensiones, aunque no todas.

Además, también se han agregado algunas extra como es el caso de las Apple AMX de las que hablaremos aquí. Por ejemplo, en los procesadores M-Series existen varias formas da acelerar ciertas tareas como las multimedia que necesitan de repetitivos cálculos vectoriales, etc. Y esas cuatro formas son:

Apple también usa una de las extensiones de ARM. Concretamente hace uso de las SIMD ARMv8 NEON de 128-bit de ancho que se pueden ejecutar hasta 4 por ciclo de reloj en los núcleos de alto rendimiento Firestorm de esta arquitectura híbrida, no es así en los Icestorm (eficientes).
También cuenta con sus propias y desconocidas extensiones Apple AMX no documentadas. Se emiten desde la CPU y se ejecutan en una unidad de aceleración especial. Todo un misterio.
Otra forma de acelerar las aplicaciones es haciendo uso de la NPU o Neural Engine, pero esta sirve solo para aplicaciones de inteligencia artificial, ya que se centra en operaciones de coma flotante simples.
Por supuesto, también puede usarse la iGPU como una GPGPU, haciendo uso de ella para acelerar ciertas tareas con Metal Compute Shaders, OpenCL, etc.

¿Qué es SIMD?

Las extensiones que suelen tener los procesadores son SIMD, mientras que las de la ISA son SISD. Es decir, estas extensiones se ejecutan haciendo uso de registros especiales y de unidades funcionales especiales donde se permite trabajar de forma vectorial o con matrices. Por eso se denominan Single Instruction Multiple Data, puesto que una misma instrucción se puede aplicar a un conjunto de datos.

Por ejemplo, imagina que tu CPU Arm ejecuta una instrucción no vectorial de la ISA, y que esa instrucción es ADD de 32-bit, es decir, una suma que representada en lenguaje ensamblador sería algo así:

ADD S2, S0, S1

Esto significa que el dado almacenado en el registro S0 se sumará al dato almacenado en el registro S1 y el resultado se almacenará en el registro S2 una vez pase por la ALU de la CPU:

S0+S1=S2

Cuando esta instrucción y esos datos sean cargados desde la memoria principal a la CPU, la unidad de control se encargará de decodificarla para su interpretación, y al ver que es una instrucción ADD enviará las señales de control oportunas para que la ALU se ponga en modo suma y sume los dos números pertenecientes a los registros indicados, finalmente guardando el resultado de la suma en el registro S2.

En cambio, usando las extensiones SIMD NEON, por ejemplo, la operación sería diferente. En este caso tendremos:

VADD V2, V0, V1

Que el equivalente sería lo siguiente:

V0 + V1 = V2

Es decir, que el resultado de V0 sumado a V1 se almacenaría en V2. Como ves, en este caso no hablamos de S, sino de registros V, que son los vectoriales de longitud 128-bit. Es decir, que correspondería a 4 registros S de 32 bits cada uno, por lo que en vez de ir de 32-bit en 32-bit, se puede ejecutar la instrucción VADD sobre todo un vector de 4 valores de 32-bit, por poner un ejemplo.

Dicho de otro modo, el equivalente a VADD V2, V0, V1 podría ser el siguiente si no se usase una instrucción SIMD o vectorial:

ADD S2, S0, S1

ADD S5, S3, S4

ADD S8, S6, S7

ADD S11 , S9, S10

Que si se pusiesen unidos S11 S8 S5 S2 (cada uno de 32-bit, recordemos), sería el equivalente al V2 (128-bit).

Es decir, en la SIMD se aprovecha un mayor paralelismo haciendo uso de operaciones simultáneas en grandes conjuntos de datos. Esto permite acelerar algunas aplicaciones como pueden ser las operaciones multimedia que necesitan gran cantidad de estas operaciones, o las aplicaciones, científicas, etc.

¿Qué es NEON?

Entre las extensiones de la arquitectura ARM existen unas extensiones SIMD denominadas NEON. Estas extensiones son ejecutadas por unos coprocesadores o FPUs integradas en la microarquitectura de la CPU y que los procesadores Apple también usan.

Cuando se desarrolla software y se optimiza el código fuente para usar NEON, y se compila para ello, esto permite al software tener una ventaja haciendo uso de las instrucciones vectoriales en vez de tenerlo que hacer con instrucciones escalares simples y tener que repetir la instrucción varias veces hasta conseguir el resultado que se busca. Esa es la ventaja de NEON.

Las instrucciones NEON se ejecutan como parte del flujo de instrucciones ARM/Thumb. Esto simplifica el desarrollo, la depuración y la integración del software en comparación con el uso de un acelerador externo. Por tanto, también pueden tener acceso a la memoria, copia entre registros de propósito general, conversión de tipos de datos, etc.

Estas extensiones tipo SIMD son de hasta 128-bit, por lo que si se comparan con las instrucciones de la ISA de AArch32 significaría que podría trabajar con 4 veces eso, o si lo comparamos con la ISA AArch64, de 64-bit, podría trabajar con el doble de eso.

Todos los procesadores Arm Cortex A-Series tienen NEON por defecto, pero también está presente en otras implementaciones que no están diseñadas por Arm, pero que usan la ISA Arm, como puede ser en los Apple Silicon.

¿Qué es Apple AMX?

Como hemos dicho anteriormente, los chips Apple Silicon, además de usar las SIMD NEON, también tienen unas instrucciones denominadas Apple AMX no documentadas y bastante misteriosas. Estas instrucciones son poco conocidas por la mayoría de usuarios, y no hay que confundirlas con las Intel AMX, aunque ambas están destinadas a lo mismo, a la multiplicación matricial desde la FPU.

Este tipo de instrucciones Apple AMX están presentes en los SoCs Apple M-Series, aunque cada generación puede tener diferentes instrucciones de este tipo. Algunas fuentes creen que el M1 tiene dos versiones diferentes de AMX, una con máscara de escritura de 7-bit y otra con máscara de escritura de 9-bit. Aunque esto solo son conjeturas.

Tipos de datos en Apple AMX

Para los datos con los que trabaja Apple AMX, se pueden usar tres formatos o tipos diferentes:

De coma flotante normalizados bajo IEEE-754 FP16, FP32 y FP64. Con el mismo ancho para los tres operandos fusionados, multiplicadores y sumadores).
Multiplicandos de coma flotante IEEE-754 FP16 y acumulando FP32.
Multiplicandos enteros de 8-bit o 16-bit, que se acumulan en 16-bit o 32-bit (con signo).

Funcionamiento Apple AMX

En la anterior figura puedes ver un ejemplo de cómo operan las instrucciones Apple AMX. Como ves, es como si fuese una cuadrícula o matriz de datos donde cada unidad puede realizar la multiplicación y acumulación. Para que esto sea posible, esta cuadrícula necesita ser alimentada por dos tipos de registros:

Registros X, cada uno de los cuales contiene 32 elementos de 16 bits (o 16 elementos de 32 bits, u 8 elementos de 64 bits).
Registros Y que contiene de manera similar 32 elementos de 16 bits (o 16 elementos de 32 bits, u 8 elementos de 64 bits).

Como una especie de SIMD que es, una sola instrucción vectorial puede realizar un producto multiplicando cada elemento de un registro X con cada elemento de un registro Y y acumularlo en un registro Z, en la posición correspondiente para obtener el resultado que se busca. De esa forma no se necesita ir procesando instrucción a instrucción escalar aplicada solo a dos valores.

Apple AMX: ¿para qué sirve?

Bueno, este tipo de matrices de multiplicación y acumulación pueden ser muy importantes para acelerar aplicaciones de:

Procesamiento de imágenes
Reconocimiento de voz y escritura
Reconocimiento facial
Compresión
Multimedia (audio y vídeo)

Diferencias entre el coprocesador matricial de Apple AMX y un motor neuronal

Pues bien, vamos a ver primero qué hace un motor neuronal, y es que también se encarga de operaciones matriciales para ayudar a las tareas de aprendizaje automático. Entonces, ¿es lo mismo que Apple AMX? ¿se necesitan ambos? Pues bien, hay que aclarar una cosa, y es que los datos entre un coprocesador y una CPU son compartidos mediante un bus de datos compartido. En cambio, los aceleradores suelen tener su propia memoria dedicada.

Es decir, mientras que el coprocesador para Apple AMX comparte ese bus con la CPU, el acelerador Neural Engine de Apple no lo hace. Dicho de otro modo, mientras que el motor AMX puede «espiar» o ver las instrucciones que la CPU extrae de la memoria, éstas resultan opacas para el acelerador.

Es ahí donde el coprocesador tiene la ventaja, ya que pueden reaccionar ante las instrucciones para las que fueron diseñados. Esto quiere decir que las instrucciones que procesa el coprocesador pueden colocarse dentro del código normal, sin tener que programar de forma especial como se haría para un acelerador. Y esto simplifica mucho las cosas en cuanto al desarrollo de software.

Diferencias entre el coprocesador matricial de Apple AMX y uno vectorial tipo SIMD

Es fácil confundir un coprocesador matricial con un motor vectorial SIMD, como podría ser NEON. Sin embargo, un motor vectorial SIMD forma parte del núcleo de un microprocesador. Al igual que la ALU (unidad aritmética lógica) y la FPU (unidad de coma flotante) forman parte de la CPU. Dentro del microprocesador hay un decodificador de instrucciones que separa una instrucción y decide qué unidad funcional activar.

Dentro de una CPU tienes la ALU, la FPU y los motores vectoriales SIMD (no mostrados) como partes separadas activadas por el decodificador de instrucciones. Un coprocesador es externo, aunque esto no implique que sea externo al propio SoC, como podrían ser los famosos Intel 8087 del pasado, que se encontraba en un chip independiente.

No tengo claro por qué el AMX no es simplemente una parte de los núcleos Firestorm de las M-Series. De todos modos, todos están en el mismo chip de silicio. Tal vez porque al ser independiente puede funcionar en paralelo con la CPU y así no restar desempeño o tal vez porque Apple haya querido respetar el estándar ARM en sus núcleos.

¿Por qué el AMX es un secreto?

Apple AMX es un secreto a voces. Es cierto que no se describe en ninguna documentación oficial, pero fue descubierto por un desarrollador llamado Dougall Johnson, que hizo un gran trabajo de ingeniería inversa con el chip M1 de Apple, descubriendo este coprocesador.

Sin embargo, aunque los desarrolladores no puedan usarlo de forma directa, lo cierto es que Apple incluye una serie de bibliotecas y frameworks especiales para acelerar tareas como:

vImage: procesamiento de imágenes de alto nivel, como conversión entre formatos, manipulación de imágenes.
BLAS: una especie de estándar industrial para el álgebra lineal (lo que llamamos las matemáticas que tratan con matrices y vectores).
BNNS: se utiliza para ejecutar y entrenar redes neuronales.
vDSP: procesamiento digital de señales, transformaciones de Fourier, convolución. Se trata de operaciones matemáticas importantes en el procesamiento de imágenes o de cualquier señal, incluido el audio.
LAPACK: funciones de álgebra lineal de alto nivel, por ejemplo, para resolver ecuaciones lineales.

Dougall Johnson descubrió que estas bibliotecas usaban el coprocesador Apple AMX para acelerar estos cálculos en programas especiales. Así fue como descubrió instrucciones especiales no documentadas en el código máquina de estas bibliotecas.

Pero, ¿por qué Apple no documenta esto y nos deja usar estas instrucciones directamente? Como se ha mencionado antes, esto es algo que ARM Ltd. querría evitar. Si las instrucciones personalizadas se utilizan ampliamente podría fragmentar el ecosistema ARM.

Sin embargo, lo que es más importante, esto es una ventaja para Apple. Al permitir que sus bibliotecas utilicen únicamente estas instrucciones especiales, Apple conserva la libertad de cambiar radicalmente el funcionamiento de este hardware más adelante. Podrían eliminar o añadir instrucciones AMX. O dejar que el motor neuronal haga el trabajo en su lugar.

En cualquier caso, al no tenerlas que usar directamente, también facilitan el trabajo a los desarrolladores. Los desarrolladores sólo tienen que utilizar el framework Accelerate y pueden ignorar cómo Apple acelera específicamente los cálculos matriciales.

Esta es una de las grandes ventajas que tiene Apple al estar integrada verticalmente. Al controlar tanto el hardware como el software, pueden hacer este tipo de trucos.

¿Cuáles son las ventajas?

Pues bien, en cuanto a las ventajas, si se compara al coprocesador de Apple AMX acelerando tareas de inteligencia artificial con hacerlo usando las instrucciones Neon, que están oficialmente soportadas por ARM, la diferencia es notable.

Apple AMX es capaz de obtener el doble de rendimiento que las NEON para operaciones matriciales. Aunque AMX no es mejor en todo a NEON, sí que lo es en trabajos de aprendizaje automático y otras cargas de trabajo similares, y ahí es donde está la diferncia…

AMD ha integrado ahora un motor neuronal en sus Ryzen 7000 Series ¿habrá conseguido algo parecido? Esto ya es otro tema, pero será interesante analizarlo…