随缘随笔
Insights Flow
0%

论文阅读[精读]-The Deep Learning Compiler: A Comprehensive Survey

发表于 更新于 分类于 阅读次数： Valine：
本文字数： 6.4k 阅读时长 ≈ 6 分钟
深度学习编译方向的一篇survey

​ 这是我读的第一篇深度学习编译方向的论文，是一篇survey工作。由于是第一篇，那基本上就得逐字逐句慢慢看了，又因为是survey，那二三十页是跑不了了。因此我为了加深印象，干脆就同步更新一下笔记，就当做”把论文翻成中文“了。

摘要

  • 学术界提出了XLA, TVM等深度学习编译技术
  • DL compiler输入模型，输出适用于不同平台的、优化后的代码
  • 本篇工作聚焦于multi-level IR，前后端的优化技术。
  • 提出了对这些技术的作者的点评，未来的改进方向
  • 这是第一篇深度学习编译优化的综述文章

Introduction

DL硬件分类：

  • 软硬件联合设计的 general-prupose 硬件
    • 线性代数库BLAS支持DL运算
    • 还有MKL-DNN，cuDNN等专门的加速库
    • 还有TensorRT等加速库
  • 专为DL设计的硬件
    • 类似的加速库
  • 从生物脑科学获得灵感的神经元硬件

库的更新经常赶不上硬件更新的速度，因此很难利用好DL硬件的算力

DL编译器

  • 优化 模型定义-> 实现代码 这一过程

  • 算子合并、层合并等方法

  • 是层次性的。

本篇工作的贡献：

  • 归纳了前端、多层IR、后端级别的优化方法
  • 分裂归纳了已有的DL编译器
  • 定量对比了不同编译器在CNN的效率提升（end-to-end和单层粒度下）
  • 提出了DL编译器的未来方向： dynamic shape and pre-/post-processing, advanced auto-tuning, polyhedral model, subgraph partitioning, quantization, unified optimizations, differentiable programming and privacy protection

Background

DL框架对比

  • tensorflow：其他语言支持最广泛：C ++, Python, Java, Go, R, and Haskell。TensorFlow Lite为Android设计。Keria作为tensorflow的前端。eager-mode类似pytorch的动态计算图
  • keria：高层神经网络库，快速构建model。兼容scikit-learn等库。由于过度封装，增加算子、获取底层信息很难。
  • Pytorch：动态框架，按行执行。合并了Caffe2。FastAI 是pytorch的高层封装，借鉴了keria。
  • Caffe/Caffe2,MXNet,CNTK,PaddlePaddle略过
  • ONNX：别的框架模型可以转换成ONNX模型，以便于交互（使用pytorchAPI）。

深度学习硬件

  • General-purpose Hardware：类似NVIDIA-GPU。同时上层有cuDNN等支持，软硬件联合开发。
  • Dedicated Hardware：A100，cloud service provider（google TPU）等等
    • TPU包含：
      • Matrix Multiplier Unit (MXU),
      • Unified Buffer (UB),
      • and Activation Unit (AU)
    • TPU由host硬件 进行CISC指令集的驱动
    • TPU可以直接把矩阵视为单元，而不是vector或者标量集
  • Neuromorphic Hardware：IBM’s TrueNorth and Intel’s Loihi.
    • 神经元可以同时存储和处理数据，存储区和运算区在一起，没有专门存储区。
    • 距离大规模应用很遥远
    • 探索rapid, life-long learning。

FPGA

cpu/gpu泛用性强，但是费电；AISC省电，但是转为某任务设计。编程FPGA是一个折中，把模型部署在FPGA上。

High-Level Synthesis (HLS) programming model帮助用户编辑FPGA，不用写很多verilog。

FGPA应用的问题：

  • AI模型用框架描述，而不是c/c++
  • DL特有的优化方式很难推广到FPGA

hardware-specific code generator：输入DL模型，输出HLS/verilog/VHDL代码，再编译成bitstream烧录:

  • The processor architecture:目标FPGA和正常处理器类似。
  • The streaming architecture：目标FPGA是个流水线。快，瓶颈是FPGA内的memory不够

DL编译器结构

正常编译器分为前端和后端。中间由IR联系起来，IR分为多个层次(multi-level IR)。

前端基于高层IR做硬件无关的优化；后端基于低层IR，做硬件专有的优化、代码生成、编译。

  • High-level IR: 定义计算和控制流，硬件无关。目标是得到程序的控制流和数据依赖性，提供图级别的优化。

  • Low-level IR:足够细粒度。需要支持后端的第三方工具链。

  • 前端： 从node-level，block-level, dataflow-level优化模型，生成 graph-IR.

  • 后端：输入high-level IR,输出 low-level IR。 可以直接把high-level IR转成LLVM IR等第三方工具链支持输入做优化，也可以通过已有的硬件结构、模型结构做自己的特殊优化

