本文主要介绍了在不同硬件平台上计算图像平均颜色的方法,包括 Intel 的 AMX 指令、Apple 的 M2 芯片的 AMX 指令以及新的 SME 指令集。
Intel 的 AMX 指令(之前的尝试):
- 作者之前使用 Intel 的 AMX 指令来计算图像平均颜色,但由于一般用户无法访问支持 AMX 的硬件,且其入门成本较高,所以这只是一个概念验证。
Apple 的 M2 芯片的 AMX 指令:
- M2 芯片也有用于 AI/ML 用例的“AMX”指令集,与 Intel 的 AMX 指令共享命名空间。
- 该指令集的状态由大型寄存器组成,
X
、Y
和Z
寄存器池,每个寄存器池可寻址 64 字节的“行”。Z
寄存器池有 64 行,单独的Z
寄存器内容为 4KiB。 - AMX 硬件可被视为与常规内核共享的“协处理器”,每个 E-Core 和 P-Core 集群有一个“AMX 核心”,这限制了并行度。
- 具体的 AMX 指令会编码如何解释行字节,实际的 AMX 硬件设计导致其是 L2 缓存的另一个“客户端”。
计算图像平均颜色的方法:
- 通用实现:将每个 RGBA8 像素解包,将其字节值添加到一个 64 位的总和中,以避免溢出问题。在
clang 18.1.0
和-O2
优化下,每 4 兆像素约需 1 毫秒,性能较好。 UDOT
/UADDW{,2}
实现:利用UDOT
和UADDW{,2}
指令,通过 SIMD 技术将字节车道累积为更大的 16/32 位总和,再将这些总和累加为更大的总和,以防止整数溢出。该实现比通用代码快约三倍,在 6 兆像素左右时可能会因内存限制或调度问题而变得不稳定。vecint
(AMX)实现:这是本文的重点。vecint
指令可对多个元素进行 8 位加法,并将结果累加到更大的数据类型中,且无需有损溢出。通过配置vecint
指令的各种位字段,可实现对输入数据的高效处理。M2 芯片支持一些额外的批量处理配置,可一次处理更多数据。在主循环中,只需三条指令即可完成加载数据和执行vecint
操作,最后将Z
中的总和写入内存并解包添加到更大的 64 位总和中,可使用成对的加宽加法指令来加速此过程。在 M2 Mac Mini 上测试,vecint
实现比UDOT
/UADDW
实现快约 3 兆像素/毫秒,但未达到理论上的 x128 倍加速,可能是多种原因导致,如加速器设计初衷、内存限制或线程调度等。
新的 SME 指令集:
- 2024 年的新 M4 芯片开始实现由 ARM 自己设计的更标准化的 SME 指令集,目前只有最新的 iPad Pro 采用了 M4 芯片。已有一些人开始分析其特性,这可能意味着 Apple 将移除或重新利用 AMX 芯片空间,因为现在有了开发者可以放心用于软件的标准指令集。
总之,本文展示了在不同硬件平台上计算图像平均颜色的多种方法和相关指令集的特点,以及它们在性能和设计上的差异。
**粗体** _斜体_ [链接](http://example.com) `代码` - 列表 > 引用
。你还可以使用@
来通知其他用户。