【TVM 教程】向 Relay 中添加 Compiler Pass
💡
原文中文,约7500字,阅读约需18分钟。
📝
内容提要
Apache TVM 是一个支持 CPU、GPU 和加速芯片的深度学习编译框架。Compiler Pass 是优化 Relay 程序的接口,用户可以通过编写 pass 来修改 AST。示例包括常量折叠 pass,展示了如何检查和替换常量表达式。用户可通过注册 Pass 实现高效的程序优化。
🎯
关键要点
- Apache TVM 是一个支持 CPU、GPU 和加速芯片的深度学习编译框架。
- Compiler Pass 是优化 Relay 程序的主要接口,用户可以通过编写 pass 修改 AST。
- 编写 pass 包括创建 C++ 类和将其与 Pass Infrastructure 交互。
- AST 遍历器的基类是 ExprFunctor,提供 VisitExpr 方法用于遍历表达式。
- ExprVisitor 用于程序分析,ExprMutator 用于转换程序。
- 常量折叠 pass 是一个示例,涉及 ConstantChecker 和 ConstantFolder 两个组件。
- ConstantChecker 用于检查表达式是否为常量,并使用 memo_ 字段缓存结果。
- ConstantFolder 执行常量折叠过程,处理 LetNode、TupleItemGetNode 和 CallNode。
- 通过 Pass Manager 注册 pass,可以将 pass 应用于 Relay 模块中的所有函数。
- FoldConstant 是一个独立函数,负责接收表达式并创建 ConstantFolder 实例。
- 可以通过 TVM API 注册 Pass 对象,并使用 Sequential 结构按顺序应用多个 pass。
🏷️
标签
➡️