一、參考鏈接

    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的一行的部分值。 拆分之后A_ip為(mc, kc), B_pj為(kc, nr), C_ij為(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。所以這種劃分更適合于行存儲優先的矩陣乘。　　　　　　　　　　　　　

（3) 拆分N，B矩陣被拆分為多列。再拆分K，A拆分成多列，B拆成多個block。拆分M，A的一列被拆分為多個小block。A的一列乘以B的每個block，得到矩陣C的一列的部分值。

　　　　　　

　　　　　　這種劃分與(2)的劃分唯一的區別是，訪問B矩陣是按列還是按行。很顯然(2)的劃分方式好于(3)。　　　　　　　

（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）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。