DL编译器的关键因素

High-level IR

High-level IR计算图的表示

图的表示法不同，优化方向也不同：

  • DAG图：有向无环图，点表示算子，边表示tensor。可以快速分析数据依赖。优点是方便优化，缺点是图上的节点、边语义不明确
  • Let-binding-based IR：对每一个变量”let“建立节点，语义明确。它对每一个let都算结果，建立一个map。每个表示通过查表找到结果

TVM等方法借鉴两种表示，博才两家之长。

tensor计算的表示法：

  • Function-based：Glow, nGraph and XLA采用，提供一些封装的算子。tensor计算视为一种函数
  • lambda expression: 用lambda函数表达，不需要声明新函数。TVM采用此方法，需要先计算输出的形状。
  • Einstein notation：比lambda表达式更简单。算子需要相关联，可交换，方便并行化。

Graph-IR 的实现

数据(tensor)的表示:

  • Placeholder:标记tensor的shape，有这个就有shape。

  • Unknown (Dynamic) shape representation：允许某一维的大小未知

  • Data layout：逻辑地址到内存分片的映射。

    • TVM and Glow把获取data layer作为一个专门的算子，这样不需要专门实现这个方法，更方便优化。
    • XLA把这个视为后端硬件的一个约束
    • Relay and MLIR要求tensor在type中描述layout

  • Bound inference：推断迭代器的上下界。TVM中iterator建立一个dag图，点代表迭代器，边代表运算。根节点的shapes of placeholders被确定，就能递归处理。

graph-IR需要支持很多算子：代数算子、tensor算子、control flow算子。举例子：

  • broadcast：放宽一般算子对形状的依赖性。
  • control flow:if或者while
  • Derivative：自动求导。
  • Customized operators：自定义算符。
    • Glow中定义算符需要实现多层的封装
    • TVM、TC只需实现一个implementation

不同编译器有不同的 graph-IR 实现，但有一定的相似之处。需要注意，graph-IR一定要是硬件无关的。

Low-level IR

实现

不同的实现

  • Halide-based IR：设计哲学是计算和规划分离。可以试多种规划，选一个最好的。原始方法需要形状确定，TVM改进了它：
    • 取消LLVM的依赖
    • 重构project module
    • 提高复用性，方便自定义算子
    • 保证每个变量只有一个定义位置
  • Polyhedral-based IR：视为多面体。loop的大小更加灵活，方便采用polyhedral transformations优化方法。有多种优化器：isl,Omega,PIP,Polylib,PPL
  • Other unique IR: 他们用一些自定义的优化方法，然后编译成LLVM IR
    • GLOW：包含declare和program两种操作。用@in,@out,@inout帮助分析内存优化的时机
    • MLIR：受到LLVM影响。用dialect提供对别的IR的抽象，包含：TensorFlow IR, XLA HLO IR, experimental polyhedral IR, LLVM IR, TensorFlow Lite。 自定义dialect很简单，方便开发者适配新硬件。
    • HLO IR of XLA：同时是high/low-level IR。足够细粒度，提供硬件级的优化，生成 LLVM IR

代码生成

绝大多数都是生成LLVM IR，接下来通过LLVM进行多种优化。需要提供几种优化：

  • 循环转换
  • 提供目标硬件的额外信息

前端的优化

前端进行计算图的优化，和硬件实现无关。前端的优化称为pass，通过多次遍历图，每轮进行不同的操作。一旦模型被import、转换成graph，前端可以获取各个地方的shape。

Node-level optimizations

  • 消除不必要的节点：比如sum（1,0），Nop Elimination
  • Zero-dim-tensor elimination：消除0维向量，或者消除某个维度为0的向量

Block-level optimizations

  • 消除单位运算+0,\(*1\)
  • 用简单运算符替代复杂运算符
  • 预计算常数

DL运算符也可以优化，比如：

  • optimization of computation order：\(A^T B^T = (BA)^T\)

  • optimization of node combination:比如把多个transpose消成一个

  • optimization of ReduceMean nodes:用AVGPool代替ReduceMean

  • Operator fusion：算子合并。难点是如何合并包含多个reshape,boardcast,reduce等节点的算子。

  • Operator sink：消除可以消除的节点

