走進(jìn)任何現(xiàn)代人工智能實(shí)驗(yàn)室、數(shù)據(jù)中心或自動(dòng)駕駛汽車開發(fā)環(huán)境，你都會(huì)聽到工程師們喋喋不休地談?wù)摳↑c(diǎn)運(yùn)算次數(shù)(FLOPS)、運(yùn)算次數(shù)(TOPS)、稀疏性、量化和模型擴(kuò)展定律。這些指標(biāo)充斥著新聞?lì)^條和產(chǎn)品數(shù)據(jù)手冊(cè)。但如果你花時(shí)間與那些實(shí)際構(gòu)建或優(yōu)化這些系統(tǒng)的人交流，你會(huì)發(fā)現(xiàn)一個(gè)截然不同的真相：原始的算術(shù)能力并非決定實(shí)際性能的關(guān)鍵。

真正重要的是數(shù)據(jù)移動(dòng)的效率。而對(duì)于當(dāng)今大多數(shù)人工智能加速器來說，數(shù)據(jù)移動(dòng)與一個(gè)鮮少在編譯器和硬件圈子之外被提及的現(xiàn)象——內(nèi)存重排(memory swizzling)——緊密相連。

內(nèi)存重排是現(xiàn)代人工智能系統(tǒng)付出的最大隱形代價(jià)之一。它既不能提高算法處理效率，也不能提高準(zhǔn)確性，更不能降低能耗，更不會(huì)帶來任何新的洞見。它的存在僅僅是為了彌補(bǔ)幾十年前設(shè)計(jì)選擇遺留下來的架構(gòu)缺陷。隨著人工智能模型規(guī)模越來越大、結(jié)構(gòu)越來越不規(guī)則，這種代價(jià)也在不斷增加。

本文探討了內(nèi)存重排存在的原因、發(fā)展歷程、代價(jià)，以及一種截然不同的架構(gòu)理念——特別是以寄存器為中心的模型——如何徹底消除對(duì)內(nèi)存重排的需求。

一個(gè)鮮為人知的問題：數(shù)據(jù)的存儲(chǔ)方式與硬件需求不符

在任何人工智能教程中，張量都被呈現(xiàn)為有序的數(shù)學(xué)對(duì)象，它們以完美的布局整齊地排列在內(nèi)存中。這些布局對(duì)程序員來說直觀易懂，并且能夠很好地融入PyTorch或TensorFlow等高級(jí)框架。

然而，硬件并非如此看待數(shù)據(jù)。

現(xiàn)代加速器——GPU、TPU和NPU——圍繞并行計(jì)算單元構(gòu)建，這些單元期望特定的數(shù)據(jù)格式：固定大小的圖塊、嚴(yán)格的對(duì)齊邊界、具有可預(yù)測(cè)步長(zhǎng)模式的序列，以及能夠無沖突地映射到內(nèi)存庫的排列方式。

遺憾的是，現(xiàn)實(shí)世界中的張量永遠(yuǎn)不會(huì)以這些格式出現(xiàn)。在處理開始之前，數(shù)據(jù)必須被重塑、重新分塊、重新排序或重新打包成硬件期望的格式。這種重塑過程稱為內(nèi)存重排。

你可以這樣理解：算法以矩陣和張量的方式思考；計(jì)算硬件則以分塊、通道和存儲(chǔ)體的方式思考。內(nèi)存重排就是轉(zhuǎn)換層——一個(gè)需要耗費(fèi)時(shí)間和能量的轉(zhuǎn)換過程。

為什么分層內(nèi)存迫使我們進(jìn)行內(nèi)存重排？

如今，幾乎所有加速器都使用分層內(nèi)存堆棧，其各層從上到下依次包含：寄存器；共享內(nèi)存或暫存內(nèi)存；L1緩存、L2緩存，有時(shí)甚至包括L3緩存；高帶寬內(nèi)存(HBM)；以及位于堆棧底部的外部動(dòng)態(tài)隨機(jī)存取存儲(chǔ)器(DRAM)。

每一層的大小、延遲、帶寬、訪問能耗以及非常重要的對(duì)齊約束都各不相同。這是CPU式架構(gòu)的遺留產(chǎn)物，在這種架構(gòu)中，緩存隱藏了內(nèi)存延遲。參見圖1和表1。

