Intel® AMX（Advanced Matrix Extensions）是x86指令集的新扩展，专为矩阵运算设计，能大幅加速AI工作负载中的矩阵乘法。这里简单介绍一下AMX指令的特性,详情参加官网：
- 基于Tile的架构
- AMX 引入了tile寄存器（TMM0–TMM7），每个 tile 是一个二维矩阵寄存器（最多 8 个）。
- 每个 tile 的大小（行数、列数）是可配置的，通过 LDTILECFG 指令进行配置。
- 支持的数据类型包括：
bf16、int8、int16、fp16（部分取决于具体的 AMX 扩展版本，如 AMX-BF16、AMX-INT8 等）。
- Tile 配置灵活
- 可动态配置 tile 的维度与行为，每个 tile 最多可配置为 16 行 × 64 字节，共 1024 字节（例如 int8 类型可容纳 16×64 = 1024 个元素）。
- Tile 配置通过 LDTILECFG 加载，可通过 STTILECFG 存储。
用户可以通过本文提供的代码学习使用AMX指令，项目地址：https://github.com/goog00/intel_amx_example，也可以通过官方示例学习。
本文通过不断提高计算访存比（计算强度），一步步优化矩阵乘法实现，最大化AMX硬件的性能来提高处理矩阵乘的效率。
在第1版中(参见matrix_mul_amx_with_policy_v1.cpp)，我们实现了一个基本的矩阵乘法，利用AMX的 _tile_loadd、_tile_dpbssd 和 _tile_stored 指令完成计算。
void MatrixMultiply(Matrix<InputType> &A, Matrix<InputType> &B, Matrix<OutputType> &C) {
_tile_loadd(2, A.Data(), A.Stride());
_tile_loadd(3, B.Data(), B.Stride());
_tile_loadd(1, C.Data(), C.Stride());
_tile_dpbssd(1, 2, 3);
_tile_stored(1, C.Data(), C.Stride());
}硬件配置： Intel(R) Xeon(R) Platinum 8458P, 4核4G
特点与瓶颈：
- 简单直接：加载矩阵A、B到tile寄存器，执行一次乘加运算（
_tile_dpbssd），结果存回C。 - 硬件利用不足：AMX支持8个tile寄存器（0-7），但此实现仅使用了3个（tile 1、2、3），大量计算资源被浪费。
- 访存效率低：每次只处理一对输入矩阵，访存带宽和计算并行性未被充分利用。
性能数据：
循环次数: 10000000
Intel Amx cost time:0.479877s, GOPS: 682.8410GOPS
这里GOPS表示每秒多少G的int8算数， 对应float的GFLOPS， 本文中的测试数据类型为int8类型。
优化方向：如何利用更多tile寄存器并增加输入数据的处理能力？
在第2版中(参见matrix_mul_amx_with_policy_v2.cpp)，我们引入了分块矩阵乘法，将输入矩阵分为2×2的子块（A0、A1、B0、B1），输出4个结果子矩阵（C00、C01、C10、C11），充分利用AMX的8个tile寄存器。
void MatrixMultiply(Matrix<InputType> &A0, Matrix<InputType> &A1, Matrix<InputType> &B0, Matrix<InputType> &B1,
Matrix<OutputType> &C00, Matrix<OutputType> &C01, Matrix<OutputType> &C10, Matrix<OutputType> &C11) {
_tile_loadd(0, A0.Data(), A0.Stride());
_tile_loadd(1, B0.Data(), B0.Stride());
_tile_loadd(2, A1.Data(), A1.Stride());
_tile_loadd(3, B1.Data(), B1.Stride());
_tile_loadd(4, C00.Data(), C00.Stride());
_tile_loadd(5, C01.Data(), C01.Stride());
_tile_loadd(6, C10.Data(), C10.Stride());
_tile_loadd(7, C11.Data(), C11.Stride());
_tile_dpbssd(4, 0, 1); // C00 += A0 * B0
_tile_stored(4, C00.Data(), C00.Stride());
_tile_dpbssd(5, 0, 3); // C01 += A0 * B1
_tile_stored(5, C01.Data(), C01.Stride());
_tile_dpbssd(6, 2, 1); // C10 += A1 * B0
_tile_stored(6, C10.Data(), C10.Stride());
_tile_dpbssd(7, 2, 3); // C11 += A1 * B1
_tile_stored(7, C11.Data(), C11.Stride());
}硬件配置： Intel(R) Xeon(R) Platinum 8458P, 4核4G
优化亮点：
- tile寄存器分块：A0/1, B0/1四个tile寄存器只读取了一次，且两次参与到了计算中去， 其计算强度得到了提升，
瓶颈：
- C的计算强度低：每次计算都需要对C进行读取，计算强度太低
性能数据：
循环次数: 10000000
Intel Amx cost time:1.03217s, GOPS: 1269.8740GOPS
优化方向：C的计算强度是否也可以得到提升？
第3版(参见matrix_mul_amx_with_policy_v3.cpp)， 根据矩阵乘法的特性 C[MxN] = A[MxK] * B[KxN]，我们可以知道， 增加K的长度并不会改变C的尺寸， 所以我们在A/B的tile（小矩阵）上增加了K纬度的长度。
void MatrixMultiply(std::vector<Matrix<InputType>> &VA, std::vector<Matrix<InputType>> &VB, Matrix<OutputType> &C) {
_tile_loadd(0, C.Data(), C.Stride());
for (int i = 0; i < VA.size(); i++) {
_tile_loadd(1, VA[i].Data(), VA[i].Stride());
_tile_loadd(2, VB[i].Data(), VB[i].Stride());
_tile_dpbssd(0, 1, 2); // C += VA[i] * VB[i]
}
_tile_stored(0, C.Data(), C.Stride());
}硬件配置： Intel(R) Xeon(R) Platinum 8458P, 4核4G
优化亮点：
- 极大的提升了C的计算强度：循环开始前读取一次C，每次在k纬度的迭代加载一对输入矩阵（VA[i]、VB[i]），结果累加到tile 0， 循环结束后再取回C， 当K取一个合适的值时，C的计算强度我会得到极大的提升
瓶颈：
- 该方案没有考虑到A和B的计算强度
性能数据：
循环次数: 10000000
Intel Amx cost time:3.94255s, GOPS: 1329.8203GOPS
优化方向：结合第2版的分块思想和第3版的累加能力，进一步提升性能。
第4版(参见matrix_mul_amx_with_policy_v4.cpp)，该版本我们融合了前两版的策略，使A，B，C的计算强度都得到了提升
void MatrixMultiply(std::vector<Matrix<InputType>> &VA0, std::vector<Matrix<InputType>> &VA1, std::vector<Matrix<InputType>> &VB0, std::vector<Matrix<InputType>> &VB1,
Matrix<OutputType> &C00, Matrix<OutputType> &C01, Matrix<OutputType> &C10, Matrix<OutputType> &C11) {
_tile_loadd(4, C00.Data(), C00.Stride());
_tile_loadd(5, C01.Data(), C01.Stride());
_tile_loadd(6, C10.Data(), C10.Stride());
_tile_loadd(7, C11.Data(), C11.Stride());
for (size_t k = 0; k < VA0.size(); ++k) {
_tile_loadd(0, VA0[k].Data(), VA0[k].Stride()); // A00(:,k)
_tile_loadd(1, VB0[k].Data(), VB0[k].Stride()); // B00(k,:)
_tile_loadd(2, VA1[k].Data(), VA1[k].Stride()); // A10(:,k)
_tile_loadd(3, VB1[k].Data(), VB1[k].Stride()); // B01(k,:)
_tile_dpbssd(4, 0, 1); // C00 += A00(:,k) * B00(k,:)
_tile_dpbssd(5, 0, 3); // C01 += A00(:,k) * B01(k,:)
_tile_dpbssd(6, 2, 1); // C10 += A10(:,k) * B00(k,:)
_tile_dpbssd(7, 2, 3); // C11 += A10(:,k) * B01(k,:)
}
// 最后一次性存回
_tile_stored(4, C00.Data(), C00.Stride());
_tile_stored(5, C01.Data(), C01.Stride());
_tile_stored(6, C10.Data(), C10.Stride());
_tile_stored(7, C11.Data(), C11.Stride());
}
硬件配置： Intel(R) Xeon(R) Platinum 8458P, 4核4G
优化亮点：
- 提升了A, B的计算强度：使用满了8个tile寄存器，最大化提升A， B的计算强度
- 提升了C的计算强度：增加了k纬度的长度， 极大提升了C的计算强度
性能数据：
循环次数: 10000000
Intel Amx cost time:9.76914s, GOPS: 2146.7103GOPS
性能提升： 相比第1版，第4版通过分块、累加和高效访存，性能从第1版的682 GOPS到2146 GOPS，翻了3倍多。
在第5版中（参考代码：matrix_mul_amx_with_policy_v5.cpp）， 我们在第4版的分块矩阵乘法基础上，通过异步线程并行执行矩阵运算任务。每个线程独立处理部分计算工作，把所有的intel amx运算单元利用起来
硬件配置： Intel(R) Xeon(R) Platinum 8458P, 4核4G
性能数据：
线程 1 - 循环次数: 5000000
执行时间: 5.6920 秒, 性能: 1842.2052 GOPS
线程 2 - 循环次数: 5000000
执行时间: 5.6874 秒, 性能: 1843.6698 GOPS
总执行时间: 5.6921 秒
时间差 (线程1 + 线程2 - 总时间): 5.6873 秒
总性能: 3684.3322 GOPS
在64核平台上，我们可以获得40T的int8算力
执行时间: 0.5042 秒, 性能: 649.8689 GOPS
线程 59 - 循环次数: 156250
执行时间: 0.4846 秒, 性能: 676.2303 GOPS
线程 60 - 循环次数: 156250
执行时间: 0.4982 秒, 性能: 657.7000 GOPS
线程 61 - 循环次数: 156250
执行时间: 0.5058 秒, 性能: 647.8485 GOPS
线程 62 - 循环次数: 156250
执行时间: 0.5047 秒, 性能: 649.2992 GOPS
线程 63 - 循环次数: 156250
执行时间: 0.5011 秒, 性能: 653.8643 GOPS
总执行时间: 0.5204 秒
总性能: 40302.2399 GOPS
从第1版到第5版的优化历程，展示了如何围绕 Intel AMX 指令的特性逐步提升矩阵乘法性能的核心策略：
- 硬件资源最大化利用：从第1版仅使用 3 个 tile 寄存器，到第2版和第4版充分利用 8 个 tile 寄存器，大幅提升了 AMX 的硬件利用率。
- 计算强度提升：
- 第2版通过分块复用 A 和 B 的 tile，增加了单次计算量。
- 第3版引入 K 维度的循环累加，显著提高了 C 的计算强度。
- 第4版融合两者，优化了 A、B、C 的整体计算强度。
- 多线程并行：第5版引入多线程，将计算任务分配到多核（最高 64 核），在 64 核配置下实现 40 TOPS 的 int8 算力，展现了 AMX 在高并发场景下的潜力。
性能成果：
- 单线程从第1版的 682 GOPS 提升到第4版的 2146 GOPS，增长超 3 倍。
- 多线程（64 核）总性能达 40302 GOPS（40.3 TOPS），凸显了并行优化的价值。
通过这一系列优化，我们不仅挖掘了 AMX 的硬件潜力，也为高性能矩阵运算提供了实用参考。欢迎读者基于本文代码实验并提出更多优化思路，一起探索 AMX 的极限性能！