背景介紹

TPU-MLIR的CodeGen是BModel生成的最后一步，該過(guò)程目的是將MLIR文件轉(zhuǎn)換成最終的Bmodel。本文介紹了CodeGen的基本原理和流程，并記錄了針對(duì)BM1684X等新架構(gòu)的CodeGen重構(gòu)過(guò)程。

與后端的關(guān)系

由于一些歷史的因素，MLIR文件中的每個(gè)OP對(duì)應(yīng)的指令并不直接在TPU-MLIR工程中生成，而是需要調(diào)用后端的函數(shù)完成最終指令的生成，這也帶來(lái)了兩個(gè)問(wèn)題

1. 如何設(shè)計(jì)編譯器與后端的接口
2. 生成指令的數(shù)據(jù)結(jié)構(gòu)存在后端還是編譯器中

關(guān)于問(wèn)題1，目前的設(shè)計(jì)是采用CodeGen與后端隔離的形式，也就是CodeGen過(guò)程不直接調(diào)用后端函數(shù)，而是將不同處理器的相應(yīng)函數(shù)全部封裝到類中，在CodeGen中調(diào)用類方法間接使用后端接口，達(dá)成解耦。

而關(guān)于問(wèn)題2，依據(jù)不同的處理器其數(shù)據(jù)結(jié)構(gòu)位置也不同，1684的數(shù)據(jù)結(jié)構(gòu)放在編譯器這邊，而BM1684X等新架構(gòu)的處理器數(shù)據(jù)結(jié)構(gòu)放在后端。無(wú)論放在哪里，其全部封裝于問(wèn)題1答案中的相應(yīng)類中，對(duì)于CodeGen過(guò)程來(lái)說(shuō)，看到的接口是一樣的。

一個(gè)OP生成指令的大致流程

代碼位置：lib/Dialect/Tpu/Transforms/Codegen/BM168xCodegen.cpp

該流程忽略CodeGen代碼內(nèi)部細(xì)節(jié)，這里只講解類似于把大象裝冰箱主要分幾步這樣的通俗介紹。

1. 在BM168XCodegen.cpp會(huì)遇到某個(gè)Op時(shí)會(huì)調(diào)用該op的codegen_local_bm168x/codegen_global_bm168x，算子的這個(gè)函數(shù)都在lib/Dialect/Tpu/Interfaces/中
2. 在具體OP中會(huì)設(shè)置一些參數(shù)，然后調(diào)用到后端的具體op的指令生成，比如Conv2d算子會(huì)調(diào)用后端函數(shù)backend_api_conv_global
3. （后端過(guò)程）直接做一系列檢查后，會(huì)直接生成指令（二進(jìn)制碼），這些二進(jìn)制碼會(huì)通過(guò)store_cmd存儲(chǔ)在指定數(shù)據(jù)結(jié)構(gòu)中，
4. 等所有op的二進(jìn)制碼全部都生成完畢后，在編譯器會(huì)調(diào)用BM1684X系列類中封裝的函數(shù)取走指令，生成Bmodel

做個(gè)形象點(diǎn)的例子：

原來(lái)裝冰箱只需要我，現(xiàn)在我嫌大象沉，我叫個(gè)張三幫我裝。
我：張三，你把這個(gè)大象給我裝冰箱里
張三吭哧吭哧幫我裝完了
我：行了，張三，你走吧；我自己把裝的運(yùn)走。

指令生成所需要的數(shù)據(jù)結(jié)構(gòu)

指令依據(jù)處理器的engine不同而有所差別，比如1684有GDMA和TIU，而新架構(gòu)的處理器SG2260會(huì)存在sdma、cdma等engine。這里拿最通用的兩種engine即BDC(后更名為TIU)和GDMA為例：

