优秀的张量编译器

这是一份关于张量计算和深度学习领域优秀编译器项目及论文的列表。

目录

• 开源项目
• 论文
  • 综述
  • 编译器与中间表示设计
  • 自动调优与自动调度
  • 代价模型
  • CPU与GPU优化
  • NPU优化
  • 图级优化
  • 动态模型
  • 图神经网络
  • 分布式计算
  • 量化
  • 稀疏性
  • 程序重写
  • 验证与测试
• 教程
• 贡献

开源项目

• TVM: 端到端机器学习编译框架
• MLIR: 多级中间表示
• XLA: 面向机器学习的优化编译器
• Halide: 用于图像和张量快速、可移植计算的语言
• Glow: 面向神经网络硬件加速器的编译器
• nnfusion: 灵活高效的深度神经网络编译器
• Hummingbird: 将训练好的机器学习模型编译为张量计算
• Triton: 用于分块神经网络计算的中间语言及编译器
• AITemplate: 将神经网络渲染为高性能CUDA/HIP C++代码的Python框架
• Hidet: 基于编译的深度学习框架
• Tiramisu: 用于表达快速且可移植代码的多面体编译器
• TensorComprehensions: 框架无关的高性能机器学习抽象
• PlaidML: 让深度学习在任何地方都能运行的平台
• BladeDISC: 面向机器学习工作负载的端到端动态形状编译器
• TACO: 张量代数编译器
• Nebulgym: 易用的AI训练加速库
• Speedster: 自动应用最先进优化技术，以在您的硬件上实现最大推理加速
• NN-512: 生成用于神经网络推理的C99代码的编译器
• DaCeML: 面向机器学习的数据中心编译器
• Mirage: 面向张量代数的多级超优化器

论文

综述

• 深度学习编译器：全面综述 由Mingzhen Li等人撰写，发表于TPDS 2020年
• 硬件上深度神经网络编译器的深入比较 由Yu Xing等人撰写，发表于ICESS 2019年

编译器与中间表示设计

• [(De/Re)-Composition of Data-Parallel Computations via Multi-Dimensional Homomorphisms]（通过多维同态进行数据并行计算的分解与重构）作者：Ari Rasch，TOPLAS 2024
• [BladeDISC：基于编译器方法优化动态形状机器学习工作负载] 作者：郑振等，SIGMOD 2024
• [Hidet：面向深度学习张量程序的任务映射编程范式] 作者：丁瑶瑶等，ASPLOS 2023
• [TensorIR：用于自动张量化程序优化的抽象] 作者：冯思源、侯博涵等，ASPLOS 2023
• [针对硬件加速器高效编程的外编译] 作者：Yuka Ikarashi、Gilbert Louis Bernstein等，PLDI 2022
• [DaCeML：面向机器学习的数据中心编译器] 作者：Oliver Rausch等，ICS 2022
• [FreeTensor：具有整体优化能力的自由形式 DSL，用于处理不规则张量程序] 作者：唐士志等，PLDI 2022
• [Roller：面向深度学习的快速高效张量编译] 作者：朱洪宇等，OSDI 2022
• [AStitch：在现代 SIMT 架构上为内存密集型 ML 训练与推理开启新的多维优化空间] 作者：郑振等，ASPLOS 2022
• [MLIR 中的可组合、模块化代码生成：一种结构化且可重定向的张量编译器构建方法] 作者：Nicolas Vasilache等，arXiv 2022
• [PET：利用部分等价变换与自动化修正优化张量程序] 作者：王浩杰等，OSDI 2021
• [MLIR：面向领域特定计算的规模化编译基础设施] 作者：Chris Lattner等，CGO 2021
• [用于统一机器学习预测服务的张量编译器] 作者：Supun Nakandala等，OSDI 2020
• [Rammer：借助 rTasks 实现深度学习编译器的整体优化] 作者：Ma Lingxiao等，OSDI 2020
• [有状态数据流多图：面向异构架构性能可移植性的数据中心模型] 作者：Tal Ben-Nun等，SC 2019
• [TASO：深度学习用张量代数超级优化器] 作者：贾志豪等，SOSP 2019
• [Tiramisu：用于表达快速且可移植代码的多面体编译器] 作者：Riyadh Baghdadi等，CGO 2019
• [Triton：面向分块神经网络计算的中间语言与编译器] 作者：Philippe Tillet等，MAPL 2019
• [Relay：深度学习用高级编译器] 作者：Jared Roesch等，arXiv 2019
• [TVM：面向深度学习的自动化端到端优化编译器] 作者：陈天奇等，OSDI 2018
• [Tensor Comprehensions：框架无关的高性能机器学习抽象] 作者：Nicolas Vasilache等，arXiv 2018
• [Intel nGraph：深度学习用中间表示、编译器和执行器] 作者：Scott Cyphers等，arXiv 2018
• [Glow：面向神经网络的图降级编译技术] 作者：Nadav Rotem等，arXiv 2018
• [DLVM：面向深度学习系统的现代编译基础设施] 作者：Richard Wei等，arXiv 2018
• [Diesel：用于 GPU 上线性代数及神经网络计算的 DSL] 作者：Venmugil Elango等，MAPL 2018
• [张量代数编译器] 作者：Fredrik Kjolstad等，OOPSLA 2017
• [Halide：用于优化图像处理流水线中并行性、局部性和重新计算的语言与编译器] 作者：Jonathan Ragan-Kelley等，PLDI 2013

