面向TPU粗粒度指令的自动张量化方法  

Automatic Tensorization for TPU Coarse-grained Instructions

作　　者：刘磊[1] 周志德 刘兴祥 车皓阳 姚雷 江贺[1] LIU Lei;ZHOU Zhide;LIU Xingxiang;CHE Haoyang;YAO Lei;JIANG He(School of Software Engineering,Dalian University of Technology,Dalian,Liaoning 116620,China;Sangfor Technologies Inc.,Shenzhen,Guangdong 518000,China;Zhejiang Zeekr Intelligent Technology Co.,Ltd.,Ningbo,Zhejiang 315800,China)

机构地区：[1]大连理工大学软件学院,辽宁大连116620 [2]深信服科技股份有限公司,广东深圳518000 [3]浙江极氪智能科技有限公司,浙江宁波315800

出　　处：《计算机科学》2024年第6期52-60,共9页Computer Science

基　　金：国家自然科学基金重点项目(62032004);CCF-深信服伏羲基金项目(2022003);中国博士后科学基金(2023M730472);国家自然科学基金(62302077)。

摘　　要：张量化是通过调用硬件特定指令对张量运算进行加速的过程。TPU支持多种粗粒度指令,可表示神经网络级别的算子,且没有明确的运算规模限制。现有张量化方法对于粗粒度指令需要手写大量的IR匹配片段,且难以实现灵活的双缓存(ping-pong buffer)形式的指令并行优化,不利于扩展至TPU场景。为此,提出了一种面向TPU粗粒度指令的自动张量化方法——Tir2TPU。首先,基于TensorIR抽象语法树的分析对运算程序进行指令替换。其次,设计了一种模拟硬件行为的并行模型以实现指令并行优化。最后,构建了基于TPU硬件特征的程序调度空间以实现快速自动调优。实验对矩阵乘法等5种机器学习模型中常用的算子进行了性能评估。实验结果表明,Tir2TPU自动优化生成的算子与TPU自有编译器相比可取得最高3.1倍、平均1.78倍的运算加速,并且可取得平均90%的手工优化性能。syntax tree.Secondly,it also utilizes a parallel model that simulates hardware behavior to generate parallel instruction flow.Finally,Tir2TPU combines a hardware-centric schedule space based on TPU features,which greatly accelerates auto-tuning process.The performance of Tir2TPU is evaluatedon 5 commonly used operators in machine learning models.Experimental results show that Tir2TPU can achieve up to 3×and an average of 1.78×speedup compared to TPU’s compiler,and consistently delivers 90%performance compared to manually optimized operators.

关 键 词：机器学习编译器 张量加速器 张量化 指令并行 算子优化 

分 类 号：TP311[自动化与计算机技术—计算机软件与理论]