圖1：當(dāng)前所有硬件處理器中典型分層內(nèi)存堆棧的容量和帶寬屬性。(來源：VSORA)

表1：當(dāng)前所有硬件處理器中典型分層內(nèi)存堆棧的容量、延遲、帶寬和訪問能耗。(來源：VSORA)

GPU繼承了這種模型，并在此基礎(chǔ)上增加了單指令多線程(SIMT)執(zhí)行。這使得它們擁有驚人的計(jì)算能力，但也對(duì)數(shù)據(jù)布局方式極其敏感。如果線程束(warp)中的相鄰線程無法訪問相鄰的內(nèi)存位置，性能就會(huì)急劇下降。如果數(shù)據(jù)塊邊界沒有對(duì)齊，張量核心就會(huì)停滯。如果共享內(nèi)存庫發(fā)生沖突，所有操作都會(huì)被暫停。

TPU也面臨著類似的限制，只是機(jī)制有所不同。它們的脈動(dòng)陣列就像精心編排的傳送帶一樣運(yùn)行。數(shù)據(jù)必須按正確的順序和時(shí)間到達(dá)。如果權(quán)重沒有按照塊優(yōu)先的格式排列，脈動(dòng)陣列就無法高效運(yùn)行。

從智能手機(jī)芯片到汽車系統(tǒng)，基于NPU的加速器也面臨著同樣的問題：多庫SRAM、固定的向量寬度以及視覺工作負(fù)載的2D局部性要求。如果沒有內(nèi)存重排，數(shù)據(jù)到達(dá)計(jì)算引擎時(shí)就會(huì)“錯(cuò)位”，導(dǎo)致性能急劇下降。

在所有這些情況下，重排都不是一種優(yōu)化——而是一種生存機(jī)制。

內(nèi)存重排的隱性成本

內(nèi)存重排耗時(shí)，有時(shí)甚至非常耗時(shí)

在實(shí)際工作負(fù)載中，內(nèi)存重排通常會(huì)消耗總運(yùn)行時(shí)間的20%到60%。這并非筆誤。在卷積神經(jīng)網(wǎng)絡(luò)中，一半的時(shí)間可能都花在了NHWCNCHW轉(zhuǎn)換上；也就是說，這是在內(nèi)存中布局4D張量的兩種不同方式。在Transformer模型中，大量時(shí)間被浪費(fèi)在重塑Q/K/V張量、拆分頭、為GEMM重新打包tile以及重組輸出上。

內(nèi)存重排會(huì)消耗能量，而能量才是真正的限制因素

一次MAC操作大約消耗四分之一皮焦耳的能量。從DRAM中移動(dòng)一個(gè)值可能需要500皮焦耳的能量。從DRAM中移動(dòng)數(shù)據(jù)所消耗的能量大約是執(zhí)行一次基本乘加運(yùn)算的1000倍。

內(nèi)存重排需要讀取大量數(shù)據(jù)塊，重新排列后再寫回。而且，每個(gè)層通常要重復(fù)多次這樣的操作。當(dāng)80%的能源預(yù)算都用于移動(dòng)數(shù)據(jù)而不是進(jìn)行計(jì)算時(shí)，內(nèi)存重排就變得不容忽視。

內(nèi)存重排會(huì)增加內(nèi)存使用量

大多數(shù)內(nèi)存重排都需要臨時(shí)緩沖區(qū)：打包的圖塊、暫存緩沖區(qū)和重塑的張量。這些額外的內(nèi)存占用可能會(huì)使模型超出L2、L3甚至HBM的極限，從而迫使模型進(jìn)行更多的數(shù)據(jù)移動(dòng)。

內(nèi)存重排使軟件更難編寫且更難移植

問問CUDA工程師，是什么讓他們夜不能寐。問問TPU編譯器工程師，為什么XLA的布局推斷代碼長(zhǎng)達(dá)數(shù)千頁。問問任何為移動(dòng)設(shè)備編寫NPU內(nèi)核的人，為什么他們害怕通道排列。

