面向深度学习编译器TVM的算子融合优化.pdfVIP

://

httwww．skx．com

pjj

:/

DOI１０．１１８９６skx．２４０１０００１８

jj

面向深度学习编译器TVM的算子融合优化

,,,

高伟１王磊１２李嘉楠１２李帅龙１２韩林１

()

国家超级计算郑州中心郑州大学郑州

１４５０００１

郑州大学计算机与人工智能学院郑州

２４５０００１

()

onwu２２＠１２６．com

yg

,,

摘要算子融合是深度学习编译器中的一种编译优化技术能够将多个算子合并为一个大的算子有效降低计算和访存的成

.,,.

本深度学习编译器TVM的算子融合方案中将算子按照功能特性进行分类并设计融合规则最后采用贪心算法进行融合

:,,,

这种融合方案存在以下问题首先功能特性的算子分类方式下的融合规则不够通用会错失算子融合机会无法实现更大粒度

;,.,,

的融合其次贪心的融合算法也无法实现算子融合的最优解针对上述问题对TVM进行改进提出按照算子输入输出映射

,;

类型的算子分类方式并设计通用的算子融合规则以扩大算子融合的粒度提出基于动态规划的融合方案搜索算法和算子融合

,,.,

代价评估模型并对搜索空间进行剪枝使得算法能够在合理时间内搜索得到优化的融合方案为评测融合方案的有效性在

以及等平台上对,,,等深度学习模型的融合比和推理时延进行测试,

CPUDCUVGGＧ１６EfficientＧB０MobileNetＧV１YOLOＧV４

,,,.

实验结果表明相较于TVM原有融合方案所提方案融合比平均提升了２７％推理时延平均获得了１．７５的加速比

:;;;;

关键词深度学习编译器TVM算子融合融合规则动态规划

中图分类号TP３１４

OeratorFusionOtimizatio