AI Overview

传统编译器里通常将编译分为三个部分:前端 (frontend)、优化器 (optimizer)、后端 (backend).编译过程以高级语言作为输入……

  • 前端 将代码转化为中间表示 (IR)
  • 优化器对 IR 进行优化.IR 是前后端连接的桥梁,也是解耦和的工具.
  • 后端负责将已经优化过的 IR 转换为底层的机器码

来到 AI 时代,我们希望模型(本质是矩阵乘法和元素操作)可以高效运行,同时需要支持多个深度学习框架(如 PyTorch, TensorFlow 之流)和多个计算硬件(如 AMD 显卡、NVIDIA 显卡、TPU 等等).这样的需要和编译器的职责是类似的.但是如果直接照搬传统编译器会有几个问题:

  1. 对 领域特定计算抽象 缺乏 深度优化能力
  2. 难以高效适配 高度异构、可配置的 AI 加速器
  3. 无法区分并满足训练与推理场景的根本性差异需求
What is AI Compiler?

AI 编译器是一种针对 AI 和机器学习应用特别优化的编译器,可以满足推理场景和训练场景的不同需求,将高级语言编写好的模型或者训练好的模型文件转换成可以在特定硬件上高效执行的程序.

  • 以 Python 为前端
  • 设计上拥有多层 IR
  • 面向深度学习优化
  • 针对不同的芯片架构设计
流程 AI Compiler 传统编译器
前端 张量表达、算子表达、计算图生成…… 词法分析、语法分析、语义分析……
优化器 图算融合、并行切分、量化…… 死代码消除、全局变量优化、常量折叠……
后端 代码生成、算子分块、内存分配…… 指令调度、指令选择、寄存器分配……
  • 前端经过多种优化(如代数简化、算子融合、操作下沉、公共子表达式消除、静态内存规划等等),得到初步优化的计算图,随后通过模式匹配和图重写等方法进一步优化,最终生成优化后的计算图
  • 后端将优化后的 IR 转换为特定平台的高效代码.后端模块通常包括
    • 硬件特定的优化器,如针对 GPU 的并行化策略、针对 TPU 的矩阵乘法优化、针对 FGPA 的流水线调度.
    • 自动调优技术,如自动搜索最优参数配置

Case Study: TVM

简单介绍一下 TVM 的工作流程:

  1. 从 Tensorflow/Pytorch 导入模型
  2. 翻译为 TVM 的高级模型语言 Relay 图级优化
  3. 将 Relay 翻译为 TE 表示 了解“计算什么”
  4. 使用 auto-tuning 模块 AutoTVM 或者 AutoScheduler 搜索最佳 schedule.
  5. 为模型编译选择最佳配置
  6. 降级为张量中间表示 (TIR)
  7. 编译为机器码

Case Study: MLIR