TPU的工作模式

在算能的产品当中,根据主控单元(主机端Host)的不同,TPU对应有两种不同的运行模式,分别被称为PCIe模式和SOC模式。

  • PCIe模式:该模式下对应的产品形态为SC系列板卡。板卡通过PCIe接口连接到主机服务器上,主机服务器作为主控单元(Host)控制板卡的运行。

  • SoC模式:该模式下对应的产品形态为SE系列边缘推理设备。推理设备上的包含一个8核的A53处理器作为设备的主控单元(Host)控制板卡的运行。


TPU编程模型

由于TPU是一个异构的架构设计,由 主机端 发送指令,设备端 接收指令,按照指令执行指定的操作。 因此,驱动TPU的进行指定计算需要分别完成主机端和设备端的两部分代码:

  • 主机端(Host): 主机端代码,运行在主机侧,发送控制TPU运行的命令。

  • 设备端(Device): 设备端代码,运行在设备侧,通常调用TPU的各种指令运行相应的运算。

主机端和设备端的代码由于目标设备不同,因此需要使用不同的编译器对代码进行编译。

PCIe模式 下:

  • 主机端代码编译器: 主机的本地C++编译器。

  • 设备端代码编译器: 用于Linux ARM A53的交叉编译器。

SoC模式 下,代码在x86平台进行编译,在ARM A53平台运行。

  • 主机端代码编译器: 用于Linux ARM A53交叉编译器。

  • 设备端代码编译器: 用于Linux ARM A53交叉编译器。

上述Linux ARM A53交叉编译器,可以通过TPUKernel工具包内scripts/prepare_toolchain.sh脚本来进行下载

对于Device端代码,编译的过程可以被看作是动态链接库更新的过程,

../_images/heterogeneous_prepare.png

如上图所示,device_*.c() 就是开发者完成的设备端代码,设备端代码编译器将设备端代码和原始的底库libbm1684x.a进行链接,形成完整的A53Lite可加载的原始动态库文件。


Host端

在完成了上述动态库的更新后,此时只需要加载动态库,然后在Host发送调用指令就可以控制TPU运行指定的计算。

typedef struct _param_func{
    ...
    ...
}__attribute__((packed)) param_func_0_1;

bm_handle_t handle;
param_func_0_1 param;
bm_dev_request(&handle, 0);
tpu_kernel_launch_sync(handle, "func_0_1", &param, sizeof(param));
bm_dev_free(&handle);
../_images/heterogeneous_run.png

如上图所示,Device端固件中已经注册了 func_0_0() , func_0_1() , func_0_2() , func_1_0() , func_1_1() , func_1_2() …等函数, 主机端现在发送 func_0_1() ,就可以调动TPU进行 func_0_1() 的计算。


主机端调用TPU进行计算有 异步同步 两种方式。

同步运行方式

同步运行方式 是指主机端在发送了调用计算命令后,不再继续下面代码的执行,直到TPU计算完成,继续进行主机端代码的执行。

主机端调用下面API后,便会等待TPU直到计算完成。

bm_status_t tpu_kernel_launch_sync(bm_handle_t handle, const char *func_name, const void *args, unsigned int size)

Launch the kernel function on device synchronously.

参数
  • handle – Handle of the device.

  • func_name – Name of the kernel function to launch on device.

  • args – Pointer to the user-discript data package.

  • size – Size of the user-discript data package in bytes.

返回

Status of launching kernel function, BM_SUCCESS means succeeded, otherwise, some errors caught.

异步运行方式

异步运行方式 是指主机端在发送了调用指令后,继续执行后续的代码,TPU和主机端异步运行。

主机端在调用了下面API后,TPU就会开始计算,同时主机端仍会继续往下执行。

bm_status_t tpu_kernel_launch_async(bm_handle_t handle, const char *func_name, const void *args, unsigned int size)

Launch the kernel function on device asynchronously.

参数
  • handle – Handle of the device.

  • func_name – Name of the kernel function to launch on device.

  • args – Pointer to the user-discript data package.

  • size – Size of the user-discript data package in bytes.

返回

Status of launching kernel function, BM_SUCCESS means succeeded, otherwise, some errors caught.

