Is x86 ready to ACE it?

CPU designs must evolve to keep up with changing workloads. Sometimes, that evolution involves extending the instruction set to efficiently represent certain types of work. Intel’s AMX extension is one such example. AMX accelerates matrix multiplication for machine learning workloads by providing a set of 2D tile registers and configuration registers. Programmers can then configure specialized execution units (“accelerators”) to target matrix data in those tile registers. AMX was first implemented on Intel’s Sapphire Rapids server CPUs with a tile matrix multiply unit (TMUL) accelerator. Now, the x86 Ecosystem Advisory Group has written up a whitepaper and specification for ACE, which introduces a second accelerator type. While ACE is an AMX accelerator alongside TMUL, I’m going to call them “AMX” and “ACE” because TMUL was the only accelerator implementation that existed at AMX’s launch, and is still the only AMX accelerator available in hardware today. Documentation also tends to call them “AMX” and “ACE”.AMX TMUL offered a highly configurable setup where code would specify matrix tile parameters for each tile register. For instance, the tile register tmm0 could be set up as a 16x64 matrix of INT8 values by specifying 16 rows and 64 bytes per row (“colsb”). TMUL matrix multiply instructions like TDPBSSD would take the tile configuration into account and perform an entire matrix multiplication operation between specified tiles. On the data type side, AMX TMUL could operate on INT8, FP16, and BF16 values. The latest iteration of TMUL, implemented on Granite Rapids-D CPUs, also supports complex numbers with FP16 real and FP16 imaginary components.ACE does away tile register configuration options, and always considers them 64 bytes by 16 rows. Complex numbers are gone, but FP8 is in. On the compute-side, ACE provides outer product instructions instead of the inner product instructions given by AMX.Arm’s Scalable Matrix Extension (SME) and its SME2 extension is an obvious comparison point. Both ISA extensions seek to accelerate matrix multiplication within the framework of a CPU ISA, providing a low latency alternative to using less integrated accelerators like GPUs. The two ISA extensions however differ in a number of areas.

ACE is built on AMX, and continues to use AMX’s 8 KB of tile registers for holding matrix values.Arm’s SME in contrast has a variable “streaming” vector length (SVL) just like SVE’s vector length (VL). The SVE and SME vector lengths do not need to be the same, and often are not. Like SVE, SME allows the vector lengths from 128 to 2048 bits, growing in powers of two.The streaming vector length defines the size of the “ZA” storage array, which is SME’s equivalent of AMX tile registers. ZA storage is a 2D array where each side matches the SME streaming vector length. ZA storage capacity therefore ranges from 256 bytes with a 128-bit streaming vector length, to 64 KB with the maximum 2048-bit vector length.While AVX512-VNNI and AMX accelerate inner products, ACE and SME accelerate outer products.An inner product (or dot product in this special case) of two vectors aand b is,\(\mathbf a \cdot \mathbf b = \mathbf a^\mathsf{T} \mathbf b = \sum_{i=1}^n a_i b_i\)Or if you are a physicist you might have learned it via the geometrical interpretation of,\(\mathbf a \cdot \mathbf b = |\mathbf a| \, |\mathbf b| \cos \theta\)Where θ is the angle between the vectors. The outer product (or tensor product in the general case) of the vectors a and b on the other hand results in a rank-1 matrix C like the following.\(\mathbf{a} \otimes \mathbf{b} = \mathbf a \mathbf b^\mathsf{T} = \mathbf C = \begin{bmatrix} a_1b_1 & a_1b_2 & \dots & a_1b_n \\ a_2b_1 & a_2b_2 & \dots & a_2b_n \\ \vdots &

\vdots & \ddots & \vdots \\ a_mb_1 & a_mb_2 & \dots & a_mb_n \end{bmatrix} \)The columns of C are all proportional to a and this is how we know that it is a rank-1 matrix, but also that matrix multiplication is just a sum of outer products like,\(\mathbf C = \mathbf AB = \sum_{i=1}^n \mathbf a_i^{row} \mathbf b_i^{col}\)In fact many operations in linear algebra can be seen as a linear combination of outer products, making outer products a natural primitive to implement in a processor, the most obvious example the singular value decomposition of a matrix,\(\mathbf M = \mathbf{U \Sigma V}^\mathsf{T}\)isn't trivial to understand when written as an inner product, but it is just the sum of the outer product of the columns of U and V multiplied by each corresponding singular value or,\(sum_{i=1}^{rank(A)} (\mathbf u_i \otimes \mathbf v_i) \sigma_i\)Even algorithms that aren’t as obviously converted to outer product forms like the FFT have been reformulated in terms of it to utilise SME as an accelerator. Arm gave the following example on how to do it, but if you search for papers you can find various more effective ways to do it.We have historically mostly used inner products because that reduces the amount of state we need to keep in registers, and since that is a precious resource that has been important to optimise for.However, as we’ve seen above almost all linear algebra can be seen through both lenses, and code can easily convert between inner and outer products. SME takes advantage of this to do matrix multiplication via outer products and ACE seeks to do the same.Model weights are often quantized to very small bit-widths to reduce memory bandwidth and capacity pressure. ACE and SME, unlike for example NVIDIA’s TensorCores, pre-process the input vectors in software and can therefore support essentially any format you can make up and not just a limited set of “native” ones.Quantized weights are converted to natively supported data types to make use of the accelerators. ACE leans on AVX-512/AVX10’s fixed 512-bit vector width to accelerate this conversion process.