自动调优与自动调度

• 面向张量计算的 GPU 内核加速自动调优，作者：Chendi Li、Yufan Xu 等，ICS 2024
• 利用张量语言模型辅助生成深度学习高性能张量程序，作者：Yi Zhai 等，OSDI 2024
• 用于内核调度的液滴搜索算法，作者：Michael Canesche 等，ACM TACO 2024
• 基于概率编程的张量程序优化，作者：Junru Shao 等，NeurIPS 2022
• 面向深度学习编译器的一次性调优器，作者：Jaehun Ryu 等，CC 2022
• 基于渐近成本模型的稀疏张量代数自动调度，作者：Peter Ahrens 等，PLDI 2022
• Bolt：弥合自动调优器与硬件原生性能之间的差距，作者：Jiarong Xing 等，MLSys 2022
• 面向领域优化深度学习加速器的全栈搜索技术，作者：Dan Zhang 等，ASPLOS 2022
• 面向 GPU 的成像与视觉流水线高效自动调度，作者：Luke Anderson 等，OOPSLA 2021
• Lorien：高效交付深度学习工作负载，作者：Cody Hao Yu 等，SoCC 2021
• 用于深度神经网络吞吐量优化的价值学习，作者：Benoit Steiner 等，MLSys 2021
• 多遍机器学习编译器自动调优的灵活方法，作者：Phitchaya Mangpo Phothilimthana 等，PACT 2021
• Ansor：为深度学习生成高性能张量程序，作者：Lianmin Zheng 等，OSDI 2020
• 针对 GPU 上 Halide 流水线的调度合成，作者：Sioutas Savvas 等，TACO 2020
• FlexTensor：面向异构系统张量计算的自动调度探索与优化框架，作者：Size Zheng 等，ASPLOS 2020
• ProTuner：使用蒙特卡洛树搜索调优程序，作者：Ameer Haj-Ali 等，arXiv 2020
• AdaTune：高效实现自适应张量程序编译，作者：Menghao Li 等，NeurIPS 2020
• 通过组合计算与数据的自动变换优化内存层次结构，作者：Jie Zhao 等，MICRO 2020
• Chameleon：用于加速深度神经网络编译的自适应代码优化，作者：Byung Hoon Ahn 等，ICLR 2020
• 面向稀疏张量代数的稀疏迭代空间变换框架，作者：Ryan Senanayake 等，OOPSLA 2020
• 使用树搜索和随机程序学习优化 Halide，作者：Andrew Adams 等，SIGGRAPH 2019
• 学习优化张量程序，作者：Tianqi Chen 等，NeurIPS 2018
• 自动调度 Halide 图像处理流水线，作者：Ravi Teja Mullapudi 等，SIGGRAPH 2016