std::vector<uint32_t>bdc_buffer;
std::vector<uint32_t>gdma_buffer;
uint32_tgdma_total_id=0;
uint32_tbdc_total_id=0;
std::vector<uint32_t>gdma_group_id;
std::vector<uint32_t>bdc_group_id;
std::vector<uint32_t>gdma_bytes;
std::vector<uint32_t>bdc_bytes;
intcmdid_groupnum=0;
CMD_ID_NODE*cmdid_node;
CMD_ID_NODE*bdc_node;
CMD_ID_NODE*gdma_node;

• bdc_buffer和gdma_buffer：放指令
• gdma_total_id和bdc_total_id：存指令總數(shù)目，因?yàn)橹噶畈灰欢ㄊ?2位的，因此使用buffer的長(zhǎng)度不能獲取到指令的總數(shù)目

• gdma_group_id和bdc_group_id：每個(gè)group中的指令數(shù)，這個(gè)group是什么意思待調(diào)查清楚，后端對(duì)其進(jìn)行控制時(shí)代碼如下所示

cmdid_overflow

• cmdid_groupnum：group的數(shù)量
• gdma_bytes和bdc_bytes：內(nèi)部是每個(gè)group中指令的字節(jié)（Bytes）數(shù)
• cmdid_node、bdc_node和gdma_node：這個(gè)node是為了并行生成GroupOp內(nèi)部所需要指令而形成指的，具體機(jī)制還有待研究

TPU-MLIR中l(wèi)ayer group和上述行為中g(shù)roup的概念區(qū)分

TPU-MLIR中的layer group是指可以存放在Lmem的一系列算子，組成一個(gè)Group Op。

而上述的group，指的是指令組。這個(gè)指令組存在的意義是防止內(nèi)存不夠用，比方說(shuō)1684只有16位尋址空間，那么大于這個(gè)數(shù)字的指令無(wú)法一次性全部搬運(yùn)到內(nèi)存，所以當(dāng)指令超出某個(gè)數(shù)時(shí)候，就會(huì)重新建一個(gè)組。

TPU-MLIR中BM168X及其相關(guān)類

這邊的相關(guān)類tpu-mlir/include/tpu_mlir/Backend該文件夾下定義的類，目的就是將不同的處理器后端封裝，從而實(shí)現(xiàn)后端于Codegen過(guò)程的隔離。

其繼承關(guān)系為：

在一次運(yùn)行中只存在一個(gè)類（設(shè)計(jì)模式中單例），該類初始化時(shí)候會(huì)經(jīng)過(guò)：讀取后端動(dòng)態(tài)鏈接庫(kù)、加載函數(shù)（設(shè)置后端的函數(shù)指針）、指令數(shù)據(jù)結(jié)構(gòu)的初始化、設(shè)置一些處理器相關(guān)的參數(shù)例如NPU_NUM、L2_SRAM起始地址等。

后端函數(shù)的加載

后端作為一個(gè)動(dòng)態(tài)庫(kù)放入了TPU-MLIR工程里，具體在third_party/nntoolchain/lib/libbackend_xxx.so。在我們要使用backend時(shí)候，先在需要函數(shù)的類中定義好函數(shù)指針，然后再將動(dòng)態(tài)庫(kù)加載后，使函數(shù)指針指向動(dòng)態(tài)庫(kù)中真正的函數(shù)。

以同步函數(shù)tpu_sync_all為例，由于之后要加上多核支持的，所以需要在相關(guān)后端Bmodel庫(kù)中定義好，

1. 注意必須和后端的函數(shù)名和參數(shù)保持一致typedef void (*tpu_sync_all)();
2. 在類內(nèi)部加入該函數(shù)成員tpu_sync_all, dl_tpu_sync_all;
3. 有成員后，在該類load_functions函數(shù)的實(shí)現(xiàn)中加入宏，CAST_FUNCTION(tpu_sync_all);該宏可以將dl_tpu_sync_all指向動(dòng)態(tài)庫(kù)中真正的函數(shù)

這時(shí)候在我們獲得到該類實(shí)例后即可使用動(dòng)態(tài)庫(kù)中的函數(shù)了。

后端store_cmd設(shè)計(jì)

