矩阵相乘优化(Gemm)
一、參考鏈接
http://tvm.d2l.ai/
https://www.cs.utexas.edu/users/pingali/CS378/2008sp/papers/gotoPaper.pdf
https://zhuanlan.zhihu.com/p/65436463
https://github.com/flame/how-to-optimize-gemm/
二、矩陣相乘優化方法
假設矩陣C = 矩陣A * 矩陣B; 矩陣A的shape為(M, K),矩陣B的shape為(K, N),矩陣C的shape為(m,n)。
普通的矩陣為 A的一行乘以B的一列,如下圖:
c/c++/python基本上是以行存儲優先的,本文將以行存儲優先作為基礎進行優化分析。
考慮兩種情況:
(1)當AB矩陣較小時,根據計算機結構可知,當從RAM中讀取AB矩陣內存,根據局部性原理可以將AB矩陣放到cache中,因為cpu訪問cache比訪問主存的快。
(2)當AB矩陣較大時,超過cache大小時,根據矩陣乘的普通方法,由于訪問“行優先存儲的B矩陣”的時候內存不連續(讀取B矩陣的一列),造成緩存cache頻繁的換入換出,從RAM讀取內存的次數大于AB矩陣的大小。因此第一種優化方法:
1. 向量化(SIMD)
向量化可以使一條指令并行的使多個相同操作數執行相同的操作,減少每次循環迭代時評估操作類型的開銷。
2. 內存對齊
內存對齊的原則:任何K字節的基本對象的地址必須都是K的倍數。
“假設 cache line 為 32B。待訪問數據大小為 64B,地址在 0x80000001,則需要占用 3 條 cache 映射表項;若地址在 0x80000000 則只需要 2 條。內存對齊變相地提高了 cache 命中率。”假定kernel一次計算執行4*4 大小的block, 根據MMult_4x4_7.c和MMult_4x4_8.c代碼,可以看出MMult_4x4_8.c使用了偏移量完成內存對齊。
3. 分塊
圖中共有6中拆分方法,依次分析:
(1) 拆分K,A矩陣拆成多列,B矩陣拆成多行。拆分M,A的一列拆分成三個block。拆分N,B的一行也可以拆分成多個小slice,每個slice放到寄存器中。遍歷A的列,A的每個block乘以B的一行,得到矩陣C的一行的部分值。 拆分之后Aip為(mc, kc), Bpj為(kc, nr), Cij為(mc, nr)。
如果kc*mc很小,那么Aip、Bpj、Cij都放入cache中,Aip(A矩陣的一個block)只需要被加載進cache一次,提高了cache命中率。對Ai和Bj進行pack使其內存連續。由于處理B矩陣是按照每個slice依次進行,所以這種劃分更適合于列存儲優先的矩陣乘。
(2) 拆分K,A拆分成多列,B拆成多行。拆分N,B的一行被拆分為三個block。拆分M,A的一列被拆分為多個slice。A的一列乘以B的每個block,得到矩陣C的一列的部分值。
類似于(1), 只是變成A的一行乘以B的一個block。所以這種劃分更適合于行存儲優先的矩陣乘?! ?/p>
(3) 拆分N,B矩陣被拆分為多列。再拆分K,A拆分成多列,B拆成多個block。拆分M,A的一列被拆分為多個小block。A的一列乘以B的每個block,得到矩陣C的一列的部分值。
這種劃分與(2)的劃分唯一的區別是,訪問B矩陣是按列還是按行。很顯然(2)的劃分方式好于(3)?! ?/p>
(4) 類似于普通矩陣乘,A的一行*B的一列。然后在K拆分,將A矩陣的每行劃分為多個block, 將B矩陣的每列劃分為多個block。
每次執行可以得到C矩陣一個block的值。當MNK非常大時,cache無法存下Cj,所以劃分方法沒有什么優勢。
(5) 類似于(4), 都是每次執行可以得到C矩陣的一個block的值。同樣當MNK非常大時,cache無妨存下Cj,所以該畫風方法沒有什么優勢。
(6) 類似于(1)。不一樣在于:對A矩陣的block遍歷是按行訪問。
和(1)相比,最外層循環是i,而非p。遍歷順序為A0*B0 + A1*B1 + .... + Ap*Bp 才能得到C一行的最終值,因此需要不斷的更新C這一行的值,可以使用一個順序的cache數組存放 pack,最后在unpack的時候累加。這與(1)中的按列訪問A矩陣的block不同,(1)中Ci += Ai*, 在最內層循環時只需要使用寄存器存儲每次Ai*Bpj的結果。
由于在循環外 cache 級別 unpack C 遠復雜于在循環內 register 級別 unpack C,因此(1)比(6)好。
總結
以上是生活随笔為你收集整理的矩阵相乘优化(Gemm)的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 良苦用心的近义词(良苦用心)
- 下一篇: 一度电等于多少万千瓦时电力(一度电等于多