主机端可以通过下面的API来实现和TPU的同步,调用下面API后,主机端会等待TPU完成计算。

bm_status_t tpu_kernel_sync(bm_handle_t handle)

Synchronize the device.

参数

handle – Handle of the device.

返回

Status of launching kernel function, BM_SUCCESS means succeeded, otherwise, some errors caught.

Device端

上一节中,介绍了用户是怎么样驱使TPU进行计算的。 在这一章节,我们将开始学习如何编写device端代码,利用TPU完成我们想要的计算,在此之前我们需要首先了解一下,TPU定义了哪些基本的指令。 在TPU架构一节,我们已经介绍了TPU的整个计算过程,回顾一下,计算过程可以划分为三部分:

1.数据在Host端内存和系统内存(System Memmory)之间的来回搬运。

2.数据在系统内存(System Memory)和Local Memory之间的来回搬运。

3.TPU对Local Memory中的数据进行相关计算。

“数据在Host端内存和系统内存之间的来回搬运”, 是主机端相关,这里不再涉及。由于数据存在于Host内存当中,因此这部分API通常由Host端来调用,关于Host端的相关命令可以参考上一小节。

指令系统

· GDMA指令

系统内存和Local-Memory中与数据搬运相关的操作都由GDMA指令完成。 与GDMA相关的指令都以 tpu_gdma_() 开头,包括数据在Local Memory之间的搬运、系统内存之间的搬运。 详细的指令说明和参数,参见“TPU_API”章节。

· BDC指令

TPU进行数据计算的相关操作都可以由BDC指令来完成。与BDC相关的指令都以 tpu_bdc_() 开头。

· HAU指令

一些不适用于并行加速计算的指令集,包括 NMS、SORT等。

内存与数据排列

· 数据的表示

Tensor

在TPU中,我们使用Tensor来描述数据。

Tensor 是一个4维数组,使用4元组(N,C,H,W)来描述一个Tensor的几何尺寸( Shape )。 Tensor(n, c, h, w)表示在(n,c,h,w)索引下的数据元素。

Stride 用于描述Tensor在实际内存当中是如何摆放,同样使用4元组(N_stride,C_stride,H_stride,W_stride)来描述, 表示 Tensor 在内存当中存放时,元素间间隔了多少元素,具体而言:

  • W_stride 描述的是从Tensor(n,c,h,w)到Tensor(n,c,h,w+1)两个元素之间,在内存存储时,间隔了多少个元素。

  • H_stride 描述的是从Tensor(n,c,h,w)到Tensor(n,c,h+1,w)两个元素之间,在内存存储时,间隔了多少个元素。

  • C_stride 描述的是从Tensor(n,c,h,w)到Tensor(n,c+X,h,w)两个元素之间,在内存存储时,间隔了多少个元素,X表示NPU的数量。

  • N_stride 描述的是从Tensor(n,c,h,w)到Tensor(n+1,c,h,w)两个元素之间,在内存存储时,间隔了多少个元素。

假设现在有一个Tensor,它的每一个元素都占1个byte,它的Shape是(4,3,2,2), 如果它的存储方式Stride是(12,4,2,1),在内存当中,它的排列方式就会如下图所示,

../_images/Tensor_shape_stride1.png

如果Tensor的存储方式Stride是(24,4,2,1),在内存当中,Tensor的排列方式就会如下所示,

../_images/Tensor_shape_stride2.png

如果Tensor的存储方式Stride是(24,8,4,2),在内存当中,Tensor的排列方式就会如下所示,

../_images/Tensor_shape_stride3.png

如果它的存储方式Stride是(24,8,4,1),在内存当中,Tensor的排列方式就会如下所示,

../_images/Tensor_shape_stride4.png

数据元素的类型

Stride 以元素个数作为计量单位。不同类型的数据元素占据不同的字节数, 在BM1684x芯片上支持如下格式的数据类型:

\[\begin{split}\begin{array}{|c|c|} \hline \textsf{数据类型} & \textsf{字节数}\\ \hline \textsf{INT8} & \textsf{1 Bytes}\\ \hline \textsf{INT16} & \textsf{2 Bytes}\\ \hline \textsf{INT32} & \textsf{4 Bytes}\\ \hline \textsf{FP16} & \textsf{2 Bytes}\\ \hline \textsf{BFP16} & \textsf{2 Bytes}\\ \hline \textsf{FP32} & \textsf{4 Bytes}\\ \hline \end{array}\end{split}\]

