如何設計一顆AI芯片？AI芯片設計入門方案

1介紹

機器學習（ML）已經在在線活動中變得無處不在。近年來，這些模型的規(guī)模和復雜性大幅增長，這有助于提高預測的準確性和有效性。然而，與此同時，這種增長給用于大規(guī)模訓練和推理這些模型的硬件平臺帶來了巨大挑戰(zhàn)?？倱碛谐杀荆═CO）是在數(shù)據(jù)中心將模型投入生產的主要制約因素之一，而功率是這些平臺TCO的重要組成部分。因此，單位TCO的性能（以及每瓦的性能）已成為針對機器學習的所有硬件平臺的重要衡量標準。

深度學習推薦模型（DLRM）已成為Meta數(shù)據(jù)中心最主要的工作負載之一。這些模型將計算密集型的傳統(tǒng)多層感知器（MLP）操作（有時稱為全連接或FC）與將稀疏特征轉換為密集表示的嵌入表相結合。這些表包含隨機索引的寬向量，并將其簡化為單個向量，然后將其與來自其他層的數(shù)據(jù)組合以產生最終結果。雖然嵌入表操作的計算要求相對較低，但由于數(shù)據(jù)訪問模式的性質和表的大小，它們的內存占用和帶寬要求依然相對較高。

圖1顯示了Meta生產數(shù)據(jù)中心中與推薦模型相關的推理工作負載的復雜性和內存占用的歷史增長和預計未來增長。虛線顯示了模型計算需求的估計增長，而實線顯示了內存占用的增加。灰色實線捕捉了用于存儲嵌入表的設備內存的占用空間，這是這些模型的重要組成部分。計算和內存需求的增長水平是一個需要解決的問題，尤其需要考慮到這些工作負載通常是如何在數(shù)據(jù)中心運行。

圖1：推理模型的增長趨勢

2動機

傳統(tǒng)上，CPU被用作Meta生產數(shù)據(jù)中心中服務推理工作負載的主要工具，但它們在滿足最新工作負載的需求方面并不具有成本效益。在某種程度上，硬件加速被認為是一種有吸引力的解決方案，它可以解決功率和性能問題，并提供一種更有效的方式來服務推理請求，同時為運行未來的模型提供足夠的計算性能余量。

圖2顯示了過去幾年中為數(shù)據(jù)中心內的推理工作負載部署的服務器的估計數(shù)量。淺色實線顯示基于CPU的服務器數(shù)量，虛線顯示配備第一代推理加速器Intel NNPI的服務器數(shù)量；深色實線顯示基于GPU的服務器數(shù)量。雖然使用NNPI加速器暫時滿足了對增加容量的最初需求，但對推理模型的要求很快超過了NNPI能力，并為使用GPU提供了動力。這帶來了利用已經用于訓練的現(xiàn)有生態(tài)系統(tǒng)的額外優(yōu)勢。因此，可以觀察到，GPU作為加速器越來越多地滿足了對模型復雜性的日益增長的需求。

雖然最近幾代GPU提供了大量的內存帶寬和計算能力，但它們在設計時沒有考慮到推理，因此處理實際推理工作負載的效率很低。開發(fā)人員使用了無數(shù)的軟件技術，如算子融合、圖轉換和內核優(yōu)化，以提高GPU的效率。但是，盡管做出了這些努力，仍然存在效率差距，這使得在實踐中部署模型具有挑戰(zhàn)性且成本高昂。

圖2：服務器對推理工作負載的需求增長

基于部署NNPI和GPU作為加速器的經驗，很明顯對于重要的推理工作負載，還有更優(yōu)化的解決方案的空間。這種最佳解決方案基于內部加速器，該加速器從頭開始構建，以滿足推理工作負載的苛刻要求，尤其側重于滿足DLRM系統(tǒng)的性能要求。然而，在關注DLRM工作負載的同時（考慮到它們正在進行的演變，以及該架構是為這些工作負載的下一代有效構建的情況），很明顯，除了性能之外，該架構還應該提供足夠的通用性和可編程性，以支持這些工作負載的未來版本以及潛在的其他類型的神經網絡模型。

雖然創(chuàng)建定制的硅解決方案為目標工作負載的充分創(chuàng)新和專業(yè)化打開了大門，但為數(shù)據(jù)中心的大規(guī)模部署創(chuàng)建加速器架構是一項艱巨的任務。因此，構建加速器時的重點和策略一直是采用和重用供應商和開源社區(qū)提供的技術以及工具和環(huán)境。這不僅縮短了上市時間，而且還利用了來自社區(qū)和供應商的支持和增強功能，減少了構建、啟用和部署此類平臺所需的資源量。

本文的其余部分解釋了Meta的第一個針對推理工作負載的加速器芯片MTIA的架構，以及隨之而來的學習。下一節(jié)詳細介紹了加速器的架構及其提供的各種功能和組件。第4節(jié)介紹了將一個示例算子映射到此體系結構，展示了如何利用各種提供的功能來高效地運行。第5節(jié)概述了加速器的軟件堆棧，第6節(jié)描述了我們的評估方法和結果。最后，第7節(jié)討論了在這個開發(fā)周期中吸取的一些重要經驗教訓。

3加速器架構

