TVM学习(四)codegen

作者:安平博,Xilinx高级工程师;来源:AI加速微信公众号

接着上一章继续深入代码,在BuildRelay中会调用Codegen函数。这个函数实现在src/relay/backend/graph_runtime_codegen.cc中。Codegen实现了内存的分配,IR节点到TIR节点的转换,tir图节点的一个调度优化。内存分配由函数relay.backend.GraphPlanMemory来实现,VisitExpr对节点进行遍历并进行节点信息的记录。LowerExternalfunctions完成ir节点到tir节点的转化以及schedule的优化。

内存分配

通过GetPackedFunc函数来获得注册到global map的内存分配函数GraphPlanMemory。我们看一下文件src/relay/backend/graph_plan_memory.cc中对内存的处理。

在处理内存分配中主要使用了StorageAllocaBaseVisitor,StorageAllocaInit,StorageAllocator这三个类。StorageAllocaBaseVisitor是一个基类,实现了对每个节点的访问,并分配token,但是token中信息是在派生类中处理的。定义了一个StorageToken的结构体,用于表示申请到内存的大小,类型等信息。在内存处理程序中,主要就是为每个节点分配这个token,同时定义token的内部信息。内存分配结果是一个节点和token的映射表。

StorageAllocator类中Plan函数为:

关键是前两行代码,第一行代码初始化了storageToken,赋予了其设备类型和数据类型信息。第二行代码遍历每个节点,并且为每个节点分配内存空间。在内存初始化函数GetInitTokenMap中,首先收集每个节点的的设备信息。调用链为CollectDeviceInfo -> GetDeviceMap(src/relay/transforms/device_annotation.cc)。在构建relay图结构的时候,每个节点是有设备号信息的,GetDeviceMap就是按照post-DFS顺序获得节点的设备号信息。当然并不是所有节点都有设备号信息,所以还需要根据节点之间的关系来推断出设备号。比如下图,add,sqrt,log节点被标注为1,2,3号设备,那么可以用两种方式来推断其它节点设备号。

1) 从一个copy节点由下而上遍历一直到遇到下一个copy,比如可以推断出add,x,y节点的设备号和copy1一样;
2) 从最后一个copy节点向下遍历,那么可以推断出substract,exp设备号和copy3一样。

设备号获得后,this->run会调用基类的run函数,基类run函数会调用派生类的CreateToken函数。CreateToken会申请StorageToken空间并且赋予设备号和数据类型,然后返回一个token_map_。和节点遍历相关函数为Run->GetToken->VisitExpr。VisitExpr会最终调用StorageAllocaInit类中定义的VisitExpr_函数来遍历节点。

节点内存初始化完成后,回到StorageAllocator类中,run会调用其定义的CreateToken函数。

分配内存空间会有两种情况,一种是can_realloc一种是不能can_realloc的。先看不can_realloc的,GetMemorySize是根据token中记录的数据类型和shape信息来获得数据的大小,Alloc函数就是为tok分配字节数量。现在看can_realloc的情况,Request中首先获取节点数据的大小。然后从free_中查询能够满足size的节点,如果有比该节点size大的就选择大的空闲区间分配,如果没有大的空间分配,选择最接近的空间分配。然后最终返回一个token_map_。

codegen

第一步是对ir节点进行遍历,转换成codegen中定义的基础节点。我们先看以下codegen中定义的节点类型,GraphNode是基础节点,GraphInputNode, GraphOpNode继承自这个基础节点。这些节点中主要提供了一些节点属性,比如name,op类型等。还提供了dmlc接口,可以实现可视化。

遍历func的parameters,将parameters转换到graph的input节点。通过AddNode添加这些input节点,并且将转换后的graphInputNode加入var_map_中,var_map_中是expr到graphNode的映射。

接下来是节点遍历,heads_=VisitExpr(func->body)。节点遍历过程中会将func中的节点转换为graphNode。对于varNode,因为已经记录在var_map_中,直接返回引用。ConstantNode会转换为GraphInputNode,tuppleNode会返回每个字段的graphNode。在遍历节点过程中,会将graphNode都添加到nodes_中。

重点看一下对CallNode的处理,只支持op是functionNode类型的。

Function生成时,走两个分支,一个是外部codegen,一个是通用分支。对应外部function codegen的处理为:

首先创建一个CCacheKey类型作为_CompileEngineLower函数的参数传入。具体CcacheKey有什么作用,以后再深入研究吧。_CompileEngineLower的实现在文件src/relay/backend/compile_engine.cc中。调用链为Lower -> LowerInternal(key)->cached_func。定义了一个cache_node并封装成cached_func返回。这块具体的操作并不是很理解,可能还需要熟悉cachedFuncNode的作用。

然后通过GraphAddCallNode将其加入nodes_中。在GraphAddCallNode中还会对op->args进行深入遍历。

内部func处理如下:

也是通过相同的pf0和pf1函数。CcacheKey的创建过程一样,但是在lowerInternal中不一样。

首先创建了一个schedule,schedule的具体实现很复杂目前还不够理解。

如果是copy节点,那么不进行lower处理,直接返回CachedFunc封装。不是copy节点,如果我们在python中自己定义了lower函数就调用python中的,如果没有就会调用TVM中的lower函数。Lower函数在src/driver/driver_api.cc文件中。在这里调用了很多tir的passes来进行一个节点转换。这块后边再详细看。

最新文章