【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 接口构建新的循环结构来实现。
➡️