【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。
➡️

继续阅读