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

延伸问答

什么是 Apache TVM?

Apache TVM 是一个支持 CPU、GPU 和加速芯片的深度学习编译框架。

Compiler Pass 在 Relay 中的作用是什么?

Compiler Pass 是优化 Relay 程序的主要接口,用户可以通过编写 pass 来修改 AST。

如何编写一个常量折叠 pass?

编写常量折叠 pass 包括创建 ConstantChecker 和 ConstantFolder 访问器与修改器,并使用 Pass Manager 注册。

AST 遍历器的基类是什么?

AST 遍历器的基类是 ExprFunctor,它提供 VisitExpr 方法用于遍历表达式。

ConstantChecker 的主要功能是什么?

ConstantChecker 用于检查表达式是否为常量,并使用 memo_ 字段缓存结果。

如何在 TVM 中注册一个 pass?

可以通过 Pass Manager 注册 pass,使用 CreateFunctionPass 函数并提供相应的优化级别和依赖项。

➡️

继续阅读