圖3顯示了加速器的高級架構，它被組織為連接在網格上的處理元素（PE）陣列。網格通過每側的交叉開關連接到一組片上存儲器塊和片外存儲控制器。有一個單獨的控制子系統(tǒng)，帶有專用處理器和外設，用于運行系統(tǒng)的控制軟件。包含PCIe接口、相關DMA引擎和安全引導處理器的主機接口單元也位于該控制子系統(tǒng)旁邊。

圖3：加速器架構

圖4顯示了PE的內部組織。PE由兩個RISC-V處理器核心和相關外設（左）、以及幾個專門執(zhí)行特定計算或數(shù)據(jù)移動的固定功能單元（右）組成。此外，每個PE都有128KB的本地存儲空間。本地互連建立了處理器、外設和自定義硬件塊之間的連接。

圖4：PE的內部組織

3.1固定功能單元

每個PE總共有五個固定功能塊和一個命令處理器，該命令處理器協(xié)調這些固定功能塊上的操作執(zhí)行。功能單元在PE中形成了一個粗粒度的管道，數(shù)據(jù)可以從一個單元傳遞到下一個單元，以執(zhí)行連續(xù)的操作。每個功能單元還可以直接訪問PE本地存儲器內的數(shù)據(jù)，執(zhí)行必要的操作，并將結果寫回，而無需將數(shù)據(jù)傳遞給其他功能單元。

3.1.1 內存布局單元（MLU）

該功能塊執(zhí)行、復制和更改本地存儲器中的數(shù)據(jù)布局有關的操作。它可以對4/8/16/32位數(shù)據(jù)類型的張量進行操作。諸如轉置、串聯(lián)或數(shù)據(jù)整形之類的操作都是使用此功能塊執(zhí)行的。輸出數(shù)據(jù)可以直接發(fā)送到下一個功能塊以立即進行操作，也可以存儲在PE的存儲器中。例如，MLU可以轉置矩陣，并將輸出直接提供給DPE塊用于矩陣乘法運算，或者它可以將數(shù)據(jù)正確格式化為深度卷積運算的一部分，并將其發(fā)送給DPE以執(zhí)行實際計算。

3.1.2 點積引擎（DPE）

該功能塊對兩個輸入張量執(zhí)行一組點積運算。首先讀取第一張量并將其存儲在DPE中，然后將第二張量流式傳輸，并對第一張量的所有行執(zhí)行點積運算。DPE每個周期可以執(zhí)行1024次INT8乘法（32×32）或512次FP16/BF16乘法（32×6）。操作完全流水線化；執(zhí)行兩個最大矩陣的乘法需要32個時鐘周期。在INT8乘法的情況下，結果輸出以INT32格式存儲，而在BF16或FP16乘法的情況中，結果以FP32格式存儲。結果總是被發(fā)送到管道中的下一個功能單元進行存儲和累積。

3.1.3 縮減引擎 (RE)

RE托管存儲元件，這些存儲元件跟蹤矩陣乘法運算的結果并在多個運算中累積它們。有四個獨立的存儲庫，可以獨立地用于存儲和累積來自DPE的結果。RE可以將初始偏置加載到這些累加器中，并且還可以通過專用縮減網絡將它們的內容發(fā)送到相鄰PE（在本節(jié)稍后討論）。在通過縮減網絡接收結果時，RE將接收到的值累積在本地存儲體之一中的值之上。然后，它可以將結果發(fā)送給相鄰的功能塊或SE，或者將其直接存儲在PE的本地存儲器中。

3.1.4 SIMD 引擎 (SE)

該功能塊執(zhí)行諸如量化/去量化和非線性函數(shù)之類的操作。在內部，該塊包含一組查找表和浮點運算單元，用于計算非線性函數(shù)的線性或三次近似，如指數(shù)、S形、tanh等。近似接受INT8或FP16數(shù)據(jù)類型作為輸入，在輸出處產生INT8或FP32結果。該單元可以直接從RE塊接收其輸入，或者從本地存儲器讀取它們。此外，該塊還能夠使用其浮點ALU執(zhí)行一組預定義的元素運算，如加法、乘法、累加等。

3.1.5 結構接口（FI）

這個功能塊充當進出PE的網關。它連接到加速器的片上網絡并通過該網絡進行通信。它制定并向片上和片外存儲器以及系統(tǒng)寄存器發(fā)送存儲器訪問請求，并接收回數(shù)據(jù)或寫入完成。它實現(xiàn)了一組類似DMA的操作，在PE的本地內存中傳輸數(shù)據(jù)。它還接收和傳輸來自處理器內核的緩存未命中和未緩存訪問，并允許其他實體（其他PE或控制子系統(tǒng)）訪問PE的內部資源。

3.1.6 命令處理器（CP）

除了托管PE的本地存儲器和寄存器外，CP塊還充當中央處理單元，同時協(xié)調固定功能塊上各種操作的執(zhí)行。它從PE中的兩個處理器核心接收指令，對這些指令執(zhí)行依賴性檢查、調度和跟蹤，并將它們分派到固定功能單元執(zhí)行。它包含兩個獨立的調度器（每個處理器內核一個）、一組命令隊列以及用于訪問本地內存和寄存器資源的仲裁邏輯。

硬件提供一組原子基元，以允許內核之間（在PE內或在多個PE之間）的同步。這些基元由處理器制定，允許對預定義寄存器進行原子更新，并且可以暫停處理器，直到外部滿足某些條件（例如，計數(shù)器達到某個值）。在更高的級別上，這些機制用于有效地實現(xiàn)軟件結構，如鎖、互斥鎖和屏障等。執(zhí)行原子操作的邏輯以及相關寄存器位于命令處理器內，并通過自定義接口與處理器內核緊密集成。

