amx

苹果的当代 M1/M2/M3/M4 机器至少有四种不同的方式供低级程序员进行大量计算: 1. CPU 核心上的标准 ARMv8 SIMD/NEON 向量指令 (宽度为 128 位,在 Firestorm 上每个周期可发出 最多四条指令) 2. CPU 发出的未记录的 AMX 指令,在专用加速器执行单元上执行 3. 神经引擎 (称为 ANE 或 NPU) 4. GPU (例如 Metal Compute Shaders)

这个仓库主要介绍第二种方法:苹果的 AMX 指令。请注意,这些指令既没有被苹果记录下来,也没有得到苹果的支持。作为一个可能会引起极大混乱的来源,苹果的 AMX 指令与 英特尔的 AMX 指令 完全不同,尽管它们都用于从 CPU 发出矩阵乘法操作。

这项研究是在 Apple M1 Max (2021) 上进行的,随后在 M2 (2023) 上进行了跟进研究,并在 M3 (2023) 和 M4 Max (2024) 上进行了进一步的跟进研究。 更旧或更新的芯片可能具有不同的 AMX 指令。一些来源 表示 M1 中包含了 AMX 指令的第 2 版,这似乎是合理的(可能使用 7 位写入掩码的所有内容来自第 1 版,而使用 9 位写入掩码的所有内容则是第 2 版中的新内容)。 M1 到 M2 的过渡 增加了 bf16 支持,以及一些其他调整。 M2 到 M3 的过渡 为 ldx、ldy 和 matint 的每个指令添加了一个额外模式。 M3 到 M4 的过渡 使 extrh、extrv、vecfp 和 vecint 的某些模式忽略 X/Y 偏移的低位。

AMX 的一个很好的图像总结如下图所示 已废弃专利 US20180074824A1。假设有一个 32x32 的计算单元网格,其中每个单元可以执行 16 位乘积累加操作,或者一个 2x2 的子网格,其中每个单元可以执行单位乘积累加操作。