答案是：內(nèi)存重排。由于硬件需要非常具體的布局，軟件必須承載巨大的復(fù)雜性。每一種新的模型架構(gòu)——無論是卷積神經(jīng)網(wǎng)絡(luò)(CNN)、長(zhǎng)短期記憶網(wǎng)絡(luò)(LSTM)、Transformer模型還是擴(kuò)散模型——都會(huì)增加新的布局模式，而這些模式必須得到支持。

最終形成了一個(gè)由布局啟發(fā)式算法、張量變換和對(duì)性能要求極高的內(nèi)存管理機(jī)制維系的生態(tài)系統(tǒng)。

主流架構(gòu)如何依賴內(nèi)存重排

1.英偉達(dá)GPU

Tensor核心需要特定的tile-main布局。共享內(nèi)存被分塊存儲(chǔ)，避免沖突需要內(nèi)存重排。線程束必須合并內(nèi)存訪問；否則，效率會(huì)大幅下降。即使是地球上優(yōu)化程度最高的GPU庫cuBLAS和cuDNN，也充滿了內(nèi)部?jī)?nèi)存重排內(nèi)核。

2.谷歌TPU

TPU依賴于脈動(dòng)陣列。數(shù)據(jù)流經(jīng)這些陣列必須完全有序。權(quán)重和激活值會(huì)不斷重新排列，以與脈動(dòng)結(jié)構(gòu)保持一致。XLA的大部分功能都是為了管理數(shù)據(jù)布局而存在的。

3.AMD CDNA、ARM Ethos、蘋果ANE和高通AI引擎

所有這些架構(gòu)都執(zhí)行內(nèi)存重排。Morton分塊、交錯(cuò)、通道堆疊等等。這是一種通用模式。每個(gè)使用分層內(nèi)存的架構(gòu)都繼承了對(duì)內(nèi)存重排的需求。

另一種理念：從根本上消除內(nèi)存重排

現(xiàn)在，想象一下，退后一步，從根本上重新思考人工智能硬件。與其接受如今復(fù)雜的內(nèi)存層次結(jié)構(gòu)——緩存層、共享內(nèi)存塊、分塊式SRAM和對(duì)齊規(guī)則——不如設(shè)想一種基于更簡(jiǎn)單前提的架構(gòu)。

如果根本沒有內(nèi)存層次結(jié)構(gòu)呢？如果整個(gè)系統(tǒng)圍繞著一個(gè)龐大而扁平的寄存器陣列運(yùn)轉(zhuǎn)呢？如果編譯器而非硬件，以確定性的精度協(xié)調(diào)每一次數(shù)據(jù)移動(dòng)呢？如果所有常見的擔(dān)憂——對(duì)齊、存儲(chǔ)體沖突、分塊策略和合并規(guī)則——都因?yàn)椴辉僦匾兀?/span>

這就是以寄存器為中心的架構(gòu)背后的理念。數(shù)據(jù)不再沿著內(nèi)存層級(jí)上下移動(dòng)，而是直接駐留在執(zhí)行計(jì)算的寄存器中。這種架構(gòu)的組織方式不是圍繞數(shù)據(jù)的移動(dòng)，而是圍繞數(shù)據(jù)的可用性。

這意味著：

無需預(yù)熱或擔(dān)心緩存未命中

無需調(diào)度線程束

無需避免內(nèi)存沖突

無需匹配數(shù)據(jù)塊大小

無需考慮張量布局

對(duì)形狀或步長(zhǎng)不敏感，因此完全無需進(jìn)行內(nèi)存重排

在這樣的系統(tǒng)中，編譯器始終精確地知道每個(gè)值所在的位置，以及它接下來應(yīng)該去往何處。它不會(huì)進(jìn)行推測(cè)、預(yù)取、分塊或依賴啟發(fā)式算法。它不會(huì)祈禱硬件運(yùn)行正常。相反，數(shù)據(jù)放置變成了一個(gè)可解決、可預(yù)測(cè)的問題。

最終，機(jī)器的吞吐量保持穩(wěn)定，延遲變得可預(yù)測(cè)，能耗大幅下降，因?yàn)椴槐匾臄?shù)據(jù)移動(dòng)已被排除在外。在這個(gè)系統(tǒng)中，性能不再受制于內(nèi)存操作——真正的計(jì)算，也就是數(shù)學(xué)運(yùn)算，終于成為了核心。

人工智能的未來：為何以寄存器為中心的架構(gòu)至關(guān)重要

