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部分，其余的与对齐排布方式相似。

审核编辑：汤梓红