TPU內(nèi)存（一）

首先我們來看一下TPU的簡要架構(gòu)。

我們可以從ppt中了解到一個TPU中包含多個NPU（Neuron Processing Unit），主要由一個local memory和多個執(zhí)行單元組成。前者用于存儲要運算的數(shù)據(jù)，后者是TPU上最小的計算單元。每個NPU一次可以驅(qū)動它的所有EU做一個MAC操作。

就整體 TPU 內(nèi)存而言，它由system memory和local memory組成。 system memory的主要部分是global memory，其實就是一塊DDR。 有時根據(jù) TPU 的特殊設(shè)計還會有其他組件，但我們不會在視頻中提及這些部分，所以現(xiàn)在了解global memory就足夠了。 而對于local memory，我們暫時只需要知道是一組Static RAM就可以了。 稍后我會進一步解釋。

通常global memory很大，用于存儲來自host端的整個數(shù)據(jù)塊。

而local memory雖然有限但在計算速度上更有優(yōu)勢。

所以有時候?qū)τ谝粋€很大的張量，我們需要把它切分成幾個部分，送到local memory中進行計算，然后把結(jié)果存回global memory。

為了在 TPU 上執(zhí)行這些操作，我們就需要用到指令。

指令主要有兩種：

1. GDMA用于system memory和local memory間或system memory內(nèi)的數(shù)據(jù)傳輸；
2. BDC用于驅(qū)動執(zhí)行單元在NPU上做計算工作；
   另外，對于那些不適合并行加速的計算，比如NMS，SORT，我們還需要HAU指令，但是這意味著我們需要額外的處理器。

對于local memory的構(gòu)成，它是由多個Static RAM組成的。每個 SRAM 稱為一個bank。此外，我們將這些 SRAM 分成多個部分給同樣數(shù)量的NPU，每個部分稱為一個lane。

而對于每個NPU，它只能訪問屬于它的那部分local memory，這使得單個NPU的執(zhí)行單元只能處理自己local memory上的那部分張量。

一旦我們調(diào)用單個 BDC 指令，所有 NPU 的執(zhí)行單元將在每個 NPU 的相同位置執(zhí)行相同的操作。 這就是 TPU 加速運算的方式。

此外，TPU 可以同時處理的數(shù)據(jù)數(shù)量取決于每個 NPU 上的執(zhí)行單元數(shù)量。

對于一個特定的TPU，EU Bytes是固定的，所以對于不同類型的數(shù)據(jù)，EU的個數(shù)會有所不同。

例如當(dāng)EU Bytes為64時，則表示一個NPU可以同時處理64個int8數(shù)據(jù)。

同理，我們可以根據(jù)數(shù)據(jù)的字節(jié)計算出對應(yīng)的EU_NUM。

對于地址分配，假設(shè)我們的local memory由16個SRAM組成，總內(nèi)存為16MB，有64個NPU，那么每個NPU的內(nèi)存為256KB。

單個通道中每個bank的內(nèi)存大小則為16KB，相當(dāng)于16x1024 字節(jié)。

所以這個塊的地址范圍是從0到16x1024 – 1。

同理，NPU0中下個bank的地址從16x1024開始到32x1024-1

按照這個規(guī)則，我們就可以得到local memory上的所有地址。