Published on

MLIR compiler paper

Authors

AI compiler相关的论文

IR表示

内存分配

1. TVM: 内存管理创新

论文标题：TVM: An Automated End-to-End Optimizing Compiler for Deep Learning

概述：

人们越来越需要将机器学习引入各种硬件设备。当前的框架依赖于供应商特定的运算符库，并针对狭窄范围的服务器级 GPU 进行优化。将工作负载部署到新平台（例如移动电话、嵌入式设备和加速器（例如 FPGA、ASIC））需要大量的手动工作。我们提出 TVM，这是一种公开图形级和操作员级优化的编译器，可为跨不同硬件后端的深度学习工作负载提供性能可移植性。 TVM 解决了深度学习特有的优化挑战，例如高级算子融合、映射到任意硬件原语以及内存延迟隐藏。它还通过采用一种新颖的、基于学习的成本建模方法来快速探索代码优化，根据硬件特性自动优化低级程序。实验结果表明，TVM 提供的跨硬件后端性能可与低功耗 CPU、移动 GPU 和服务器级 GPU 的最先进的手动调整库相媲美。我们还展示了 TVM 针对新加速器后端的能力，例如基于 FPGA 的通用深度学习加速器。该系统是开源的，并在几家大公司内进行生产使用。

关键内存优化点：

  • 提出运行时内存重用策略

  • 实现动态内存分配与管理

  • 跨硬件平台的内存优化技术

  • abstract

图形级和运算符级优化，为跨不同硬件后端的深度学习工作负载提供性能可移植性。TVM 解决了深度学习特有的优化挑战，例如高级运算符融合、映射到任意硬件基元和内存延迟隐藏。它还通过采用一种新颖的、基于学习的成本建模方法来快速探索代码优化，从而自动根据硬件特性优化低级程序。

  • introduction

当前的 DL 框架，如 TensorFlow、MXNet、Caffe 和 PyTorch，都依赖于计算图中间表示来实现优化，例如自动微分和动态内存管理 .

框架也必须在以下两者之间做出艰难的选择：（1） 避免产生不在预定义算子库中的新算子的图形优化，以及 （2） 使用这些新算子的未优化实现。

TVM从现有框架中获取深度学习程序的高级规范，并为各种硬件后端生成低级优化代码。为了吸引用户，TVM 需要提供与跨不同硬件后端的大量手动优化算子库相比的性能。要实现这一目标，需要解决下面描述的关键挑战。

  1. 利用特定的硬件功能和抽象。
    1. 硬件指令的输入是多维的，具有固定或可变长度;它们规定了不同的数据布局;并且它们对内存层次结构有特殊要求。
    2. 加速器设计通常也有利于更精简的控制 并将大多数调度复杂性卸载到编译器堆栈中。对于专用加速器，系统现在需要生成明确控制管道依赖关系的代码，以隐藏内存访问延迟
  2. 用于优化的大搜索空间
    1. 内存访问、线程模式和新颖硬件原语的组合选择为生成的代码（例如，循环平铺和排序、缓存、展开）创造了一个巨大的配置空间，如果我们实现黑盒自动调优，这将产生大量的搜索成本。可以采用预定义的成本模型来指导搜索，但由于现代硬件的复杂性日益增加，构建准确的成本模型很困难。此外，这种方法需要我们为每种硬件类型构建单独的成本模型。

TVM 通过三个关键模块来应对这些挑战。

（1） 我们引入了一种张量表达式语言来构建运算符，并提供程序转换原语，这些原语通过各种优化生成不同版本的程序。该层扩展了 Halide [ 32] 的计算/调度分离概念，还将目标硬件内部函数与转换原语分开，从而支持新型加速器及其相应的新内部函数。此外，我们还引入了新的转换基元来解决与 GPU 相关的挑战，并支持部署到专用加速器。然后，我们可以应用不同的程序转换序列，为给定的运算符声明形成一个丰富的有效程序空间。

（2） 我们引入了一个自动化程序优化框架来寻找优化的张量算子。优化器以基于 ML 的成本模型为指导，该模型会随着我们从硬件后端收集更多数据而进行调整和改进。

（3） 在自动代码生成器的基础上，我们引入了一个图形重写器，它充分利用了高级和运算符级优化。

图优化：

  1. 算子融合
  2. 常量折叠
  3. 静态内存规划传递：预分配内存，保存每一个中间变量。
  4. datalayout转化
    1. 有多种方法可以在计算图中存储给定的张量。最常见的数据布局选项是 column major 和 row major。在实践中，我们可能更喜欢使用更复杂的数据布局。例如，DL 加速器可能会利用 4×4 矩阵运算，要求将数据平铺成 4×4 块以优化访问位置。
    2. 数据布局优化将计算图转换为可以使用更好的内部数据布局在目标硬件上执行的计算图。它首先为每个运算符指定首选数据布局，给定内存层次结构规定的 constraints 。然后，如果生产者和消费者的首选数据布局不匹配，我们会在他们之间执行适当的布局转换。

2. DLVM: 深度学习专用内存管理

论文标题：DLVM: Dynamic Memory Management for Deep Learning

核心创新：

  • 动态内存预测分配
  • 基于计算图的内存生命周期分析
  • 减少内存碎片

3. Halide: 异构系统内存优化

论文标题：Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines

内存管理特点：

  • 显式内存调度
  • 计算与存储分离
  • 降低内存访问延迟

4. XLA (Accelerated Linear Algebra)：张量内存优化

论文标题：XLA: Domain-Specific Just-In-Time Compilation for Machine Learning

内存优化技术：

  • 静态内存分配
  • 张量内存压缩
  • JIT编译时内存重用

5. Tensor Comprehensions: 自动内存管理

