Relay #5
Description
Relay IR
文章, Relay: A High-Level IR for Deep Learning
文档: https://docs.tvm.ai/dev/relay_intro.html
tvm 上的讨论 issue apache/tvm#1673
是 TVM 的继 NNVM 之后的新的高级中间表示.
强调 DL IR 在可扩展性上的三个挑战
- 表达性, 支持 tree, graph, list
- 可组合, 能够轻松替换某一实现, 比如 quantization, operator fusion, automatic differentiation
- 可移植, 支持多种硬件后端
以往, operator 的 C++ 实现 与 dataflow graph 耦合了
贡献
- Relay, 语言特点: 看起来就是 ML-style 的函数式编程语言, tensor-oriented, statically typed, functional IR.
- 将 DL 框架中的问题变为编译问题, 应用现有编程语言技巧.
- backpropagation 变成 source code transformation
- shape inference 变为 type inference
- graph rewriting, 变成 program optimization
- executor 变成 interpreter, VM, AOT copmiler
- Relay 的类型系统. 加上对形状的推理, HM type system 加上 type relation 扩展.
- 平台无关的 operator language, 和 compiler pass manager.
Relay 表达能力: 能够 import TensorFlow, PyTorch, MxNet, NNVM, ONNX 的模型.
源到源的 AD 算法, 支持高阶函数与高阶导数, 不需要delimited continuations?
2. 背景知识
DL 如何训练模型, 现有 DL 的两种类型, 动态神经网络, 多种硬件 backends
3.1. 设计:
Relay pipeline 可以分为三个经典部分：
- 前端，输入格式转换为Relay;
-
- 支持内存中构建 AST, (C++/Python API)
- Relay text format
- on-disk serialization format.
- 编译器，其类型检查中继AST，应用优化，并编译运算符;
- 后端，其中选择执行机制并使用可用的硬件加速器。
- Relay 是高级 IR, 需要依赖 TVM 或 Halide 来生成底层代码
3.2. Relay IR, 高阶的函数式可微分语言
计算图, 有向无环图, 包含多个输入, 一个输出.
- let, 共享变量. let 引入 lexical scope, 解决顺序问题.
- 控制流, 使用 recursion 和 pattern matching
- First-Class 函数, 计算图服用
- ADT, 表达 lists, trees, graphs 多种数据类型
基本上就是一个 ML-style 的语言, 本来计算图的表达, 能用过程式语言写, 自然也能用函数式语言写.
其中的好处是什么?
困难: 把 shape inference 和 automatic differentiation adapt 到 这个 new IR 上
3.3. 类型系统.
shape of a tensor, tuple of integers, 可以是 symbolic 的, 支持 polymorphic over shaped, Any 表示在某一维度上没有静态信息.
operator 的类型, 能表达 output 与 input 的关系. 可能 output 是 input 的函数, 也可能是两者满足某种约束.
4. Case studies.
基于 relay 去实现如下功能
- model importer , 能够 import TensorFlow, PyTorch, MxNet, NNVM, ONNX 的模型.
- 通用自动微分算法
- partial evaluator
- generic operator fusion algorithm
- automatic quantization framework
- Ahead-of-Time Compiler
默认执行方式是 一个基于 AST 遍历的 interpreter, 使用 JIT 执行 operators.
5. 实验
- Relay 使能可组合的优化.
- 优化分为多个passes, 比较 -O0, -O1, -O2, -O3
- Relay 性能好
- 与 NNVM, PyTorch, TensorFlow, TF-XLA, MXNet 性能对比
- Relay 能支持多种 Challenging Backends
- 支持树莓派, FPGA, 能够轻松进行 quantization