TVM编译入门概念 一

最近阅读到陈天奇的一篇文章，链接分享
21年底，他说在三年前也就是2018年左右就完成了基于多层优化的解决方案。
当把全栈解决方案搭建起来并且不断实践之后发现有两种隔阂阻碍整个行业的发展。做一些总结。

• 本文合适的受众是ai编译器相关的入门级选手或供大家回忆知识点。
• 我的风格是突出重点字，句。
• 如果大家有好的建议，请不吝批判和提供建议

ai编译器目前有一些问题存在，打算从这些问题入手，去理解ai编译器相关的基本概念。

• 竖向隔阂
  • 阻碍手工优化方案和自动编译方案
  • 两种流派：手工算子优化，自动优化
  • 大部分当前框架基本都只为两种方案之一去做设计
  • 打破竖向墙，让手工优化和自动优化做整合，是行业面临大问题
• 层级隔阂(横向隔离)
  • 多层优化生态问题
  • 每层优化的抽象设计是分开的
  • 优点是抽象内部的优化非常灵活
  • 麻烦是抽象到另一个抽象的转化往往需要通过translator或lowering的批量转换
  • 许多困难都集中在一个抽象到另一个抽象的边界上
  • 而且这些转换往往是单向的，在高阶做优化 (如计算图) 然后传递给张量表达层
  • 但张量表达层以及硬件层级的信息往往难以反馈给更高的表达层
  • 重点：张量优化本身可以反过来指导计算图层级的算子融合和数据排布，但是当前的单向架构，比较难自然地利用这一类的反馈信息
• 总结
  • 深度学习编译和优化不是分为多个层级的优化就可以完全优化的问题
  • 解决优化问题需要各层级抽象之间的联动
  • 目前可以做的事情
    • 搭建某个层级的抽象或者dialect
    • 并且让这些抽象通过多层优化的方式从高到低的进行逐层变换和转换
  • 目前的困难所在
    • 自底向上反馈迭代的优化
    • 在边界引入更多可模块化的整合变换
    • 仅有多层优化的概念是不够的
  • 由于横向和竖向隔离这两个问题，依然难以做好端到端的整体优化，这些隔阂和问题的存在和基础架构的选择无关，不论是MLIR或者TVM，一旦采用了多层优化的方式，都难以避免这个问题。希望接下来完能快速打通这个流程，真正理解到这个问题。

---

Tensor IR (TIR)

• 介绍
  • 旧Relay和新Relax IR 最后都会被转换到TIR
  • TIR源码封装在tvm.tir中
  • TIR的数据结构就是一颗抽象语法树(AST)
  • 语法树可以表示的内容就是 变量的声明，变量的初始化，变量的计算，函数调用，控制流等。
  • 只需要采用观察者模式去遍历这颗AST树就可以实现一对一的硬件码级别的翻译。也就是codegen模块
  • TVM前端都会被封装到IRModule中进行编译，
  • linux下IRModule就是一个.so动态链接库
  • PrimFunc叫做张量函数
  • 内部封装了对应动态库的入口函数
  • 一个IRModule可以有多个PrimFunc
  • Codegen实际上就是对TIR AST树进行遍历然后一对一的将AST Node翻译为TIR Node 对应的数据结构并发回给调用函数VisitExpr_和 VisitStmt

tvm.IR基础设施

• 介绍
  • tvm.ir 抽象
  • Type和Expr是关键概念
  • Type中包含张量的基础数据类型、自定义的复杂类型的函数类型
  • Expr既包含lowering IR PrimExpr，也包含 RelayExpr
  • 具体的部分可以去看源码中的定义，分别在tvm/include/tvm/ir这个文件下的头文件里
  • 另外说一点，开发者们对代码进行了多次重构，路径结构可能有改变，我下的版本是2023年10月最新版本的tvm

Relay IR介绍

• 清楚一个概念，onnx模型文件可以直接转换到Relay IR这一层
  • onnx模型文件可以通过https://netron.app/这个网站可视化展开哦
  • 代码角度看，Relay的基类Expr 就是 tvm.ir基础设施中的RelayIR
  • 在tvm/relay/expr.h中可以看到 relay命名空间中Expr指的就是 RelayExpr
    
  • VarNode 是Relay里面的变量
  • VarNode 也会继承自 ExprNode这个类
  • VarNode当中 有vid,type_annotatiton两个成员变量
  • 如Relay IR看到的以@和以%开头的变量