3.2處理器內核

每個PE包含兩個RISC-V內核，它們運行應用程序的代碼并向CP發(fā)出命令，用于將各種計算加載到固定功能單元。內核按序執(zhí)行指令，且單周期只執(zhí)行一個指令，具有五級流水線（AX25-V100，來自Andes Technology），并經過大量定制以適應所需的功能。該組自定義包括自定義接口、自定義寄存器、自定義指令和自定義異常。自定義接口將內核連接到CP，以向固定功能單元發(fā)出命令，并在內核和本地存儲器之間來回移動數(shù)據(jù)。自定義寄存器存儲在發(fā)出命令時發(fā)送給CP的命令信息，添加自定義指令以啟動每個固定功能單元上的所需操作。最后，自定義異常確保了向CP發(fā)出的每個命令的正確性，并在命令中出現(xiàn)非法值時引發(fā)異常。

其中一個處理器內核配備了RISC-V矢量擴展，這為PE增加了額外的靈活性，并允許實現(xiàn)無法很好地映射到現(xiàn)有固定功能單元的操作。矢量處理單元包含32個矢量寄存器，每個寄存器的寬度為64B，所有矢量功能單元的寬度相同。它實現(xiàn)了RISC-V矢量擴展的0.8.1版本。

3.3本地存儲 (LS)

每個PE總共有128KB的本地內存供處理器和功能單元使用。CP實現(xiàn)存儲器組的仲裁方案，并協(xié)調來自內核和固定功能單元的訪問。本地存儲器被映射到系統(tǒng)的地址空間，并且可以由內核通過常規(guī)加載/存儲指令來訪問。

在本地存儲器之上引入了一個抽象層，以簡化它們的操作之間的使用和依賴性檢查。每個PE可以定義映射到現(xiàn)有本地存儲器的循環(huán)緩沖區(qū)（CB）。每個CB都指定有一個ID，并有一對寄存器，用于指定其大小（深度）和本地存儲器中的起始地址。此外，每個CB還實現(xiàn)一組讀寫指針，以實現(xiàn)硬件FIFO。

在CB中，讀取操作總是從讀取指針開始讀取數(shù)據(jù)，而寫入操作總是從寫入指針開始寫入數(shù)據(jù)。讀寫操作攜帶一個偏移量，這允許它們訪問緩沖區(qū)當前頭或尾以外的位置（圖5）。固定功能單元使用CB ID作為其輸入/輸出操作數(shù)；例如，矩陣乘法運算使用兩個CB作為其輸入操作數(shù)。在允許操作開始之前，命令處理器檢查輸入CB中數(shù)據(jù)的可用性和輸出CB中的空間。它只允許在必要的元素和空間檢查通過時啟動操作。因此，可以保證操作具有完成所需的資源，并且不會在執(zhí)行過程中暫停功能單元。

圖5：從環(huán)形緩沖區(qū)讀取

命令處理器還使用CB ID來強制執(zhí)行不同自定義指令之間的依賴性檢查和互鎖。它確保訪問和修改特定CB的操作始終按程序順序執(zhí)行，而在不同CB或同一CB的不同區(qū)域上的操作可以并行執(zhí)行。與使用絕對本地存儲器地址來實施這種互鎖相比，這顯著簡化了依賴性檢查。

CB還簡化了不同操作之間生產者-消費者執(zhí)行模型的實現(xiàn)。這些操作可以由不同的內核或不同的固定功能單元發(fā)起。例如，程序可以向硬件發(fā)出一系列DMA操作（將數(shù)據(jù)從外部存儲器移動到CB），然后進行一組使用該數(shù)據(jù)的自定義計算操作（例如MATMUL），而不需要兩者之間的顯式同步。MATMUL指令由命令處理器自動停止，直到之前的DMA操作將足夠的數(shù)據(jù)帶入CB，然后立即啟動，從而使程序不再明確檢查數(shù)據(jù)的可用性。

雖然一些指令（如DMA操作）會自動調整讀和寫指針（因為它們會將數(shù)據(jù)移入和移出CB，從而產生或消耗元素），但其他自定義指令不會移動指針。這允許CB內部的數(shù)據(jù)在被明確標記為已消耗之前被不同的操作多次重用。硬件提供了額外的自定義指令，可以調整每個CB中的讀指針和寫指針，允許在必要時將數(shù)據(jù)元素明確標記為生成或消耗的數(shù)據(jù)元素。

3.4存儲器子系統(tǒng)和互連

除了PE內的本地存儲器外，加速器還具有128MB的片上SRAM，這些SRAM被組織為網格周圍的切片。這種片上存儲器可以用作可尋址的暫存存儲器，也可以用作通用、共享的存儲器側緩存。網格的每一側都有四個LPDDR5控制器，總共提供176 GB/s（理論上）的片外帶寬。該加速器總共可以支持128GB的片外存儲器容量。存儲器地址分布在這些控制器之間，以及片上SRAM片之間。當片上SRAM被配置為高速緩存時，每四個高速緩存片都與單個存儲器控制器相關聯(lián)，并高速緩存其地址。

將所有PE和存儲器連接在一起的片上網絡基于具有特殊增強功能的AXI互連?；ミB由兩個網絡組成，用于分別承載存儲器和寄存器訪問。存儲器訪問網絡配備了多播功能，允許將來自多個PE的請求合并為一個PE（如果它們被發(fā)送到同一組地址）。然后向存儲器塊發(fā)送單個請求以檢索數(shù)據(jù)并將其返回給所有請求PE。然而，多播僅支持位于網格中同一行或列的PE，不能用于任意一組PE。

