英特尔最近公布了高级矩阵扩展(AMX)的详细信息,AMX是一种新的x86扩展,专为矩阵操作而设计,目标是加速机器学习工作负载。AMX是英特尔在DL Boost技术品牌下销售的一系列特定于AI的扩展中的第三个。第一个扩展AVX512_VNNI是在Cascade Lake中引入的。VNNI旨在通过将值对的乘法和加法与累加写回操作相结合来加速CNNS内核(特别是8位和16位值)。库珀湖紧随其后,推出了下一个扩展-AVX512_BF16。AVX512_BF16提供了许多指令,用于将单精度浮点值转换为bfloat16,以及对bfloat16对执行点积并累加结果。AMX将是第三种DL Boost技术,英特尔表示将在2021年推出基于蓝宝石Rapids微体系结构的第4代Xeon Scalable。
Advanced Matrix Extension或AMX是一个新的x86扩展。从复杂性的角度来看,它比其他两个DL Boost扩展要复杂得多。_VNNI和_BF16构建在AVX512基础之上,而AMX是一个独立的扩展,具有自己的存储和操作。AMX引入了一个新的矩阵寄存器堆,它具有8个秩2张量(矩阵)寄存器,称为“块”。它还介绍了能够在这些瓷砖上操作的加速器的概念。该扩展的实现方式与添加AVX/2和AVX512的方式类似,不需要对整体架构进行任何特殊更改。AMX指令在指令流中与通过与主机的存储器访问相一致的瓦片进行的存储器加载/存储操作同步。与其他扩展一样,AMX可以与其他x86代码以及其他扩展(如AVX512)交错。
矩阵寄存器堆包括八个瓦片(命名为TMM0…。TMM7),每个具有16行乘64字节列的最大大小,用于总大小为1KiB/寄存器和用于整个寄存器堆的8KiB。通过瓦片控制寄存器(TILECFG),程序员能够配置那些瓦片的大小(根据行和字节_每_行)。根据所实现的算法,可以更改拼贴的大小以更自然地表示该算法。
AMX目前只包含12条指令。它们可以松散地分成三类:配置指令、平铺处理(例如,加载、存储、零)和平铺操作(即,点积)。基本流程是设置如何在磁贴上操作的配置。这只做一次。您将在内存中加载一个代表较大图像中的一小部分的平铺,对该平铺进行操作,然后对表示图像下一部分的下一个平铺重复操作。完成后,您将把生成的瓷砖存储回内存。算法可能会在矩阵乘法之后使用AVX-512进行后处理。
目前的规范只定义了一个加速器--平铺矩阵乘法(TMUL)单元,它包括能够在平铺上操作的融合乘加单元网格。TMUL单元仅执行一个操作Tile C[M][N]+=Tile A[M][K]*Tile B[K][N],这些矩阵的大小由程序员预先配置在TILECFG中。
AMX的定义方式类似于具有各种子扩展的AVX512。有底部的瓷砖建筑和附带的说明。英特尔称其为AMX-Tiles。AMX-瓦片包含用于配置瓦片寄存器文件的所有必要指令(即,加载/存储TILECFG)以及用于处理瓦片的指令(即,加载/存储瓦片、零/复位瓦片)。TMUL单元实际上是AMX-INT8和AMX-BF16的一部分。似乎对需要支持的硬件没有任何要求,也不排除硬件只支持INT8或只支持BF16子扩展的可能性。事实上,当前的ISA扩展参考手册让Sapphire Rapids的支持变得模棱两可,说它支持AMX,但实际上并没有透露支持的级别。AMX-BF16子扩展添加了bfloat16点积瓦片操作支持,而AMX-INT8子扩展添加了字节点积瓦片操作支持。由于AMX-INT8支持有符号/无符号字节操作的所有四种排列,因此在AMX-INT8中有四条指令。
英特尔最近报告说,蓝宝石快速硅片已经通电。随着该公司接近发布,有关该芯片及其新功能的更多细节可能会浮出水面。除了最近推出的库珀湖至强处理器和计划于明年推出的蓝宝石急流处理器之外,英特尔预计也将在今年晚些时候推出冰湖至强处理器。
完整的高级矩阵扩展(AMX)规范可在英特尔®体系结构指令集扩展编程参考(319433-0,2020年6月)的第3章中找到