如何對GPU中的矩陣乘法（GEMM）進行優化

本篇文章主要是介紹如何對GPU中的矩陣乘法（GEMM）進行優化。目前針對GEMM的優化，網絡上已經有非常多的教程和示例了。大部分的重要資料我都看了看。但總的來說，還是不夠接地氣，然后理解起來還是會比較費解。所以希望寫這么一篇文章，盡可能地去把GPU的GEMM優化說清楚，說明白。然后讓小白讀者也能通過這么一兩篇文章去更好地了解GEMM優化的相關技術。

不像上次的reduce優化一樣，能一篇文章說完。這次的GEMM優化會分為三個部分。第一個部分只說優化思路和分析，沒有任何代碼，這么做考慮也是為了減輕讀者的負擔，看代碼太累，盡可能地讓讀者先明白原理，為什么要這么做。第二個部分是對代碼的詳細解析，這個里面就是一行一行地去分析代碼。因為之前的很多博客進行了分析，但是代碼本身并沒有開源，或者說開源了代碼，但沒有解析，看起來太累了。我希望提供一個盡可能詳細的代碼解析，讀者看完之后能明白相關優化技巧，并且可以直接把代碼拿去驗證使用。第三個部分主要涉及到匯編器，最重要的是說明在NV的卡上，怎么去解決寄存器的bank沖突來獲取極致的性能。

本篇文章是GEMM優化的第一個部分，在這篇文章中，只說優化思路和分析。

前言

在高性能領域，對于矩陣乘（GEMM）的優化是一個非常重要的課題。GEMM可以非常廣泛地應用于航空航天、流體力學等科學計算領域，這也是之前HPC的主要應用場景。后來深度學習開展地如火如荼，由于對高算力的需要，也成為HPC的主要應用場景之一。這些年涌現了一系列的深度學習模型。模型里面最耗時的東西，包括卷積、全連接層、attention，都可以轉換成GEMM操作。所以說，GEMM優化的重要性，怎么突出都不過分。

目前網上能找到的針對GEMM優化的資料主要有這么幾個方面：（1）論文，目前針對GPU進行GEMM優化的論文非常多，這里主要推薦Understanding the GPU Microarchitecture和Fast implementation of dgemm on fermi gpu以及Dissecting the NVIDIA Volta GPU Architecture via Microbenchmarking。這幾篇論文在業界都比較有影響力，就是代碼開源方面做的不算太好。（2）官方博客，主要是CUTLASS和NervanaSystems-SGEMM優化。還有前段時間曠視發的文章CUDA矩陣乘法優化，寫的都很詳細。（3）github的一些demo，代碼量不大，看起來比較舒服。我是看了這兩個：

demo1 ：

https://github.com/Cjkkkk/CUDA_gemm

demo2 ：

https://github.com/yzhaiustc/Optimizing-SGEMM-on-NVIDIA-Turing-GPUs

demo1代碼寫的好理解一些，但是優化工作沒做完全，沒有做到prefetch。demo2是效果很好，11個優化技巧，不斷逼近cublas。但是代碼真的看起來比較難受，最重要的很多參數寫死了，不好去調。 總而言之，目前列舉的上述資料存在著這么兩個問題：（1）文檔方面，讀起來還是比較費勁，對于小白來說，還是不夠簡單不夠傻，看起來太累了；（2）代碼方面，要么是沒公開代碼，要么是代碼太多了，看不下去；還有的就是代碼可讀性很強，但是優化工作還不是特別深，或者就是代碼優化做的很好，但是可讀性差了。方方面面總是有點欠缺，所以希望能夠寫一篇盡可能地在文檔上簡單明了，在代碼上詳細且可讀性好的文章。當然，這是一個逐步迭代的過程，所以這篇文章也會持續進行更新哈。 本篇文章主要是采納了cutlass的行文思路，主要介紹GEMM中的數據分塊和如何在多級存儲進行數據搬運。這也是HPC優化的核心思想，怎么樣讓數據放在更近的存儲上來掩蓋計算的延時，從而減少存儲墻的影響。文章分為四個方面進行敘述，首先介紹在global memory層面如何進行分塊以及數據搬運，隨后介紹在shared memory層面如何進行分塊以及數據搬運，而后介紹在register層面如何進行分塊以及避免bank沖突，最后介紹如何進行prefetch以更好地掩蓋訪存時延。

一、從global memory到shared memory