· Tensor在gloabl memory的排列方式

global memory 由一块DDR内存组成。

一个Shape为(N,C,H,W)的Tensor,在 global memory 排列,对应的Stride为:

  • W_Stride = 1,

  • H_Stride = W,

  • C_Stride = H*W,

  • N_Stride = C*H*W。

这种排列方式被称为 连续存储方式

举例: 一个Shape(N=2, C=2,H=3,W=2)的Tensor在global memory排列方式

../_images/global_mem.png

上图中,n0c0h0w0 代表 Tensor(0,0,0,0) 位置的元素。

· Tensor在local memory的排列方式

local memory的物理组成和地址分配

Local Memory共由多片SRAM(静态随机存取存储器)构成,每一片SRAM都被称为一个 Bank

bm1684x芯片一共由16个Bank构成。16个Bank组成整个 local memory

整个Local Memory同时被划分为了64个 lane (对应64个NPU,以下用NPU指代lane),地址分配如下图所示:

../_images/local_mem.png

BM1684x的Local Memory大小为 256KB * 64,地址按照NPU进行分配, 其中,NPU0对应 0~256*1024-1 的地址,NPU1对应 256*1024~2*256*1024-1 的地址,依次类推。


Tensor在Local memory上排列的基本规则

Tensor在Local Memory上的排布方式与global Memory的排布方式不同,主要区别在于 C维度的数据排布方式。 一个Shape为(N,C,H,W)的Tensor,Tensor(N,c,H,W)代表:当C = c时,Tensor的数据切片。 对于不同的c, Tensor(N,c,H,W)分配在不同的NPU上。

举例,Tensor的Shape(N=2,C=3,H=2,W=3),Stride(N_stride = 9,C_stride = 9, H_stride = 3, W_stride = 1) 那么Tensor在Local Memory上的数据排列方式如下所示。

../_images/local_mem_hw_arrange.png

不同大小的C影响着实际存储方式,假设现在一共由X个NPU,考虑下面几种存储情形:

情形1: 当Tensor的Shape的维度C = X-1时,当Tensor从NPU0开始存储时,那么Tensor在Local Memory上的排布方式如下所示。

../_images/scatter_0.png

在NPU X-1的地址空间,没有数据分布。

情形2: 当Tensor的Shape的维度C = X-1时,当Tensor从NPU1开始存储时,那么Tensor在Local Memory上的排布方式如下所示。

../_images/scatter_1.png

在NPU0的地址空间,没有数据分布。

情形3: 当Tensor的Shape的维度C = X + 2,当Tensor从NPU0开始存储时,那么Tensor在Local Memory上的排布方式如下所示。

../_images/scatter_2.png

C=X+1维度的数据被排布到NPU0,而C=X+2维度的数据被排布到NPU1。并且注意到,下一个N维度仍然从NPU0开始排布。

情形4: 当Tensor的Shape的维度C= X + 2, 而Tensor从NPU X-1 开始存储时,那么Tensor在Local Memory上的排布方式如下所示。

../_images/scatter_3.png

可以看到, C=0维度的数据被排布到了NPU X-1, C=1维度的数据被排布到了NPU0上,依次类推,C= X+2维度的数据被排布到了NPU0上。


Local Memory上几种常用数据排布方式

上面介绍了Tensor在Local Memory上存储的基本原则,现在介绍1684x指令集常用的几种数据排布方式:

1. 64-Bytes对齐存储方式

“64-Bytes对齐存储方式”是最常用的Tensor存储方式,它是指Tensor排放存储要满足以下几个约束:

  • Tensor的起始地址是64的整数倍

  • W_stride = 1

  • H_stride = W

  • C_stride = ceil(H*W, 16) * 16, 如果数据元素是 32-bits ; ceil(H*W, 32) * 32, 如果数据元素是 16-bits ; ceil(H*W, 64) * 64, 如果数据元素是 8-bits

  • N_stride = C_stride * (单个NPU上channel的个数)

