Google的TPU芯片的發(fā)展歷史和硬件架構

Google在高性能處理器與AI芯片主要有兩個系列：1)針對服務器端AI模型訓練和推理的TPU系列，主要用于Goggle云計算和數(shù)據(jù)中心；2)針對手機端AI模型推理的Tensor系列，主要用于Pixel智能手機。

結合最近幾年Google在HotChips、ISCA、ISSCC發(fā)布的論文和報告，總結了Google的TPU芯片的發(fā)展歷史和硬件架構，可作為學習、研發(fā)高性能處理器與AI芯片的參考資料。

1. TPUv1

Google第一代TPU芯片，服務器端推理芯片。

硬件架構 功能特性

1).TPU指令通過PCIe Gen3 x16總線從主機發(fā)送到指令緩沖區(qū)。矩陣乘法單元是TPU的核心，包含256x256個MAC，可以對有符號或無符號整數(shù)執(zhí)行8位乘法和加法。16位乘積被收集在矩陣單元下方的32位累加器的4 MiB中。4MiB表示4096256個元素的32位累加器。矩陣單元在每個時鐘周期產(chǎn)生一個256元素的部分和。

2).當混合使用 8 位權重和 16 位激活時（反之亦然），矩陣單元以半速計算，而當兩者都是 16 位時，它以四分之一速度計算。

3).省略了稀疏架構支持。稀疏性將在未來的設計中占據(jù)高度優(yōu)先地位。

4).TPU 指令遵循 CISC 傳統(tǒng)，包括重復字段。這些 CISC 指令的平均每條指令時鐘周期 (CPI) 通常為 10 到 20?？偣布s有 12 條指令，但以下 5 條是關鍵指令：Read_Host_Memory、Read_Weights、MatrixMultiply/Convolve、Activate、Write_Host_Memory。其他指令是備用主機內(nèi)存讀/寫、設置配置、兩個版本的同步、中斷主機、調(diào)試標記、nop 和暫停。

2. TPUv2

Google的第二代TPU，定位是服務端AI推理和訓練芯片。

硬件架構 TPUv2改變

單個向量存儲器，而不是固定功能單元之間的緩沖區(qū)。

通用向量單元，而不是固定功能激活管道。

連接矩陣單元作為向量單元的卸載。

將 DRAM 連接到內(nèi)存系統(tǒng)而不是直接連接到矩陣單元。

轉向 HBM 以獲得帶寬。

添加互連以實現(xiàn)高帶寬擴展。

TPUv2 Core

超長指令字架構：利用已知的編譯器技術。

線性代數(shù)ISA：標量、向量和矩陣，為通用性而構建。

TPU 核心：標量單元 322b VLIW 捆綁包：

2 個標量槽

4 個向量槽（2 個用于加載/存儲）

2 個矩陣插槽（推入、彈出）、

1 個雜項插槽

6 個立即數(shù)

存儲系統(tǒng)

針對 SRAM 暫存器進行加載和存儲

在核心內(nèi)提供可預測的調(diào)度

可能會因同步標志而停止

可通過異步 DMA 訪問

在同步標志中指示完成

互連器

具有 4 個鏈路的片上路由器

每個鏈路 500 Gbps

組裝成2D環(huán)面

軟件視圖：使用 DMA，就像 HBM 一樣；限制推送 DMA；只需定位另一個芯片 ID

3. TPUv3

TPU3是是對TPU2的溫和重新設計，采用相同的技術，MXU和HBM容量增加了兩倍，時鐘速率、內(nèi)存帶寬和ICI帶寬增加了1.3倍。TPU3超級計算機還可以擴展到1024個芯片。

硬件架構 功能特性

協(xié)同設計：具有軟件可預測性的簡化硬件（例如，VLIW、暫存器）。

使用 bfloat16 脈動陣列計算密度：HBM 為計算提供支持，XLA編譯器。

具有原則性線性代數(shù)框架的靈活大數(shù)據(jù)核心。

4. Edge TPU

Google發(fā)布的嵌入式TPU芯片，用于在邊緣設備上運行推理。

5. TPUv4i

TPUv4i：Google于2020年發(fā)布，定位是服務器端推理芯片.

硬件架構 功能特性 1).單核TPUv4i 用于推理，雙核 TPUv4（可擴展至 4096 個芯片）用于訓練。 2).選擇編譯器兼容性，而不是二進制兼容性。 3).通過通用內(nèi)存 (CMEM)增加了片上 SRAM 存儲。 4).四維張量 DMA 引擎充當協(xié)處理器，可完全解碼和執(zhí)行 TensorCore DMA 指令。 5).添加了一個共享片上互連 (OCI)，用于連接芯片上的所有組件。 6).引入了四輸入加法器運算單元。 7).時鐘頻率達到 1.05 GHz。 8).2個ICI鏈路鏈接板端4 個芯片。 9).具有廣泛的跟蹤和性能計數(shù)器等硬件功能。

6. TPUv4

谷歌2020年發(fā)布，服務器推理和訓練芯片，芯片數(shù)量是TPUv3的四倍。

硬件架構 功能特性 1).通過引入具有光學數(shù)據(jù)鏈路的光路交換機（OCS）來解決規(guī)模和可靠性障礙，允許 4K 節(jié)點超級計算機通過重新配置來容忍 1K CPU 主機在 0.1%–1.0% 的時間內(nèi)不可用。 2).公開了 DLRM（SparseCore 或 SC）中嵌入的硬件支持，DLRM 是自 TPU v2 以來 TPU 的一部分。 3).結合了前兩種功能，為超級計算機規(guī)?；ミB的需求添加了全對全通信模式。 編輯：黃飛

?