背景介紹

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

與后端的關系

由于一些歷史的因素，MLIR文件中的每個OP對應的指令并不直接在TPU-MLIR工程中生成，而是需要調用后端的函數完成最終指令的生成，這也帶來了兩個問題

1. 如何設計編譯器與后端的接口
2. 生成指令的數據結構存在后端還是編譯器中

關于問題1，目前的設計是采用CodeGen與后端隔離的形式，也就是CodeGen過程不直接調用后端函數，而是將不同處理器的相應函數全部封裝到類中，在CodeGen中調用類方法間接使用后端接口，達成解耦。

而關于問題2，依據不同的處理器其數據結構位置也不同，1684的數據結構放在編譯器這邊，而BM1684X等新架構的處理器數據結構放在后端。無論放在哪里，其全部封裝于問題1答案中的相應類中，對于CodeGen過程來說，看到的接口是一樣的。

一個OP生成指令的大致流程

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

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

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

做個形象點的例子：

原來裝冰箱只需要我，現在我嫌大象沉，我叫個張三幫我裝。
我：張三，你把這個大象給我裝冰箱里
張三吭哧吭哧幫我裝完了
我：行了，張三，你走吧；我自己把裝的運走。

指令生成所需要的數據結構

指令依據處理器的engine不同而有所差別，比如1684有GDMA和TIU，而新架構的處理器SG2260會存在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：存指令總數目，因為指令不一定是32位的，因此使用buffer的長度不能獲取到指令的總數目

• gdma_group_id和bdc_group_id：每個group中的指令數，這個group是什么意思待調查清楚，后端對其進行控制時代碼如下所示

cmdid_overflow

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

TPU-MLIR中layer group和上述行為中group的概念區分

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

而上述的group，指的是指令組。這個指令組存在的意義是防止內存不夠用，比方說1684只有16位尋址空間，那么大于這個數字的指令無法一次性全部搬運到內存，所以當指令超出某個數時候，就會重新建一個組。

TPU-MLIR中BM168X及其相關類

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

其繼承關系為：

在一次運行中只存在一個類（設計模式中單例），該類初始化時候會經過：讀取后端動態鏈接庫、加載函數（設置后端的函數指針）、指令數據結構的初始化、設置一些處理器相關的參數例如NPU_NUM、L2_SRAM起始地址等。

后端函數的加載

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

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

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

這時候在我們獲得到該類實例后即可使用動態庫中的函數了。

后端store_cmd設計

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

注：store_cmd類設計的非常復雜，參雜各種設計模式在里面，只大概梳理一下類之間關系

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

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

關于類之間的關系與邏輯

1. 使用單例設計模式，在store_cmd中只存在一個ConcretCmdStorer類，該類中會存所有EngineStorer的類，當調用不同的engine時，會調用不同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套上其裝飾器的殼子（目的實現其他功能），以VectorDumpStorerDecorator為例，會使用宏為每個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
}

每個具體的EngineStorer，注意其功能并非把命令存下來，他只干解析命令，比方說拿到一條320位的命令（瞎說的），EngineStorer會將其解析成長度為10的32位數組（std::vector）。

真正存命令是使用VectorDumpEngineStorerDecorator，裝飾器的作用是：執行被裝飾類的特定函數時，進行更多的操作，具體可以《設計模式》的書。這點對于理解store_cmd非常重要，作者在設計store_cmd時，使用了很多裝飾器、為每個EngineStorer賦予了額外的功能，其中把指令儲存也看作一個裝飾器。VectorDumpEngineStorerDecorator該裝飾器執行EngineStorer類中的store函數后，會追加執行take_cmds函數，該函數將所有指令存儲到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中類與暴露給編譯器接口的關系

實際上上述的各種類不能直接暴露給編譯器，因為必須傳的是c函數的函數接口，因此必須將類中各種函數封裝進c語言函數形式，以store_cmd為例，get_storer會獲得唯一的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中的重構修改

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

對于BM168X派生類來說，后端工程中添加了很多新的函數，這些函數主要是將存指令的數據結構放入了后端管理，涉及的后端有1684X之后架構的處理器，而1684并不適配新的函數。這意味著：

存儲指令的數據結構需要發生改變，許多數據結構已經不需要。如下圖所示：

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

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

對于BM168XCodegen，之前是需要在上面的Code結構體獲取相關的數據，而修改后須使用新接口，

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

并且對于后續架構處理器的指令生成來說，目前需要存儲sdma和hau的指令，所以相關指令也需要添加入Bmodel。如下圖所示(這里主要用到了FlatBuffer操作)：

codegen_save

總結

從中可以看出，TPU-MLIR雖然能夠滿足當前TPU上的基本需求，但隨著應用場景的擴展和TPU架構的不斷演進，其需要滿足很多新的要求。這就需要開發者不斷思考和挖掘新的接口和架構，使其具有一定的擴展性和適應性。歡迎并感謝各位有識之士為TPU-MLIR多提建議，貢獻代碼！