Hello SME!使用可扩展矩阵扩展生成快速矩阵乘法内核
Hello SME! Generating Fast Matrix Multiplication Kernels Using the Scalable Matrix Extension
AI总结:
本文深入研究Apple M4芯片的SME扩展性能,基于其高吞吐量特性设计即时编译代码生成器,生成的小矩阵乘法内核在几乎所有测试配置中均优于厂商优化的BLAS实现。
摘要:
现代中央处理器(CPU)采用单指令多数据流水线来加速计算密集型浮点和定点工作负载。传统上,这些流水线及相应的指令集架构(ISA)是为向量并行性而设计的。近年来,主要硬件供应商通过引入带有相应 ISA 扩展的矩阵单元,进一步提高了其 CPU 的吞吐量。可扩展矩阵扩展(SME)于 2021 年针对 Arm 架构发布,Apple 的 M4 芯片是首个支持 SME 的芯片。本文对 M4 上的 SME 进行了深入研究。我们的微基准测试确定了 M4 的 SME 加速的最大浮点和定点吞吐量,并研究了向矩阵寄存器传输及从中传输的可实现带宽。此外,我们利用所获得的见解,设计了一个基于 SME 的小矩阵乘法即时编译代码生成器。所展示的结果表明,M4 的 SME 支持以 FP32 为中心,可实现超过 2.3 FP32 TFLOPS 的吞吐量。为了最大化读写带宽,向矩阵寄存器加载和存储必须分两步完成。我们即时编译生成的小矩阵乘法内核在几乎所有测试配置中都优于供应商优化的 BLAS 实现。
英文摘要:
Modern central processing units (CPUs) feature single-instruction, multiple-data pipelines to accelerate compute-intensive floating-point and fixed-point workloads. Traditionally, these pipelines and corresponding instruction set architectures (ISAs) were designed for vector parallelism. In recent years, major hardware vendors have further increased the throughput of their CPUs by introducing matrix units with corresponding ISA extensions. The Scalable Matrix Extension (SME) has been announced for the Arm architecture in 2021 and Apple's M4 chip is the first to support SME. This paper presents an in-depth study of SME on M4. Our microbenchmarks determine the maximum floating-point and fixed-point throughput of M4's SME acceleration and study the achievable bandwidth for transfers to and from the matrix registers. Furthermore, we used the insights gained to design a just-in-time code generator for SME-based small matrix multiplications. The results presented show that M4's SME support is FP32-centric, with an achievable throughput of over 2.3 FP32 TFLOPS. To maximize read and write bandwidth, loading and storing to and from the matrix registers must be done in two steps. Our just-in-time generated small matrix multiplication kernels outperform the vendor-optimized BLAS implementation in almost all tested configurations.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。