成本模型

• TLP：用于张量程序调优的基于深度学习的成本模型，作者：Yi Zhai 等，ASPLOS 2023
• 面向稀疏张量程序自动调度的渐近成本模型，作者：Peter Ahrens 等，PLDI 2022
• TenSet：用于学习型张量编译器的大规模程序性能数据集，作者：Lianmin Zheng 等，NeurIPS 2021
• 用于自动代码优化的基于深度学习的成本模型，作者：Riyadh Baghdadi 等，MLSys 2021
• 面向张量处理单元的学习型性能模型，作者：Samuel J. Kaufman 等，MLSys 2021
• DYNATUNE：深度神经网络编译中的动态张量程序优化，作者：Minjia Zhang 等，ICLR 2021
• MetaTune：基于元学习的成本模型，用于快速高效的自动调优框架，作者：Jaehun Ryu 等，arXiv 2021
• 基于 LightGBM 的加速张量程序编译，作者：Gonghan Liu1 等，JPCS 2021

CPU 和 GPU 优化

• DeepCuts：适用于多种 GPU 工作负载的深度学习优化框架，作者：Wookeun Jung 等，PLDI 2021
• CNN 优化的解析表征与设计空间探索，作者：Rui Li 等，ASPLOS 2021
• UNIT：统一张量化指令编译，作者：Jian Weng 等，CGO 2021
• PolyDL：用于创建高性能 DL 原语的多面体优化，作者：Sanket Tavarageri 等，arXiv 2020
• Fireiron：面向 GPU 的数据移动感知调度语言，作者：Bastian Hagedorn 等，PACT 2020
• Volta 张量核心的自动内核生成，作者：Somashekaracharya G. Bhaskaracharya 等，arXiv 2020
• Swizzle Inventor：GPU 内核的数据移动合成，作者：Phitchaya Mangpo Phothilimthana 等，ASPLOS 2019
• 在 CPU 上优化 CNN 模型推理，作者：Yizhi Liu 等，ATC 2019
• 张量收缩的缓存解析建模与分块大小优化，作者：Rui Li 等，SC 19

NPU优化

• [Heron：面向深度学习加速器的自动约束高性能库生成]（Jun Bi 等人，ASPLOS 2023）
• [AMOS：通过硬件抽象实现空间加速器上张量计算的自动映射]（Size Zheng 等人，ISCA 2022）
• [迈向神经网络与加速器的协同设计]（Yanqi Zhou 等人，MLSys 2022）
• [AKG：基于多面体变换的神经处理单元自动内核生成]（Jie Zhao 等人，PLDI 2021）

图级优化

• [POET：在小型设备上训练神经网络，集成重计算与分页技术]（Shishir G. Patil 等人，ICML 2022）
• [Collage：深度学习后端的无缝集成与自动放置]（Byungsoo Jeon 等人，PACT 2022）
• [Apollo：通过逐层优化实现基于分区的算子融合]（Jie Zhao 等人，MLSys 2022）
• [张量图超优化中的等式饱和]（Yichen Yang 等人，MLSys 2021）
• [IOS：用于 CNN 加速的算子间调度器]（Yaoyao Ding 等人，MLSys 2021）
• [利用图替换优化 DNN 计算图]（Jingzhi Fang 等人，VLDB 2020）
• [适用于 ML 编译器的可迁移图优化器]（Yanqi Zhou 等人，NeurIPS 2020）
• [FusionStitching：提升深度学习工作负载中的内存密集型计算性能]（Zhen Zheng 等人，arXiv 2020）
• [Nimble：面向深度学习的轻量级并行 GPU 任务调度]（Woosuk Kwon 等人，Neurips 2020）

动态模型

