特斯拉超算芯片，可超越GPGPU？
作者：陳巍 博士存算一體/GPU架構和AI專家，高級職稱。中關村云計算產業聯盟，中國光學工程學會專家，國際計算機學會（ACM）會員，中國計算機學會（CCF）專業會員。曾任AI企業首席科學家、存儲芯片大廠3D NAND設計負責人，主要成就包括國內首個大算力可重構存算處理器產品架構（已在互聯網大廠完成原型內測），首個醫療領域專用AI處理器（已落地應用），首個RISC-V/x86/ARM平臺兼容的AI加速編譯器（與阿里平頭哥/芯來合作，已應用），國內首個3D NAND芯片架構與設計團隊建立（與三星對標），國內首個嵌入式閃存編譯器（與臺積電對標，已平臺級應用）。
2022年9月最后一天，特斯拉的人工智能日，特斯拉“擎天柱”機器人正式登臺亮相。按照特斯拉工程師的說法，2022 年人工智能日是特斯拉機器人第一次在沒有任何外部支持的情況下被“放出“。“他“步姿端莊，大方向場上觀眾打招呼。除了動作稍微遲緩之外，其他都很自然。
1 特斯拉機器人強大在于“內芯”？
特斯拉展示了機器人在辦公室周圍“工作”的視頻。名為擎天柱的機器人搬運物品，給植物澆水，甚至自主的在工廠工作了一段時間?！蔽覀兊哪繕耸潜M快制造出有用的人形機器人”，特斯拉表明，他們的目標是讓機器人的價格低于 2 萬美元，或者比特斯拉的電動汽車便宜。
特斯拉機器人之所以這么強，除了特斯拉本身在AI技術的積累外，更主要得益于特斯拉強勁的自研AI芯片。這顆AI芯片，不是傳統上的CPU，更不是GPU，是一種更適合復雜AI計算的形態。
D1處理器與其他自動駕駛/機器人處理器的對比
1.1 基于數據流近存架構打造一顆超越GPU的通用AI芯片
特斯拉打造自有芯片的原因是，GPU 并不是專門為處理深度學習訓練而設計的，這使得GPU在計算任務中的效率相對較低。特斯拉與 Dojo（Dojo既是訓練模組的名稱，又是內核架構名稱） 的目標是“實現最佳的 AI 訓練性能。啟用更大、更復雜的神經網絡模型，實現高能效且經濟高效的計算?！?特斯拉的標準是制造一臺比其他任何計算機都更擅長人工智能計算的計算機，從而他們將來不需要再使用 GPU。
構建超級計算機一個關鍵點是如何在擴展計算能力同時保持高帶寬（困難）和低延遲（非常困難）。特斯拉給出的解決方案是強大的芯片和獨特的網格結構組成的分布式 2D 架構（平面），或者說是數據流近存計算架構。
特斯拉算力單元的層級劃分
按照層次劃分的話，每354個Dojo核心組成一塊D1芯片，而每25顆芯片組成一個訓練模組。最后120個訓練模組組成一組ExaPOD計算集群，共計3000顆D1芯片。
一個特斯拉Dojo芯片訓練模組可以達到6組GPU服務器的性能，成本卻少于單組GPU服務器。單臺Dojo服務器算力甚至達到了54PFLOPS。只用 4 個 Dojo 機柜就能取代由 4000 顆 GPU 組成的 72 組 GPU 機架。Dojo 將通常需要幾個月的AI計算（訓練）工作減少到了1 周。這樣的“大算力出奇跡”，與特斯拉自動駕駛的風格一脈相承。顯然芯片也會大大加速特斯拉AI技術的進步速度。
當然，這一芯片模組還沒有到達“完美”的程度，盡管采用了數據流近存計算的思路，其算力能效比并沒有超過GPU。單個服務器的功耗巨大，電流達到了2000A，需要特殊定制的電源供電。特斯拉D1芯片已經是近存計算架構的結構極限了。如果特斯拉采用“存內計算”或者“存內邏輯”架構，或許芯片性能或能效比還會有大幅度提升。
特斯拉Dojo芯片服務器由12個Dojo訓練模組組成（2層，每層6個）
1.2 特斯拉的Dojo架構設計哲學
Dojo核心是一個8路譯碼的內核，具有較高吞吐量和4路矩陣計算單元（8x8）以及 1.25 MB 的本地 SRAM。但是Dojo核心的尺寸卻不大，相比之下，富士通的A64FX在同一工藝節點上占據的面積是其兩倍以上。
通過Dojo核心的結構，我們可以看出特斯拉在通用AI處理器上的設計哲學：
- 面積精簡：特斯拉通過將大量計算內核集成到芯片中，以最大限度提高AI計算的吞吐量，因此需要在保障算力的情況下使單個內核的面積盡可能小，更好的折衷超算系統中算力堆疊和延遲的矛盾。
- 緩存與延遲精簡：為了實現其區域計算效率最大化，Dojo內核以相對保守的 2 GHz 運行（保守時鐘電路往往占用較少的面積），只使用基本的分支預測器和小的指令緩存，在如此精簡只保留必要部件的架構下。其余面積盡可能留給向量計算和矩陣計算單元。當然，如果內核程序的代碼占用量很大，或分支較多時，這種策略可能會犧牲一些性能。
- 功能精簡：通過削減對運行內部計算不是必須的處理器功能來進一步減少功耗和面積使用。Dojo核心不進行數據端緩存，不支持虛擬內存，也不支持精確異常。
對于特斯拉和馬斯克而言，Dojo不僅僅形狀布局像道場，其設計哲學也與道場的精神息息相關，充分體現了“少即是多”的處理器設計美學。
2 D1核心是RISC-V架構嗎？
我們先來看看每個Dojo的結構和特點。
每個Dojo核心是帶有向量計算/矩陣計算能力的處理器，具有完整的取指、譯碼、執行部件。Dojo核心具有類似CPU的 風格，似乎比GPU 更能適應不同的算法和分支代碼。D1的指令集類似于 RISC-V，處理器運行在2GHz，具有4組8x8矩陣乘法計算單元。同時具有一組自定義向量指令，專注于加速AI計算。
對RISC-V領域熟悉的大概能看出，特斯拉Dojo架構圖的配色方案像是在致敬伯克利的BOOM處理器架構圖，上黃中綠下紫。
特斯拉Dojo核心與伯克利BOOM/ IBM Cell核心對比
2.1 D1核心整體架構
D1核心結構（藍色部分為添加/修改的細節）
從目前的架構圖來看，Dojo核心由前端、執行單元、SRAM和NoC路由4部分組成，比CPU和GPU的控制部件都更少，具有類似CPU的AGU和思路類似GPU張量核心（Tensor core）的矩陣計算單元。
Dojo核心結構比BOOM更加精簡，沒有Rename這些改善執行部件利用率的組件，同時也難于支持虛擬內存。但這樣設計的好處是減少了控制部分占用的面積，可以把芯片上更多的面積劃分給計算執行單元。每個Dojo核心提供了1.024TFLOPS的算力?？梢钥吹?，每個幾乎所有的算力都由矩陣計算單元提供。因而矩陣計算單元和SRAM共同決定了D1處理器的計算能效比。
Dojo核心的主要參數
- 分支預測：相對GPU這類SIMT架構，Dojo核心也沒有SIMT堆棧核心來進行多線程分支任務的分配。但Dojo核心具有 BTB（分支目標緩沖區），因此D1可以通過簡單的分支預測來提升性能。
BTB將分支成功的分支指令的地址和它的分支目標地址都放到一個緩沖區中保存起來，緩沖區以分支指令的地址作為標識。可以通過預測分支的路徑和緩存分支使用的信息來減少流水線處理器中分支的性能損失。 - 指令緩存：較小的L1指令緩存直接與核心中的SRAM相連獲取計算指令。
- 取指：每個Dojo內核具有 32 B 的取指窗口，最多可容納 8 條指令。
- 譯碼：一個8路解碼器每個周期可以處理兩個線程。譯碼階段從取指緩沖獲取指令并譯碼，并根據每條指令的要求分配必要的執行資源。
- 線程調度：在較寬的8路譯碼之后，則是向量的調度器（Scheduler）和寄存器堆（Register File）。貌似這里沒有分支聚合的掩碼判斷，實際的分支執行效率可能會比GPU略低。希望特斯拉有一個強大的編譯器吧。
- 執行單元：具有2路ALU和2路AGU，以及針對向量/矩陣計算的512位SIMD和矩陣計算單元（分別執行512位向量計算和4路8x8矩陣乘法）。其中矩陣計算單元是D1芯片的算力主體。（在下一節具體介紹）
- ALU和AGU主要負責矩陣計算之外的少量邏輯計算。其中AGU是地址生成單元，主要用于生成操作SRAM所需的地址和訪問其他核心的地址。通過由與 CPU 的其余部分并行運行地址計算。
普通CPU 在執行各種操作時，需要計算從內存（或SRAM）中取數據所需的內存地址。例如，必須先計算數組元素的內存位置，然后 CPU內核才能從實際內存位置獲取數據。這些地址生成計算涉及不同的整數算術運算，例如加法、減法、模運算或位移。計算內存地址可以編譯多個通用機器指令，也可以類似特斯拉Dojo這樣通過AGU的硬件電路直接執行。這樣各種地址生成計算可以從ALU卸載，減少執行AI計算所需等待的CPU 周期數，從而提高計算性能。 - SIMD主要負責激活等特殊功能計算和數據的累加。
- 矩陣計算單元是Dojo的主要算力原件，負責二維矩陣計算，進而實現卷積、Transformer等計算。
Intel Nehalem架構中使用AGU來提升單周期地址訪問效率
Dojo內核的連接方式比較像 IBM 的 Cell處理器中的SPE內核連接方式。主要的相似點包括：
- D1或 SPE 上運行的代碼都不能直接訪問系統內存，應用程序主要在本地 SRAM 中工作；
- 如果需要來自主存儲器（DDR或HBM）的數據，須使用 DMA 操作進行讀入
- D1 和 Cell 的 SPE 都不支持虛擬內存。
下面將介紹計算與矩陣乘法模塊與內核的存儲。
2.2 算力核心矩陣計算單元與片內存儲
Dojo架構算力增強的核心是矩陣計算單元。矩陣計算單元與核心SRAM的數據交互構成了主要的內核數據搬運功耗。
特斯拉矩陣計算單元相應的專利如下圖。該模塊關鍵部件是一個8x8矩陣-矩陣乘法單元（圖中稱為矩陣計算器）。輸入為數據輸入陣列和權重輸入陣列，計算矩陣乘法后直接在輸出進行累加。每個Dojo核心包括4路8x8矩陣乘法單元。
特斯拉矩陣計算單元專利
由于架構圖上只有一個L1 緩存和SRAM，大膽猜測特斯拉精簡了RISC-V的緩存結構，目的是節約緩存面積并減少延遲。每個核心1.25MB的SRAM塊可以為SIMD和矩陣計算單元提供2x512位的讀（對應AI計算的權重和數據）和512位的寫帶寬，以及面向整數寄存器堆的64位讀寫能力。計算的主要數據流是從SRAM到SIMD和矩陣乘法單元。
矩陣計算單元的主要處理流程為：
通過多路選擇器（Mux）從SRAM中加載權重到權重輸入陣列（Weight input array），同時SRAM中加載數據到數據輸入陣列（Data input array）。
輸入的數據與權重在矩陣計算器（Matrix computation Unit）中進行乘法計算（內積或外積？）
乘法計算結果輸出到輸出累加（Output accumulator）中進行累加。這里計算時可以通過矩陣劃分拼接的方式進行超過8x8的矩陣計算。
累加后的輸出傳入后處理器寄存器堆進行緩存，隨后進行后處理（可執行例如激活、池化、Padding等操作）。
整個計算流程由控制單元（Control unit）直接控制，無需CPU干預。
執行單元與SRAM/NoC的數據交互
Dojo核心內的SRAM具有非常大的讀寫帶寬，可以以 400 GB/秒的速度加載并以 270 GB/秒的速度寫入。Dojo核心指令集具有專用的網絡傳輸指令，通過NoC路由，可以直接將數據移入或移出 D1 芯片中甚至Dojo訓練模塊中其他內核的SRAM 存儲器。
與普通的SRAM不同，Dojo的SRAM包括列表解析引擎（list parser engine）和一個收集引擎（gather engine）。列表解析功能是 D1芯片的關鍵特性之一，通過列表解析引擎可以將復雜的不同數據類型的傳輸序列進行打包，提升傳輸效率。
列表解析功能
為了進一步減少操作延遲、面積和復雜度，D1 并不支持虛擬內存。在通常的處理器中，程序使用的內存地址不是直接訪問物理內存地址，而是由 CPU 使用操作系統設置的分頁結構轉換為物理地址。
在 D1內核中， 4 路 SMT 功能讓計算具備顯式并行性，簡化 AGU 和尋址計算方式，以讓特斯拉以足夠低的延遲訪問 SRAM，其優勢是可避免中間L1 數據緩存的延遲。
2.3 Dojo指令集
D1處理器指令集
D1參考了RISC-V 架構的指令，并且自定義了一些指令，特別是矢量計算相關的指令。
D1指令集支持 64 位標量指令和 64字節 SIMD 指令，網絡傳輸與同步原語和機器學習/深度學習相關的專用原語（例如8x8矩陣計算）。
在網絡數據傳輸和同步原語方面，支持從本地存儲（SRAM）到遠程存儲傳輸數據的指令原語（Primitives），以及信號量（Semaphore）和屏障約束（ Barrier constraints）。這可以使D1支持多線程，其存儲操作指令可以在多個 D1 內核中運行。
針對機器學習和深度學習，特斯拉定義了包括 shuffle、transpose 和 convert 等數學操作的指令，以及隨機舍入（ stochastic rounding ），padding相關的指令。
2.4 數據格式
D1核心具備FP32和FP16這兩個標準的計算格式，同時還具備更適合Inference的BFP16格式。為了達到混合精度計算提升性能的目的， D1還采用了用于較低精度和更高吞吐量的 8 位 CFP8 格式。
采用CFP8的優勢在于可以節約更多的乘法器空間實現幾乎同樣的算力，這對提升D1的算力密度非常有幫助。
Dojo 編譯器可以在尾數精度附近滑動，以涵蓋更廣泛的范圍和精度。在任何給定時間，最多可以使用 16 種不同的矢量格式，靈活提升算力。
D1處理器的數據格式
根據特斯拉提供的信息，在矩陣乘法單元內部可使用CFP8來進行計算（存儲為CFP16格式）。
3 Dojo架構處理器能否超過GPU？
D1處理器由臺積電制造，采用7納米制造工藝，擁有500億個晶體管，芯片面積為645mm2，小于英偉達的A100（826 mm2）和AMD Arcturus（750 mm2）。
3.1 Dojo數據流近存計算架構
D1處理器結構
每個D1處理器由 18 x 20 的Dojo核心拼接構成。每個D1處理器中有354個Dojo核心可用。（之所以只使用360個核心中的354個是出于良率和每處理器核心穩定考慮）由臺積電制造，采用7nm制造工藝，擁有500億個晶體管，芯片面積為645mm2。
每個Dojo核心有一塊1.25MB的SRAM作為主要的權重和數據存儲。不同的Dojo核心通過片上網絡路由（NoC路由）進行連接，不同的Dojo內核通過復雜的NoC網絡進行數據同步，而不是共享數據緩存。NoC 可以處理跨節點邊界4個方向（東南西北）的 8 個數據包，每個方向 64 B/每個時鐘周期，即在所有四個方向上一個數據包輸入和一個數據包輸出到網格中每個相鄰的Dojo核心。該NoC路由還可以在每個周期對核心內的 SRAM 進行一次 64 B 雙向讀寫。
跨處理器傳輸和D1處理器內部的任務劃分
每個Dojo核心都是一個相對完整的帶矩陣計算能力的類CPU（由于每個核心具備單獨的矩陣計算單元，且前端相對較小，所以這里稱為類CPU）其數據流架構則有點類似于SambaNova的二維數據流網格結構，數據直接在各個處理核心之間流轉，無需回到內存。
D1芯片運行在2GHz，擁有巨大的440MB SRAM。特斯拉將設計重心放在計算網格中的分布式SRAM，通過大量更快更近的片上存儲和片上存儲之間的流轉減少對內存的訪問頻度，來提升整個系統的性能，具有明顯的數據流存算一體架構（數據流近存計算）特征。
每顆D1 芯片有 576 個雙向 SerDes 通道，分布在四周，可連接到其他 D1 芯片，單邊帶寬為 4 TB/秒。
D1處理器芯片主要參數
3.2 Dojo訓練模組的Chiplet封裝互連技術
每個D1訓練模塊由5x5的 D1芯片陣列排布而成，以二維Mesh結構互連。片上跨內核SRAM達到驚人的11GB，當然耗電量也達到了15kW的驚人指標。能效比為0.6TFLOPS/W@BF16/CFP8。（希望是我算錯了，否則這個能效比確實不是太理想）。外部32GB共享HBM內存。（HBM2e或HBM3）
特斯拉D1處理器的散熱結構專利
特斯拉使用了專用的電源調節模塊（VRM）和散熱結構來進行功耗管理。在這里功耗管理的主要目的有2個：
減少不必要的功耗損失，提升能效比。
減少散熱形變造成的處理器模組失效。
根據特斯拉的專利，我們可以看到電源調節模塊與芯片本身垂直，極大的減少了對處理器平面的面積占用，且可以通過液冷來迅速平衡處理器的溫度。
特斯拉D1處理器的散熱和封裝結構專利
訓練模組在封裝上采用InFO_SoW（Silicon on Wafer）封裝來提高芯片間的互連密度。該封裝除了TSMC的INFO_SoW技術之外，也采用了特斯拉自己的機械封裝結構，以減小處理器模組的失效。
每個訓練模塊外部邊緣的 40 個 I/O 芯片達到了 36 TB/s的聚合帶寬，或者10TB/s的橫跨帶寬。每層訓練模塊都連接著超高速存儲系統：640GB 運行內存可以提供超過 18TB/s的帶寬，另外還有超過 1TB/s的網絡交換帶寬。
數據傳輸方向與芯片平面平行，供電及液冷方向與芯片平面垂直。這是一個非常優美的結構設計，不同的訓練模塊之間還可以互連。通過立體結構，節約了芯片模組的供電面積，盡可能減少計算芯片間的距離。
一個 Dojo POD 機柜由兩層計算托盤和存儲系統組成。每一層托盤都有 6 個 D1 訓練模組。兩層共 12個訓練模組組成的一個機柜，可提供 108PFLOPS 的深度學習算力。
Dojo模組與Dojo POD機柜
3.3 電源管理與散熱控制
超算平臺的散熱，一直是衡量超算系統水平的重要維度。
D1 芯片的熱設計功率
特斯拉在 Dojo POD 上使用了全自研的 VRM（電壓調節模組），單個 VRM可以在不足 25 美分硬幣面積的電路上，提供52V電壓和超過 1000A 的巨大電流，電流目的為0.86A每平方毫米，共計12個獨立供電相位。
特斯拉的電源調節模組
對高密度芯片散熱而言，其重點是控制熱膨脹系數（CTE）。Dojo系統的芯片密度極高，如果CTE稍微失控，都可能導致結構變形/失效，進而出現連接故障。
特斯拉這套自研 VRM 在過去2年內迭代了 14 個版本，采用了MEMS振蕩器（MO）來感知電源調節模組的熱形變，最終才完全符合內部對 CTE 指標的要求。這種通過MEMS技術主動調節電源功率的方式，與控制火箭箭身振動的主動調節方式類似。
3.4 Dojo架構處理器的編譯生態
D1處理器軟件棧
對于D1這類AI芯片來說，編譯生態的重要性不低于芯片本身。
在D1處理器平面上，D1被劃分為矩陣式的計算單元。編譯工具鏈負責任務的劃分和配置數據存儲，并且通過多種方式進行細粒度的并行計算，并減少存儲占用。
Dojo編譯器支持的并行方式包括數據并行、模型并行和圖并行。支持的存儲分配方式包括分布式張量、重算分配和分割填充。
編譯器本身可以處理各種CPU中常用的動態控制流，包括循環和圖優化算法。借助Dojo編譯器，用戶可將Dojo大型分布式系統視作一個加速器進行整體設計和訓練。
整個軟件生態的頂層基于PyTorch，底層基于Dojo驅動，中間使用Dojo編譯器和LLVM形成編譯層。這里加入LLVM后，可以使特斯拉更好的利用LLVM上已有的各種編譯生態進行編譯優化。
特斯拉Dojo 編譯器
4 結語
通過特斯拉AI日，我們看到了特斯拉機器人的真身，并且對其強大的“內芯”有了更多的認識。
特斯拉的Dojo核心與以往的CPU和GPU架構特點都有差別，可以說是結合了CPU特點的精簡GPU，相信其在編譯上也會與CPU和GPU有較大的差異。為了提升計算密度，特斯拉做了極致精簡的優化，并且提供了主動調節的電源管理機制。
特斯拉Dojo架構不止名為道場，其設計也確實以簡為道，以少為多。那這種架構會不會成為繼CPU和GPU之后的另一算力芯片架構典型形態呢？讓我們拭目以待。