假設有矩陣A、B，需要計算矩陣A和B的乘，即矩陣C。A、B、C三個矩陣的維度分別為，，，且三個矩陣中的數據都是單精度浮點數。對于C中每一個元素，C[i][j]，可以看作是A的一行和B的一列進行一次歸約操作。采用最naive的GEMM算法，在GPU中，一共開啟個線程，每個線程需要讀取矩陣A的一行與矩陣B的一列，而后將計算結果寫回至矩陣C中。因而，完成計算一共需要從global memory中進行次讀操作和次寫操作。大量的訪存操作使得GEMM效率難以提高，因而考慮global memory中進行分塊，并將矩陣塊放置到shared memory中。其示意圖如下： 對global memory進行分塊的GEMM算法示意圖見上圖右側。首先將A、B、C三個矩陣劃分為多個維度為?，，?的小矩陣塊。三個矩陣形成?，，?的小矩陣網格。其中，，，。隨后在GPU中開啟??個block，每個block負責C中一個維度為??的小矩陣塊的計算。計算中一共有K次迭代，每一次迭代都需要讀取A中一個維度為??的小矩陣塊和B中一個維度為??的小矩陣塊，并將其放置在shared memory中。因而，完成C中所有元素的計算一共需要從global memory中讀取?，即??個單精度浮點數。相比于naive的GEMM算法，訪存量減少為原來的?。通過global memory中分塊算法極大地減少了對global memory的訪存量。并且，相比于naive算法，對global進行分塊可以更充分地利用數據局部性。在naive算法中，每一個線程都需要直接從global memory中取數，其時延非常長，計算性能非常差。而進行分塊后，將維度為?，?的小矩陣塊先存儲到shared memory之中。而后計算單元進行計算時可以直接從shared memory中取數，大大減少了訪存所需要的時延。

二、從shared memory到register

隨后，我們進一步考慮從shared memory到register的過程。在這里，只分析一個block中的計算。當進行K輪迭代中某一輪迭代時，GPU將維度為，的小矩陣塊存儲到shared memory中，而后各個線程將shared memory中的數據存入register中進行計算。 在不對shared memory分塊時，一個block中含有個線程，每一個線程負責C中一個元素的計算。則一個block一共需要對shared memory進行次讀操作。而后考慮對shared memory進行分塊，對的小矩陣進行再一次劃分，將其劃分為多個維度為的子矩陣。則一個block需要負責個子矩陣。其中，，。隨后，在一個block中開啟個線程，每個線程負責一個維度為的子矩陣的計算。在計算中，一個block一共需要從shared memory讀取，即個單精度浮點數。相比于未分塊的算法，對于shared memory中的訪存量減少為原來的。并且，由于將數據放入register中，可以直接對數據進行運算，減少了從shared memory中取數的時延。

三、register分塊

在這里，我們考慮最后一層，即register中的計算，并且只分析一個thread。在完成以上的過程后，對于一個線程而言，它現在擁有：個A矩陣的寄存器值，個B矩陣的寄存器值，以及個C矩陣的寄存器值。通過這些寄存器的值，需要計算個數。這需要條FFMA指令。 這個時候會涉及到寄存器的bank conflict。在NV的GPU中，每個SM不僅會產生shared memroy之間的bank 沖突，也會產生寄存器之間的bank沖突。這一點對于計算密集型的算子十分重要。像shared memory一樣，寄存器的Register File也會被分為幾個bank，如果一條指令的的源寄存器有2個以上來自同一bank，就會產生沖突。指令會重發射，浪費一個cycle。PS：這個地方是從曠視的博客中看的。然后對于maxwell架構的GPU而言，bank數為4，寄存器id%4即所屬bank。 我們假設對這個thread來說，、。并且計算C的寄存器以一種非常naive的情況分配，如下圖左側所示。則需要產生16條FFMA指令，列舉如下：

FFMA R0, R16, R20, R0 FFMA R1, R16, R21, R1 ……

可以從中看出，這會產生大量的register bank沖突，所以需要對參與計算的寄存器重新進行分配和排布,如上圖右側所示。在有些地方，這種方式也可以叫做register分塊。

四、數據的prefetch

最后，我們來講講如何通過對數據進行prefetch來減少訪存的latency。我們再來回顧GEMM的過程，并且仔細地看看這個訪存的latency到底是怎么導致的。對于一個block而言，需要計算一個的矩陣塊，這個時候需要進行K次迭代，每次迭代都需要先將來自A和B的兩個小塊送到shared memory中再進行計算。而從global中訪存實際上是非常慢的，所以導致了latency。雖然GPU中可以通過block的切換來掩蓋這種latency，但是由于分配的shared memory比較多，活躍的block并不太多，這種延時很難被掩蓋。對于一個thread，需要計算一個的小矩陣，但是必須先將數據從shared memory傳到寄存器上，才能開始進行計算。所以導致了每進行一次迭代，計算單元就需要停下來等待，計算單元不能被喂飽。

為此，需要進行數據的Prefetch來盡可能地掩蓋這種latency。思想也比較簡單，需要多開一個buffer，進行讀寫分離。示意圖如下。當block進行第2輪迭代時，需要對A2和B2進行計算，在計算單元進行計算的同時，我們將A3和B3提前放置到shared memory。而后，在進行第3輪迭代時，就可以直接對shared memory中的A3和B3進行計算，而不需要等待從global memory搬運到shared memory的時間。寄存器上的Prefetch也是同理。

總結

GEMM的優化思想，基本上就是這么幾方面的內容。希望大家通過介紹能夠對GEMM的優化有一個比較直觀且具體的理解。