除了基于AXI的主要互連外，PE還通過一個稱為縮減網絡的專用網絡相互連接。這是一個單向網絡，只從北向南和從西向東傳播。它將一個PE的RE塊中的累加器的部分和傳送到另一個PE。使用該網絡，PE可以方便地累積其計算結果，而不必將其保存和恢復在存儲器中。在行或列中的最后一個PE可以在累積所有部分值之后將最終結果存儲在存儲器中。

3.5并行性與數(shù)據(jù)復用

在任何深度學習加速器中，并行性、局部性和數(shù)據(jù)復用在有效利用有限的硬件資源方面都發(fā)揮著重要作用。MTIA體系結構提供了一組功能，以允許在神經網絡模型和算子中實現(xiàn)多個并行度和最大限度地利用時間和空間數(shù)據(jù)復用，如下所述。

并行性：該體系結構支持多種級別的并行性和各種操作的重疊。數(shù)據(jù)級并行（DLP）是通過在固定功能單元中使用寬矢量以及矢量處理器來實現(xiàn)的。多個PE也可以以數(shù)據(jù)并行的方式對同一任務進行操作。命令處理器利用指令級并行性，允許不同的固定功能塊同時處理多個未完成的操作。內存級并行（MLP）是通過允許每個PE對片上和片外存儲器發(fā)出許多未完成的請求來實現(xiàn)的。最后，線程級并行（TLP）可以通過利用多個PE（或PE組）運行并行線程以及在每個PE內有兩個獨立的線程來實現(xiàn)。PE內的線程可以協(xié)作執(zhí)行給定的任務，通過一個線程編排數(shù)據(jù)移動，而另一個線程則編排計算。

緩存：在硬件的各個功能塊中有多個級別的緩存，以提高局部性并減少內存帶寬消耗。這包括處理器核心中的指令和數(shù)據(jù)高速緩存、大型片上末級高速緩存以及DPE塊中輸入操作數(shù)的高速緩存。DPE級別的緩存允許引擎保存操作數(shù)A和操作數(shù)B的數(shù)據(jù)，并在命中時保存對本地內存的訪問。

循環(huán)緩沖區(qū)/本地存儲器：當PE執(zhí)行計算時，循環(huán)緩沖區(qū)為保存輸入操作數(shù)提供存儲。調整指針以及偏移到循環(huán)緩沖區(qū)內任何位置的靈活性來允許程序在決定將每行數(shù)據(jù)標記為已消耗之前多次訪問每行數(shù)據(jù)。

專用縮減：擁有專用縮減網絡不僅可以從系統(tǒng)的主要片上網絡中避免大部分數(shù)據(jù)傳輸，而且還提供了一種將PE分組在一起并以聚合形式使用其本地存儲器的方法。這又允許在PE中存儲更大部分的輸入操作數(shù)，并減少從片外存儲器加載它們的帶寬要求。此外，DPE塊利用歸約樹（空間和）來計算乘法運算的輸出，這是已知的更節(jié)能的方法。

多播：如前所述，當多個PE訪問內存中的同一組地址時，系統(tǒng)的NoC允許合并來自多個PE的請求。這減少了內存帶寬，并通過允許共享數(shù)據(jù)來提高數(shù)據(jù)移動的能效，同時只從內存中讀取一次數(shù)據(jù)并將其傳遞給所有請求者.

圖6顯示了由片上SRAM和片外DDR控制器包圍的PE網格的芯片示意圖，而表I列出了芯片特征和參數(shù)的摘要。

圖6:MTIA die示意圖

表I MTIA特征和參數(shù)概述。

4映射FC層

為了演示上述所有特征是如何協(xié)同工作的，讓我們考慮一個FC算子，它以CT=a×BT的形式執(zhí)行矩陣乘法運算，并看看它是如何映射到PE的子網格的。以轉置方式執(zhí)行操作的原因是保持k作為兩個張量的內部維度，以提高存儲器訪問的效率。矩陣A假設為m×k，矩陣B假設為k×n（因此BT將為n×k），產生輸出C，輸出C將為m×n矩陣（或CT為n×m矩陣）。假設輸入具有行主存儲器布局。當內部維度（k）不是32B的倍數(shù)時，外部維度（m或n）步幅與32B邊界對齊，以實現(xiàn)有效的數(shù)據(jù)移動。為了簡單起見，我們將假設所有元素都是INT8數(shù)據(jù)類型。

如前所述，DPE對32（m）×32（k）×32個（n）輸入的塊進行操作，生成RE中累積的32（n）×32（m）個部分結果。此操作需要32個時鐘周期。為了饋送DPE的流水線，必須在32個周期內將32（m）×32（k）個塊矩陣A和32（n）×32（k）個塊矩陣BT從外部存儲器引入PE的本地存儲器，需要64B/周期的帶寬。為了緩解這種帶寬壓力，RE塊中的四個累加器用于累加2×2個部分結果塊，總共保持輸出矩陣的64（n）×64（m）個元素。通過以這種方式使用累加器，每個32×32輸入塊我們使用兩次，從而將外部帶寬需求減少到32B/周期。

