作者：安平博，Xilinx高級(jí)工程師；來(lái)源：AI加速微信公眾號(hào)

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

內(nèi)存分配

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

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

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

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

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

設(shè)備號(hào)獲得后，this->run會(huì)調(diào)用基類(lèi)的run函數(shù)，基類(lèi)run函數(shù)會(huì)調(diào)用派生類(lèi)的CreateToken函數(shù)。CreateToken會(huì)申請(qǐng)StorageToken空間并且賦予設(shè)備號(hào)和數(shù)據(jù)類(lèi)型，然后返回一個(gè)token_map_。和節(jié)點(diǎn)遍歷相關(guān)函數(shù)為Run->GetToken->VisitExpr。VisitExpr會(huì)最終調(diào)用StorageAllocaInit類(lèi)中定義的VisitExpr_函數(shù)來(lái)遍歷節(jié)點(diǎn)。

節(jié)點(diǎn)內(nèi)存初始化完成后，回到StorageAllocator類(lèi)中，run會(huì)調(diào)用其定義的CreateToken函數(shù)。

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

codegen

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

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

接下來(lái)是節(jié)點(diǎn)遍歷，heads_=VisitExpr(func->body)。節(jié)點(diǎn)遍歷過(guò)程中會(huì)將func中的節(jié)點(diǎn)轉(zhuǎn)換為graphNode。對(duì)于varNode，因?yàn)橐呀?jīng)記錄在var_map_中，直接返回引用。ConstantNode會(huì)轉(zhuǎn)換為GraphInputNode，tuppleNode會(huì)返回每個(gè)字段的graphNode。在遍歷節(jié)點(diǎn)過(guò)程中，會(huì)將graphNode都添加到nodes_中。

重點(diǎn)看一下對(duì)CallNode的處理，只支持op是functionNode類(lèi)型的。

Function生成時(shí)，走兩個(gè)分支，一個(gè)是外部codegen，一個(gè)是通用分支。對(duì)應(yīng)外部function codegen的處理為：

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

然后通過(guò)GraphAddCallNode將其加入nodes_中。在GraphAddCallNode中還會(huì)對(duì)op->args進(jìn)行深入遍歷。

內(nèi)部func處理如下：

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

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

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