Dataflow-level optimizations

  • Common sub-expression elimination (CSE)：预计算可以计算的控制流。
  • Dead code elimination (DCE)：如果结果不用，就是dead，不需要存中间结果。经常在上面的优化以后出现
  • Static memory planning，对于内存受限的机器很重要：
    • in-place memory sharing：复用输入和输出的memory
    • Standard memory sharing：在不覆盖的情况下，复用前面用到的存储空间
  • Layout transformation：
    • 计算出tensor最好的存储形式，然后在计算图增加layout transformation节点。
    • 最优方式随tensor计算公式不同，硬件不同而不同
    • 速度提升明显

后端的优化

Hardware-specific Optimization

一种方式是编译到 LLVM IR，另一种方式是定义自己的优化方式，用模型的信息。举5个例子：

  • Hardware intrinsic mapping：编译成硬件已经优化过的一些kernel。

  • Memory allocation and fetching：针对硬件上不同存储的延迟不同做优化

  • Memory latency hiding：等内存延迟的时候干别的，好多硬件自己实现了，但TPU没有。TVM通过虚拟线程解决这事

  • Loop oriented optimizations：Halide和LLVM实现过了。

    • Loop fusion: 把多个边界一样、没有依赖的循环放到一起
    • Sliding windows：直到需要数据时再计算
    • Tiling：把一个循环拆成多重循环，更好的局部性
    • Loop reordering：重新排列循环顺序。让空间局部的循环们放在一起
    • Loop unrolling：循环展开，更好地做指令级并行。

  • Parallelization：更好的支持线程级并行等，提高硬件利用率。需要额外的、模型的知识，更难开发。

Auto-tuning

由于各种优化很多，搜索空间很大，需要自动搜索优化方式。TVM, TC,XLA都支持这个。包含四个要素：

  • Parameterization：搜索参数时需要知道的一些已有参数
    • Data and target：tensor的形状，GPU上各个memory的延迟、大小
    • Optimization options：优化方法有什么，对应什么超参。TC，XLA支持把超参参数化，比如batch_size。
  • Cost model：评价优化方法的模型
    • Black-box model：只关心运行时间，对内部情况不关心。TC采用
    • ML-based cost model: 用ML模型衡量现在情况有多好。TVM，XLA采用
    • Pre-defined cost model:预先定义好。搜索更快，但是定义很难，和DL model有关。
  • Searching technique：如何搜索
    • Initialization and searching space determination：随机/特定初始化。TVM允许用户决定搜索空间
    • Genetic algorithm (GA) ：类似蚁群算法，TC采用
    • Simulated annealing algorithm (SA)：模拟退火。TVM采用
    • Reinforcement learning (RL): Chameleon采用（基于TVM开发）
  • Acceleration：如何加速搜索过程

    • Parallelization：

    • Configuration reuse：复用上次的结果，如果某些局部的参数和上次一样。

      Optimized Kernel Libraries

硬件预定义好了一些快速的计算库(比如cuDNN)。后端可以调用这些库

  • 调用库带来巨大的性能提升
  • 但调用库需要符合库的调用约定，可能破坏最佳控制流
  • 库函数对编译器是黑盒，可能影响优化，不能做operator fusion等。

已有编译器的分类

通过上面已经讲过的不同方面的优化方法，可以梳理、归纳、分类一下现在比较火的DL编译器。总体结果如下图：

评测

这一部分评测了一下已有的一些编译器的性能，具体实验就略过了，对学习DL编译优化知识没什么帮助。

结论和未来方向

这一部分归纳梳理了一下DL编译未来的发展空间：

  • Dynamic shape and pre/post processing：

    • 一方面类似NLP，只有运行时才能知道输入的形状。另一方面，模型结构本身可能也会变化。
    • 随着模型增大，模型的加载时间可能成为瓶颈。已有还没有触及

  • Advanced auto-tuning：受限于时间，现有优化都是在找一些局部最优。但局部最优的组合大概率不是全局最优

    • 目前ML方法还有潜力，可以考虑在auto-tuning中进一步应用，而不只是cost model。

  • Polyhedral model：在auto-tuning中引入多面体模型

    • 可以复用之前结果
    • 可以减少搜索空间
    • 挑战是如何在稀疏情况下应用之。

  • Subgraph partitioning：把计算图拆分成不同子图。可以把不同子图分布并行到异质的设备上。

  • Quantization：增强量子化可以在编译时进一步提升优化的空间，挑战是：

    • 如何简单的增量开发
    • 量子化操作如何和其他的优化步骤交互。

  • Unified optimizations：如何同时采用不同编译器的优化，已由编译器大多聚焦于一些方面的优化。MLIR通过dialect某种程度上可以复用不同编译器的优化。

  • Differentiable programming：让编译器支持可微程序

  • Privacy protection：能不能在中间输出层加噪音，保护隐私。

  • Training support：已有编译方法专注于部署。

相关文章