【TVM 教程】编写自定义 Pass

💡 原文中文,约4300字,阅读约需11分钟。
📝

内容提要

Apache TVM是一个支持CPU和GPU等加速芯片的深度学习编译框架。本文介绍如何在TVM中编写自定义IR转换pass,以满足专用硬件的需求,包含向量加法示例及IR节点的分析与转换方法。

🎯

关键要点

  • Apache TVM是一个支持CPU和GPU等加速芯片的深度学习编译框架。
  • 用户可以自定义分析和IR转换以适应专用硬件。
  • 教程假设读者熟悉TVM算法编写、HalideIR结构和IR节点属性。
  • 示例中展示了如何编写简单的向量加法并构建默认调度。
  • IR转换pass是将语句映射到新语句的函数。
  • TVM提供两个类用于分析和转换IR。
  • 使用post_order_visit收集IR节点信息,需存储结果于数组中。
  • 转换接口支持前序和后序回调,返回None保留原始节点。
  • 编写向量化函数以处理可被8整除的循环。
  • 通过元组列表将自定义pass注入TVM的标准较低级pass。
  • 转换pass适合放在第1阶段之后,完成原始IR生成和循环级别。
  • 总结编写自定义IR转换pass的步骤,包括信息收集和节点转换。

延伸问答

Apache TVM 是什么?

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

如何在 TVM 中编写自定义 IR 转换 pass?

在 TVM 中编写自定义 IR 转换 pass 需要使用 tvm.tir.stmt_functor.post_order_visit 收集 IR 节点信息,并使用 tvm.tir.stmt_functor.ir_transform 转换这些节点。

什么是 IR 转换 pass?

IR 转换 pass 是将语句映射到新语句的函数,用于处理和优化 IR 节点。

如何收集 IR 节点的信息?

可以使用 tvm.tir.stmt_functor.post_order_visit 函数来收集 IR 节点的信息,并在回调函数中存储结果。

自定义的 IR 转换 pass 应该放在哪个阶段?

自定义的 IR 转换 pass 适合放在第 1 阶段之后,完成原始 IR 生成和循环级别。

如何实现向量化函数?

向量化函数可以通过查找可被 8 整除的循环,并使用 TVM IR maker 接口构建新的循环结构来实现。

➡️

继续阅读