TVM代码流程分析

TVM代码流程分析

TVM - 代码生成流程

本节主要介绍TVM的代码生成流程，即调用relay.build或tvm.build之后发生了什么，将深入到TVM的源代码进行剖析。（这里采用的依然是TVM v0.6）

首先区分两个build的区别：tvm.build主要针对单一算子（参照Tensor Expression一文），而relay.build是针对整个模型进行编译（参照GCN优化一文），而Relay最后也会调用到tvm::build做代码生成。

relay.build

通常的模型编译由以下两条语句完成。

# Build with Relay
 


    

 

 跟踪细节

 

那么对relay.build进行跟踪，跳转进来是python/tvm/relay/build_module.py（这里是因为在relay/__init__.py中将build函数直接import到relay的命名空间，因此跳过了build_module这一层），其中的build函数是build_module内的全局函数(helper)。

 
 
 
 


    # do somthing
 
    
 
 


        
 
 


    


        
 
 


 
    
 


        
 
 


        
 
 
 
 
 
 
 


    

首先是寻找AutoTVM是否有预先tune好的参数记录，然后构造tophub_context，在其内部构建了BuildModule之后，才跳转到BuildModule.build，然后返回BuildModule.__init__中的内容。

 


    """Build a Relay function to run on TVM graph runtime. This class is used
    to expose the `RelayBuildModule` APIs implemented in C++.
    """
    
 


        
 
 


        
 
 


        
 
 


        
 
 


        
 
 


        
 
 


        
 
 


 
    
 
 
 
 
 


        
 
 


 
        # Setup the params.
        
 


            


        # Build the function
        
 
 


        # Get artifacts
        
 
 


        
 
 


        
 
 


 
        

而_build_module._BuildModule()又通过FFI在python/tvm/relay/_build_module.py中与C++函数建立联系（tvm._ffi._cytpes.function.Function.__call__）。

对应的C++函数在src/relay/backend/build_module.cc

 
 


  
 
 
 


  
 


}
 
TVM_REGISTER_GLOBAL("relay.build_module._BuildModule")
 
 
 
 


  
 
 


});

也就是注册了一个RelayBuildModule供调用，由于主要用的是build函数，因此到RelayBuildModule中找对应的函数。这里TVM又用PackedFunc做了一层封装，见下。

 
 
 


                         
 
 
 
 


      


      
 
 
 
 


      
 
 
 
 
 
 


        
 


        
 
 


      


      


}

也就是调用的是this->Build，再跳转过去会指向BuildRelay。

  
 


      
 


      
 
 
 
 


    


    
 
 
 
 
 


    


    
 
 


 
    


    
 
 
 


    
 


    


 
    
 
 


    
 
 


 
    
 
 
 


    
 
 
 
 


      
 
 


    
 
 


      
 
 


        


        


        


    


  

经过多番跳转，终于到达build的核心模块，再来看TVM逐步做的工作。

  1. 优化
  2. 计算图生成
  3. 后端代码生成

优化

先是优化Optimize，可以看到这里的优化主要是设备无关的优化，是graph-level的针对tensor运算的优化。（这里的优化pass都已经在C++中实现，先前版本的NNVM似乎还是在Python中调用）

  
 


      
 


      
 
 


      
 
 
 
 


    


 
    


    
 
 
 


 
    
 


    


    


    


    


    


    


    


    


    


    


    


    


    


    


    


 
    


    
 
 
 


    


    
 
 


 
    


    
 
 
 


    
 
 
 
 


      
 


          
 


    


 
    


    
 
 


    
 
 


    


 
    
 


  

计算图生成

对应GraphCodegen类，以同样的方式调用src/relay/backend/build_module.cc中的relay.build_module._GraphRuntimeCodegen（一样是FFI），然后跳转至src/relay/backend/graph_runtime_codegen.cc，其中已经用TVM_REGISTER_GLOBAL注册了对应函数，即用GraphRuntimeCodegenModule生成对应Object。

因此实际graph_codegen_->Codegen的函数是一个PackedFunc，定义在GraphRuntimeCodegen.Codegen，用来将relay::Function func进行遍历，然后生成计算图。

后端代码生成

Relay得到lower后的函数，最后一步则是交给tvm::build做代码生成，跳转到src/codegen/build_module.cc中的build函数（注意这里重载了几个版本），然后跳转到核心build，注意这里的build函数支持异构编译，只要再inputs划分好不同硬件设施即可。

