TVM学习资料TensorIR快速入门

　　本篇文章译自英文文档BlitzCoursetoTensorI作者是冯思远。更多TVM中文文档可访问TVM中文站https：hyper。ai
　　TensorIR是深度学习领域的特定语言，主要有两个作用：在各种硬件后端转换和优化程序。自动tensorized程序优化的抽象。importtvmfromtvm。ir。moduleimportIRModulefromtvm。scriptimporttirasTimportnumpyasnpIRModule
　　IRModule是TVM的核心数据结构，它包含深度学习程序，并且是IR转换和模型构建的基础。
　　编辑切换为居中
　　添加图片注释，不超过140字（可选）
　　上图展示的是IRModule的生命周期，它可由TVMScript创建。转换IRModule的两种主要方法是TensorIR的schedule原语转换和pass转换。此外，也可直接对IRModule进行一系列转换。注意，可以在任何阶段将IRModule打印到TVMScript。完成所有转换和优化后，可将IRModule构建为可运行模块，从而部署在目标设备上。
　　基于TensorIR和IRModule的设计，可创建一种新的编程方法：基于PythonAST语法，用TVMScript编写程序。使用PythonAPI转换和优化程序。使用命令式转换API交互检查和提高性能。创建IRModule
　　IRModule是TVMIR的一种可往返语法，可通过编写TVMScript来创建。
　　与通过张量表达式创建计算表达式（使用张量表达式操作算子）不同，TensorIR允许用户通过TVMScript（一种嵌在PythonAST中的语言）进行编程。新方法可以编写复杂的程序并进一步调度和优化。
　　下面是向量加法的示例：tvm。script。irmoduleclassMyModule：T。primfuncdefmain（a：T。handle，b：T。handle）：我们通过T。handle进行数据交换，类似于内存指针T。funcattr（｛globalsymbol：main，tir。noalias：True｝）通过handle创建BufferAT。matchbuffer（a，（8，），dtypefloat32）BT。matchbuffer（b，（8，），dtypefloat32）foriinrange（8）：block是针对计算的抽象withT。block（B）：定义一个空间（可并行）block迭代器，并且将它的值绑定成iviT。axis。spatial（8，i）B〔vi〕A〔vi〕1。0irmoduleMyModuleprint（type（irmodule））print（irmodule。script（））
　　输出结果：classtvm。ir。module。IRModulefromtvm。scriptimporttirasTtvm。script。irmoduleclassModule：T。primfuncdefmain（A：T。Buffer〔8，float32〕，B：T。Buffer〔8，float32〕）None：functionattrdictT。funcattr（｛globalsymbol：main，tir。noalias：True｝）bodywithT。block（root）foriinT。serial（8）：withT。block（B）：viT。axis。spatial（8，i）T。reads（A〔vi〕）T。writes（B〔vi〕）B〔vi〕A〔vi〕T。float32（1）
　　此外，我们还可以使用张量表达式DSL编写简单的运算符，并将它们转换为IRModule。fromtvmimportteAte。placeholder（（8，），dtypefloat32，nameA）Bte。compute（（8，），lambdai：A（i）1。0，nameB）functe。createprimfunc（〔A，B〕）irmodulefromteIRModule（｛main：func｝）print（irmodulefromte。script（））
　　输出结果：fromtvm。scriptimporttirasTtvm。script。irmoduleclassModule：T。primfuncdefmain（A：T。Buffer〔8，float32〕，B：T。Buffer〔8，float32〕）None：functionattrdictT。funcattr（｛globalsymbol：main，tir。noalias：True｝）bodywithT。block（root）fori0inT。serial（8）：withT。block（B）：i01T。axis。spatial（8，i0）T。reads（A〔i01〕）T。writes（B〔i01〕）B〔i01〕A〔i01〕T。float32（1）构建并运行IRModule
　　可将IRModule构建为特定target后端的可运行模块。modtvm。build（irmodule，targetllvm）CPU后端的模块print（type（mod））
　　输出结果：classtvm。driver。buildmodule。OperatorModule
　　准备输入数组和输出数组，然后运行模块：atvm。nd。array（np。arange（8）。astype（float32））btvm。nd。array（np。zeros（（8，））。astype（float32））mod（a，b）print（a）print（b）
　　输出结果：〔0。1。2。3。4。5。6。7。〕〔1。2。3。4。5。6。7。8。〕转换IRModule
　　IRModule是程序优化的核心数据结构，可通过Schedule进行转换。schedule包含多个primitive方法来交互地转换程序。每个primitive都以特定方式对程序进行转换，从而优化性能。
　　上图是优化张量程序的典型工作流程。首先，用TVMScript或张量表达式创建一个初始IRModule，然后在这个初始IRModule上创建schedule。接下来，使用一系列调度原语来提高性能。最后，我们可以将其降低并构建成一个可运行模块。
　　上面只演示了一个简单的转换。首先，在输入irmodule上创建schedule：schtvm。tir。Schedule（irmodule）print（type（sch））
　　输出结果：classtvm。tir。schedule。schedule。Schedule
　　将嵌套循环展开成3个循环，并打印结果：通过名字获取blockblockbsch。getblock（B）获取包围block的循环（i，）sch。getloops（blockb）展开嵌套循环i0，i1，i2sch。split（i，factors〔2，2，2〕）print（sch。mod。script（））
　　输出结果：fromtvm。scriptimporttirasTtvm。script。irmoduleclassModule：T。primfuncdefmain（A：T。Buffer〔8，float32〕，B：T。Buffer〔8，float32〕）None：functionattrdictT。funcattr（｛globalsymbol：main，tir。noalias：True｝）bodywithT。block（root）fori0，i1，i2inT。grid（2，2，2）：withT。block（B）：viT。axis。spatial（8，i04i12i2）T。reads（A〔vi〕）T。writes（B〔vi〕）B〔vi〕A〔vi〕T。float32（1）
　　还可对循环重新排序。例如，将循环i2移到i1之外：sch。reorder（i0，i2，i1）print（sch。mod。script（））
　　输出结果fromtvm。scriptimporttirasTtvm。script。irmoduleclassModule：T。primfuncdefmain（A：T。Buffer〔8，float32〕，B：T。Buffer〔8，float32〕）None：functionattrdictT。funcattr（｛globalsymbol：main，tir。noalias：True｝）bodywithT。block（root）fori0，i2，i1inT。grid（2，2，2）：withT。block（B）：viT。axis。spatial（8，i04i12i2）T。reads（A〔vi〕）T。writes（B〔vi〕）B〔vi〕A〔vi〕T。float32（1）转换为GPU程序
　　要在GPU上部署模型必须进行线程绑定。幸运的是，也可以用原语来增量转换。sch。bind（i0，blockIdx。x）sch。bind（i2，threadIdx。x）print（sch。mod。script（））
　　输出结果：fromtvm。scriptimporttirasTtvm。script。irmoduleclassModule：T。primfuncdefmain（A：T。Buffer〔8，float32〕，B：T。Buffer〔8，float32〕）None：functionattrdictT。funcattr（｛globalsymbol：main，tir。noalias：True｝）bodywithT。block（root）fori0inT。threadbinding（2，threadblockIdx。x）：fori2inT。threadbinding（2，threadthreadIdx。x）：fori1inT。serial（2）：withT。block（B）：viT。axis。spatial（8，i04i12i2）T。reads（A〔vi〕）T。writes（B〔vi〕）B〔vi〕A〔vi〕T。float32（1）
　　绑定线程后，用cuda后端来构建IRModule：ctxtvm。cuda（0）cudamodtvm。build（sch。mod，targetcuda）cudaatvm。nd。array（np。arange（8）。astype（float32），ctx）cudabtvm。nd。array（np。zeros（（8，））。astype（float32），ctx）cudamod（cudaa，cudab）print（cudaa）print（cudab）
　　输出结果：〔0。1。2。3。4。5。6。7。〕〔1。2。3。4。5。6。7。8。〕
　　下载Python源代码：tensorirblitzcourse。py
　　下载JupyterNotebook：tensorirblitzcourse。ipynb
　　以上就是该文档的全部内容，查看更多TVM中文文档，请访问