TPU內存（一）

首先我們來看一下TPU的簡要架構。

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

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

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

而local memory雖然有限但在計算速度上更有優勢。

所以有時候對于一個很大的張量，我們需要把它切分成幾個部分，送到local memory中進行計算，然后把結果存回global memory。

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

指令主要有兩種：

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

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

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

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

此外，TPU 可以同時處理的數據數量取決于每個 NPU 上的執行單元數量。

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

例如當EU Bytes為64時，則表示一個NPU可以同時處理64個int8數據。

同理，我們可以根據數據的字節計算出對應的EU_NUM。

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

單個通道中每個bank的內存大小則為16KB，相當于16x1024 字節。

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

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

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