中文

LoopStack:一个轻量级张量代数编译器栈

机器学习 2022-05-03 v1 性能 符号计算

摘要

我们提出 LoopStack,一个面向张量操作的领域专用编译器栈,由前端 LoopTool 和高效的优化代码生成器 LoopNest 组成。该栈使我们能够编译完整的神经网络并生成面向 AVX2、AVX512、NEON 和 NEONfp16 指令集的代码,同时融入了其他机器学习编译器后端常缺失的优化。我们在一系列完整神经网络、常用网络模块以及单个算子上对栈进行了评估,结果表明 LoopStack 生成的机器码在两类情况下均匹配且经常超越最先进机器学习框架的性能。我们还表明,对于大量调度,LoopNest 的编译比 LLVM 快数个数量级,同时达到持平或更好的运行时间性能。此外,LoopStack 具有极小的内存占用——二进制大小 245KB,有效代码不足 30K 行,使其非常适用于移动和嵌入式设备。

关键词

引用

@article{arxiv.2205.00618,
  title  = {LoopStack: a Lightweight Tensor Algebra Compiler Stack},
  author = {Bram Wasti and José Pablo Cambronero and Benoit Steiner and Hugh Leather and Aleksandar Zlateski},
  journal= {arXiv preprint arXiv:2205.00618},
  year   = {2022}
}