• [Axon：一种用于深度学习图中动态形状的语言]（Alexander Collins 等人，arXiv 2022）
• [DietCode：动态张量程序的自动优化]（Bojian Zheng 等人，MLSys 2022）
• [The CoRa 张量编译器：针对稀疏张量的零填充最小化编译]（Pratik Fegade 等人，MLSys 2022）
• [Nimble：高效编译动态神经网络以支持模型推理]（Haichen Shen 等人，MLSys 2021）
• [DISC：面向机器学习工作loads的动态形状编译器]（Kai Zhu 等人，EuroMLSys 2021）
• [Cortex：递归深度学习模型的编译器]（Pratik Fegade 等人，MLSys 2021）

图神经网络

• [Graphiler：利用消息传递数据流图优化图神经网络]（Zhiqiang Xie 等人，MLSys 2022）
• [Seastar：面向图神经网络的顶点中心编程]（Yidi Wu 等人，Eurosys 2021）
• [FeatGraph：灵活高效的图神经网络系统后端]（Yuwei Hu 等人，SC 2020）

分布式计算

• [SpDISTAL：编译分布式稀疏张量计算]（Rohan Yadav 等人，SC 2022）
• [Alpa：自动化分布式深度学习中的算子间及算子内并行化]（Lianmin Zheng、Zhuohan Li、Hao Zhang 等人，OSDI 2022）
• [Unity：通过代数变换与并行化的联合优化加速 DNN 训练]（Colin Unger、Zhihao Jia 等人，OSDI 2022）
• [在层次化系统上为深度学习合成最优并行化布局与规约策略]（Ningning Xie、Tamara Norman、Diminik Grewe、Dimitrios Vytiniotis 等人，MLSys 2022）
• [DISTAL：分布式张量代数编译器]（Rohan Yadav 等人，PLDI 2022）
• [GSPMD：面向 ML 计算图的一般化可扩展并行化]（Yuanzhong Xu 等人，arXiv 2021）
• [打破分布式机器学习工作loads中的计算与通信抽象壁垒]（Abhinav Jangda 等人，ASPLOS 2022）
• [OneFlow：从头重新设计分布式深度学习框架]（Jinhui Yuan 等人，arXiv 2021）
• [超越深度神经网络的数据与模型并行化]（Zhihao 等人，MLSys 2019）
• [利用自动数据流图划分支持超大规模模型]（Minjie Wang 等人，EuroSys 2019）
• [分布式 Halide]（Tyler Denniston 等人，PPoPP 2016）

量化

• [自动化后端感知的训练后量化]（Ziheng Jiang 等人，arXiv 2021）
• [量化深度学习模型的高效执行：一种编译器方法]（Animesh Jain 等人，arXiv 2020）
• [高性能量化机器学习内核的自动生成]（Meghan Cowan 等人，CGO 2020）

稀疏