張量維度m、n和k分別以64、64和32的倍數(shù)分布在PE網格上。因此，每個PE以數(shù)據(jù)并行方式在較大結果矩陣的不同子塊上工作。縮減維度（k）沿著行（或列）分布在多個PE上。這便于在乘法完成之后使用歸約網絡來累積部分結果。PE將計算出的部分結果相互傳遞，以累積并傳遞給下一個PE。當沿給定行或列的兩個或多個PE使用來自任一輸入矩陣的相同輸入數(shù)據(jù)塊時，芯片網絡的多播功能用于合并來自多個PE的請求，并向內存發(fā)送單個請求，從而進一步降低內存帶寬需求。

圖7顯示了在4×4 PE子網格上分布尺寸為512（m）、1024（k）和256（n）的FC算子的示例?？s減尺寸（k）分布在沿著同一行的兩個PE上，尺寸m分布在四行上。列0和2中的PE以及列1和3中的PE參與矩陣A的行多播讀取。類似地，每列中的所有PE都參與矩陣BT的列多播讀取。

圖7：將FC操作映射到子網格

在PE中，操作以生產者-消費者的方式在兩個內核之間進行劃分。圖8顯示了對應于PE中每個內核的偽代碼。Core0發(fā)出一組DMA操作，將數(shù)據(jù)從主存儲器移動到CB_A和CB_B，用于本地存儲矩陣A和B。在并行線程中，Core1發(fā)出一組矩陣乘法（MML）指令，分別從CB_A和CB_B讀取數(shù)據(jù)，并將結果存儲在累加器寄存器中。可以觀察到，每個數(shù)據(jù)塊被使用兩次，以在每個累加器寄存器中產生部分結果。如果操作是最后一次迭代，則通過發(fā)出POP指令將數(shù)據(jù)標記為在CB中消耗，否則相應的CB偏移量將增加，以在下一次迭代中移動到下一個數(shù)據(jù)塊。最后，調用歸約運算（REDUCE）來累積PE之間的所有部分和?？s減鏈中的最后一個PE使用DMA操作將數(shù)據(jù)發(fā)送回主存儲器。

圖8:PE中運行的FC運算的偽代碼

PE中的兩個內核必須在操作開始時同步，因為它們中只有一個執(zhí)行必要的初始化任務（例如，設置要使用的CB）。但之后，沒有顯式的每次迭代同步；生產者-消費者同步由硬件負責：如果消費者（MML操作）試圖使用沒有足夠數(shù)據(jù)的CB，硬件會暫停操作，直到生產者（DMA操作）在CB中放置足夠的數(shù)據(jù)，此時它允許進行矩陣乘法。這種異步性使生產者和消費者線程解耦，并允許生產者繼續(xù)前進，為以后的迭代引入更多數(shù)據(jù)。

5軟件棧

MTIA的軟件堆棧在設計時考慮到了兩個主要目標：高效生產，這意味著實現(xiàn)比其他同類最佳解決方案更高的性能/TOC；同時，使用起來簡單明了，甚至比現(xiàn)有的替代方案更簡單。MTIA的軟件堆棧是圍繞PyTorch設計和構建的，以受益于其功能，并實現(xiàn)與生產環(huán)境中可用的ML的其他基礎組件的無縫集成。本節(jié)的其余部分提供了軟件堆棧的每個組件的概述，如圖9所示。

圖9：MTIA的軟件棧

ML服務平臺：在軟件堆棧的頂部，我們有專用于ML模型服務平臺（如圖9所示的應用層）。這些服務平臺在PyTorch之上運行，大多與硬件無關，支持在異構硬件系統(tǒng)上執(zhí)行，包括CPU、GPU和MTIA等加速器。

PyTorch Runtime：為MTIA開發(fā)了PyTorch Run集成，提供了必要的功能和特性，包括MTIA張量、主機端內存分配器和類似CUDA的流式API，用于在設備上調度所需的操作。運行時支持不同的模型執(zhí)行模式，包括熱切模式，以及全圖編譯和執(zhí)行，以最大限度地提高性能。它還支持將模型拆分為跨多個卡的分區(qū)，從而在它們之間提供必要的同步和通信通道。

編譯器：軟件堆棧中的下一個重要組件是一組編譯器，它由多個部分組成：

一個基于PyTorch FX的ML模型編譯器，它對表示為FX IR的PyTorch圖進行了多次轉換和模型級優(yōu)化，并逐漸將其轉換為LLVM IR。它負責利用PE網格和MTIA的內存子系統(tǒng)進行圖優(yōu)化。它實現(xiàn)了一個張量放置方案，該方案采用盡最大努力的方法將生產者-消費者數(shù)據(jù)保存在片上存儲器中。它還可以將模型拆分為子圖，這些子圖旨在跨多張卡運行，甚至跨同一芯片內的子網格運行。

用于ML內核開發(fā)的基于DSL的編譯器（代號KNYFE），它對ML內核進行了簡短的高級描述，并生成了低級優(yōu)化的C++代碼。它使用底層硬件特定的API來實現(xiàn)ML算子，并被廣泛用于開發(fā)MTIA中使用的許多ML內核。

基于LLVM的編譯工具鏈將LLVM IR轉換為外設的可執(zhí)行文件。LLVM的使用主要是因為它提供了RISC-V支持，并負責最低級別的優(yōu)化，如寄存器分配、排隊和代碼生成。大多數(shù)主要的優(yōu)化，如PE之間的工作和數(shù)據(jù)的平鋪或調度，都是由前面提到的高級編譯器執(zhí)行的。

