深入學(xué)習(xí)和掌握TPU硬件架構(gòu)有困難？TDB助力你快速上手！

TDB介紹

TDB(TPU DeBugger)是針對(duì)TPU-MLIR編譯出來(lái)的BModel設(shè)計(jì)的一系列調(diào)試工具集合, 可以支持對(duì)BModel反匯編、結(jié)構(gòu)可視化、單步執(zhí)行仿真等功能，使用方法靈活。能夠快速定位BModel與原始模型推理結(jié)果不一致的問(wèn)題，進(jìn)而修復(fù)TPU-MLIR的編譯或模型出錯(cuò)點(diǎn)。

下圖是TDB工具集的框架。

TDB主要是基于Python開(kāi)發(fā)的，其核心功能模塊包括TPU架構(gòu)相關(guān)的指令解析、指令運(yùn)行、數(shù)據(jù)IO功能，以及通用的BModel反匯編框架及MLIR的索引機(jī)制。在此基礎(chǔ)上，形成了TDB調(diào)試、BModel_checker, BModel_dis, MLIR2graph等工具。TDB目前支持的處理器包括BM1684、BM1684X。

下表是TDB工具的具體功能

本文重點(diǎn)介紹tdb.py的使用方法和實(shí)現(xiàn)細(xì)節(jié)。

tdb.py的使用

tdb.py提供了一個(gè)和pdb、gdb界面類(lèi)似的調(diào)試窗口，用于分析BModel運(yùn)行，支持添加斷點(diǎn)，單步執(zhí)行，查看內(nèi)存數(shù)據(jù)（修改內(nèi)存數(shù)據(jù)），數(shù)據(jù)比對(duì)。相關(guān)差異可參考下圖：

進(jìn)入調(diào)試

usage:tdb.py[-h][--inputs[INPUTS]][--ref_data[REF_DATA...]][--plugins[PLUGINS]][--ddr_size[DDR_SIZE]][-v][context_dir]

TPUDebugger.

positionalarguments:
context_dirThepathofBModel.

options:
-h,--helpshowthishelpmessageandexit
--inputs[INPUTS]TheinputsdataoftheBModel.
--ref_data[REF_DATA...]
ThereferencedataoftheBModel.
--plugins[PLUGINS]Theextrapluginstobeadded.
--ddr_size[DDR_SIZE]
Theddr_sizeofcmodel.
-v,--verboseuseprogressbar

在context_dir直接運(yùn)行tdb.py即可調(diào)試當(dāng)前目錄下的compilation.bmodel文件；可通過(guò)--ref_data來(lái)傳入?yún)⒖紨?shù)據(jù)，示例如下：

tdb.py--ref_data../yolov5s_1o_bm1684x_f32_tpu_outputs.npz

調(diào)試命令

在進(jìn)入TDB后，可以使用多個(gè)快捷命令來(lái)控制運(yùn)行狀態(tài)。按下兩次tab或者鍵入 ? 獲取命令提示。

基本命令

插件系統(tǒng)

為了減少代碼的冗余程度，設(shè)計(jì)上TDB本體只負(fù)責(zé)對(duì)指令執(zhí)行過(guò)程的基本控制，而額外通過(guò)插件系統(tǒng)對(duì)TDB的功能進(jìn)行擴(kuò)展。插件提供的功能包括：

• 在TDB執(zhí)行指令的不同生命周期接收回調(diào)信息
• 對(duì)TDB本身的指令進(jìn)行擴(kuò)展

目前內(nèi)置的“插件” 有 breakpoints、display、info、data-check 等，可以根據(jù)需求實(shí)現(xiàn)自己的TDB插件。

breakpoint

breakpoint插件用于實(shí)現(xiàn)斷點(diǎn)功能，在設(shè)計(jì)之初就考慮了可以通過(guò)多種形式添加斷點(diǎn)，包括指令名、指令地址、指令id、mlir的Operation名、value-id、location...等，同時(shí)可以根據(jù)需求靈活建立斷點(diǎn)。

目前支持的斷點(diǎn)類(lèi)型包括：

display

display目前支持的較為簡(jiǎn)陋，是通過(guò)對(duì)info函數(shù)的python調(diào)用封裝實(shí)現(xiàn)的：

display self.info("asm 5")

data-checker

默認(rèn)在通過(guò)--ref_data傳入數(shù)據(jù)的情況下會(huì)直接開(kāi)啟。也可以通過(guò)Bmodel_checker.py 使用

print

提供了內(nèi)置的一些打印功能，主要聚焦于打印當(dāng)前指令以及指令對(duì)應(yīng)的數(shù)據(jù)，并會(huì)在TDB每一個(gè)stop時(shí)打印待執(zhí)行的下一個(gè)atomic command

info

提供了內(nèi)置的一些打印功能，主要聚焦于打印出不同格式的當(dāng)前指令的上下文。

實(shí)現(xiàn)細(xì)節(jié)

文件結(jié)構(gòu)

debugger首先提供了完整的指令解析和運(yùn)行功能，這一部分的核心代碼位于 debugger/target_{產(chǎn)品型號(hào)/產(chǎn)品型號(hào)/common}。其中不同處理器在指令解析、cModel運(yùn)行等方式的實(shí)現(xiàn)均有差異，這些差異分別在各自的文件夾中實(shí)現(xiàn)。而一部分相同的內(nèi)容和對(duì)差異行為的抽象的聲明則在 python/debugger/target_common/ 目錄下提供。這部分內(nèi)容比較復(fù)雜，在后面單獨(dú)介紹。

此外，debugger下還有一些其他文件，包括

