Intel AMX Programming Model Lands In LLVM Compiler

One of the big features to look forward to with Intel's Xeon "Sapphire Rapids" is the introduction of AMX as the Advanced Matrix Extensions. While Sapphire Rapids looks to be at least one year out still, the company's open-source compiler engineers have already been hard at work on the software infrastructure support.

AMX is Intel's new programming paradigm with a focus on better AI performance both for training and inference. AMX is built around the concept of "tiles" as a set of two-dimensional registers for representing a larger memory image and accelerators that can operate on said tiles. Initial AMX features are for BFloat16, TILE, and INT8 while the design is extensible for new accelerators to be added later.

Since July we've seen AMX support beginning to appear for LLVM and GCC compilers. New instructions around AMX were initially added along with related pieces like the assembler support.

Being merged today into LLVM 12 though is the initial AMX programming model for this open-source compiler. The AMX programming model was discussed for a while on the LLVM mailing list and consists of new intrinsics for the LLVM intermediate representation, the C interface exposed to the user (developer), lowering from AMX intrinsics to pseudo instructions, register allocation for the AMX tile registers, transforming the appropriate load/store sizes to AMX intrinsics, and support for morphing AMX psueudo instructions into real AMX instructions.

The AMX programming model was merged today while more details can be found via the mailing list discussion that's been taking place the past few months around all aspects of the model from the C usage to the internal LLVM representation.

Separately, Intel engineers have also been working on the Linux kernel support around AMX and other changes for ensuring Advanced Matrix Extensions are well supported by the open-source software ecosystem by the time Sapphire Rapids is ready to launch.