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部分，其余的与对齐排布方式相似。
-
DDR
+关注
关注
11文章
712浏览量
65317 -
内存
+关注
关注
8文章
3019浏览量
73999 -
数据结构
+关注
关注
3文章
573浏览量
40121 -
NPU
+关注
关注
2文章
279浏览量
18582
发布评论请先 登录
相关推荐
工商网监
评论