TPU內(nèi)存（二）

我們知道，TPU中的數(shù)據(jù)結(jié)構(gòu)是張量，可以看做是一個(gè)四維數(shù)組，形狀為（N,C,H,W）。

要描述一個(gè)張量在算能的TPU上是如何排列的，我們首先要知道一個(gè)概念。 那是Stride。

它用于衡量同一 NPU 中張量的兩個(gè)元素之間的距離。

例如，W_Stride 表示張量 n,c,h,w 和 n,c,h,w+1 之間有多少個(gè)元素

而H_stride表示從n,c,h,w到n,c,h+1,w我們需要經(jīng)過的元素個(gè)數(shù)

同樣，我們可以得到C_stride和N_stride在Global memory上的含義。

但是對(duì)于local memory，我們可以看到是有所不同的，C_stride指的是從n,c,h,w到n,c+X,h,w的元素個(gè)數(shù)，其中X表示 NPU 的數(shù)量。

而在N_stride中我們還需要考慮我們開始存儲(chǔ)數(shù)據(jù)的local memory的索引。

稍后我將進(jìn)一步解釋這一點(diǎn)。

有了tensor的shape和stride，我們基本上就可以得到這個(gè)tensor的每個(gè)元素在內(nèi)存上的地址

但是步長的單位是張量中的單個(gè)元素，所以對(duì)于不同的數(shù)據(jù)類型，要計(jì)算它們的地址，我們還必須將它們的字節(jié)數(shù)考慮進(jìn)去。

例如，在一個(gè) F32 張量中，w 和 w+1 個(gè)元素之間的實(shí)際距離是 1 * 4。

在global memory中，數(shù)據(jù)以連續(xù)的方式存儲(chǔ)，

這很容易理解。 由于global memory是一個(gè)完整的DDR，我們把tensor的每個(gè)元素挨個(gè)存儲(chǔ)，所以w_stride等于1，h_stride等于w，對(duì)于c_stride來說，就是w的h倍，n_stride則是c_stride的c倍。

例如，對(duì)于形狀為 (2,2,3,2) 的張量，
w_stride為1，每2個(gè)元素后開始一個(gè)新的h，所以h_stride為2，每個(gè)通道包含3 * 2個(gè)元素，所以c_stride應(yīng)該為6，同理，我們可以很容易地得到n_stride, 12。

但是對(duì)于本地內(nèi)存，就變得有點(diǎn)復(fù)雜了，

首先，張量的不同通道會(huì)被放到不同的NPU上。 如果通道大于 NPU 的數(shù)量，它將返回到第一個(gè) NPU開始存放。

這就是為什么local memory的C_stride是n,c,h,w到n,c+X,h,w。 Stride僅衡量同一memory中的距離。

例如，我們使用 X個(gè)NPU 來存儲(chǔ)具有 X + 2 個(gè)通道的張量，我們將從第一個(gè) NPU 到最后一個(gè) NPU 放入每個(gè)通道的元素。 然后其余的通道再次從第一個(gè) NPU 開始存放。

對(duì)于張量的每個(gè)batch，我們將從同一 NPU 上新的一行開始存儲(chǔ)。

像這個(gè)例子中，當(dāng)我們完成第一批的存儲(chǔ)后，即使同一個(gè)bank中的剩余內(nèi)存為空，我們也不會(huì)存儲(chǔ)任何東西，而是重新從NPU0開始。

這就解釋了為什么我們?cè)谟?jì)算 N_stride 時(shí)需要考慮local memory的起始索引。

基于上述原則，local memory中的張量以多種不同方式排布。

最常用的一種是對(duì)齊排布。

這意味著張量的起始地址應(yīng)該可以被 EU_BYTE 整除。

另外，對(duì)于不同通道的數(shù)據(jù)，用于保存的區(qū)域大小應(yīng)該是EU_NUM的倍數(shù)，

從數(shù)學(xué)角度看，C_stride的計(jì)算應(yīng)該是這樣的（看PPT）。 當(dāng)H * W小于EU_NUM時(shí)，C_stride為EU_NUM。 當(dāng)大于EU_NUM但小于2倍EU_NUM時(shí)，C_stride應(yīng)為EU_NUM的2倍。

關(guān)于N_stride，由于有時(shí)通道數(shù)大于NPU_NUM或者local memory的起始索引不為零，可能會(huì)導(dǎo)致不同通道的數(shù)據(jù)存儲(chǔ)在同一個(gè)NPU中，N_stride的公式也應(yīng)該做round- up 操作，如PPT中所示。

例如，我們將在具有 64 EU_BYTE 的 TPU 上處理形狀為 (2,3,4,5) 的fp16張量。其中包含了4個(gè)NPU，而本地內(nèi)存的起始索引設(shè)置為 0。

所以我們從NPU0開始存儲(chǔ)張量，W_stride和H_stride顯然是1和5。

對(duì)于C_stride，由于H * W小于EU_NUM，所以C_stride為32。

另外，因?yàn)檫@個(gè)張量的通道小于 NPU_NUM，所以N_stride 也是 32。

但是當(dāng)起始索引設(shè)置為2時(shí)，情況會(huì)有點(diǎn)不同，C步長仍然是32，但是由于張量的最后一個(gè)通道被存到第一個(gè)NPU，下一batch中的數(shù)據(jù)應(yīng)該從NPU2的下一行開始存儲(chǔ)，則 N_stride 應(yīng)為 64。

另一個(gè)常見的排布類型就是緊密排布，除了C_stride部分，其余的與對(duì)齊排布方式相似。