通过 MAC 位权维度的变换探索张量处理引擎的性能提升
硬件体系结构
2025-03-11 v1 机器学习
摘要
通用矩阵乘法（GEMM）是 AI 计算的基石，使得张量处理引擎（TPE）在 GPU 和领域特定架构中日益关键。现有架构主要优化数据流或操作数复用策略。然而，考虑矩阵乘法与乘加器（MAC）之间的交互作用可提供更大的优化潜力。本研究引入了一种关于矩阵乘法的新颖硬件视角，重点关注 MAC 的位权维度。我们以矩阵三重循环为例，提出了一种更细粒度的 TPE 表示法，引入了设计和优化 PE 微架构的新方法。基于该表示法及其变换，我们提出了四种优化技术，以改善时序、面积和功耗。我们使用 SMIC-28nm 工艺在 RTL 中实现了我们的设计，并在四种经典 TPE 架构中评估了其有效性：脉动阵列、3D-Cube、乘法加法树和 2D-Matrix。我们的技术分别实现了 1.27 倍、1.28 倍、1.56 倍和 1.44 倍的面积效率提升，以及 1.04 倍、1.56 倍、1.49 倍和 1.20 倍的能效提升。应用于位切片架构时，与 Laconic 相比，我们的方法实现了 12.10 倍的能效提升和 2.85 倍的面积效率提升。我们的 Verilog HDL 代码以及时序、面积和功耗报告可在 https://github.com/wqzustc/High-Performance-Tensor-Processing-Engines 获取
引用
@article{arxiv.2503.06342,
title = {Exploring the Performance Improvement of Tensor Processing Engines through Transformation in the Bit-weight Dimension of MACs},
author = {Qizhe Wu and Huawen Liang and Yuchen Gui and Zhichen Zeng and Zerong He and Linfeng Tao and Xiaotian Wang and Letian Zhao and Zhaoxi Zeng and Wei Yuan and Wei Wu and Xi Jin},
journal= {arXiv preprint arXiv:2503.06342},
year = {2025}
}