其中 ceil 是向上取整的意思。可通过 tpu_aligned_stride() 计算 stride。

为举例简便,假设NPU个数=4

举例1: Tensor的Shape(.N=2,.C=3,.H=4,.W=5),数据类型为float16, NPU0开始存储

../_images/align_fp32_start_0_2_3_4_5.png

举例2: Tensor的Shape(.N=2,.C=3,.H=4,.W=5),数据类型为float16,NPU2开始存储

../_images/align_fp32_start_2_2_3_4_5.png

2. 紧凑存储方式

“紧凑存储方式”也是较为常用的Tensor存储方式。

假设Tensor的Shape为(N,C,H,W),按照“紧凑存储方式存储”要满足以下约束:

  • Tensor的起始地址是4的整数倍。

  • W_stride = 1

  • H_stride = W

  • C_stride = H * W

  • N_stride = C_stride * (单个NPU上channel的个数)

可通过 tpu_compact_stride() 计算 stride。

为举例简便,假设NPU个数=4

举例1: Tensor的Shape(.N=2,.C=3,.H=4,.W=5),数据类型为float16, NPU0开始存储

../_images/compact_fp32_start_0_2_3_4_5.png

举例2: Tensor的Shape(.N=2,.C=3,.H=4,.W=5),数据类型为float16,NPU2开始存储

../_images/compact_fp32_start_2_2_3_4_5.png

3. 矩阵存储方式

“矩阵存储方式”是 矩阵运算指令 用的数据存储方式。

对于一个n x m的矩阵,可以用Tensor的4维数组的形式来进行表示, 这个Tensor的Shape为(N=n,C=ceil(m/w),H=1,W=w),其中w可以为(1,m)之间的任意值。

为举例简便,假设NPU个数=4

举例1:矩阵的形状为(2x40),数据类型为float16, w = 40

../_images/matrix_w_40.png

举例2:矩阵的形状为(2x40),数据类型为float16, w = 20

../_images/matrix_w_20.png

举例3:矩阵的形状为(2x40),数据类型为float16, w = 15

../_images/matrix_w_10.png

举例4:矩阵的形状为(2x40),数据类型为float16, w = 8

../_images/matrix_w_8.png

举例5:矩阵的形状为(2x40),数据类型为float16, w = 6

../_images/matrix_w_6.png

4. 向量存储

对于一个1 x m的向量,可以用Tensor的4维数组的形式来进行表示, 这个Tensor的Shape为(N=1,C=ceil(m/w),H=1,W=w),其中w可以为(1,m)之间的任意值。


5. 行64字节对齐存储

一种 4D 张量在 local memory 中的存储格式。张量的 shape 是 (N, C, H, W),满足

  • 地址被 64 整除

  • W-stride 是 1

  • H-stride 是 ceil(W / 16) * 16,如果元素的数据类型的位宽是 32-bit, ceil(W / 32) * 32,如果是 16-bit,ceil(W / 64) * 64,如果是 8-bit

  • C-stride 是 H * H-stride

  • N-stride 是 C-stride 乘以每个 NPU 的 channel 数

可通过 tpu_line_aligned_stride() 计算 stride。


6. 64IC/32IC 存储

64IC/32IC 存储是卷积核在 local memory 中的特殊存储方式,仅用于卷积计算过程。其中 INT8 kernel 以 64IC 格式存放, FP16/BFP16 kernel 以 32IC 格式存放。

假设卷积核的 shape 是 (ic, oc, kh, kw),分别表示 input channel、 output channel、 卷积核的高度以及卷积核的宽度。

64IC 存储满足:

  • W_stride = 64

  • H_stride = 64 * kw

  • C_stride = 64 * kw * kh * ceil(ic/64)

  • N_stride = 64 * kw * kh * ceil(ic/64)

其中,每 64 个 input channel 作为一组进行存储,每组之间的 stride 为 64 * kw * kh。

32IC 存储满足:

  • W_stride = 32

  • H_stride = 32 * kw

  • C_stride = 32 * kw * kh * ceil(ic/32)

  • N_stride = 32 * kw * kh * ceil(ic/32)

其中,每 32 个 input channel 作为一组进行存储,每组之间的 stride 为 32 * kw * kh。