• 稀疏抽象机 由 Olivia Hsu 等人撰写，ASPLOS 2023
• SparseTIR：用于深度学习中稀疏编译的可组合抽象 由 Zihao Ye 等人撰写，ASPLOS 2023
• WACO：学习工作负载感知的稀疏张量程序格式与调度协同优化 由 Jaeyeon Won 等人撰写，ASPLOS 2023
• Looplets：一种用于结构化协同迭代的语言 由 Willow Ahrens 等人撰写，CGO 2023
• 稀疏张量格式转换与优化的代码合成 由 Tobi Popoola 等人撰写，CGO 2023
• Stardust：将稀疏张量代数编译为可重构数据流架构 由 Olivia Hsu 等人撰写，arXiv 2022
• 稀疏抽象机 由 Olivia Hsu 等人撰写，arXiv 2022
• 无损压缩与稀疏计算的统一编译 由 Daniel Donenfeld 等人撰写，CGO 2022
• SparseLNR：利用循环嵌套重排加速稀疏张量计算 由 Adhitha Dias 等人撰写，ICS 2022
• SparTA：通过带有稀疏属性的张量实现深度学习模型稀疏化 由 Ningxin Zheng 等人撰写，OSDI 2022
• MLIR 中对稀疏张量计算的编译器支持 由 Aart J.C. Bik 等人撰写，TACO 2022
• 稀疏数组编程模型的编译 由 Rawn Henry 和 Olivia Hsu 等人撰写，OOPSLA 2021
• MLIR 中的高性能稀疏张量代数编译器 由 Ruiqin Tian 等人撰写，LLVM-HPC 2021
• 动态稀疏张量代数编译 由 Stephen Chou 等人撰写，arXiv 2021
• 高效稀疏张量格式转换例程的自动生成 由 Stephen Chou 等人撰写，PLDI 2020
• TIRAMISU：面向稠密与稀疏深度学习的多面体编译器 由 Riyadh Baghdadi 等人撰写，arXiv 2020
• 带工作区的张量代数编译 由 Fredrik Kjolstad 等人撰写，CGO 2019
• 稀疏计算数据依赖关系简化以生成高效的编译器检查器 由 Mahdi Soltan Mohammadi 等人撰写，PLDI 2019
• Taichi：一种用于空间稀疏数据结构上高性能计算的语言 由 Yuanming Hu 等人撰写，ACM ToG 2019
• 稀疏多面体框架：编译器生成的检查器-执行器代码组合 由 Michelle Mills Strout 等人撰写，IEEE 2018 年会议论文集
• 稀疏张量代数编译器的格式抽象 由 Stephen Chou 等人撰写，OOPSLA 2018
• ParSy：针对并行化的稀疏矩阵计算检查与变换 由 Kazem Cheshmi 等人撰写，SC 2018
• Sympiler：通过解耦符号分析来转换稀疏矩阵代码 由 Kazem Cheshmi 等人撰写，SC 2017
• 张量代数编译器 由 Fredrik Kjolstad 等人撰写，OOPSLA 2017
• 下一代泛型编程及其在稀疏矩阵计算中的应用 由 Nikolay Mateev 等人撰写，ICS 2000
• 基于高级规格的稀疏矩阵代码合成框架 由 Nawaaz Ahmed 等人撰写，SC 2000
• 自动非零元素结构分析 由 Aart Bik 等人撰写，SIAM 计算期刊 1999
• SIPR：生成稀疏矩阵计算高效代码的新框架 由 William Pugh 等人撰写，LCPC 1998
• 稀疏矩阵计算的自动数据结构选择与转换 由 Aart Bik 等人撰写，TPDS 1996
• 稀疏矩阵计算的编译技术 由 Aart Bik 等人撰写，ICS 1993

程序重写

• 通过高级调度重写进行验证的张量程序优化 由 Amanda Liu 等人撰写，POPL 2022
• 基于访问模式的纯张量程序重写（表示珍珠） 由 Gus Smith 等人撰写，MAPL 2021
• 等式饱和用于张量图超优化 由 Yichen Yang 等人撰写，MLSys 2021

验证与测试

• NNSmith：为深度学习编译器生成多样且有效的测试用例 由 Jiawei Liu 等人撰写，ASPLOS 2023
• 基于覆盖率的张量编译器模糊测试，结合 IR 传递突变 由 Jiawei Liu 等人撰写，OOPSLA 2022
• Halide 语言的端到端翻译验证 由 Basile Clément 等人撰写，OOPSLA 2022
• 深度学习编译器缺陷的全面研究 由 Qingchao Shen 等人撰写，ESEC/FSE 2021
• 利用程序合成验证并改进 Halide 的项重写系统 由 Julie L. Newcomb 等人撰写，OOPSLA 2020

教程

• 机器学习编译
• 深入深度学习编译器

贡献

我们鼓励大家为本仓库做出贡献。请打开一个问题或发送一个拉取请求。

关于链接格式的说明

我们更倾向于使用指向信息更丰富页面的链接，而不是单独的 PDF 文件。例如，对于 arXiv 论文，我们更喜欢 https://arxiv.org/abs/1802.04799，而不是 https://arxiv.org/pdf/1802.04799.pdf。对于 USENIX 论文（OSDI/ATC），我们更倾向于 https://www.usenix.org/conference/osdi18/presentation/chen，而非 https://www.usenix.org/system/files/osdi18-chen.pdf。对于 ACM 论文（ASPLOS/PLDI/Eurosys），我们更偏好 https://dl.acm.org/doi/abs/10.1145/3519939.3523446，而不是 https://dl.acm.org/doi/pdf/10.1145/3519939.3523446。