ML內核庫：另一個重要組件是內核庫和ML算子庫，用于構建在設備上執(zhí)行的ML模型。這些內核中的許多都是使用前面提到的DSL編譯器開發(fā)的，但一些對性能要求最高的內核，例如全連接（FC）層和嵌入包（EB）層，是由專家直接在低級C++中使用公開的內部函數(shù)開發(fā)的，以確保它們可以在硬件上實現(xiàn)最高級別的性能。

主機驅動程序和固件接口：MTIA平臺軟件使主機能夠訪問加速器。它管理設備的生命周期和資源，并幫助啟動和跟蹤設備上的運行時操作。堆棧的這一部分大致分為兩部分：主機軟件和固件。主機軟件包括Linux驅動程序、用于提供統(tǒng)一設備接口的設備訪問庫、用于與PyTorch接口的流式API，以及用于管理和監(jiān)控設備的軟件工具和實用程序。

設備固件：設備固件包括基于ROM的預引導固件、在其自己的處理器上運行的安全引導固件、運行在執(zhí)行運行時和管理操作的控制子系統(tǒng)上的控制核心處理器固件，以及最后運行在計算網格中的PE上的PE監(jiān)視器，該監(jiān)視器調度和監(jiān)視在PE上運行的工作負載。主控制固件是基于Zephyr實時操作系統(tǒng)工作的。

6結論

我們通過將MTIA與基線加速器（NNPI）和最近部署的GPU進行比較來評估其性能。需要注意的是，我們報告了使用正在開發(fā)的軟件堆棧收集的結果，因為我們認為這反映了端到端的性能，并且代表了生產環(huán)境。然而，這個堆棧目前并不像GPU的軟件堆棧那樣優(yōu)化。因此，在某些情況下，GPU更高效，但我們希望隨著時間的推移能夠縮小這一差距，并讓MTIA軟件堆棧在所有DLRM工作負載空間中發(fā)揮架構的全部優(yōu)勢。我們評估了復雜性、規(guī)模和準確性方面不同的完整DLRM模型。由于這些加速器都基于不同的硬件平臺，我們首先比較了它們的系統(tǒng)級硬件規(guī)范（表II）。這些平臺如下：Yosemite V2服務器（帶有六個NNPI加速器卡），Zion4S服務器（帶有八個Nvidia A100 GPU），Yosemite V3服務器（帶有十二個MTIA加速器卡）。

表II: 測試所用的推理硬件平臺

雖然我們可以比較MTIA、NNPI和GPU的絕對性能，但每個設備在計算吞吐量、內存帶寬和內存容量方面都有不同的能力，它們也在不同的電力預算下運行。因此，在我們的研究中，我們報告了性能/W（來代表性能/TCO，充分考慮到TCO的敏感性），因為功率是數(shù)據(jù)中心部署資源調配的一個重要因素。我們使用總平臺功率除以加速器卡的數(shù)量來確定為每個加速器提供的功率，而不是使用卡的最大TDP。

6.1基準測試表現(xiàn)

我們首先評估了幾個重要的操作符和內核的性能，這些操作符和內核突破了體系結構的極限，是生產DLRM中主要組件的代表。表III顯示了batch大小為64和256的代表性DLRM中請求的延遲細分。該模型有大約750層，其中近550層由EB算子組成。對于batch大小為64的情況，F(xiàn)C主導執(zhí)行時間，其次是EB，而對于batch大小為256的情況，EB略微主導FC，兩者合計占執(zhí)行時間的62%。應該注意的是，對于較大的輸入，內核能夠更好地分攤設置成本，并更多地重用數(shù)據(jù)，從而實現(xiàn)硬件中固定功能單元的更高利用率。

表III-中等復雜度DLRM中各算子細分

基于細分，我們使用一組基準來評估MTIA硬件的效率。這些基準測試雖然不是全面的工作負載，但允許重要算子使用各種形狀和大小，并揭示硬件中可能存在的潛在缺陷。GemmBench用于評估密集計算；它創(chuàng)建了一個由FC層鏈組成的模型。在我們的基準測試中，我們關注FP16和INT8（量化）數(shù)據(jù)，這需要額外的量化和去量化層。TBEBench用于評估稀疏計算，并允許我們配置TBE算子的batch大小、表數(shù)、每個表的行數(shù)、嵌入維度和池化因子。BatchGEMMBench、ConcatBench和TransposeBbench用于有效覆蓋推薦模型中常見的其他重要算子。我們還評估了幾個元素內核，包括量化、反量化和tanh。

密集計算：我們評估了INT8和FP16全連接（FC）層（圖10和圖11）。當精度足夠時，INT8量化可以使FC吞吐量提高2倍。對于我們評估的一組形狀，趨勢線大致跟蹤INT8和FP16上的MTIA和GPU，表明軟件實現(xiàn)在一系列算術強度上得到了很好的優(yōu)化。在許多情況下，MTIA實現(xiàn)了每瓦2倍或更高的性能，并且對于低batch特別有效，這有助于在嚴格的延遲要求下為請求提供服務。對于大batch，GPU能夠隨著工作量的增加而實現(xiàn)更高的利用率，因此MTIA的性能/W增益更低。請注意，當張量可以直接從SRAM流式傳輸時，MTIA是最有效的，這意味著圖優(yōu)化和管理數(shù)據(jù)局部性對于模型級別的良好性能非常重要。

