基于自适应静态数据布局策略的深度学习张量程序自动生成框架  

A deep learning tensor program automatic generation framework based on adaptive layout of static data

作　　者：樊哲 南子渊 郝一帆 杜子东 陈云霁[1,2] FAN Zhe;NAN Ziyuan;HAO Yifan;DU Zidong;CHEN Yunji(State Key Laboratory of Computer Architecture,Institute of Computing Technology,Chinese Academy of Sciences,Beijing 100190;University of Chinese Academy of Sciences,Beijing 100049)

机构地区：[1]中国科学院计算技术研究所计算机体系结构国家重点实验室,北京100190 [2]中国科学院大学,北京100049

出　　处：《高技术通讯》2023年第11期1160-1171,共12页Chinese High Technology Letters

基　　金：国家重点研发计划(2020AAA0103802);国家自然科学基金(61925208,61732020,U19B2019);中国科学院战略性先导科技专项(XDB32050200);中国科学院稳定支持基础研究领域青年团队计划(YSBR-029);中国科学院青年创新促进会和科学探索奖资助项目;北京智源人工智能研究院以及北京市科技新星计划(Z191100001119093)。

摘　　要：如何确定静态数据布局是深度学习张量程序自动生成框架面临的重大挑战。Ansor作为目前应用最广泛、最具前景的此类框架,其根据预先指定的单一静态数据布局策略,训练性能预测模型,依据该模型搜索最佳性能的张量程序。但其存在单一策略非最优和性能预测模型不准确的问题。为此,本文提出基于自适应静态数据布局(AL)策略的深度学习张量程序自动生成框架AL-Ansor。AL-Ansor在搜索过程中自适应地选取多种静态数据布局策略,共同训练性能预测模型,从而搜索得到性能更高的张量程序。本文以32核Intel Xeon CPU为目标硬件平台,在多个卷积层上进行实验,结果表明,在同样的搜索次数下,相较于基于3种指定静态数据布局策略的Ansor,AL-Ansor生成的张量程序分别有13.81%、12.41%和16.59%的平均性能提升。How to determine the layout of static/const data is a big challenge faced by tensor program automatic generation frameworks.Ansor,the most broadly-used and promising framework among them,solves this issue by training a performance cost model according to a layout strategy specified in advance,then searching the tensor program with the optimal performance based on the cost model.However,there are two problems:a single strategy cannot be suitable for all tasks,and the performance cost model is not accurate.In order to solve these problems,AL-Ansor,a tensor program automatic generation framework based on the adaptive layout(AL)strategy of static data,is proposed.It adaptively chooses multiple layout strategies during the search process,and trains the performance cost model according to them.In this way,AL-Ansor can find a tensor program with higher performance.Taking convolutional layers as workloads,this work evaluates Ansor and AL-Ansor in a target server with a 32-core Intel Xeon CPU.The experimental results show that AL-Ansor improves the execution performance by 13.81%,12.41%,and 16.59%,respectively,on average,compared against Ansor with three specified layout strategies.

关 键 词：深度学习 张量程序自动生成框架 静态数据布局策略 自适应策略 性能预测模型 

分 类 号：TP18[自动化与计算机技术—控制理论与控制工程]