后端的store_cmd功能是指編譯器調(diào)用算子的過(guò)程中，把配置的指令保存到約定空間的過(guò)程。（以下是后端代碼，以后會(huì)選擇性開放）。后端的重點(diǎn)函數(shù)在store_cmd.cpp中，以cmodel/src/store_cmd.cpp；cmodel/include/store_cmd.h為例

注：store_cmd類設(shè)計(jì)的非常復(fù)雜，參雜各種設(shè)計(jì)模式在里面，只大概梳理一下類之間關(guān)系

store_cmd分別有EngineStorer系列類和CmdStorer系列類：

• EngineStoreInterface（接口類）、繼承于EngineStoreInterface接口的GDMAEngineStorer、BDEngineStorer等具體類、EngineStorerDecorator（裝飾類接口）、繼承于EngineStorerDecorator的VectorDumpEngineStorerDecorator等具體裝飾類
• CmdStorerInterface（接口）、繼承于接口的ConcretCmdStorer、StorerDecorator：裝飾接口、VectorDumpStorerDecorator具體裝飾類。

關(guān)于類之間的關(guān)系與邏輯

1. 使用單例設(shè)計(jì)模式，在store_cmd中只存在一個(gè)ConcretCmdStorer類，該類中會(huì)存所有EngineStorer的類，當(dāng)調(diào)用不同的engine時(shí)，會(huì)調(diào)用不同EengineStorer，如下代碼