隨著人工智能系統(tǒng)的演進(jìn)，整齊劃一的張量和完美的矩形計(jì)算單元的世界正在逐漸消逝。現(xiàn)代模型不再是可預(yù)測(cè)的、步調(diào)一致的密集層堆疊。相反，它們向各個(gè)方向擴(kuò)展：它們能夠接收多模態(tài)輸入，整合稀疏和不規(guī)則結(jié)構(gòu)，進(jìn)行自適應(yīng)推理，并處理越來越長(zhǎng)的序列。它們還必須實(shí)時(shí)響應(yīng)安全關(guān)鍵型應(yīng)用，并且必須在嚴(yán)格的能耗預(yù)算內(nèi)完成這些操作——從汽車到邊緣設(shè)備，無一例外。

換句話說，塑造GPU和TPU架構(gòu)的那些假設(shè)——對(duì)規(guī)則性、密集網(wǎng)格和整齊計(jì)算單元的預(yù)期——正在瓦解。未來的工作負(fù)載不再符合硬件的預(yù)期。

以寄存器為中心的架構(gòu)提供了一條截然不同的道路。因?yàn)樗苯釉跀?shù)據(jù)所在的位置進(jìn)行操作，而不是強(qiáng)制將數(shù)據(jù)轉(zhuǎn)換為適合計(jì)算單元的格式，所以它繞過了內(nèi)存重排的整個(gè)機(jī)制。它不依賴于固定的張量形狀。

當(dāng)訪問模式變得不規(guī)則或動(dòng)態(tài)時(shí)，它不會(huì)出現(xiàn)故障。它避免了為了滿足計(jì)算單元的需求而進(jìn)行代價(jià)高昂的內(nèi)存重排。隨著模型變得越來越異構(gòu)和復(fù)雜，這種架構(gòu)能夠隨著其復(fù)雜性的增加而擴(kuò)展，而不是與之對(duì)抗。

這不僅僅是漸進(jìn)式的改進(jìn)，它代表了我們對(duì)人工智能計(jì)算思維方式的轉(zhuǎn)變。通過消除不必要的數(shù)據(jù)移動(dòng)——現(xiàn)代加速器中最大的瓶頸和能耗來源——以寄存器為中心的架構(gòu)使硬件能夠適應(yīng)人工智能本身復(fù)雜多變的現(xiàn)實(shí)。

內(nèi)存重排是每個(gè)分層內(nèi)存加速器都必須付出的隱形代價(jià)。它是GPU、TPU、NPU以及幾乎所有人工智能芯片運(yùn)行的基礎(chǔ)。同時(shí)，它也日益成為一個(gè)負(fù)擔(dān)。內(nèi)存重排會(huì)引入延遲、消耗能源、增加內(nèi)存使用量并使軟件復(fù)雜化——而對(duì)實(shí)際的數(shù)學(xué)運(yùn)算卻沒有任何貢獻(xiàn)。

一種以寄存器為中心的架構(gòu)通過移除導(dǎo)致內(nèi)存重排的層級(jí)結(jié)構(gòu)，從根本上消除了內(nèi)存重排。它用確定性的數(shù)據(jù)流取代了猜測(cè)和啟發(fā)式方法。它優(yōu)先考慮局部性，而無需重新排列數(shù)據(jù)。它讓算法驅(qū)動(dòng)硬件，而不是反過來。

隨著人工智能工作負(fù)載變得更加不規(guī)則、動(dòng)態(tài)和對(duì)功耗敏感，能夠保持?jǐn)?shù)據(jù)穩(wěn)定和可預(yù)測(cè)(而不是無休止地重新排列數(shù)據(jù))的架構(gòu)將定義下一代計(jì)算。

內(nèi)存重排是上一個(gè)硬件時(shí)代的必要補(bǔ)救措施，它不應(yīng)該定義下一個(gè)時(shí)代。

Lauro Rizzatti是VSORA的商業(yè)顧問，VSORA是一家提供重新定義性能的硅半導(dǎo)體解決方案的技術(shù)公司。他是一位知名的芯片設(shè)計(jì)驗(yàn)證顧問，也是硬件仿真領(lǐng)域的行業(yè)專家。

(責(zé)編：Franklin)