圖10:INT8 FC性能

圖 11: FP16 FC 性能

稀疏計算：雖然一個典型的推薦模型可能包括數(shù)百個EmbeddingBag（EB）算子，但它們可以合并到一個或多個TableBatchedEmbedding（TBE）算子中，以分攤內核啟動開銷，并增加可以在設備上并行化的工作。圖12顯示了在一組具有代表性的算子的應用中，在MTIA和GPU上運行的TBE基準測試的性能（以GB/s/W為單位）。請注意，我們在這里以GB/s為單位報告性能，因為此基準測試主要是內存綁定的，而測量帶寬（而不是查找/秒）可以更好地了解硬件利用率。在這里，我們利用片上SRAM的緩存配置來利用跨batch和batch內的局部性。在這些示例中，所有表條目都使用8位量化，圖中顯示的三元組描述了算子的池化因子、表中的行數(shù)和嵌入維度（每行元素）。MTIA實現(xiàn)了當前內核實現(xiàn)的GPU性能/W的0.6倍至1.5倍。

圖 12: TBE 性能表現(xiàn)

考慮到軟件堆棧的演變性質，我們觀察到存在顯著的改進空間：MTIA僅達到其內存帶寬的10-20%，而GPU實現(xiàn)了其HBM帶寬的60%左右。為了確保硬件沒有缺陷，我們使用了為RTL驗證而開發(fā)的手寫內核，并且在SRAM中具有足夠的位置時，可以觀察到高達500 GB/s（超過最大限度的60%）或6 GB/s/W的性能水平。我們希望通過改進TBE內核的軟件流水線和指令調度來彌補這一差距。

其他算子：雖然FC和TBE往往主導執(zhí)行時間，但我們發(fā)現(xiàn)其他算子也同樣重要，特別是考慮到優(yōu)化前者所花費的精力。我們評估了BatchMatMul、Concat、Transpose和M=256、K=128、N=32的幾個元素內核，張量數(shù)據(jù)放在SRAM和DRAM中（圖13）。這些算子往往是內存綁定的，例如BatchMatMul和Tanh，它們分別達到SRAM帶寬的90%和80%以上。當數(shù)據(jù)被放置在DRAM中時，由于更難隱藏額外的內存延遲，效率平均下降到40%左右。我們相信，實現(xiàn)數(shù)據(jù)布局優(yōu)化、算子融合和最大限度地減少DRAM中的數(shù)據(jù)提取可能會緩解這一問題。

圖13：其它算子的性能表現(xiàn)

6.2模型性能

我們測試了表IV中的五種不同的具有代表性DRLM的性能，它們的復雜性從低到高不等。MTIA可以運行與NNPI和GPU上運行的推薦模型相同的模型。以目前軟件堆棧的成熟度水平，MTIA實現(xiàn)了與GPU接近的性能/W，并超過了NNPI的性能/W，而建模表明，隨著軟件堆棧的進一步成熟，還有很大的改進空間。

表IV: 用于評估的DLRM模型。

圖14顯示了上述一組DLRM的性能（單位為TFLOPS/s/W）。與NNPI相比，MTIA實現(xiàn)了1.6倍的效率，而與GPU相比，它達到了0.9倍的效率。在這些結果中有兩個重要因素需要考慮：模型特性和實現(xiàn)中的軟件優(yōu)化水平。對于低復雜度模型，MTIA比GPU具有顯著優(yōu)勢，因為這些模型由具有較小輸入形狀的FC層主導，并且MTIA非常有效地處理了這一點，例如，LC2顯示出近3倍的改進。對于中等復雜度的模型，MTIA仍然比GPU更有效率增益，但它更低，因為FC不太占主導地位，并且GPU軟件堆棧提供了其他算子的更高效實現(xiàn)（具有TBE和積極的算子融合）。對于高復雜度模型，我們看到GPU軟件堆棧針對大型輸入進行了更好的優(yōu)化，MTIA需要類似的優(yōu)化才能實現(xiàn)相同或更高級別的效率。這些初步結果使我們深入了解了軟件堆棧中未來應該考慮關注的領域（例如，大型FC、TBE優(yōu)化、算子融合等），并為我們接下來討論的下一代架構提供了重要的參考。

圖14：DLRM的性能

7討論

構建芯片總是一個困難、漫長和耗時的過程，尤其是當初次嘗試時。對于MTIA，所產生的芯片需要實現(xiàn)高性能，處理廣泛的推薦模型，并提供一定程度的可編程性，以允許在生產中快速部署模型。本節(jié)重點介紹了我們對體系結構選擇的重要觀察和思考，以及它們如何影響軟件堆棧、性能和開發(fā)人員效率。這些經驗教訓也為改進和加強未來幾代架構提供了指導。

雙核PE：選擇在PE內有兩個獨立的處理器內核、并允許兩者控制固定的功能單元，在線程級別提供了很大程度的并行性和靈活性，它能夠允許計算與數(shù)據(jù)傳輸解耦。雖然這種解耦簡化了編程，并減輕了特定算子受指令約束時的性能問題（通過提供兩倍于總指令的吞吐量），但在軟件中正確有效地使用這兩個內核需要付出一些努力。在第一次運行之前，兩個內核之間用于初始化和清理的同步等細節(jié)很難正確處理，但之后通過軟件堆棧中的適當集成，在所有工作負載中都能夠得到很好的利用。