virtualvoidstore_cmd(intengine_id,void*cmd,CMD_ID_NODE*cur_id_node,
intport)override
{
switch(engine_id)
{
caseENGINE_BD:
caseENGINE_GDMA:
caseENGINE_HAU:
caseENGINE_SDMA:
port=0;
break;
caseENGINE_CDMA:
ASSERT(port break;
caseENGINE_VSDMA:
engine_id=ENGINE_SDMA;
break;
default:
ASSERT(0);
break;
}
returnthis->get(engine_id,port)->store(cmd,cur_id_node);
}

1. Cmd裝飾類的作用是將所有的EngineStorer套上其裝飾器的殼子（目的實(shí)現(xiàn)其他功能），以VectorDumpStorerDecorator為例，會(huì)使用宏為每個(gè)EngineStorer、套上VectorDumpEngineStorerDecorator的殼子。

voiddecorate_engines()
{
#defineDECOR_STORER(name,idx)\
if(outputs_[ENGINE_##name][idx])\
{\
autoname##_str=std::make_shared(\
StorerDecorator::get(ENGINE_##name,idx),\
&(outputs_[ENGINE_##name][idx]));\
StorerDecorator::get(ENGINE_##name,idx)=name##_str;\
engine_decorators_.push_back(name##_str);\
}
DECOR_STORER(BD,0)
DECOR_STORER(GDMA,0)
DECOR_STORER(HAU,0)
DECOR_STORER(SDMA,0)
for(inti=0;i{
DECOR_STORER(CDMA,i)
}

#undefDECOR_STORER
}

每個(gè)具體的EngineStorer，注意其功能并非把命令存下來(lái)，他只干解析命令，比方說(shuō)拿到一條320位的命令（瞎說(shuō)的），EngineStorer會(huì)將其解析成長(zhǎng)度為10的32位數(shù)組（std::vector）。

真正存命令是使用VectorDumpEngineStorerDecorator，裝飾器的作用是：執(zhí)行被裝飾類的特定函數(shù)時(shí)，進(jìn)行更多的操作，具體可以《設(shè)計(jì)模式》的書。這點(diǎn)對(duì)于理解store_cmd非常重要，作者在設(shè)計(jì)store_cmd時(shí)，使用了很多裝飾器、為每個(gè)EngineStorer賦予了額外的功能，其中把指令儲(chǔ)存也看作一個(gè)裝飾器。VectorDumpEngineStorerDecorator該裝飾器執(zhí)行EngineStorer類中的store函數(shù)后，會(huì)追加執(zhí)行take_cmds函數(shù)，該函數(shù)將所有指令存儲(chǔ)到output_中。

classVectorDumpEngineStorerDecorator:publicEngineStorerDecorator
{
private:
std::vector<uint32_t>*&output_;

voidtake_cmds()
{
autocmds=EngineStorerDecorator::get_cmds();
(*output_).insert((*output_).end(),cmds.begin(),cmds.end());
}

public:
VectorDumpEngineStorerDecorator(ComponentPtrcomponent,
std::vector<uint32_t>**output)
:EngineStorerDecorator(component),output_(*output){}

virtualvoidstore(void*cmd,CMD_ID_NODE*cur_id_node)override
{
EngineStorerDecorator::store(cmd,cur_id_node);
if(!enabled_)
return;
this->take_cmds();
}

virtualvoidstore_cmd_end(unsigneddep)override
{
EngineStorerDecorator::store_cmd_end(dep);
this->take_cmds();
}
};

store_cmd中類與暴露給編譯器接口的關(guān)系

實(shí)際上上述的各種類不能直接暴露給編譯器，因?yàn)楸仨殏鞯氖莄函數(shù)的函數(shù)接口，因此必須將類中各種函數(shù)封裝進(jìn)c語(yǔ)言函數(shù)形式，以store_cmd為例，get_storer會(huì)獲得唯一的ConcretCmdStorer類

voidstore_cmd(void*cmd,intengine_id,CMD_ID_NODE*cur_id_node,intport,
intthread_id)
{
get_storer()->store_cmd(engine_id,cmd,cur_id_node,port);
}

TPU-MLIR中的重構(gòu)修改

分為三部分：BM168X及派生類、BM168XCodeGen

對(duì)于BM168X派生類來(lái)說(shuō)，后端工程中添加了很多新的函數(shù)，這些函數(shù)主要是將存指令的數(shù)據(jù)結(jié)構(gòu)放入了后端管理，涉及的后端有1684X之后架構(gòu)的處理器，而1684并不適配新的函數(shù)。這意味著：

存儲(chǔ)指令的數(shù)據(jù)結(jié)構(gòu)需要發(fā)生改變，許多數(shù)據(jù)結(jié)構(gòu)已經(jīng)不需要。如下圖所示：

需要添加新的接口函數(shù)，即使獲取指令的方式不同，但是在Codegen過(guò)程看到的應(yīng)該是一樣的行為。

這里傳入的參數(shù)是const char*是為了簡(jiǎn)化參數(shù)定義，可以用特定格式字符串來(lái)指定后端engine。如gdma1，這里gdma表示GDMA Engine, 0表示第0個(gè)GDMA Engine（一個(gè)TPU內(nèi)可能有多個(gè)相同的engine), 1表示第0號(hào)GDMA Engine的第1個(gè)線程（每個(gè)Engine可能支持多線程）。

對(duì)于BM168XCodegen，之前是需要在上面的Code結(jié)構(gòu)體獲取相關(guān)的數(shù)據(jù)，而修改后須使用新接口，

修改前：auto gdma_ptr = (uint8_t *)(*bm168x)->gdma_buffer.data();修改后：auto gdma_ptr = (uint8_t *)(*bm168x).get_inst_data("gdma0");

并且對(duì)于后續(xù)架構(gòu)處理器的指令生成來(lái)說(shuō)，目前需要存儲(chǔ)sdma和hau的指令，所以相關(guān)指令也需要添加入Bmodel。如下圖所示(這里主要用到了FlatBuffer操作)：

codegen_save

總結(jié)

從中可以看出，TPU-MLIR雖然能夠滿足當(dāng)前TPU上的基本需求，但隨著應(yīng)用場(chǎng)景的擴(kuò)展和TPU架構(gòu)的不斷演進(jìn)，其需要滿足很多新的要求。這就需要開發(fā)者不斷思考和挖掘新的接口和架構(gòu)，使其具有一定的擴(kuò)展性和適應(yīng)性。歡迎并感謝各位有識(shí)之士為TPU-MLIR多提建議，貢獻(xiàn)代碼！