040、AMX Dialect:Intel高级矩阵扩展指令集
AMX Dialect:Intel高级矩阵扩展指令集
从一次诡异的性能回退说起
去年秋天,团队在优化一个深度学习推理管线。模型是BERT-large,batch size=16,跑在Sapphire Rapids服务器上。一切看起来都很正常——直到我们对比了AVX-512和AMX两种实现的性能数据。
AVX-512版本:单次推理12.3ms。AMX版本:单次推理14.1ms。慢了15%。
当时我盯着这个数字看了十分钟。AMX的峰值吞吐量是AVX-512的4倍以上,理论上不应该出现这种情况。更诡异的是,用perf stat看,AMX版本的L1缓存缺失率比AVX-512版本高了3倍,而指令数反而少了40%。
后来扒开MLIR生成的LLVM IR,才发现问题出在tile布局上。AMX的tile寄存器是2D结构,但MLIR默认的tile分配策略把矩阵切成了不连续的内存块,导致每次加载tile都要做大量地址计算和内存搬运。这个坑,踩得我至今记忆犹新。
AMX到底是什么
Intel Advanced Matrix Extensions(AMX)是Sapphire Rapids引入的矩阵计算加速单元。它不像AVX-512那样是向量指令的扩展,而是一个独立的协处理器,专门处理8位和16位整数的矩阵乘加运算。
AMX的核心是8个tile寄存器,每个tile可以配置为最大16行×64字节。对于int8数据,这意味着每个tile可以容纳16×64=1024个元素;对于bf16,则是16×32=512个元素。这些tile寄存器