通用計算：以RISC-V矢量支持的形式添加通用計算是一個合適的選擇：有些算子是在架構定義階段之后開發(fā)或被發(fā)現(xiàn)重要性的，因此架構不包括對它們的任何卸載支持。像LayerNorm和BatchedReduceAdd這樣的算子很容易用向量實現(xiàn)，而且這些實現(xiàn)被證明優(yōu)于使用標量核和固定函數(shù)單元的版本。

自動代碼生成：關于如何在PE中集成和操作固定功能單元的一些架構選擇使編譯器的自動代碼生成變得困難。處理器必須安裝并發(fā)出明確的命令來操作任何固定的功能塊。雖然這是通過向處理器添加自定義指令和寄存器來完成的，但它仍然需要組裝許多參數(shù)并將其傳遞給每個目標引擎，以指定操作的細節(jié)。從程序中控制一組異構的固定函數(shù)單元并平衡它們之間的數(shù)據(jù)流對編譯器來說是一個具有挑戰(zhàn)性的工作。在各種輸入形狀和大小的固定功能塊上實現(xiàn)期望的利用水平也是很困難的。雖然我們基于DSL的KNYFE編譯器使編寫內核變得更容易并能夠自動處理其中許多問題，但它仍然需要學習一種新的DSL。

沖區(qū)：添加循環(huán)緩沖區(qū)抽象極大地簡化了在同一內存區(qū)域上工作的自定義操作之間的依賴性檢查，因為循環(huán)緩沖區(qū)ID被用作依賴性檢查的單元（類似于處理器核中的寄存器ID）。它們還簡化了固定功能塊和處理器之間生產者-消費者關系的實現(xiàn)，因為硬件會推遲操作，直到循環(huán)緩沖區(qū)中有足夠的數(shù)據(jù)（或空間）可用，而不需要在軟件級別進行任何顯式同步。靈活的尋址機制還允許對循環(huán)緩沖區(qū)內的任何位置進行任意訪問，這簡化了數(shù)據(jù)重用，因為不同的操作可以多次訪問循環(huán)緩沖區(qū)中的不同段。然而，這需要軟件明確管理緩沖區(qū)內的空間，并決定何時將數(shù)據(jù)標記為已消耗，如果執(zhí)行不當，這可能會產生難以調試的問題。

內存延遲：PE和片上SRAM內存的訪問延遲都比典型的要長。讓許多客戶端訪問PE存儲器使仲裁方案復雜化并增加了延遲周期。對于固定功能塊，這種延遲會被計入操作的延遲中，但當處理器試圖使用本地內存時，軟件必須采用軟件流水線等技術來隱藏延遲，這會增加寄存器壓力。當在矢量處理器上開發(fā)內核時，這種情況會加劇，因為它限制了可以執(zhí)行的寄存器分組的數(shù)量，從而增加了動態(tài)指令數(shù)。

在外圍放置片上SRAM，同時在多個片上均勻地分配請求，造成了內存訪問延遲的很大程度的不均勻性。由于請求總是在最后一段數(shù)據(jù)到達后完成，因此當發(fā)出大于最小請求時，總體延遲會受到影響。雖然在大多數(shù)情況下，這種延遲可以通過將數(shù)據(jù)預取到PE的內存中來隱藏，但在具有小池化組的EmbeddingBag算子等情況下，延遲會凸顯出來，因為內存訪問模式事先未知，并且沒有足夠的未處理請求來隱藏延遲。

緩存一致性：雖然系統(tǒng)實現(xiàn)了共享內存范例，但沒有硬件支持緩存一致性。在這個共享存儲器系統(tǒng)中，不需要PE間的一致性，因為不同的PE以數(shù)據(jù)并行的方式對數(shù)據(jù)集的專用部分進行操作。然而，PE內部的一致性有時會導致準確性問題：如果兩個處理器內核、固定的功能單元和一個內核接觸到同一組內存地址，或者同一個內核在不同的算子之間重用地址，則必須從緩存中顯式清除這些地址的緩存副本，否則可能會使用過時的數(shù)據(jù)。

架構層次：加速器的推薦模型在其使用的層和算子的大小和復雜性方面有很大差異。雖然映射到PE網格上的大層可以從可用硬件資源中提取所需的利用率水平，并分攤創(chuàng)建和調度的開銷，但較小的層或較低的batch處理大小必須采用諸如利用子圖并行性和融合等技術來獲得相同的利用率。盡管在軟件級別有足夠的空間來執(zhí)行此類優(yōu)化或減少部署作業(yè)的開銷，但我們相信，在架構級別進行一些資源調配會使解決此問題變得更容易。這是因為對于較小的任務，網格必須劃分為較小的子網格，以便每個子網格都能處理較小的任務。而設置和拆除這些子網格的任務是系統(tǒng)固件的一部分。我們認為，在體系結構本身中擁有另一個層次，例如PE集群，可能會使這個問題更容易解決，因為與PE的單片網格相比，它提供了自然的隔離和管理單元。

在第一代MTIA中，我們構建了一種架構，與GPU和其他加速器相比，該架構可以為DLRM的工作顯著提升效率。隨著軟件堆棧的成熟，我們希望隨著時間的推移繼續(xù)提高這種效率。為這種體系結構編寫內核、構建編譯器和優(yōu)化模型的經驗讓我們深入了解了哪些功能有更大的影響，我們希望這些項目硬件和軟件方面的經驗教訓能夠在未來幾代架構中得到整合和利用。

審核編輯：湯梓紅