Awesome Tensor Compilers 快速上手指南

awesome-tensor-compilers 并非单一的可安装软件，而是一个精选的张量编译器与深度学习编译技术资源列表。它汇集了开源项目、学术论文和教程。

要开始使用其中的技术，您需要根据具体需求选择列表中的某个编译器项目（如 TVM, MLIR, Triton 等）进行安装和使用。本指南将以生态最成熟、应用最广泛的 Apache TVM 为例，演示如何从零开始体验张量编译技术。其他项目（如 MLIR, XLA）的安装逻辑类似，请参考其各自官方文档。

环境准备

在开始之前，请确保您的开发环境满足以下基本要求：

• 操作系统: Linux (推荐 Ubuntu 18.04/20.04/22.04) 或 macOS。Windows 用户建议使用 WSL2。
• Python: 版本 3.8 或更高。
• 构建工具: gcc, g++, cmake, git, llvm-dev (LLVM 是多数编译器的后端依赖)。
• 深度学习框架: 已安装 PyTorch 或 TensorFlow (用于模型导入)。

前置依赖安装命令 (Ubuntu 示例):

sudo apt-get update
sudo apt-get install -y python3 python3-pip python3-setuptools gcc g++ cmake git llvm-dev libedit-dev libssl-dev zlib1g-dev libboost-all-dev libncurses5-dev libsqlite3-dev

安装步骤

推荐使用 Python pip 直接安装预编译包，这是最快的方式。如果需要针对特定硬件（如特定的 GPU 或 NPU）进行深度优化，则需从源码编译（此处仅展示快速安装）。

1. 安装 Apache TVM:

pip install apache-tvm==1.0.0

(注：版本号可根据实际需求调整，建议查看 PyPI 获取最新版)

2. 验证安装:

python3 -c "import tvm; print(tvm.__version__)"

若输出版本号且无报错，则安装成功。

基本使用

以下示例展示如何使用 TVM 将一个简单的 PyTorch 模型编译并优化为可在 CPU 上高效运行的张量程序。

1. 创建测试脚本 (quick_start.py):

import torch
import tvm
from tvm import relay
from tvm.contrib.download import download_testdata
from tvm.relay.backend import Runtime

# 1. 定义一个简单的 PyTorch 模型
class SimpleNet(torch.nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.linear = torch.nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

model = SimpleNet()
model.eval()

# 2. 创建虚拟输入
input_shape = (1, 10)
input_data = torch.randn(input_shape)

# 3. 将 PyTorch 模型转换为 Relay IR (TVM 的中间表示)
scripted_model = torch.jit.trace(model, input_data).eval()
mod, params = relay.frontend.from_pytorch(scripted_model, [("input0", input_shape)])

# 4. 配置编译目标 (以 LLVM/CPU 为例)
target = tvm.target.Target("llvm")
dev = tvm.cpu(0)

# 5. 编译模型
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(mod, target=target, params=params)

# 6. 创建运行时并执行
tvm_model = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
tvm_model.set_input("input0", tvm.nd.array(input_data.numpy()))
tvm_model.run()
output_tvm = tvm_model.get_output(0).numpy()

# 7. 验证结果 (对比 PyTorch 原生输出)
output_pytorch = model(input_data).detach().numpy()
print("TVM Output shape:", output_tvm.shape)
print("Results match:", torch.allclose(torch.tensor(output_tvm), torch.tensor(output_pytorch), atol=1e-5))

2. 运行脚本:

python3 quick_start.py

预期输出: 如果配置正确，脚本将输出模型的形状以及 Results match: True，表明张量编译器已成功接管并执行了计算任务。

---

提示：若要探索列表中其他项目（如 mlir, triton, hidet），请访问 awesome-tensor-compilers 仓库中对应的 "Open Source Projects" 链接获取其专属安装指南。