TVM學習（四）codegen

作者：安平博，Xilinx高級工程師；來源：AI加速微信公眾號

接著上一章繼續(xù)深入代碼，在BuildRelay中會調用Codegen函數(shù)。這個函數(shù)實現(xiàn)在src/relay/backend/graph_runtime_codegen.cc中。Codegen實現(xiàn)了內存的分配，IR節(jié)點到TIR節(jié)點的轉換，tir圖節(jié)點的一個調度優(yōu)化。內存分配由函數(shù)relay.backend.GraphPlanMemory來實現(xiàn)，VisitExpr對節(jié)點進行遍歷并進行節(jié)點信息的記錄。LowerExternalfunctions完成ir節(jié)點到tir節(jié)點的轉化以及schedule的優(yōu)化。

內存分配

通過GetPackedFunc函數(shù)來獲得注冊到global map的內存分配函數(shù)GraphPlanMemory。我們看一下文件src/relay/backend/graph_plan_memory.cc中對內存的處理。

在處理內存分配中主要使用了StorageAllocaBaseVisitor，StorageAllocaInit，StorageAllocator這三個類。StorageAllocaBaseVisitor是一個基類，實現(xiàn)了對每個節(jié)點的訪問，并分配token，但是token中信息是在派生類中處理的。定義了一個StorageToken的結構體，用于表示申請到內存的大小，類型等信息。在內存處理程序中，主要就是為每個節(jié)點分配這個token，同時定義token的內部信息。內存分配結果是一個節(jié)點和token的映射表。

StorageAllocator類中Plan函數(shù)為：

關鍵是前兩行代碼，第一行代碼初始化了storageToken，賦予了其設備類型和數(shù)據(jù)類型信息。第二行代碼遍歷每個節(jié)點，并且為每個節(jié)點分配內存空間。在內存初始化函數(shù)GetInitTokenMap中，首先收集每個節(jié)點的的設備信息。調用鏈為CollectDeviceInfo -> GetDeviceMap(src/relay/transforms/device_annotation.cc)。在構建relay圖結構的時候，每個節(jié)點是有設備號信息的，GetDeviceMap就是按照post-DFS順序獲得節(jié)點的設備號信息。當然并不是所有節(jié)點都有設備號信息，所以還需要根據(jù)節(jié)點之間的關系來推斷出設備號。比如下圖，add，sqrt，log節(jié)點被標注為1，2，3號設備，那么可以用兩種方式來推斷其它節(jié)點設備號。

1） 從一個copy節(jié)點由下而上遍歷一直到遇到下一個copy，比如可以推斷出add，x，y節(jié)點的設備號和copy1一樣；
2） 從最后一個copy節(jié)點向下遍歷，那么可以推斷出substract，exp設備號和copy3一樣。

設備號獲得后，this->run會調用基類的run函數(shù)，基類run函數(shù)會調用派生類的CreateToken函數(shù)。CreateToken會申請StorageToken空間并且賦予設備號和數(shù)據(jù)類型，然后返回一個token_map_。和節(jié)點遍歷相關函數(shù)為Run->GetToken->VisitExpr。VisitExpr會最終調用StorageAllocaInit類中定義的VisitExpr_函數(shù)來遍歷節(jié)點。

節(jié)點內存初始化完成后，回到StorageAllocator類中，run會調用其定義的CreateToken函數(shù)。

分配內存空間會有兩種情況，一種是can_realloc一種是不能can_realloc的。先看不can_realloc的，GetMemorySize是根據(jù)token中記錄的數(shù)據(jù)類型和shape信息來獲得數(shù)據(jù)的大小，Alloc函數(shù)就是為tok分配字節(jié)數(shù)量?，F(xiàn)在看can_realloc的情況，Request中首先獲取節(jié)點數(shù)據(jù)的大小。然后從free_中查詢能夠滿足size的節(jié)點，如果有比該節(jié)點size大的就選擇大的空閑區(qū)間分配，如果沒有大的空間分配，選擇最接近的空間分配。然后最終返回一個token_map_。

codegen

第一步是對ir節(jié)點進行遍歷，轉換成codegen中定義的基礎節(jié)點。我們先看以下codegen中定義的節(jié)點類型，GraphNode是基礎節(jié)點，GraphInputNode, GraphOpNode繼承自這個基礎節(jié)點。這些節(jié)點中主要提供了一些節(jié)點屬性，比如name，op類型等。還提供了dmlc接口，可以實現(xiàn)可視化。

遍歷func的parameters，將parameters轉換到graph的input節(jié)點。通過AddNode添加這些input節(jié)點，并且將轉換后的graphInputNode加入var_map_中，var_map_中是expr到graphNode的映射。

接下來是節(jié)點遍歷，heads_=VisitExpr(func->body)。節(jié)點遍歷過程中會將func中的節(jié)點轉換為graphNode。對于varNode，因為已經記錄在var_map_中，直接返回引用。ConstantNode會轉換為GraphInputNode，tuppleNode會返回每個字段的graphNode。在遍歷節(jié)點過程中，會將graphNode都添加到nodes_中。

重點看一下對CallNode的處理，只支持op是functionNode類型的。

Function生成時，走兩個分支，一個是外部codegen，一個是通用分支。對應外部function codegen的處理為：

首先創(chuàng)建一個CCacheKey類型作為_CompileEngineLower函數(shù)的參數(shù)傳入。具體CcacheKey有什么作用，以后再深入研究吧。_CompileEngineLower的實現(xiàn)在文件src/relay/backend/compile_engine.cc中。調用鏈為Lower -> LowerInternal(key)->cached_func。定義了一個cache_node并封裝成cached_func返回。這塊具體的操作并不是很理解，可能還需要熟悉cachedFuncNode的作用。

然后通過GraphAddCallNode將其加入nodes_中。在GraphAddCallNode中還會對op->args進行深入遍歷。

內部func處理如下：

也是通過相同的pf0和pf1函數(shù)。CcacheKey的創(chuàng)建過程一樣，但是在lowerInternal中不一樣。

首先創(chuàng)建了一個schedule，schedule的具體實現(xiàn)很復雜目前還不夠理解。

如果是copy節(jié)點，那么不進行l(wèi)ower處理，直接返回CachedFunc封裝。不是copy節(jié)點，如果我們在python中自己定義了lower函數(shù)就調用python中的，如果沒有就會調用TVM中的lower函數(shù)。Lower函數(shù)在src/driver/driver_api.cc文件中。在這里調用了很多tir的passes來進行一個節(jié)點轉換。這塊后邊再詳細看。