Advanced Matrix Extensions

Advanced Matrix Extensions (AMX), also known as Intel Advanced Matrix Extensions (Intel AMX), are extensions to the x86 instruction set architecture (ISA) for microprocessors from Intel designed to work on matrices to accelerate artificial intelligence (AI) and machine learning (ML) workloads.^[1] Especially they perform matrix multiplication at the hardware level, making them apt for problems and algorithms that use matrix multiplication as their core.

AMX was introduced by Intel in June 2020 and first supported by Intel with the Sapphire Rapids microarchitecture for Xeon servers, released in January 2023.^[2][3] It introduced 2-dimensional registers called tiles upon which accelerators can perform operations. It is intended as an extensible architecture; the first accelerator implemented is called tile matrix multiply unit (TMUL).^[4][5]

In Intel Architecture Instruction Set Extensions and Future Features revision 46, published in September 2022, a new AMX-FP16 extension was documented. This extension adds support for half-precision floating-point numbers. In revision 48 from March 2023, AMX-COMPLEX was documented, adding support for half-precision floating-point complex numbers. Both extensions are available in the Granite Rapids set of server processors (with AMX-COMPLEX support only being available in Granite Rapids-D ^[6]).

TMUL unit supports BF16 and INT8 input types.^[7] AMX-FP16 and AMX-COMPLEX also add support for real and complex FP16 numbers. The register file consists of 8 tiles, each with 16 rows of size of 64 bytes (32 BF16/FP16 or 64 INT8 elements). The only supported operation is matrix multiply and accumulate (MMA), which is the extension of the fused multiply–add (FMA) operation for scalar values as applied to matrix operands:^[8]

C n m = C n m + ∑ j = 1 J A n j B j m . {\displaystyle C_{nm}=C_{nm}+\sum _{j=1}^{J}A_{nj}B_{jm}.} ^[4]

4th Gen Intel Xeon Scalable processor core can perform 2048 INT8 or 1024 BF16 operations per cycle:^[9][10] the maximal input sizes are 16 × J {\textstyle 16\times J} for A and J × 16 {\textstyle J\times 16} for B, where J is 64 for INT8 and 32 for BF16. The matrix multiplication requires 256 J {\textstyle 256J} multiplication and 256 J {\textstyle 256J} additions, thus performing 512 J {\textstyle 512J} operations in 16 cycles.^[10]

• Compiler and assembler support
  • LLVM 13^{[11][12][13][14]}
  • GCC 11^[15][16][17]
  • GNU Assembler (GAS) initial support committed on 25 June 2020^[18][14]
• Operating system support
  • glibc support for detecting AMX feature in CPUs committed on 25 June 2020^[19]
  • Linux kernel support since version 5.16^[20]
  • VMware vSphere support for AMX in virtual machines released in ESXi version 8.0u1 for VMs using Hardware Version 20^[21]

• Intel Intrinsics Guide
• Wikichip: Advanced Matrix Extension (AMX) - x86