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来进行一个节点转换。这块后边再详细看。