深入学习和掌握TPU硬件架构有困难？TDB助力你快速上手！

 

TDB介绍

TDB(TPU DeBugger)是针对TPU-MLIR编译出来的BModel设计的一系列调试工具集合, 可以支持对BModel反汇编、结构可视化、单步执行仿真等功能，使用方法灵活。能够快速定位BModel与原始模型推理结果不一致的问题，进而修复TPU-MLIR的编译或模型出错点。

下图是TDB工具集的框架。

TDB主要是基于Python开发的，其核心功能模块包括TPU架构相关的指令解析、指令运行、数据IO功能，以及通用的BModel反汇编框架及MLIR的索引机制。在此基础上，形成了TDB调试、BModel_checker, BModel_dis, MLIR2graph等工具。TDB目前支持的处理器包括BM1684、BM1684X。

下表是TDB工具的具体功能

本文重点介绍tdb.py的使用方法和实现细节。

tdb.py的使用

tdb.py提供了一个和pdb、gdb界面类似的调试窗口，用于分析BModel运行，支持添加断点，单步执行，查看内存数据（修改内存数据），数据比对。相关差异可参考下图：

进入调试

usage: tdb.py [-h] [--inputs [INPUTS]] [--ref_data [REF_DATA ...]] [--plugins [PLUGINS]] [--ddr_size [DDR_SIZE]] [-v] [context_dir]
 
TPU Debugger.
 
positional arguments:
  context_dir           The path of BModel.
 
options:
  -h, --help            show this help message and exit
  --inputs [INPUTS]     The inputs data of the BModel.
  --ref_data [REF_DATA ...]
                        The reference data of the BModel.
  --plugins [PLUGINS]   The extra plugins to be added.
  --ddr_size [DDR_SIZE]
                        The ddr_size of cmodel.
  -v, --verbose         use progress bar
 

在context_dir直接运行tdb.py即可调试当前目录下的compilation.bmodel文件；可通过--ref_data来传入参考数据，示例如下：

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

调试命令

在进入TDB后，可以使用多个快捷命令来控制运行状态。按下两次tab或者键入 ? 获取命令提示。

基本命令

命令	说明
r/run	从头执行到结束，run指令包含重新初始化
s/start	加载 bmodel 并初始化cmodel，cmodel 为单例模式，不会重复加载，只会清空内存；bmodel 每次初始化均会重新加载。
n/next	执行下一条指令，可以使用n 3来执行接下来的三条指令
c/continue	继续执行指令
q/quit	退出
py	直接在环境中执行python命令，集成了pdb的代码补全功能，连按两下 tab 键获取提示

插件系统

为了减少代码的冗余程度，设计上TDB本体只负责对指令执行过程的基本控制，而额外通过插件系统对TDB的功能进行扩展。插件提供的功能包括：

• 在TDB执行指令的不同生命周期接收回调信息
• 对TDB本身的指令进行扩展

目前内置的“插件” 有 breakpoints、display、info、data-check 等，可以根据需求实现自己的TDB插件。

breakpoint

 

breakpoint插件用于实现断点功能，在设计之初就考虑了可以通过多种形式添加断点，包括指令名、指令地址、指令id、mlir的Operation名、value-id、location...等，同时可以根据需求灵活建立断点。

命令	介绍	示例
b/break	添加断点，具体查看下方的	b %1
enable num	允许在index为num 的断点停下	enable 1   enable 1,2,3
disable num	停止使用index为num的断点	disable 1  disable 1,2,3
delete	删除断点	delete 1
info b	查看断点信息	info b

目前支持的断点类型包括：

