【TVM 教程】使用元组输入(Tuple Inputs)进行计算和归约
💡
原文中文,约5800字,阅读约需14分钟。
📝
内容提要
本文介绍了Apache TVM中元组输入的用法,适用于批量计算和归约操作。通过示例展示了如何使用元组输入进行多个输出的调度和计算,强调调度应基于操作而非张量。
🎯
关键要点
-
Apache TVM 是一个深度学习编译框架,支持 CPU、GPU 和各种加速芯片。
-
元组输入用于在单个循环中计算多个输出或执行归约操作,如 argmax。
-
通过将相同 shape 的算子作为 te.compute 的输入,可以进行批量计算。
-
使用协同输入描述归约操作,能够比较操作数的值并保留索引。
-
调度操作时,多个输出只能一起调度,调度应基于操作而非张量。
❓
延伸问答
Apache TVM 是什么?
Apache TVM 是一个端到端的深度学习编译框架,支持 CPU、GPU 和各种机器学习加速芯片。
元组输入在 TVM 中的主要用途是什么?
元组输入用于在单个循环中计算多个输出或执行归约操作,如 argmax。
如何使用元组输入进行批量计算?
可以将相同 shape 的算子作为 te.compute 的输入,进行批量计算。
什么是协同输入,它在归约操作中有什么作用?
协同输入用于描述归约操作,能够比较操作数的值并保留索引。
在调度操作时,元组输入有什么限制?
多个输出只能一起调度,调度应基于操作而非张量。
如何在 TVM 中实现 argmax 的归约操作?
可以使用 te.comm_reducer() 来描述 argmax 的归约操作,比较操作数的值并保留索引。
➡️