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寄存器