A 512-bit vector is large enough to hold a lookup table mapping a data type up to 6 bits to an 8-bit output using VPERMB, and for a 7-bit input VPERMI2B can use two 512-bit vector registers together as the lookup table.A new VUNPACKB instruction added in ACE/AVX10.3 can extract 2 to 7-bit elements into byte-aligned positions and then the previously mentioned vector permute instructions can perform the data type conversion. ACE thus provides a degree of future proofing by being able to handle any data type between 2 and 7 bits with just these three instructions, and you can with software implement essentially any method of your choosing.The x86-64 EAG also hopes this flexibility will let ACE hardware apply to applications beyond de-quantizing model weights, such as codebooks for compression of data.Data conversion using VUNPACKB and VPERM(2I)B is a two-step process, but is flexible and can handle any input data type from 2 to 7 bitsArm can’t rely on having vector registers wide enough to work as lookup tables for data conversions, because SVE lets implementations define streaming vector lengths from 128 to 2048 bits. SME2 therefore adds a 512-bit ZT0 register specifically designed to work as a 16 x 4B lookup table. LUTI2 and LUTI4 instructions perform data conversions by decompressing 2-bit or 4-bit index values from a vector register, looking up their values in ZT0, and placing output values into a destination vector register.Adding the fixed width ZT0 register lets Arm accelerate data conversions within the variable vector length SVE/SME framework, but isn’t as flexible as ACE’s VUNPACKB + permute combination. Model weights quantized down to anything other than 2 or 4 bits wouldn’t benefit from SME2’s lookup table acceleration and more complicated codebook approaches that require multiple different lookup tables cannot be supported.Data conversion using SME2’s ZT0 lookup table register. Less flexible, but Arm’s LUT instructions don’t need an intermediate register to unpack into, and don’t depend on a fixed vector length.On the other hand, SME2’s mechanism lets Arm express a data type conversion with a single instruction instead of two.

It also reduces pressure on vector registers, because the lookup table is stored in a separate register and LUT instructions don’t need a separate vector register to hold intermediate unpacked values. On the flexibility front, Arm could continue extending the ISA and add LUT instruction variants to support more data type widths.Low precision data types like FP8 suffer from low dynamic range. The Open Compute Project’s Microscaling Formats specification tackles this with scaling factors. One scaling value applies across a block of values, increasing the dynamic range of multiplication results without increasing the bit width of each data element. ACE supports this with a new 1024-bit Block Scale Register named BSR0. BSR0 is divided into two 512-bit halves, one for each of the two inputs of an outer product operation. Each half is further divided into four groups of 16 8-bit scale values, selected by an immediate operand in outer product instructions.From the ACE v1 public specification, showing an outer product with scalingArm’s SME supports scaling via 8-bit LSCALE and LSCALE2 fields in the Floating Point Mode Register (FPMR). BF1CVT uses LSCALE, and BF2CVT uses LSCALE2. Although Arm documents two instructions, the scheme is actually similar to how ACE selects BSR groups using an immediate. There’s a 1-bit opcode difference between BF1CVT and BF2CVT, which a creative decoder could treat as an immediate. Using scale factors is a two-step process with SME, where BF1CVT/BF2CVT apply a scale and a separate outer product instruction carries out the math. Update instructions for both ACE’s BSR0 and SME’s FPMR will overwrite the entire register, signaling that ISA designers on both sides don’t expect scale values to change often.There’s a distinct possibility that neither register gets renamed in hardware implementations, so any scale value update would create a serialization barrier for FP8 instructions. If scale values do change often enough to warrant renaming, then both ISAs are in a difficult position. All of ACE’s FP8 outer product instructions reference BSR0, and therefore depend on the last BSR0 write. SME has separate instructions for scaling and outer product math, but FPMR also includes fields that control FP8 format and overflow behavior.