TPU内存(二)

我们知道,TPU中的数据结构是张量,可以看做是一个四维数组,形状为(N,C,H,W)。

要描述一个张量在算能的TPU上是如何排列的,我们首先要知道一个概念。 那是Stride。

它用于衡量同一 NPU 中张量的两个元素之间的距离。

例如,W_Stride 表示张量 n,c,h,w 和 n,c,h,w+1 之间有多少个元素

而H_stride表示从n,c,h,w到n,c,h+1,w我们需要经过的元素个数

同样,我们可以得到C_stride和N_stride在Global memory上的含义。

但是对于local memory,我们可以看到是有所不同的,C_stride指的是从n,c,h,w到n,c+X,h,w的元素个数,其中X表示 NPU 的数量。

而在N_stride中我们还需要考虑我们开始存储数据的local memory的索引。

稍后我将进一步解释这一点。

有了tensor的shape和stride,我们基本上就可以得到这个tensor的每个元素在内存上的地址

但是步长的单位是张量中的单个元素,所以对于不同的数据类型,要计算它们的地址,我们还必须将它们的字节数考虑进去。

例如,在一个 F32 张量中,w 和 w+1 个元素之间的实际距离是 1 * 4。

在global memory中,数据以连续的方式存储,

这很容易理解。 由于global memory是一个完整的DDR,我们把tensor的每个元素挨个存储,所以w_stride等于1,h_stride等于w,对于c_stride来说,就是w的h倍,n_stride则是c_stride的c倍。

例如,对于形状为 (2,2,3,2) 的张量,
w_stride为1,每2个元素后开始一个新的h,所以h_stride为2,每个通道包含3 * 2个元素,所以c_stride应该为6,同理,我们可以很容易地得到n_stride, 12。

但是对于本地内存,就变得有点复杂了,

首先,张量的不同通道会被放到不同的NPU上。 如果通道大于 NPU 的数量,它将返回到第一个 NPU开始存放。

这就是为什么local memory的C_stride是n,c,h,w到n,c+X,h,w。 Stride仅衡量同一memory中的距离。

例如,我们使用 X个NPU 来存储具有 X + 2 个通道的张量,我们将从第一个 NPU 到最后一个 NPU 放入每个通道的元素。 然后其余的通道再次从第一个 NPU 开始存放。

对于张量的每个batch,我们将从同一 NPU 上新的一行开始存储。

像这个例子中,当我们完成第一批的存储后,即使同一个bank中的剩余内存为空,我们也不会存储任何东西,而是重新从NPU0开始。

这就解释了为什么我们在计算 N_stride 时需要考虑local memory的起始索引。

基于上述原则,local memory中的张量以多种不同方式排布。

最常用的一种是对齐排布。

这意味着张量的起始地址应该可以被 EU_BYTE 整除。

另外,对于不同通道的数据,用于保存的区域大小应该是EU_NUM的倍数,

从数学角度看,C_stride的计算应该是这样的(看PPT)。 当H * W小于EU_NUM时,C_stride为EU_NUM。 当大于EU_NUM但小于2倍EU_NUM时,C_stride应为EU_NUM的2倍。

关于N_stride,由于有时通道数大于NPU_NUM或者local memory的起始索引不为零,可能会导致不同通道的数据存储在同一个NPU中,N_stride的公式也应该做round- up 操作,如PPT中所示。

例如,我们将在具有 64 EU_BYTE 的 TPU 上处理形状为 (2,3,4,5) 的fp16张量。其中包含了4个NPU,而本地内存的起始索引设置为 0。

所以我们从NPU0开始存储张量,W_stride和H_stride显然是1和5。

对于C_stride,由于H * W小于EU_NUM,所以C_stride为32。

另外,因为这个张量的通道小于 NPU_NUM,所以N_stride 也是 32。

但是当起始索引设置为2时,情况会有点不同,C步长仍然是32,但是由于张量的最后一个通道被存到第一个NPU,下一batch中的数据应该从NPU2的下一行开始存储,则 N_stride 应为 64。

另一个常见的排布类型就是紧密排布,除了C_stride部分,其余的与对齐排布方式相似。