特斯拉自研芯片架構(gòu)帶來的啟示
如果我們說特斯拉只對機器學習感興趣是輕描淡寫的。事實上這家電動汽車制造商建造了一臺名為 Dojo 的內(nèi)部超級計算機,針對訓練其機器學習模型進行了優(yōu)化。
與許多其他超級計算機不同,Dojo 沒有使用現(xiàn)成的 CPU 和 GPU,例如 AMD、Intel 或 Nvidia 的。特斯拉根據(jù)他們的需求設(shè)計了自己的微架構(gòu),讓他們做出更通用架構(gòu)無法做出的權(quán)衡。
在本文中,我們將根據(jù)特斯拉在 Hot Chips 上的演示來了解該架構(gòu)。該架構(gòu)沒有單獨的名稱,因此為簡單起見,每當我們在后面提到 Dojo 時,我們都是在談?wù)摷軜?gòu)。
框圖
從高層次上看,Dojo 是一個 8 寬內(nèi)核,具有四路 SMT,運行在保守的 2 GHz,具有 CPU 風格的pipeline,使其比 GPU 更能容忍不同的算法和分支代碼。Dojo 的指令集在標量方面類似于 RISC-V,但 Tesla 的工程師添加了一組自定義向量指令,專注于加速機器學習。
特斯拉 Dojo 核心的簡化框圖
Tesla 將 Dojo 描述為“高吞吐量、通用 CPU”。從性能的角度來看,這肯定有一些道理。但為了提高計算密度,特斯拉做出了犧牲,與我們在臺式機、筆記本電腦和智能手機中熟悉的 CPU 相比,Dojo 內(nèi)核非常難以使用。在某些方面,Dojo 內(nèi)核的處理方式更像 IBM 的 Cell 中的 SPE,而不是傳統(tǒng)的通用 CPU 內(nèi)核。
像Cell的 SPE?
在 2000 年代中期推出的IBM Cell 處理器具有八個“協(xié)同處理元件”(Synergistic Processing Elements)或 SPE,由一個功能齊全的 CPU 內(nèi)核(“電源處理元件”或 PPE:Power Processing Elemen)控制。乍一看,Dojo 與 SPE 有很多相似之處。
Dojo 和 SPE 都針對矢量處理進行了優(yōu)化,并且依賴于單獨的主機處理器進行工作分配。在 Dojo 或 SPE 上運行的代碼不能直接訪問系統(tǒng)內(nèi)存。相反,應(yīng)用程序預(yù)計主要在一小部分本地 SRAM 中工作。此本地 SRAM 由軟件管理,不能用作緩存。如果需要來自主存儲器的數(shù)據(jù),則必須使用 DMA 操作將其引入。
最后,Dojo 和 Cell 的 SPE 都缺乏對虛擬內(nèi)存的支持。稍后我們將詳細介紹這意味著什么,但簡而言之,它使多任務(wù)處理變得非常困難。
Tesla Dojo 和 IBM Cell SPE 之間的非常簡化的比較
Dojo 在幾個重要方面有所不同。由于 Dojo 在設(shè)計時并未考慮小規(guī)模部署,因此主機處理器駐留在單獨的主機系統(tǒng)上。這些主機系統(tǒng)具有帶接口處理器的 PCIe 卡,然后通過高速網(wǎng)絡(luò)鏈路連接到 Dojo 芯片。相比之下,Cell 的主處理器駐留在同一個芯片上。這使得單獨部署單個 Cell 芯片成為可能——這在 Dojo 中是不可能的。
Dojo 的 1.25 MB 本地塊是 SRAM,它比只有 256 KB SRAM 的 Cell SPE 大得多,帶寬也更高。Cell 的 256 KB SRAM 只有一個端口,每個周期能夠提供 128B。Dojo 的 SRAM 有五個 64B 端口。當然,架構(gòu)目標是非常不同的。Dojo 是寬時鐘和低時鐘,而 Cell SPE 具有專為高時鐘設(shè)計的窄而深的pipelines。
前端:CPU Comforts等
讓我們從前端開始簡要介紹一下 Dojo 的pipelines。有某種分支預(yù)測器,因為特斯拉的圖表顯示了 BTB(branch target buffer:分支目標緩沖區(qū))。它的預(yù)測能力可能無法達到我們在 AMD、ARM 和 Intel 的高性能內(nèi)核上看到的水平,因為 Dojo 需要優(yōu)先考慮在矢量執(zhí)行上花費裸片區(qū)域。但即使是一個基本的分支預(yù)測器也比沒有預(yù)測器有了很大的進步,Dojo 的分支預(yù)測能力在處理分支代碼或更大的指令占用空間時應(yīng)該比 GPU 提供更好的性能。
一旦分支預(yù)測器生成了下一條指令提取指針,Dojo 可以每個周期從“小”指令緩存中提取 32 個字節(jié)到每個線程的提取緩沖區(qū)中(per-thread fetch buffers)。該指令緩存可能有助于降低本地 SRAM 的指令帶寬壓力,確保數(shù)據(jù)端可以訪問 SRAM 時盡可能少地爭用。此外,指令高速緩存不連貫。如果將新代碼加載到本地 SRAM 中,則必須在分支到該新代碼之前刷新指令緩存。
來自特斯拉的 Hot Chips 演示文稿
從提取緩沖區(qū)中,Dojo 的解碼器每個周期可以處理來自兩個線程的八條指令。我對“每個周期兩個線程”的含義有點困惑,因為具有 SMT 的 CPU 通常每個周期處理一個線程并在周期邊界上切換線程。也許 Dojo 將解碼器分成兩個集群,并選擇兩個線程在每個循環(huán)中為它們提供數(shù)據(jù)。這可能會減少所采用分支的解碼吞吐量損失。
在解碼時,某些指令(如branches, predicated operations和immediate loads (“l(fā)ist parsing”) )可以在前端執(zhí)行并從pipeline中刪除。這有點像更新的 x86 CPU 在重命名器(renamer)中消除了寄存器到寄存器的副本。但是您沒聽錯——Dojo 不會通過pipeline跟蹤“已消除”的指令以維持in-order retirement。其他處理器跟蹤所有內(nèi)容以退出,以便它們可以在任何指令邊界處停止,并保持恢復(fù)執(zhí)行所需的所有狀態(tài)。這種能力被稱為“precise exceptions”,現(xiàn)代操作系統(tǒng)使用它來提供各種好東西,比如paging to disk,或者告訴你你的代碼到底在哪里搞砸了。
特斯拉并不關(guān)心precise exceptions。Dojo 確實具有調(diào)試模式,其中更多指令通過pipeline以提供“更精確”的exceptions,但沒有像正常無序 CPU 那樣的有序退出邏輯。
Dojo 的執(zhí)行引擎
在看到一個寬前端之后,看到一個只有兩個 ALU 和兩個 AGU 的 4 寬整數(shù)執(zhí)行引擎可能會有點奇怪。但是這個漏斗形的流水線是有意義的,因為一些指令是在前端執(zhí)行和丟棄的。
Dojo 也不會進入客戶端系統(tǒng),其中標量整數(shù)性能很重要。因此,整數(shù)端提供了足夠的吞吐量來處理控制流和地址生成,以保持向量和矩陣單元的饋送。
Dojo 的向量和矩陣執(zhí)行端放置在核心pipeline中的標量執(zhí)行引擎之后,并且有兩個執(zhí)行pipeline。兩個pipeline可能聽起來不多,但 Dojo 在這些pipeline后面有非常寬的執(zhí)行單元。一個pipeline可以執(zhí)行 512 位向量執(zhí)行,而另一個pipeline執(zhí)行 8x8x4 矩陣乘法。因此,只要指令公開足夠的顯式并行性,Dojo 就能夠?qū)崿F(xiàn)非常高的吞吐量——尤其是在使用矩陣單元的情況下。Tesla 聲稱具有 354 個 Dojo 內(nèi)核的芯片在 2 GHz 時可以達到 362 BF16 TFLOPS,這表明每個內(nèi)核每個周期可以執(zhí)行 512 BF16 FLOPS。
我們不確定 Dojo 是否可以完全亂序執(zhí)行。但特斯拉確實表示,整數(shù)端可以在向量端之前運行很遠,這表明它可以執(zhí)行過去停滯的指令,直到其中一個調(diào)度程序被填滿。缺乏按序引退也指向無序執(zhí)行能力。
通常,實現(xiàn)無序執(zhí)行會帶來很多復(fù)雜性。這是因為 CPU 必須按順序執(zhí)行指令。AMD、ARM 和 Intel 的高性能 CPU 使用大型重新排序緩沖區(qū)(和其他結(jié)構(gòu))跟蹤指令,以便它們的結(jié)果可以按程序順序提交。這意味著如果一個程序做了一些愚蠢的事情,比如除以零,這些內(nèi)核可以準確地顯示出哪條指令出錯了。而且,它們可以顯示 CPU 狀態(tài),該狀態(tài)反映了故障之前所有指令的動作,但之后沒有。這意味著您可以修復(fù)導(dǎo)致指令錯誤的任何原因,并恢復(fù)執(zhí)行。Dojo 放棄了這種能力。作為交換,Dojo 可以避免與通過其pipeline跟蹤每條指令相關(guān)的功率和面積開銷,以確保按程序順序提交結(jié)果。
SRAM 訪問
通常,我們會在這里討論緩存。但是 Dojo 不能直接訪問系統(tǒng)內(nèi)存,所以我們將討論 1.25 MB 的 SRAM 塊。它每個周期可以處理兩個 512 位負載,使其每個周期的帶寬與支持 AVX-512 的英特爾 CPU 相匹配。Tesla 表示 SRAM 有五個 512 位端口(2 個加載端口、1 個存儲端口和兩個到網(wǎng)格停止端口)。但是標量端只有兩個 AGU,這可能意味著內(nèi)核無法在每個周期內(nèi)承受兩個 512 位負載和一個 512 位存儲。
因為 Dojo 的本地 SRAM 塊不是緩存,所以它不需要與數(shù)據(jù)一起存儲的標記和狀態(tài)位。SRAM 前面也沒有 L1D 緩存,因此它必須足夠快以處理所有加載和存儲指令而不會造成瓶頸,即使它的大小更接近 L2 緩存。不將 SRAM 實現(xiàn)為緩存可能是特斯拉保持低延遲的方式。如果我們回顧很久以前的 AMD 的 Hammer 架構(gòu),我們可以看到在檢測到 L1D 未命中后訪問 1 MB L2 需要 8 個周期(總延遲為 12 個周期)。如果直接尋址 1 MB 的 SRAM 而不是用作緩存,則可能會刪除三個階段,從而將延遲降低到 5 個周期:
來自較早的 Hot Chips 演示文稿的幻燈片,顯示了 L2 緩存訪問所涉及的pipeline階段。如果 L2 不是緩存,則可以跳過的階段標記為紅色。
考慮到幾十年的工藝節(jié)點改進,甚至比 Athlon 更低的時鐘速度目標,很容易看出特斯拉如何以類似 L1 的延遲訪問 L2 大小的 SRAM 塊。跳過一級緩存當然可以節(jié)省面積和功耗。
為了進一步減少延遲、面積和核心復(fù)雜性,Dojo 沒有虛擬內(nèi)存支持。因此,它沒有 TLB 或頁面遍歷機制(page walk mechanisms)?,F(xiàn)代操作系統(tǒng)利用虛擬內(nèi)存為每個進程提供自己的內(nèi)存視圖。程序使用的內(nèi)存地址不是直接訪問物理內(nèi)存地址,而是由 CPU 使用操作系統(tǒng)設(shè)置的分頁結(jié)構(gòu)轉(zhuǎn)換為物理地址。這就是現(xiàn)代操作系統(tǒng)如何使程序相互隔離,并防止一個行為不端的應(yīng)用程序?qū)е抡麄€系統(tǒng)癱瘓。
虛擬內(nèi)存也是您可以運行比物理內(nèi)存更多的程序的方式。當您用完實際內(nèi)存時,操作系統(tǒng)會取消映射頁面,將其寫入磁盤,并為您的程序提供所需的內(nèi)存。當其他糟糕的程序嘗試訪問該內(nèi)存時,CPU 會嘗試將虛擬地址轉(zhuǎn)換為物理地址,但發(fā)現(xiàn)轉(zhuǎn)換不存在。CPU 拋出一個頁面錯誤異常,操作系統(tǒng)通過將被驅(qū)逐的頁面讀回物理內(nèi)存并填寫頁表條目來處理該異常。
在 Dojo 上這些都不可行。內(nèi)核的 4 路 SMT 功能更多的是讓單個應(yīng)用程序公開顯式并行性,而不是提高多任務(wù)處理性能。例如,一個線程可以執(zhí)行向量計算,而另一個線程將數(shù)據(jù)從系統(tǒng)內(nèi)存異步加載到 SRAM(通過 DMA)。
為了進一步簡化設(shè)計,Dojo 可以只用 21 個地址位尋址 SRAM,這可以簡化 Dojo 的 AGU 和尋址總線。這些權(quán)衡可能讓特斯拉以足夠低的延遲訪問該 SRAM,以避免在其前面實現(xiàn)單獨的 L1 數(shù)據(jù)緩存。
內(nèi)存存取
說到系統(tǒng)內(nèi)存,Dojo 芯片并不直接連接到內(nèi)存。相反,它們連接到配備 HBM 的接口處理器。這些接口處理器還負責與主機系統(tǒng)通信。
具有 25 個獨立芯片的 Dojo tile 可以訪問 160 GB 的 HBM 內(nèi)存 Tesla 表示,他們可以跨 tile 邊界從每個芯片邊緣傳輸 900 GB/s,這意味著可以以 4.5 TB/s 的鏈路帶寬訪問接口處理器及其 HBM。因為訪問 HBM 需要通過單獨的芯片,所以訪問延遲可能非常高。
是小,如何做到的?
Dojo 是一個 8 寬的內(nèi)核,具有至少一些 OoO 執(zhí)行能力、不錯的向量吞吐量和一個矩陣乘法單元。但即使有 1.25 MB 的本地 SRAM,它最終還是一個非常小的內(nèi)核。相比之下,富士通的A64FX在同一工藝節(jié)點上占據(jù)的面積是其兩倍以上。處理器設(shè)計就是要做出正確的權(quán)衡。特斯拉希望通過將大量內(nèi)核封裝到芯片上來最大限度地提高機器學習的吞吐量,因此單個內(nèi)核必須很小。為了實現(xiàn)其區(qū)域效率,Dojo 使用了一些熟悉的技術(shù)。它以保守的 2 GHz 運行。較低的時鐘電路往往占用較少的面積。它可能有一個基本的分支預(yù)測器和一個小的指令緩存。如果程序的代碼占用量很大或有很多分支,這會犧牲一些性能。
但特斯拉還通過削減運行內(nèi)部工作負載所不需要的功能來進一步減少功耗和面積使用。它們不進行數(shù)據(jù)端緩存,不支持虛擬內(nèi)存,也不支持precise exceptions。
結(jié)果是一個處理器內(nèi)核提供了現(xiàn)代 CPU 內(nèi)核的性能靈活性,同時在許多方面不如 Intel 8086 對用戶和程序員友好。除了內(nèi)核之外,Tesla 還通過設(shè)計專門用于部署的 Dojo 芯片來節(jié)省芯片面積大規(guī)模。
物理實現(xiàn)
縮小,Dojo 核心是在一個非常大的 645 平方毫米的裸片上實現(xiàn)的,稱為 D1。與我們熟悉的其他芯片不同,單個 Dojo 芯片不能自給自足。它沒有 DDR 或 PCIe 控制器。die 邊緣周圍有 IO 接口,可以讓 die 與相鄰的 die 進行通信,延遲約為 100 ns。
要訪問系統(tǒng)內(nèi)存,Dojo D1 芯片必須與具有板載 HBM 的接口處理器通信。該接口處理器又通過 PCIe 連接到主機系統(tǒng)(接口處理器安裝在 PCIe 卡上)。理論上,最小的功能性 Dojo 部署將涉及一個 Dojo 芯片、一個接口處理器卡和一個主機系統(tǒng)。但特斯拉將 Dojo die部署在每個包含 25 個die的模塊中,以提供規(guī)模感。Dojo D1 die專門設(shè)計用于超級計算機的構(gòu)建塊,僅此而已。
這種專業(yè)化可以節(jié)省更多的裸片面積。Dojo D1 不會在 DDR 和 PCIe 控制器上花費空間。大部分裸片都被大量 Dojo 核心占據(jù),除了外部設(shè)計用于與相鄰裸片接口的定制 IO 連接器。
相比之下,設(shè)計時考慮到更多部署靈活性的芯片在 IO 上花費了大量空間。AMD 的 Zen 1 “Zeppelin” 芯片很好地展示了這一點。Zeppelin 可以直接連接到 DDR4 內(nèi)存、PCIe 設(shè)備、SATA 驅(qū)動器和 USB 設(shè)備——非常適合滿足客戶要求。在服務(wù)器中,IFOP 接口讓它與相鄰的芯片進行通信。位于 IFOP 附近的大塊 SRAM 可能是窺探過濾器(snoop filters),有助于在高核心數(shù)設(shè)置中有效地保持緩存一致性。Dojo 不會嘗試跨內(nèi)核保持緩存一致性,并且不會將任何 SRAM 用于窺探過濾。
來自 AMD 在 ISSCC 2021 上的小芯片演示
AMD 通過將大約 44% 的“Zeppelin”裸片區(qū)域用于除內(nèi)核和緩存之外的邏輯來為這種靈活性付出代價。Dojo 僅將 28.9% 的裸片面積用于 SRAM 和內(nèi)核以外的東西。
最后的話
特斯拉 Dojo 超級計算機背后的微架構(gòu)展示了如何實現(xiàn)非常高的計算密度,同時仍保持 CPU 處理分支代碼的能力。要到達那里,您需要放棄定義我們現(xiàn)代計算體驗的大部分舒適設(shè)施。如果您可以假設(shè)圍繞 Dojo 核心構(gòu)建您的桌面,那么任何使用過 MS-DOS 的人都可能會覺得這種體驗很熟悉。您不能同時運行多個應(yīng)用程序。單個行為不端的應(yīng)用程序可能會迫使您重新啟動系統(tǒng)。如果您沒有足夠的 RAM 來運行某個程序,您可以完全忘記運行它(無需分頁到磁盤)。
但這些權(quán)衡在特斯拉的超級計算機中非常有意義。Tesla 不需要 Dojo 內(nèi)核來同時處理多個正在運行的應(yīng)用程序,Dojo 只需運行內(nèi)部的可信代碼。所以,特斯拉并不關(guān)心虛擬內(nèi)存支持。同樣,在 Dojo 上運行的機器學習程序?qū)⒃诰帉憰r考慮到該特定系統(tǒng)。您不會有一批任意程序可能要求比可用內(nèi)存更多的內(nèi)存。這意味著您不需要precise exceptions(和虛擬內(nèi)存)來允許諸如過度使用內(nèi)存、內(nèi)存壓縮或在磁盤之間交換內(nèi)存頁面等技術(shù)。precise exceptions對于調(diào)試也很有用,但 Tesla 通過單獨的調(diào)試模式以更便宜的方式進行調(diào)試。
可以肯定的是,特斯拉為獲得高計算密度而采用的權(quán)衡取舍在消費者或服務(wù)器 CPU 中是不可能的。但他們在行動中非常有趣,我們必須感謝特斯拉花時間在 Hot Chips 上展示。 在過去的二十年里,進程節(jié)點的改進一直在放緩,導(dǎo)致單線程性能的進步更慢。在過去的五年中,電源和冷卻限制一直在降低多線程性能。但對更多計算能力的需求并沒有放緩,因此公司正在轉(zhuǎn)向更專業(yè)的硬件以跟上步伐。
Tesla 的 Dojo 超級計算機中的架構(gòu)是一個很好的例子,說明了如何進行權(quán)衡以提高計算密度,以及當前趨勢如何有利于為吞吐量受限的應(yīng)用程序引入專用硬件。