機器之心發布
作者：XiaoyuWang
九大章節，一萬餘字，這篇文章可能是目前為止Maxas彙編器工作原理最全面、最細緻的解析。
在從事深度學習框架的實現工作時，了解到 Nervana 有一個稱為 Maxas 的彙編代碼生成器項目，可以生成性能超過 nVidia 官方版本的矩陣相乘的 GPU 機器碼，由此對其工作原理產生興趣。
項目地址：https://github.com/NervanaSystems/maxas
其作者 Scott Gray 在代碼外提供了詳細的文檔（https://github.com/NervanaSystems/maxas/wiki/SGEMM），但由於該算法的複雜性，行文晦澀，邏輯跳躍，尤其是對一些方法的動機沒有交待，很容易迷失在細節中。
本文可以看作按作者對該文檔的理解進行的重寫，但求在細節上不厭其煩，對其代碼的前因後果作儘可能完整的交待，不過大體結構還是按照 maxas 文檔來安排，文中圖片也全部出自該文檔。
值得說明的是 Maxas 使用的算法完全依賴於 Maxwell 架構的一些特性， 隨著新一代 GPU 的架構的演進這個項目本身已經完全過時了，但其解決問題的思路仍然值得借鑑。
背景
單精度矩陣相乘（SGEMM）應該是每一個學習 CUDA 編程的程式設計師最熟悉的案例了，從 nVidia 推出 CUDA 的第一版起這就是其官方教程裡唯一的例子，這不僅因 SGEMM 是幾乎所有計算密集型軟體最重要的底層模塊，更是因為通過這個例子可以很好地展示 GPU 中的各種優化技巧，特別是對 GPU 內存層次的利用。
對於兩個 NxN 的矩陣 A 和 B 的相乘，一個最簡單的並行方法是對於其輸出矩陣 C（大小同為）的每一個元素開一個線程，該線程載入 A 的一行和 B 的一列，然後對其做一次向量的內積。但問題是在 GPU 上訪問顯存的延時相當的大（~100 時鐘周期），如果 A 的一行因為在內存中是連續的還能夠利用 GPU 的超大顯存帶寬一次載入多個元素平攤其載入時間以及緩存來降低延時，對於 N 上千的大矩陣來說 B 的一個列中的元素的內存地址也要相隔上千，意味著載入一次中除了需要的那個列元素外大部分數據都是無用的，同時這種訪問模式幾乎不可能有緩存線命中，總而言之這個並行方法的內存訪問效率低到令人髮指。
對其的優化就要用到共享內存了，共享內存是位於 GPU 上的片上緩存，速度可與一級緩存相當，而且同一個線程塊中的線程可以通過共享內存交換數據，唯一的缺點是容量有限。為了利用這一小塊高速內存，原來的並行算法必須有所改進：矩陣將在每個維度被劃分成個小片，輸出矩陣 C 可以寫為：
A 和 B 也做同樣的劃分，其中不是元素而是小片矩陣，當然小片大小為 1 時小片矩陣就退化為單個元素。顯然矩陣乘法的定義依然在此適用：。
如果把小片看作一個元素，整個矩陣的規模相當於被縮小了 倍。對於每個小片的結果可以由一組線程負責，其中每個線程對應小片中的一個元素。這個線程組將 A 的行小片和 B 的列小片一一載入共享內存，在共享內存上對其做矩陣相乘，然後疊加在原有結果上。這樣對小片中的元素每載入一次可以被在高速的共享內存中被訪問 K 次，由此得到遠高於原始方法的內存存取效率。
網上找的分片算法示意圖，圖示比 CUDA 官方教程中的更加清楚一點。
以上只是對這種算法的一個簡單描述，經過這樣的優化整個算法已經可以達到相當高的效率了，而且是進一步進行優化的基礎。為了理解後文中的內容，讀者需要細讀 CUDA 的官方教程以對這個分片算法非常熟悉，並且對 nVidia GPU 的編程模型有比較深入的理解。
基本思想
如上節所述，分片算法在利用了片上高速緩存之後，不但小片矩陣的乘法速度可以大大加快，還可以利用計算小片矩陣相乘的時間將下一個小片從主內存傳送至片上共享內存，換句話說此時整個矩陣相乘的時間已經完全由小片矩陣相乘所決定，如果要進一步提高性能就要在小片矩陣相乘上做文章了。
在共享內存內部做矩陣相乘雖然已經很快了，但距離硬體性能的極限還是有距離，主要瓶頸是兩個。首先共享內存的延時終究還是比不過寄存器，在 Maxwell/Pascal 上寄存器延遲時 6 個時鐘周期，在共享內存上達到 23 個周期。此外，GPU 的運算單元無法直接操作共享內存的數據，需要有一個傳輸指令將其送到寄存器上，而這個 mov 指令會佔用和實際計算指令幾乎相當的時間，可謂相當大的開銷。要達到硬體的峰值性能，我們希望 GPU 的每個運算單元的流水線都被運算指令佔滿，每個時鐘周期都有結果被計算出來。為了在現有基礎上達到這一目標，maxas 和之前的一些研究提出了如下方法：
1. 利用新加入的向量指令，一個指令可以傳輸四個連續的浮點數，大大減少傳輸指令的數量，並且有利於用計算指令隱藏傳輸指令消耗的時間；
2. 通過交叉布置計算指令和傳輸指令，實現數據預讀取和計算的流水線；
3. 分片算法利用高速的共享內存緩存主顯存上需要多次存取的數據，那麼把這個思路發展下去，在小片矩陣內部作進一步分片，利用寄存器去緩存共享內存的數據，得到進一步的加速。但是這個新的分片算法和之前的有所不同，也帶來了額外的困難。
為了實現這些方法需要對 GPU 指令和寄存器的精確控制，已經不在 CUDA 語言表達能力的範圍之內，所以其實現必須由 GPU 原生彙編語言完成（並非 PTX 這樣的偽彙編語言），但不妨礙用表達能力更強的類似 C 的偽代碼來說明這個實現。從偽代碼到實際的彙編代碼有相當直接的轉換方法，在 maxas 中用 perl 實現了這一轉換。
maxas 算法概要
只考慮兩個矩陣相乘，在之前的直觀算法中，計算一個 C 矩陣的元素是按照矩陣乘法的定義，取 A 中的一行和 B 中的一列做內積。A 中的一行和 B 中的一列都要被用到 64 次。如果要充分利用寄存器的優勢三個的矩陣（每個矩陣佔 16KB）都要放在寄存器中對寄存器文件（每個 SM 64K）是巨大的壓力，更嚴重的問題是和共享內存不同，寄存器在線程間是不能共享的，如果每個線程要在自己的寄存器中保存所負責的 C 矩陣的那部分，它也必須在寄存器中保存所用到的 A 的行和 B 的列。結果是大量寄存器用於保存重複的數據，而且對共享內存的訪問很可能造成 bank 衝突，進一步惡化延時。
如果換一個思路，不從輸出矩陣 C 的角度，而從輸入矩陣的角度，不難發現 A 的第 k 列僅被用於和 B 的第 k 行的元素相乘，也就是說如果取 A 的第 k 列和 B 的第 k 行，將其中所有元素對兩兩相乘並加到其所貢獻的輸出矩陣元素上：
這些行和列就完成了其在矩陣相乘中的使命，可以被扔掉了。這種算法可以大大減少輸入矩陣對寄存器的佔用，而且載入個數據就可以進次加乘運算，完全符合利用寄存器進一步緩存共享內存數據的要求。不難看出該方法在 A 的列和 B 的行大小不一樣時依然可以適用，只要它們的列指標和行指標相同。
maxas 對於小片矩陣乘法是用 64 個線程來並行實現的，其中每個線程負責計算個矩陣的乘積，64 個線程按照布局，這樣就確定了小片的大小為一個邊長個元素的矩陣（每線程 8 元素 x8 線程）。這一點區別於原始分片算法中每個線程計算矩陣中的一個元素，也是充分利用寄存器的超低延遲的關鍵。
圖2. maxas 計算兩個 64x64 矩陣相乘的示意圖，綠色的 4x4 小片是線程 0 負責的那部分元素，黃色是其他線程負責那部分的左上角元素。圖中只標出了左上角 4x4 矩陣的線程號，其他 15 個只是其重複。每個黑框中包含 32 線程，代表一個 warp。這塊 32x32 矩陣計算完成後，線程 0 以及其他線程保持相對位置，依次平移到另外三個綠色小片標出的地方，完成這個 64x64 矩陣的計算。左邊的向量是 A 矩陣的一個列，上方的向量是 B 矩陣中與之對應的行，其中標為綠色的數據（各 8 個浮點數）是線程 0 所需要用到的，其他線程需要的不難類推。
maxas 的整個實現就是為了這個算法服務的，後文中所要解決的問題也是該算法在實現過程中所遇到的。以上的那些參數選擇，比如為什麼選擇 64 個線程，都是根據 GPU 硬體資源決定的，以便在滿足每個線程所需的寄存器資源基礎上，創建儘可能多的線程 warp，以便調度器在某些 warp 等待數據時將別的 warp 切換進來進行計算。
將輸入矩陣載入共享內存
為了實現以上算法，這 64 個線程首先被用來將兩個輸入矩陣所需要的部分從主顯存載入共享內存。這裡需要指出一個原文中沒有提到的假設，即在 maxas 中默認矩陣是按照行優先儲存的，即矩陣的每一列在內存中是連續的，否則無法解釋後面的一系列算法。由於算法的不同載入的方法也有所不同，並且在原方法基礎上增加了一些優化：
1. B 矩陣用到的是行數據，而默認的行優先矩陣儲存中行數據是不連續的，需要做轉
置，行變為列，這樣 A 和 B 的載入方法可以完全相同，以降低代碼複雜度；
2. A 和 B 矩陣被作為一維紋理載入，這樣不但可以利用紋理數據的緩存，而且不用耗費時間進行數組越界檢查，因為紋理載入會自動將越界的數據置為 0，對於矩陣相乘不會有任何影響。
了解以上預處理可以更方便地理解後面的偽代碼。創建紋理和轉置的工作應該是在 GPU 內核執行前完成的，不影響內核執行的性能。
紋理內存中的數據也是分段被載入共享內存的，不過按照原方法每段載入的應該是一個個的小片，但為了充分利用寄存器資源，maxas 採用了完全不同的計算方法。如果線程塊計算的是，首先將矩陣 A 按每 64 行一條劃分為的條，所需的輸入數據全部在第條上，當然這一條數據還還是很大，需要進一步分次載入，maxas 中每次載入並消耗，分次完成。對於矩陣 B 方法類似，只不過是按列劃分為的條，轉置後載入方法和 A 完全相同。其內存布局如下圖所示：
圖 3. 每次循環中被一個 warp 載入共享內存的一段紋理，可以看作 Bj 或轉置後的 Ai，這樣這個矩陣其實又回到了常規的列優先儲存。這個圖轉置後看其實更直觀。
圖中每一個格子代表一個線程負責載入的數據單元，綠色格子是線程 0 要先後載入的四個單元，黃色格子是其餘 31 個線程第一次載入的那部分數據。整個 warp 每次載入兩行，如此重複四次完成。在 GPU 上執行單位是 32 個線程組成的 warp，所以 64 個線程是分為兩個 warp 執行。其中一個 warp（線程 0-31）載入 A 另一個（線程 32-63）載入 B。此圖有一個容易造成困惑的地方是圖中的矩陣形狀為而不是， 這是因為後面每個線程會用到向量指令一次載入 4 個浮點數，即每個格子本身就是四個浮點數。在後面的代碼中會看到在紋理內存上使用向量指令時偏移量會相對實際元素的數量除以 4。
這個加載方法顯然不是唯一的，我的理解是因為 A 和 B 的載入方法完全一樣，只是所用的紋理不同，所以相比一個線程同時加載 A 和 B 可以減少與計算無關的指令的代碼量。
載入的數據被暫時在寄存器上，等待被儲存到共享內存。共享內存中的數據排布如下圖：
圖 4. 矩陣 A 被載入的數據在共享內存中的排布。
因為共享內存的的偏移單位是 1 字節，終於又回到了直觀的表示，由此可見以上兩圖的數據儲存方式其實是完全一樣的，均為的列優先儲存，唯一的區別是在共享內存中 B 的數據會直接跟在 A 後面。
圖 5. 載入輸入矩陣 A 和 B 的示意圖，注意圖中 lda, ldb, ldc, bx, by, k 的意義，這些變量在後面的代碼中都會被用到。
以下偽代碼是整個過程的描述，英文注釋為 maxas 文檔中所帶，額外的注釋為中文。
tid = threadId.x;bx = blockId.x; // 可以看作C_ij的iby = blockId.y; // 可以看作C_ij的j
blk = tid >= 32 ? by : bx; ldx = tid >= 32 ? ldb/4 : lda/4; // lda和ldb為A的行數或B的列數，ldx為其在向量載入中的表示（每次4個浮點數故除以4），可以看成每一行的跨度tex = tid >= 32 ? texB : texA; // 64個線程分為2個warp，第一個載入紋理A，第二個載入紋理B
tid2 = (tid >> 4) & 1; // 此處將32個線程分為兩組，tid2=0為第一組載入圖三中的第一行，tid2=1載入第二行tid15 = tid & 15; // 本線程在每一行中列的位置
// 這些track變量表示本線程需要載入的數據在tex中的偏移，乘以4即在$A_i$或$B_j_T$中的偏移track0 = blk*64/4 + tid15 + (ldx * tid2); // 本線程載入數據的起始位置，解釋見後文track2 = track0 + ldx*2; // 本線程載入的第二部分數據，在第一部分兩行後，以此類推track4 = track0 + ldx*4;track6 = track0 + ldx*6;
end = track0 + (k-8)*ldx; // 載入結束的標誌，其中k為A的列數和B的行數，即兩個相乘矩陣的公共維度，對於NxN的矩陣， k=N。因為每次載入8行所以減去8作為最後一次的載入的標記
writeS = tid15*4*4 + tid2*64*4; // 本線程載入數據在共享內存中的偏移，注意其相當於track0的第二項的16倍，因為向量載入的偏移1代表（4個數x每個浮點數4位元組）writeS += tid >= 32 ? 2048 : 0; // 如果載入的是矩陣B的數據，放在矩陣A的數據之後，而矩陣A佔用（64列x8行x每元素4位元組）=2048位元組
while (track0 < end){ tex.1d.v4.f32.s32 loadX0, [tex, track0]; tex.1d.v4.f32.s32 loadX2, [tex, track2]; tex.1d.v4.f32.s32 loadX4, [tex, track4]; tex.1d.v4.f32.s32 loadX6, [tex, track6];
st.shared.v4.f32 [writeS + 4*0*64], loadX0; st.shared.v4.f32 [writeS + 4*2*64], loadX2; st.shared.v4.f32 [writeS + 4*4*64], loadX4; st.shared.v4.f32 [writeS + 4*6*64], loadX6;
// our loop needs one bar sync after share is loaded bar.sync 0;
// Increment the track variables and swap shared buffers after the sync. // We know at this point that these registers are not tied up with any in flight memory op. track0 += ldx*8; track2 += ldx*8; track4 += ldx*8; track6 += ldx*8; writeS ^= 4*16*64; // 又見magic number！其意義是A和B在共享內存中一共佔4*16*64=4096位元組，但是共享內存一共分配了8192位元組的兩組，每次載入後用這個操作切換到另外一組，其目的是實現一個流水線，在一個warp載入數據進一組時另一個warp就可以操作另一組已經完成載入的數據了
// Additional loop code omitted for clarity.}
以上代碼中還有兩個問題需要用一定篇幅來澄清：
1. track0 = blk*64/4 + tid15 + (ldx * tid2); 中的第一項 blk*64/4 是用來從紋理中選擇輸入矩陣 A 或轉置矩陣 B 中第 blk 條左上角相對整個輸入矩陣左上角的一維偏移。由於所有條的左上角都在輸入矩陣的第一列中，而行優先儲存中第一列中任一點的偏移就是其行數，對於第 blk 條左上角就是 blk*64，而 / 4 來自向量因子。tid15 + (ldx * tid2) 的意義比較清楚，即本線程在圖 3 中對應的黃色格點相對綠色格點的位置，tid15 是列坐標，tid2 是行坐標，在一維表示時需要乘以該方向的跨度 ldx。
2. 代碼中用了四個 track 變量來記錄四次載入的偏移，很容易想到只用一個 track 變量，載入一次後對偏移加兩行再做一次載入來完成同樣的工作:
tex.1d.v4.f32.s32 loadX0, [tex, track];track += ldx*2;tex.1d.v4.f32.s32 loadX2, [tex, track];track += ldx*2;tex.1d.v4.f32.s32 loadX4, [tex, track];track += ldx*2;tex.1d.v4.f32.s32 loadX6, [tex, track];track += ldx*2;
這樣做的問題是 tex.1d.v4.f32.s32 指令發出後其 track 操作數是不會被保存的，為了保證其不被下一個增量指令修改，必須要插入控制代碼等待前一條載入指令完成，而最壞情況下該指令可能花去上百個時鐘周期。用四個偏移變量就可以不用等待傳輸完成就可以將四條載入指令一一發射出去，也是起到了流水線的作用。其代價是每個線程需要額外佔用三個寄存器，所幸 GPU 上有足夠的寄存器可以提供。
將共享內存中的數據載入寄存器
上節的工作完成後，共享內存中就有 A 和 B 的數據各 8 行，每行 64 個浮點數。將其各取出一行就可以將其中的元素進行前述的加乘操作，完成後各再取出一行直到共享內存中的 8 行數據被用完，此時其他 warp 應該已經在共享內存的另一組完成了從紋理內存的傳輸，計算線程只需切換到另一組進行計算即可。
如圖 2 所示，對於每個線程，其實只需要 64 個浮點數中的 8 個，其在 A 和 B 向量中位置可以根據圖上的計算出，具體計算過程在代碼中是通過一頓騷位操作實現的，在此可以提前做一說明：
readAs = ((tid >> 1) & 7)
圖中線程 2*i 和 2*i+1 會用到同一段 A，可以寫作 (i/2) % 8。這段位操作就是這個表達式的位操作實現。
readBs = (((tid & 0x30) >> 3) | (tid & 1))
B 的行向量中的選擇更複雜一點，首先觀察到對於偶數線程每隔 16 個線程 B 方向就要差 2 段（8 個浮點數），所以對於 6 個比特位表示的 64 線程，tid & 0x30 表示其中代表 tid mod 16 的後四位可以被遮蓋掉，只有前兩位對選擇 B 有意義。其後的 >>3 其實是先 >>4 將前兩位拉到個位數，再 *2 來表達相差的兩段。| (tid & 1) 等價於 +(tid & 1)，表達的線程 2*i+1 永遠會選擇線程 2*i 後的那段數據，這也補上了之前相差兩段中缺失的部分。
在圖 2 中可能很早就有人注意到其中的線程排布順序非常彆扭，並沒有按照線程號按行或列一個個排下來，其原因是為了避免共享內存訪問的 bank 衝突。bank 衝突的定義和發生的條件在 CUDA 官方文檔中有詳細說明，簡單地說就是共享內存的訪問按地址被分成若干個 bank（最簡單的方法是做求餘數的操作），如果兩個線程要訪問位於同一 bank 的共享內存，其訪問無法同時完成，訪存時間成倍增加，增加的倍數由一個 bank 最多被多少個線程同時訪問決定。當然這是最一般的情況，GPU 中提供了一些機制，比如廣播，儘量減輕 bank 衝突對訪問時間的影響。圖 2 所示的線程順序就是為了消除 bank 衝突所作的調整後的最佳排序。另一個彆扭的地方是每個線程計算 4 個而不是直接計算，這也是為了避免 bank 衝突的技巧，在每個線程的實際計算中 4 個完全等價於 1 個矩陣。
在實現代碼中還用到了一個技巧，雖然每線程只需要輸入 16 個輸入數據，實際分配的寄存器是這個數字的兩倍，目的和前述的類似，是為了用兩組寄存器實現流水線，即每個線程在用一行數據作計算時預先讀取讀取下一行的數據。
readAs = ((tid >> 1) & 7) readBs = (((tid & 0x30) >> 3) | (tid & 1))
while (track0 < end){ // Process each of our 8 lines from shared for (j = 0; j < 8; j++) { // We fetch one line ahead while calculating the current line. // Wrap the last line around to the first. prefetch = (j + 1) % 8; // +1代表預讀取 // Use even/odd rows to implement our double buffer. if (j & 1) // 在兩組寄存器j0和j1直接切換 { // 共享內存到寄存器的傳輸還是使用向量指令 // 在maxas的代碼中可見j0Ax是4個連續的寄存器，一個指令就可以完成到這4個寄存器的傳輸而無需一一寫出 ld.shared.v4.f32 j0Ax00, [readAs + 4*(prefetch*64 + 0)]; ld.shared.v4.f32 j0By00, [readBs + 4*(prefetch*64 + 0)]; ld.shared.v4.f32 j0Ax32, [readAs + 4*(prefetch*64 + 32)]; ld.shared.v4.f32 j0By32, [readBs + 4*(prefetch*64 + 32)]; } else { ld.shared.v4.f32 j1Ax00, [readAs + 4*(prefetch*64 + 0)]; ld.shared.v4.f32 j1By00, [readBs + 4*(prefetch*64 + 0)]; ld.shared.v4.f32 j1Ax32, [readAs + 4*(prefetch*64 + 32)]; ld.shared.v4.f32 j1By32, [readBs + 4*(prefetch*64 + 32)]; } } // swap our shared memory buffers after reading out 8 lines readAs ^= 4*16*64; readBs ^= 4*16*64;
// Additional loop code omitted for clarity.}
C 矩陣的計算：寄存器分配和計算順序
現在所需要的數據已經被儘可能高效地被送到寄存器了，似乎可以直接使用 FFMA 加乘指令對它們直接進行操作了，畢竟這才是矩陣相乘內核應該做的事。不幸的是在此之前還要解決一個可能是整個項目中最大的麻煩，就是寄存器訪問的 bank 衝突。
為了在一個流計算單位內容納大量線程，GPU 準備了多達 32K 的寄存器文件，顯然其訪問無法和 CPU 一樣直接，而是和共享內存一樣要通過 bank 進行，因此 bank 衝突也是難免的，而一旦出現會對性能造成很大影響。Maxwell 上的寄存器文件有 4 個 32 位的 bank，每個寄存器通過對其編號的 mod 4 操作被映射到對應的 bank 上。如果在 C 矩陣的計算中出現以下指令就會出現寄存器 bank 衝突：
FFMA R0, R4, R5, R0; # R0, R4 在 bank 0，R5 在bank 1，R0和R4產生bank衝突
後來每一代 GPU 架構的寄存器 bank 都會有變動，比如 Volta 架構就是分為 2 個 64 位的 bank，這也是 maxas 無法在現在的主流 GPU 上發揮性能的主要原因。
直接使用彙編語言的一大優勢就是可以通過手動分配寄存器儘可能減少 bank 衝突：
0-63 號分配給 C 矩陣；
64-71 和 80-87 分配給 A 矩陣，72-79 和 88-95 分配給 B 矩陣（分配兩倍於實際大小用於流水線預讀取）。
因為是用向量指令載入，分配給 A 和 B 的每四位寄存器編號必須是連續的，也就是所有四個 bank 都會連續出現，因此在 A 和 B 的寄存器選擇上並沒有優化空間，maxas 能做到的是儘量調整分配給 C 的 63 個寄存器的順序，其所採用的編號方案如下圖：
圖 6. 每個線程內部所用數據的寄存器編號，每個寄存器所在 bank 用不同顏色標出，如果某個 C 元素的顏色和其對應的 A 或 B 元素相同就會發生 bank 衝突，這種元素在圖中用黑框標出。
顯然這是調整寄存器編號能得到的最好結果，圖中黑框標出的 bank 衝突不管如何調整 C 矩陣的編號是無法避免的，因為其來源是 A 和 B 用到了同一個 bank，而 A 和 B 中的操作數既需要佔據所有四個 bank（每個 bank 兩個數），又需要與另一個矩陣中的其他所有操作數配對，A 的每一寄存器必然會和 B 中的兩個寄存器產生 bank 衝突。事實上如果 C 使用最簡單的寄存器編號方式，比如在第一行佔據 0~7，那麼其中每一個寄存器都會和對應的 B 操作數發生 bank 衝突，就是非常壞的一種編號方法。
要進一步消除通過寄存器分配所不能消除的那部分 bank 衝突，需要用到 Maxwell 提供的一種操作數重用特性。在發出一條指令時，可以將其的一些操作數設為重用，硬體將把該操作數送往一個重用緩存。送如果後一條指令在同一位置要用到同一個操作數，則該指令可以直接從緩存而不用通過 bank 去取這個操作數，從而避免 bank 衝突。
FFMA R2, R4.reuse, R5, R2; # 此處指定R4將會被重用，將其放入緩存FFMA R0, R4.reuse, R5, R0; # R0和R4本來產生bank衝突，但因為上一條指令緩存了第二個操作數R4，只要到bank中取R0即可，從而避免了bank衝突FFMA R0, R5, R4, R0; # R0和R4的bank衝突依然會發生，因為所緩存的R4在第2個操作數，但本指令中R4落在第3個操作數上。
如果在線程中簡單地一行行或一列列遍歷圖 6 中 C 矩陣的 64 個寄存器，並且將 A 的寄存器設為重用，因為就可以解決 16 個 A 和 B 的寄存器 bank 衝突中的 14 個，不能解決的是寄存器 R3 和 R35，因為它們是該行的第一個用到該 A 操作數的指令，之前沒有指令將其送入重用緩存。知道原因後這兩個 bank 衝突也可以被輕鬆化解，只要在遍歷每行時偶數行從右到左（0 行從 26 到 3）奇數行從左到右（1 行從 7 到 30）。但是 maxas 即使對此還是不滿足，它在前述的來回遍歷基礎上又加上一個漩渦提出了一個更詭異的遍歷方法：
1, 0, 2, 3, 5, 4, 6, 7, 33, 32, 34, 35, 37, 36, 38, 39,45, 44, 46, 47, 41, 40, 42, 43, 13, 12, 14, 15, 9, 8, 10, 11,17, 16, 18, 19, 21, 20, 22, 23, 49, 48, 50, 51, 53, 52, 54, 55,61, 60, 62, 63, 57, 56, 58, 59, 29, 28, 30, 31, 25, 24, 26, 27
按照文檔中的說法，每個操作數有 8 字節的重用緩存可以緩存兩個 4 字節寄存器，而逐行遍歷只用了其中一個用於緩存 A 寄存器的數據，所以緩存使用率偏低。我推測考慮到有 B 操作數也有重用緩存但沒有利用，逐行遍歷的重用緩存利用率為 4/8/2=25%。對於來回遍歷的利用率估算不是那麼直觀，文檔直接給出了其利用率為 39%，而漩渦遍歷的利用率可以達到 49%。從 maxas 最後生成的彙編代碼來看其中有的指令確實對 A 和 B 同時使用了重用緩存，同時在對每個操作數緩存了兩個寄存器：
--:-:-:-:1 FFMA R37, R71.reuse, R72.reuse, R37;--:-:-:-:1 FFMA R36, R71.reuse, R73.reuse, R36;#前2條指令對第3個操作數緩存了R72和R73，它們被接下來的2條指令用到--:-:-:-:1 FFMA R38, R69.reuse, R73, R38; --:-:-:-:1 FFMA R39, R69.reuse, R72, R39;#前4條指令對第2個操作數緩存了R69和R71，它們被接下來的4條指令用到--:-:-:-:1 FFMA R45, R71.reuse, R74.reuse, R45;--:-:-:-:1 FFMA R44, R71, R75.reuse, R44;--:-:-:-:1 FFMA R46, R69.reuse, R75.reuse, R46;--:-:-:-:1 FFMA R47, R69, R74.reuse, R47;
不過，在來回遍歷可以完全解決 bank 衝突的情況下依然試圖提高重用緩存的使用率的目的並不在於提高重用率，而且因為 FFMA 指令之間插入的從共享內存到寄存器的載入指令。這樣做的目的是為了載入指令的延遲可以被不依賴於其數據的計算指令所填充。遍歷 C 矩陣寄存器的前八條指令和其間插入的載入指令如下：
01:-:-:-:0 FFMA R1, R66.reuse, R72.reuse, R1;--:-:-:-:1 LDS.U.128 R80, [R106+0x200];--:-:-:-:1 FFMA R0, R66, R73.reuse, R0;--:-:-:-:0 FFMA R2, R64.reuse, R73.reuse, R2;--:-:-:-:1 LDS.U.128 R88, [R107+0x200];--:-:-:-:1 FFMA R3, R64, R72.reuse, R3;--:-:-:-:0 FFMA R5, R67.reuse, R72.reuse, R5;--:-:-:-:1 LDS.U.128 R84, [R106+0x300];--:-:-:-:1 FFMA R4, R67, R73.reuse, R4;--:-:-:-:0 FFMA R6, R65.reuse, R73.reuse, R6;--:-:1:-:1 LDS.U.128 R92, [R107+0x300];--:-:-:-:1 FFMA R7, R65, R72.reuse, R7;
由於載入指令的運行時間需要 20 多個時鐘周期（對應於共享內存的延遲），在這段時間內其第一個的操作數都有可能被通過 bank 訪問到，此時有可能其後的計算指令也被發射並需要訪問同一個 bank，這就造成了一個延遲的 bank 衝突。不過這只是一個基本原理，maxas 的遍歷順序是如何具體避免這樣的 bank 衝突目前還沒有完全搞清楚。
從本節可以看出即使所有計算全部在寄存器中進行，還是要用到兩個技巧來得到最佳性能：
1. 最優化的寄存器編號
2. 最優化的遍歷順序
至於計算本身已經顯得如此簡單以至於 maxas 文檔都懶得提了。
傳輸 C 矩陣到主顯存
每個線程塊完成所負責的那個小片矩陣的計算後，最後一個任務就是將其從寄存器傳輸到主顯存。由於寄存器無法在線程間共享（其實有_shfl_sync() 指令可以但是此處不適用），所以每個線程必須將所計算的 4 個矩陣先傳輸至共享內存。之所以不從寄存器直接傳輸到主顯存是因為按照現有的線程布局無法利用 GPU 的超大帶寬。要充分利用 GPU 的帶寬我們希望每個 warp 的 32 個線程所同時傳輸的數據是連續的，這樣一個時鐘周期裡就可以一下子傳輸 128 字節的數據，如果這些數據離得太遠，在最壞可能下需要分 32 次才能傳輸出去。
根據圖 2 的線程布局，每一列連續的 64 字節數據分布在 8 個線程中，比如第 1 列前 4 行的結果都保存在線程 0,2,4,6,8,10,12,14 所控制的寄存器中，每個線程在該行有 8 個寄存器，而且為了避免 bank 衝突這 8 個寄存器都不是連續的，因此不能使用向量傳輸指令，所以需要分 8 次才能完成一個 warp 的傳輸，而此時 warp 中的其他 24 個線程將無所事事，因為其數據都不在這一列上。為了解決這個問題可以首先利用共享內存暫存所有線程的數據，然後用令一種線程布局將共享內存中的連續的數據同時傳輸出去。
首先還是要先將寄存器中的數據寫入共享內存。每個線程的寄存器中所保存的四個矩陣是分成在列上對齊的兩對。按圖 6 所示的 maxas 的寄存器分配，每一列上有八個寄存器。比如第一列有寄存器 3,7,1,5，屬於同一個矩陣的一列，以及 35,39,33,37，屬於令一個矩陣的一列。由於結果矩陣 C 也是按照行優先儲存的，如果將寄存器 3,7,1,5 拷貝到 4 個連續的寄存器（maxas 中命名為 cs），35,39,33,37 拷貝到 cs，就可以用向量儲存指令在兩個指令內將 8 個數拷貝到共享內存中對應的位置。圖 7 中的左圖是這個過程的示意圖，可以看作將圖 2. 的矩陣每隔四列抽出一條來拼在一起。完成後在共享內存中得到一個的矩陣，其中每一列都是連續的且對應於 C 矩陣中的一列。這時候改變一下線程次序，令 warp 中一個線程傳輸該列上一個字節的數據，就可以完成一次傳輸 32 個連續的浮點數。這個共享內存中的緩衝區可以利用之前為載入 A 和 B 所分配的空間，在完成 C 的計算後 A 和 B 的數據已經沒有用了。
圖 7. 左圖為寄存器寫入共享內存的線程布局，右圖為此後從同一塊共享內存讀取的線程布局。本圖中每一列是圖 2 中矩陣 C 的一列，相鄰的 2 列在矩陣 C 中間隔 4 列。
該方法的實現代碼如下。雖然這個方法需要反覆在寄存器和共享內存之間搬運數據，共享內存的延遲可以通過在 2 個 warp 間切換任務而得到掩蓋，畢竟比多次訪問主顯存要快多了。其中值得注意的是雖然這個方法明顯是分步完成的，但是代碼中沒有任何同步指令，這是因為所有的數據交換都是在同一個 warp 的 32 個線程直接完成的，warp 之間的數據保持獨立。GPU 的執行模型可以嚴格保證同一 warp 內的同一指令可以同時完成。能夠省去同步指令也是圖 2 中並行方法的一個優勢。
tid31 = tid & 31;tid32 = tid & 32; // 只可能取兩個值，0 為第一個warp，32為第二個warp
// Remove the high bits if present from the last loop's xor. // Also remove the 2048 added onto readBs.// 之前對A和B在共享內存中分配了兩倍於所需的容量（4096位元組），一塊用於已載入數據的計算，一塊用於載入下一段A和B，每一塊的前2048位元組存放A，後2048位元組存放B//這個AND操作等價於取第一塊存放A的內存來存放CreadAs &= 0x7ff; // 本線程左上角數據在64x64矩陣中的行坐標readBs &= 0x7ff; // 本線程左上角數據在64x64矩陣中的列坐標
// Write to shared using almost the same shared mapping as before but collapse readBs down to stride one.writeCs = (readBs / 4) * 64 + readAs; // 本線程左上角數據在64x64矩陣中的相對左上角的一維偏移，根據行坐標和列坐標計算出。64/4是行優先儲存矩陣進行向量傳輸的跨度。對於線程0這就是圖7左圖中的綠色格子中最上面的一個。
// Read out with a mapping amenable to coalesced global writesreadCs = ((tid32
ldc4 = ldc * 4; // ldc是矩陣C在行優先儲存中列方向的跨度，因子4代表其單位為字節而不是浮點數。
cx = bx*64 + tid31; // cx可看作所要寫入主顯存的那一整列的在整個矩陣C中所對應的行數cy = by*64 + (tid32 >> 1); // cy可看作所要寫入主顯存的那一整列的在整個矩陣C中所對應的列數，顯然對於同一個warp列數是一樣的
// Cy00, Cy04, Cy08, Cy12 是圖7右圖中上面那四個綠色格點在整個矩陣C中的偏移// 雖然它們在共享內存中相隔1列，在矩陣C中它們之間的間隔是4列，所有它們之間的偏移是ldc4*4Cy00 = (cy*ldc + cx) * 4 + C;Cy04 = Cy00 + ldc4 * 4;Cy08 = Cy00 + ldc4 * 8;Cy12 = Cy00 + ldc4 * 12;
foreach copy vertical line of 8 registers from C into .v4.f32 cs0 and cs4{ // 步驟1. 從不連續的寄存器到共享內存 // Feed the 8 registers through the warp shuffle before storing to global // 這裡有一個缺失的步驟，每個線程首先將其上下對其的兩個4x4矩陣中取出同一列的各四個元素，此時它們為了避免bank衝突 // 不得不位於不連續的寄存器上，這個步驟將其複製到8個連續的額外的寄存器cs0-cs7，上面的矩陣使用cs0-3，下面的使用cs4-7 st.shared.v4.f32 [writeCs + 4*00], cs0; // 將連續的寄存器cs0,cs1,cs2,cs3用向量指令傳輸到共享內存中該4個數對應的位置 st.shared.v4.f32 [writeCs + 4*32], cs4; // 和上一行同樣的操作，因為上下兩個4x4矩陣間隔32個數，需要對寫入位置增加4*32位元組的偏移
// 步驟2. 從共享內存讀取到寄存器，重用cs寄存器，不過此時並不用到其連續的特性 // cs0, cs2, cs4, cs6位於同一行上，在行優先儲存中相差64個元素，4*64位元組 ld.shared.f32 cs0, [readCs + 4*(0*64 + 00)]; // cs1,cs3,cs5,cs7位於另一行上，由圖7右圖可見與上一行相差32個元素 ld.shared.f32 cs1, [readCs + 4*(0*64 + 32)]; ld.shared.f32 cs2, [readCs + 4*(1*64 + 00)]; ld.shared.f32 cs3, [readCs + 4*(1*64 + 32)]; ld.shared.f32 cs4, [readCs + 4*(2*64 + 00)]; ld.shared.f32 cs5, [readCs + 4*(2*64 + 32)]; ld.shared.f32 cs6, [readCs + 4*(3*64 + 00)]; ld.shared.f32 cs7, [readCs + 4*(3*64 + 32)]; // 步驟3. 將cs寄存器的數寫入主顯存，對於整個warp相當於將一列連續的32個浮點數寫入主顯存。邏輯上可以看作是步驟2的反過程，除了改列的位置在共享內存和主顯存中有所不同。 st.global.f32 [Cy00 + 4*00], cs0; st.global.f32 [Cy00 + 4*32], cs1; st.global.f32 [Cy04 + 4*00], cs2; st.global.f32 [Cy04 + 4*32], cs3; st.global.f32 [Cy08 + 4*00], cs4; st.global.f32 [Cy08 + 4*32], cs5; st.global.f32 [Cy12 + 4*00], cs6; st.global.f32 [Cy12 + 4*32], cs7; // 在下一次循環中輸出本線程所計算的4個4x4矩陣的下一列，對應於C矩陣中的下一列，注意不要和圖7中共享內存中的下一列混淆。 Cy00 += ldc4; Cy04 += ldc4; Cy08 += ldc4; Cy12 += ldc4;
// After processing forth set shift over to the stride 32 registers // 補充說明，4次循環後4個4x4矩陣的左邊兩個已經傳輸到主顯存了，接下去要傳輸右邊的兩個。 // 左邊和右邊兩對4x4矩陣在C矩陣中對應的位置可以通過平移32列而重合，考慮到矩陣本身寬度有4列（在之前4次循環中已經通過 += ldc4 4次得到實現） // 實際需要額外平移的是左右兩對4x4矩陣的間距32-4=28列，這是這就是28這個magic number的來由 if (4th iteration) { Cy00 += ldc4 * 28; Cy04 += ldc4 * 28; Cy08 += ldc4 * 28; Cy12 += ldc4 * 28; }}
maxas 文檔中另有一張圖表達這個過程，但可能由於未能對該充分理解，感覺其意義不大反而容易造成混淆故沒有在此引用。代碼本身已經足夠描述這一過程了。
完成到主顯存的傳輸後，maxas 所生成的 GEMM 內核的任務就完成了
256 線程實現
在以上所描寫的每線程塊 64 線程的基礎上，可以將其擴展 4 倍到 256 線程，每個線程所做的工作不變。這樣每個線程塊所計算的小片矩陣的兩個維度各擴大 2 倍，達到，此時輸入矩陣的載入和結果的輸出會有相應的變化，但理解了 64 線程實現後這些變化就非常簡單，在此無需贅述。對於比較大的矩陣 256 線程實現有一些性能上的優勢，詳細測試結果參見 maxas 文檔。
結語
本文雖然儘可能詳盡地對原文檔中的偽代碼進行了注釋，但這還是相對高層的實現，具體到 GPU 機器碼還有一個重要的課題，即控制碼沒有在本文中涉及。考慮到本文的目的僅是介紹一些 GPU 優化的思路和實現方法，對此 maxas 文檔中涉及控制碼的部分沒有進行解讀。
總的來說，maxas 所用的優化思路還是比較清晰的，按其說法之前已經有文獻提出了，其最困難的地方在於 nVidia 不願意透露其硬體的實現細節，以至於都需要其作者經過艱苦的反向工程猜測出來的才能達到硬體性能的極限。可能作者自己搭建了一個測試平臺來快速驗證某些指令的細微差別所帶來的性能的影響。無論如何這是一個偉大的工作，值得任何一位有志於衝擊硬體性能極限的工程師深入研究。