// Build for heterogeneous execution.
 
 
 
 


                      
 
 


                      
 
 
 


  
 


  
 


 
  
 
 
 


  
 
 


    
 
 
 
 
 
 


      
 
 
 
 


        
 
 


        


      


    


  


 
  
 
 


    
 
 


  


 
  
 
 
 
 
 
 


    
 
 


        
 
 
 


    
 
 
 


    
 
 
 


    


    
 
 
 
 


    
 
 
 
 
 
 


      


    


    


  


 
  
 
 
 
 


  


  
 
 
 
 
 
 


    
 
 


      


    


  


  
 


}

当中最最核心的则是mhost = codegen::Build，最后跳转过去就开始调用代码生成模块了（src/codegen/codegen.cc）。

 
 
 


                      
 
 
 


  


 
  
 
 
 
 
 


  


  
 
 
 
 


  
 
 
 
 


                      
 
 


                      
 


  
 


}

以生成LLVM IR为例，codegen.build_llvm会在src/codegen/llvm/llvm_module.cc注册，然后调用同个文件中的LLVMModuleNode->Init。这时会跳转到src/codegen/llvm/codegen_llvm.cc中的CodeGenLLVM类进行代码生成。

tvm.build

用tvm.build对算子进行编译则是按照以下方式进行调用，例子来自Tensor Expression。

调用tvm.build后首先跳转到python/tvm/build_module.py，其中的build函数主要做两个步骤：

  1. lower高层次代码
  2. 后端代码生成

代码变换

lower高层次代码对应的是

而lower函数同样在python/tvm/build_module.py中，类似于relay.build中的Optimize，但这里执行的是operator-level的优化，主要针对循环变换。

 


          


          


          


          


 
    # initialization
 
    # Phase 0
    
 
 


        
 
 


 
    
 
 
 


        
 
 


 
    
 
 


    
 
 
 
 
 


 
    # Phase 1
    
 
 
 
 


    
 
 
 
 
 


    
 
 


    
 
 
 


        
 
 


 
    # Phase 2
    
 
 


        
 
 
 


    
 


        
 
 


    


        
 
 


    
 
 


    
 
 
 


    
 
 


    
 
 


        


        


        


        


        


    
 
 
 


        
 
 


 
    # Phase 3
    
 
 


    
 
 


    
 
 


        
 
 


    
 
 
 


        
 
 


    # Instrument BoundCheckers
    
 


        
 
 


    
 


        
 


 
    

优化Pass的主体实施都在src/api/api_pass.cc中，以tvm.ir_pass进行注册（注意由于C++函数中已经在tvm的命名空间里，故搜索时直接搜ir_pass才会出来对应的API）。

代码生成

lower完之后就进入到后端代码生成，对应build函数中的

同样的原理，跳转至tvm/codegen.py，初始化tvm.codegen的API codegen._Build，调用FFI，跳转至src/api/api_codegen.cc，最后跳转至src/codegen/codegen.cc中的tvm::Build，之后的后端代码生成则与relay.build相同。

TVM - Tensor Expression

本节以向量加法为例，记录TVM最最基本的Tensor Expression的使用，以及简单的编译运行流程。

下面的代码为简单的向量加法，参考自Tensor Expression官方教程，在TVM v0.6下执行（注意与v0.7dev的模块有区别)。

 


 
 
 


 
# Tensor Expression
# args: (shape, label)
 
 
 


 
 
 


# args: (shape, function, label)
# function represented in lambda expression (element-wise)
#     lambda axis1, axis2, ... : f(axis1, axis2, ...)
 
 
 
 
 
 
 
 


 
# generate schedule
 
 


# print low level codes
print(tvm.lower(s,[A,B,C],simple_mode=True))

其中placeholder代表特定维度的张量，最后生成的代码会要求用户输入两个tensor，如果是C++代码，则要求用户输入两个float*。注意，会发现这个过程实际上是没有计算发生的，而只是定义了计算如何进行。

输出的low-level代码如下所示，还是相当好理解的，即i从0到10循环，循环内每次计算C[i]的值。

 
 


  
 
 
 
 


    
 
 
 
 


  


}

一些常用的循环优化API可以在这里找到。这里使用循环分割split作为尝试。

split(parent[, factor, nparts])

Split the stage either by factor providing outer scope, or both. Return outer, innervaiable of iteration.

 
 
 


