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

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

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

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

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

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

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

指令主要有兩種：

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

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

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

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

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

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

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

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

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

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

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

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

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