一句话:提出端到端神经算子合成框架 Syno,自动发现精度或速度更优的新算子。
神经网络对精度和性能的追求永无止境,但现有 NAS 与张量编译器多局限于组合或优化人工设计算子,难以自动发现全新设计。本文研究神经算子合成,提出端到端框架 Syno。Syno 在张量维度上定义细粒度原语,保证训练所需性质,并通过表达式规范化避免搜索冗余;采用引导式合成流程生成与指定输入/输出尺寸匹配的合法算子,并用高效随机树搜索探索设计空间。实验表明,Syno 发现的算子在不同硬件和编译器选择下平均加速 1.37× 至 2.06×,即使在 NAS 优化模型上精度损失也低于 1%。
The desires for better prediction accuracy and higher execution performance in neural networks never end. Neural architecture search (NAS) and tensor compilers are two popular techniques to optimize these two goals, but they are both limited to composing or optimizing existing manually designed operators rather than coming up with completely new designs. In this work, we explore the less studied direction of neural operator synthesis, which aims to automatically and efficiently discover novel neural operators with better accuracy and/or speed. We develop an end-to-end framework Syno, to realize practical neural operator synthesis. Syno makes use of a novel set of fine-grained primitives defined on tensor dimensions, which ensure various desired properties to ease model training, and also enable expression canonicalization techniques to avoid redundant candidates during search. Syno further adopts a novel guided synthesis flow to obtain valid operators matched with the specified input/output dimension sizes, and leverages efficient stochastic tree search algorithms to quickly explore the design space. We demonstrate that Syno discovers better operators with average speedups of 1.37× to 2.06× on various hardware and compiler choices, while keeping less than 1% accuracy loss even on NAS-optimized models.