[论文解读] A matrix math facility for Power ISA(TM) processors

José E. Moreira, Kit Barton+23 more|arXiv (Cornell University)|Apr 7, 2021

Parallel Computing and Optimization Techniques参考文献 12被引用 6

一句话总结

本文介紹了 Power ISA™ 版本 3.1 中的 Matrix-Multiply Assist (MMA) 功能，這是為 IBM POWER10 處理器引入的一套新指令，用於加速矩陣運算。透過使用 512 位累加器暫存器以 32 位元、64 位元和 bfloat16 精度對小矩陣執行 rank-k 更新，MMA 功能在矩陣乘法上的每核心性能相比 POWER9 提升了 4 倍，功耗僅增加 8%，實現顯著的能源效率提升。

ABSTRACT

Power ISA(TM) Version 3.1 has introduced a new family of matrix math instructions, collectively known as the Matrix-Multiply Assist (MMA) facility. The instructions in this facility implement numerical linear algebra operations on small matrices and are meant to accelerate computation-intensive kernels, such as matrix multiplication, convolution and discrete Fourier transform. These instructions have led to a power- and area-efficient implementation of a high throughput math engine in the future POWER10 processor. Performance per core is 4 times better, at constant frequency, than the previous generation POWER9 processor. We also advocate the use of compiler built-ins as the preferred way of leveraging these instructions, which we illustrate through case studies covering matrix multiplication and convolution.

研究动机与目标

• 為解決 POWER10 上 BLAS3 工作負載的性能差距，特別是針對需要高吞吐量密集線性代數的商業分析與機器學習工作負載。
• 設計硬體與軟體協同優化的解決方案，加速矩陣乘法、卷積及其他核心運算，而不擴展架構暫存器檔案或需要新的編譯工具鏈。
• 透過編譯器內建函數實現對新 MMA 指令的高效使用，確保向後相容性與可維護性，同時達成高效率。
• 證明 MMA 功能在性能與能源效率方面相較於 POWER9 和先前僅支援向量的指令集有顯著提升。

提出的方法

• MMA 功能引入 8 個新的 512 位累加器暫存器，用於儲存 rank-k 矩陣更新的中間結果，與用於輸入的 64×128 位向量-標量暫存器分離。
• Matrix Math Engine (MME) 功能單元可與其他 Power ISA 指令並行執行 MMA 指令，實現與一般運算的完全重疊。
• MMA 指令作用於小矩陣（例如 32 位浮點數的 4×4 矩陣），使用向量暫存器作為輸入，累加器作為輸出，從而高效計算 GEMM 和卷積核心運算。
• 設計採用電源門控的 MME 單元，並初始將累加器設為非架構狀態，允許未來以向後相容的方式提升為完整架構狀態。
• 在 GCC 和 LLVM 中引入編譯器內建函數，支援高階程式碼生成並提供細粒度控制，避免手寫組語彙碼的同時維持性能。
• 透過預矽核心模型的模擬驗證實作，量測 DGEMM 和 HPL 工作負載下的性能與功耗。

实验结果

研究问题

• RQ1 在功耗受限且沿用舊有架構指令集的環境下，新的指令集擴展如何提升密集線性代數核心的性能？
• RQ2 專用矩陣數學引擎在不擴展暫存器檔案或不需新編譯技術的情況下，能在多大程度上提升機器學習與分析工作負載的性能？
• RQ3 在實際處理器核心中，使用 MMA 指令相比傳統僅支援向量的方案，在性能與功耗效率方面有何提升？
• RQ4 編譯器內建函數是否能有效暴露 MMA 功能的潛力，同時維持程式碼的可維護性與性能？
• RQ5 MMA 功能在生產工作負載中，於不同精度格式（32 位元、64 位元、bfloat16）下的可擴展性如何？

主要发现

• POWER10 處理器在矩陣乘法核心運算中，每核心性能相比 POWER9 提升 4 倍，且時鐘頻率不變。
• 與 POWER10 上僅使用 VSX 的程式碼相比，MMA 功能實現 2.5 倍性能提升，功耗僅增加 8%，顯示出強勁的能源效率。
• 與 POWER9 相比，MMA 加速的 POWER10 核心在功耗降低 24% 的情況下實現 5 倍性能提升，使每次運算的能耗降低近 7 倍。
• MMA 指令已整合至 OpenBLAS 和 Eigen，支援雙精度、單精度與 bfloat16 精度，可立即用於高性能線性代數函式庫。
• 使用編譯器內建函數可實現高效程式碼生成，且開銷極小，提供可維護性高於手寫組語彙碼的替代方案。
• MMA 功能可高效實現卷積與離散傅立葉變換核心運算，並具備在稀疏運算與訊號處理工作負載中廣泛應用的潛力。