TensorIR：面向自动张量化程序优化的抽象
机器学习
2022-10-31 v2 人工智能
编程语言
摘要
将深度学习模型部署到各类设备已成为重要课题。硬件专用化的浪潮带来了用于多维张量计算的多样化加速原语。这些新的加速原语，连同新兴的机器学习模型，带来了巨大的工程挑战。本文提出TensorIR，一种用于优化带有这些张量计算原语的程序的编译器抽象。TensorIR将现有机器学习编译器中使用的循环嵌套表示泛化，使张量计算成为一等公民。最后，我们在此抽象之上构建了端到端框架，以针对给定张量计算原语自动优化深度学习模型。实验结果表明，TensorIR编译能自动利用给定硬件后端的张量计算原语，并交付跨平台可与最先进手工优化系统竞争的性能。
引用
@article{arxiv.2207.04296,
title = {TensorIR: An Abstraction for Automatic Tensorized Program Optimization},
author = {Siyuan Feng and Bohan Hou and Hongyi Jin and Wuwei Lin and Junru Shao and Ruihang Lai and Zihao Ye and Lianmin Zheng and Cody Hao Yu and Yong Yu and Tianqi Chen},
journal= {arXiv preprint arXiv:2207.04296},
year = {2022}
}
备注
Accepted to ASPLOS 2023