断点类型名称	断点命令使用示例	断点来源说明
asm-name	b dma.tensor  b arith.cast	解析出的asm的指令名称，如： %R8, %B2420 = "arith.cast"(%R0, %D288) {round_mode = 1} : (memref<16x64x1x512xf16, strides: [512, 512, 512, 1]>, none) -> (memref<16x64x1x512xf32, strides: [512, 512, 512, 1]>, none)
address	b R0  b G33386496	解析出的asm的 memref，如： %G33386496, %D293 = "dma.matrix"(%R8, %B2422) {decompress = False} ...  目前对Local Memory仅支持offset，不支持通过 NPU_OFFSET 等进行索引
cmd-id	b B23  b D465	解析出的 asm 的 cmd_id，如： %R0, %D291 = "dma.matrix"(%G35471360, %B2420) {decompress = False} ...
value-id	b %173	final.mlir 的 Operation 的前缀，如： %173 = "tpu.Cast"(%169) {ginfo ... 目前仅匹配字符串，只要字符串中包含对应的 value-id 就视为匹配
location	b #loc("conv1_glow")  b #loc(3)	final.mlir的Operation 的loc，目前仅支持 name
op-name	b tpu.Load	final.mlir中的Operation 的名字，如：tpu.Load
file-line	b mlir:312 b asm:168	在final.mlir和asm的行号上打断点

display

 

display目前支持的较为简陋，是通过对info函数的python调用封装实现的：

display self.info("asm 5")

data-checker

 

默认在通过--ref_data传入数据的情况下会直接开启。也可以通过Bmodel_checker.py 使用

print

 

提供了内置的一些打印功能，主要聚焦于打印当前指令以及指令对应的数据，并会在TDB每一个stop时打印待执行的下一个atomic command

info

 

提供了内置的一些打印功能，主要聚焦于打印出不同格式的当前指令的上下文。

实现细节

文件结构

debugger首先提供了完整的指令解析和运行功能，这一部分的核心代码位于 debugger/target_{产品型号/产品型号/common}。其中不同处理器在指令解析、cModel运行等方式的实现均有差异，这些差异分别在各自的文件夹中实现。而一部分相同的内容和对差异行为的抽象的声明则在 python/debugger/target_common/ 目录下提供。这部分内容比较复杂，在后面单独介绍。

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

• debugger/disassembler.py：提供对Bmodel 文件的基本解析
• debugger/atomic_dialect.py：基于disassembler.py，将 Bmodel的内容转换为 mlir 的类型系统
• debugger/final_mlir.py：解析 `final.mlir`` 文件并建立索引供 TDB使用
• debugger/tdb_support.py：提供了基于cmd.Cmd的基本的TDB功能的实现，同时加载 bmodel、日志、断点的基类、插件系统等。

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

• context.py：提供了CModelContext 这一基类以及相关方法的声明
• cmodel.py：提供了CModelRunner 和 Memory 两个基类，提供了一些方法的声明
• decoder.py：提供了 Decoder 类定义以及解析tiu dma指令方法的声明
• op_support.py：提供了一些类型定义、Layout、MemoryType 等枚举类型的声明、一些公共方法、以及一些Operation 基类声明

集成上下文环境

为了更方便的获取不同target的相同功能（如指令解析或执行），每个target 需要继承CModelContext并实现自己的TargetContext来主动暴露自己提供功能。每个TargetContext需要以类变量的方式定义：

• MemRef：描述了一块内存区域的view（address、dtype、layout等）
• device：描述了该target的类型
• memmap：描述了处理器上的内存地址信息
• dma_sys、tiu_sys（部份处理器）：描述了停止指令的类型
• get_memory_type：根据输入地址获取地址类型（global、local、l2等 ）
• get_runner：获取该target运行指令的Runner实例
• 其他，如local_layout_to_stride，
• merge_instruction：在每个subnet中，atomic指令满足DAG，所以可以进行 topological sort（这一功能实际应该由 Decoder 提供），从而可以被顺序执行

每个TDB实例在加载BModel后，会根据BModel的处理器来获取到对应target的Context实例（是实例而不是类），从而能实现指令的解析和执行。

指令解析

每个target都需要继承DecoderBase并实现自己的Decoder来提供各个target的指令解析功能，每个Decoder都需要实现以下接口：

• decode_tiu_cmd
• decode_tiu_cmds
• decode_dma_cmd
• decode_dma_cmds

BModel中，指令以CmdGroup（对多核实现，通过CoreCmdGroup）的形式存储，在python/debugger/atomic_dialect.py中定义了decode_cmdgroup，可以看做是调用Decoder方法的入口

def decode_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))
    # hack injection subnet_id
    for cmd in cmdgroup.all:
        cmd.subnet_id = subnet_id
    return cmdgroup
 

指令运行

每个target都需要继承MemoryBase和CModelRunner并实现自己的 Memory和TargetRunner来提供各个target的指令运行功能，其中：

• 每个Memory都需要实现set_data、get_data、clear_memory方法
• 每个TargetRunner都需要实现tiu_compute、dma_compute、init_memory、clear_memory、dynamic_compute 等方法

其中，Runner实例默认需要包含对应target的Memory实例。

插件系统

插件系统同时支持了指令执行回调和TDB功能扩展，相关的代码位置包括：

• 插件基类：tdb_support.py:TdbPlugin
• 指令扩展基类：tdb_support.py:TdbPluginCmd
• 回调注册功能实现：tdb_support.py:add_callback
• tdb 中实现相关逻辑：tdb_support.py:TdbCmdBackend.add_plugin
• tdb 中注册生命周期：tdb.py中各个地方添加的 add_callback 装饰器
• 各个 plugin 的具体实现：plugins/xxx.py

总结

BModel调试不是一件容易的事情，TDB提供了一种手段，可以看到BModel内部的推理过程并进行干预。本文介绍了TDB的使用方法和实现细节, 一方面让大家熟悉BModel的调试方法，能够更快地定位模型转换过程中的问题并修复；另一方面，可以，并分析指令的格式和内容，增加开发的背景知识。