论文标题：Tensor Comprehensions: Framework-Agnostic Deep Learning Performance Optimization

内存管理亮点：

  • 自动内存分配策略
  • 基于机器学习的内存优化
  • 跨框架内存管理

6. Apache TensorFlow: 内存分配算法

论文标题：TensorFlow: A System for Large-Scale Machine Learning

内存分配创新：

  • 记忆池(Memory Pool)技术
  • 图级别内存重用
  • 设备间内存协调

7. CUDA深度学习内存优化

论文标题：Efficient Memory Management for Deep Learning on GPU

关键技术：

  • GPU显存动态分配
  • 零拷贝内存传输
  • 内存碎片整理算法

论文: "Zero-Copy Memory Management in Deep Learning Frameworks"

  • 会议: EuroSys 2021
  • 摘要: 提出一种减少内存拷贝次数的框架，显著提升运行效率。
  • 亮点: 对实时推理系统有重要意义。

论文: "Optimizing Memory Placement and Movement for Deep Learning"

  • 会议: NeurIPS 2020
  • 摘要: 提出一种内存调度算法，通过静态规划优化深度学习计算图中的内存分配和数据移动，显著降低内存峰值。
  • 亮点: 使用动态规划解决跨层内存复用问题，适合多种DL框架。

论文: "Dynamic Memory Management for TensorFlow"

  • 会议: USENIX ATC 2017
  • 摘要: 介绍了TensorFlow如何在运行时动态分配和回收张量内存，提升了多任务的内存利用率。
  • 亮点: 提供了张量生命周期的管理思路，可推广到其他编译器。

论文: "Naiad: Efficient and Transparent Memory Management for Deep Learning"

  • 会议: OSDI 2016
  • 摘要: 提出基于张量生命周期的内存优化方法，通过显式地追踪数据的使用范围，减少冗余分配。
  • 亮点: 兼顾训练和推理任务，支持复杂的深度学习模型。

论文: "Memory Optimization of Deep Learning Models via Efficient Lifetime Analysis"

  • 会议: PLDI 2019
  • 摘要: 使用静态分析技术，结合模型的计算图结构，进行内存块复用优化。
  • 亮点: 针对不同内存需求的操作进行定制化复用策略。

论文: "MLIR: A Compiler Infrastructure for the End of Moore's Law"

  • 会议: arXiv 2020
  • 摘要: MLIR引入了多级中间表示（IR），在降低内存使用和跨层优化方面表现出色。
  • 亮点: 提供了一个统一的编译框架，便于不同深度学习任务的内存优化。

论文: "Deep Learning Inference with Limited Memory"

  • 会议: CVPR 2019
  • 摘要: 提出一种针对推理阶段的内存复用技术，尤其适用于嵌入式设备和移动设备。
  • 亮点: 专注于设备受限场景的优化策略。

论文: "Unified Memory Architecture for Accelerators"

  • 会议: HPCA 2020
  • 摘要: 在异构硬件（如GPU+TPU）的环境中，通过统一内存池的设计减少内存瓶颈。
  • 亮点: 提供了对多种硬件架构友好的内存管理方法。

代码生成 & schedule

AI Powered Compiler Techniques for DL Code Optimization

在 CPU 上创建深度学习基元的高性能实现是一项具有挑战性的任务。多个考虑因素（包括多级高速缓存层次结构和 CPU 平台的宽 SIMD 单元）会影响选择程序转换以应用性能优化。在本文中，我们提出了机器学习驱动的编译器技术来优化循环嵌套。我们采用双管齐下的方法进行代码优化：首先，我们应用高级优化来优化代码，以充分利用高速缓存内存。然后，我们执行低级、特定于目标的优化，以有效地将代码矢量化，以便在机器的 SIMD 单元上良好运行。对于高级优化，我们使用多面体编译技术和深度学习方法。对于低级优化，我们使用特定于目标的代码生成器，该生成器使用向量内联函数和强化学习 （RL） 技术生成代码，以找到代码生成器的最佳参数。我们对流行的深度学习工作负载中出现的各种矩阵乘法对开发的技术进行了实验评估。实验结果表明，本文中介绍的编译器技术在顺序和并行运行的基线上分别实现了 7.6 倍和 8.2 倍的加速。

  • 发表时间： 2021 年 4 月 12 日星期一

CompilerGym: Robust, Performant Compiler Optimization Environments for AI Research

人们将人工智能 (AI) 技术应用于编译器优化的兴趣正在迅速增加，但编译器研究的进入门槛很高。与其他领域不同，编译器和人工智能研究人员无法访问能够快速迭代和开发想法的数据集和框架，并且开始需要大量的工程投资。我们需要的是一个简单的、可重用的实验基础设施，用于现实世界的编译器优化任务，它可以作为比较技术的通用基准，并作为加速该领域进展的平台。 我们推出了 CompilerGym，这是一组用于现实世界编译器优化任务的环境，以及一个用于向编译器研究人员展示新优化任务的工具包。 CompilerGym 使任何人都可以通过易于使用的软件包来试验生产编译器优化问题，无论他们使用编译器的经验如何。我们以流行的 OpenAI Gym 界面为基础，使研究人员能够使用 Python 和熟悉的 API 与编译器进行交互。 我们描述了 CompilerGym 架构和实现，描述了三个包含的编译器环境的优化空间和计算效率，并提供了广泛的实证评估。与之前的作品相比，CompilerGym 提供了更大的数据集和优化空间，计算效率提高了 27 倍，具有容错能力，并且能够检测底层编译器中的再现性错误。 为了让任何人都能轻松地尝试编译器（无论其背景如何），我们的目标是加速人工智能和编译器研究领域的进展。

  • 发表时间： 2021 年 9 月 17 日星期五
Discuss on Twitter • View on GitHub