arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2207.04296cs.LGcs.AIcs.PL
收藏

TensorIR:一种用于自动张量化程序优化的抽象

TensorIR: An Abstraction for Automatic Tensorized Program Optimization

Shanghai Jiao Tong University(上海交通大学);Carnegie Mellon University(卡内基梅隆大学);OctoML(OctoML公司);University of Washington(华盛顿大学);UC Berkeley(加州大学伯克利分校);Amazon Web Services(亚马逊云科技)

Siyuan Feng, Bohan Hou, Hongyi Jin, Wuwei Lin, Junru Shao, Ruihang Lai, Zihao Ye, Lianmin Zheng, Cody Hao Yu, Yong Yu, Tianqi Chen

摘要:

在各种设备上部署深度学习模型已经成为一个重要的课题。硬件专业化的浪潮为多维张量计算带来了各种各样的加速原语。这些新的加速原语,以及新兴的机器学习模型,带来了巨大的工程挑战。在本文中,我们提出了TensorIR,一个编译器抽象优化程序与这些张量计算原语。TensorIR概括了现有机器学习编译器中使用的循环嵌套表示,将张量计算作为一等公民。最后,我们在我们的抽象之上构建了一个端到端框架,以自动优化给定张量计算原语的深度学习模型。实验结果表明,TensorIR编译自动使用给定硬件后端的张量计算原语,并提供与跨平台的最先进手动优化系统竞争的性能。

英文摘要:

Deploying deep learning models on various devices has become an important topic. The wave of hardware specialization brings a diverse set of acceleration primitives for multi-dimensional tensor computations. These new acceleration primitives, along with the emerging machine learning models, bring tremendous engineering challenges. In this paper, we present TensorIR, a compiler abstraction for optimizing programs with these tensor computation primitives. TensorIR generalizes the loop nest representation used in existing machine learning compilers to bring tensor computation as the first-class citizen. Finally, we build an end-to-end framework on top of our abstraction to automatically optimize deep learning models for given tensor computation primitives. Experimental results show that TensorIR compilation automatically uses the tensor computation primitives for given hardware backends and delivers performance that is competitive to state-of-art hand-optimized systems across platforms.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