print(tvm.lower(s,[A,B,C],simple_mode=True))

由于对schedule的操作是原地变换，因此可以直接输出lower后的代码，发现确实已经改变了，原来的循环体变成5*2的循环。

 
 


  
 
 
 
 


    
 
 
 
 


      
 
 
 
 
 
 
 
 
 
 


    


  


}

当然这一个schedule变换并没有带来任何好处，只是为了说明Tensor Expression应该怎么用。

之后就可以调用build生成目标代码了，可以设置target和target_host。

然后可以创造运行时环境，进行运行测试。

 
 


 
 


 
 
 


 
 
 


 
 
 


fadd(a,b,c) # run
 
# test
 
 


print(fadd.get_source())

生成的C代码如下

 
 
 
 
 
 
 
 
 


  
 
 
 
 
 
 
 
 
 


    
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 


  


}

 

 生成的myadd.c完整代码如下

 

最后通过fadd.save("myadd.c")保存文件。

TVM - Relay IR Pass

本节介绍Relay IR Pass的构造。

Relay IR Pass核心依然是在C++中实现，但提供了Python接口，方便上层直接调用并对计算流图进行变换优化。

Pass管理器在include/tvm/relay/transform.h中，里面包含所有Pass的声明，希望做到

  • 管理调度不同的优化pass
  • 收集需要的分析信息，并且保持是最新的
  • 减少程序员实现新pass的麻烦

Python的接口函数声明在python/tvm/relay/transform.py中，在python/tvm/relay/_transform.py中通过FFI对C++函数进行调用，命名空间为relay._transform。

具体C++的实现则分为两个部分：

  • 高层IR图变换，源码在src/relay/pass中，集中变换则是在src/relay/backend/build_module.cc中的relay::Module Optimize
  • 后端代码的图变换，源码在src/relay/backend/vm中，集中变换在python/tvm/build_module.py中的lower函数

Pass的构造

  • PassInfo
·       
  
 
 
 
 
 


·       
  
 


·       
  
 


·       
  
 


·       
  
  • PassContext
·       
  
 
 
 
 
 


·       
  


·       
  
 


·       
  
 


·       
  
 


·       
  
 


·       
  
 


·       
  


·       
 

·       
  
 
 
 
 
 


·       
  


·       
  
 
 
 


·       
  
 
 
 


·       
  


·       
 

·       
  


·       
  


·       
  
 
 


·       
  


·       
  
 
 


·       
 

·       
  


·       
  
 
 


·       
  


·       
 

·       
  
 
 


·       
  


·       
  
 


·       
  


·       
  
 


·       
  
 


·       
      
 
 


·       
  


·       
  


·       
 

·       
  


·       
  
 


·       
      
  • Pass Constructs：提供基类
·       
  
 
 
 
 


·       
  
 
 
 
 
 


·       
  
 
 
 
 


·       
                              
 
 
 
 
 


·       
  

也就是说，一个Pass一定是作用在特定context下的IRModule，所有Pass都设计成Module到Module的映射，完整Pass的定义在src/relay/ir/transform.cc和src/ir/transform.cc中。

Module-Level

 
 
 
 


  
 


  
 
 


  
 
 
 
 
 
 
 
 


  


};

Function-Level

 
 
 
 


  
 


  
 
 
 


  
 
 
 
 
 
 
 
 


  
 
 
 
 


  


};

Sequential

类似于PyTorch中的nn.Sequential，顺序执行多个Pass

 
 
 
 


  
 


  


  
 


  
 
 
 
 


  
 
 
 
 
 
 
 
 


};

 

References

 

 

posted @   吴建明wujianming  阅读(708)  评论(1)  编辑  收藏  举报
编辑推荐：
· 记一次.NET内存居高不下排查解决与启示
· 探究高空视频全景AR技术的实现原理
· 理解Rust引用及其生命周期标识（上）
· 浏览器原生「磁吸」效果！Anchor Positioning 锚点定位神器解析
· 没有源码，如何修改代码逻辑？
阅读排行：
· 全程不用写代码，我用AI程序员写了一个飞机大战
· DeepSeek 开源周回顾「GitHub 热点速览」
· 记一次.NET内存居高不下排查解决与启示
· MongoDB 8.0这个新功能碉堡了，比商业数据库还牛
· .NET10 - 预览版1新功能体验（一）
点击右上角即可分享
微信分享提示