• debugger/disassembler.py：提供對(duì)Bmodel 文件的基本解析
• debugger/atomic_dialect.py：基于disassembler.py，將 Bmodel的內(nèi)容轉(zhuǎn)換為 mlir 的類(lèi)型系統(tǒng)
• debugger/final_mlir.py：解析 `final.mlir`` 文件并建立索引供 TDB使用
• debugger/tdb_support.py：提供了基于cmd.Cmd的基本的TDB功能的實(shí)現(xiàn)，同時(shí)加載 bmodel、日志、斷點(diǎn)的基類(lèi)、插件系統(tǒng)等。

各自的target_{device} 下提供的功能在后面介紹。而target_common包括

• context.py：提供了CModelContext 這一基類(lèi)以及相關(guān)方法的聲明
• cmodel.py：提供了CModelRunner 和 Memory 兩個(gè)基類(lèi)，提供了一些方法的聲明
• decoder.py：提供了 Decoder 類(lèi)定義以及解析tiu dma指令方法的聲明
• op_support.py：提供了一些類(lèi)型定義、Layout、MemoryType 等枚舉類(lèi)型的聲明、一些公共方法、以及一些Operation 基類(lèi)聲明

集成上下文環(huán)境

為了更方便的獲取不同target的相同功能（如指令解析或執(zhí)行），每個(gè)target 需要繼承CModelContext并實(shí)現(xiàn)自己的TargetContext來(lái)主動(dòng)暴露自己提供功能。每個(gè)TargetContext需要以類(lèi)變量的方式定義：

• MemRef：描述了一塊內(nèi)存區(qū)域的view（address、dtype、layout等）
• device：描述了該target的類(lèi)型
• memmap：描述了處理器上的內(nèi)存地址信息
• dma_sys、tiu_sys（部份處理器）：描述了停止指令的類(lèi)型
• get_memory_type：根據(jù)輸入地址獲取地址類(lèi)型（global、local、l2等 ）
• get_runner：獲取該target運(yùn)行指令的Runner實(shí)例
• 其他，如local_layout_to_stride，
• merge_instruction：在每個(gè)subnet中，atomic指令滿(mǎn)足DAG，所以可以進(jìn)行 topological sort（這一功能實(shí)際應(yīng)該由 Decoder 提供），從而可以被順序執(zhí)行

每個(gè)TDB實(shí)例在加載BModel后，會(huì)根據(jù)BModel的處理器來(lái)獲取到對(duì)應(yīng)target的Context實(shí)例（是實(shí)例而不是類(lèi)），從而能實(shí)現(xiàn)指令的解析和執(zhí)行。

指令解析

每個(gè)target都需要繼承DecoderBase并實(shí)現(xiàn)自己的Decoder來(lái)提供各個(gè)target的指令解析功能，每個(gè)Decoder都需要實(shí)現(xiàn)以下接口：

• decode_tiu_cmd
• decode_tiu_cmds
• decode_dma_cmd
• decode_dma_cmds

BModel中，指令以CmdGroup（對(duì)多核實(shí)現(xiàn)，通過(guò)CoreCmdGroup）的形式存儲(chǔ)，在python/debugger/atomic_dialect.py中定義了decode_cmdgroup，可以看做是調(diào)用Decoder方法的入口

defdecode_cmdgroup(
context:CModelContext,cmd_group:CmdGroup,subnet_id:int,core_id=0
)->StaticCmdGroup:
context=context
decoder=context.decoder

tiu=decoder.decode_tiu_cmds(cmd_group.tiu_cmd,core_id=core_id)
dma=decoder.decode_dma_cmds(cmd_group.dma_cmd,core_id=core_id)

cmdgroup=StaticCmdGroup(tiu,dma,context.merge_instruction(tiu,dma))
#hackinjectionsubnet_id
forcmdincmdgroup.all:
cmd.subnet_id=subnet_id
returncmdgroup

指令運(yùn)行

每個(gè)target都需要繼承MemoryBase和CModelRunner并實(shí)現(xiàn)自己的 Memory和TargetRunner來(lái)提供各個(gè)target的指令運(yùn)行功能，其中：

• 每個(gè)Memory都需要實(shí)現(xiàn)set_data、get_data、clear_memory方法
• 每個(gè)TargetRunner都需要實(shí)現(xiàn)tiu_compute、dma_compute、init_memory、clear_memory、dynamic_compute 等方法

其中，Runner實(shí)例默認(rèn)需要包含對(duì)應(yīng)target的Memory實(shí)例。

插件系統(tǒng)

插件系統(tǒng)同時(shí)支持了指令執(zhí)行回調(diào)和TDB功能擴(kuò)展，相關(guān)的代碼位置包括：

• 插件基類(lèi)：tdb_support.py:TdbPlugin
• 指令擴(kuò)展基類(lèi)：tdb_support.py:TdbPluginCmd
• 回調(diào)注冊(cè)功能實(shí)現(xiàn)：tdb_support.py:add_callback
• tdb 中實(shí)現(xiàn)相關(guān)邏輯：tdb_support.py:TdbCmdBackend.add_plugin
• tdb 中注冊(cè)生命周期：tdb.py中各個(gè)地方添加的 add_callback 裝飾器
• 各個(gè) plugin 的具體實(shí)現(xiàn)：plugins/xxx.py

總結(jié)

BModel調(diào)試不是一件容易的事情，TDB提供了一種手段，可以看到BModel內(nèi)部的推理過(guò)程并進(jìn)行干預(yù)。本文介紹了TDB的使用方法和實(shí)現(xiàn)細(xì)節(jié), 一方面讓大家熟悉BModel的調(diào)試方法，能夠更快地定位模型轉(zhuǎn)換過(guò)程中的問(wèn)題并修復(fù)；另一方面，可以，并分析指令的格式和內(nèi)容，增加開(kāi)發(fā)的背景知識(shí)。