# Kernel_matmul contributed by < [vestata](https://github.com/vestata/matmul.git) > ## 介紹 此專案源自於 jserv 的 Linux 核心實作 [2024q1 第 9 週測驗題](https://hackmd.io/@sysprog/linux2024-quiz9#%E6%B8%AC%E9%A9%97-2) ,有關以下之探討是否可以更加發揮硬體性能已加速矩陣相乘。 ## 目標 預計會完成以下討論: - [ ] 以 CMWQ 重寫,並針對批次的矩陣乘法運算,提出有效的存取模型 - [x] 在 GitHub 找出矩陣乘法相關專案,如 matmul, matmul-bench, matmul-cpu, libxsmm, Matrix_Multiply_using_Arm_Neon_and_Avx,並進行效能比較和實作分析,從而歸納出提升矩陣乘法的手法 - [ ] 嘗試在 Linux 核心模組使用 SSE/AVX/NEON 等 SIMD 指令集並降低資料存取的延遲 - [x] 研讀 [LLaMA Now Goes Faster on CPUs](https://justine.lol/matmul/) 並歸納加速矩陣乘法的手段 ## 紀錄 ### [matmul](https://github.com/attractivechaos/matmul) #### 程式 在初始矩陣的時候使用的手法,它直接宣告 row * col 的連續記憶空間,再 for 迴圈分配到每一 row,此方法在宣告與釋放記憶體很整潔: ```c m = (float**)malloc(n_rows * sizeof(float*)); m[0] = (float*)calloc(n_rows * n_cols, sizeof(float)); for (i = 1; i < n_rows; ++i) m[i] = m[i-1] + n_cols; ``` ##### `s_dot()` `s_dot()`:用來計算兩個向量的內積,`s_dot_1()` 使用一個一個計算線性的方式,`s_dot_8()`:使用的是一次計算 8 個的的線性方式。 所以先計算出 8 個為一次計算,使用 `n8 = n>>3<<3` 便可以獲得以 bitwise 獲得小於 n 的最大 8 的倍數。 ##### `sdot_sse()` 以下進入 SSE 指令的內積操作,一樣一次處理 8 個資料: ```c #ifdef __SSE__ #include <xmmintrin.h> float sdot_sse(int n, const float *x, const float *y) { int i, n8 = n>>3<<3; __m128 vs1, vs2; float s, t[4]; vs1 = _mm_setzero_ps(); vs2 = _mm_setzero_ps(); for (i = 0; i < n8; i += 8) { __m128 vx1, vx2, vy1, vy2; vx1 = _mm_loadu_ps(&x[i]); vx2 = _mm_loadu_ps(&x[i+4]); vy1 = _mm_loadu_ps(&y[i]); vy2 = _mm_loadu_ps(&y[i+4]); vs1 = _mm_add_ps(vs1, _mm_mul_ps(vx1, vy1)); vs2 = _mm_add_ps(vs2, _mm_mul_ps(vx2, vy2)); } for (s = 0.0f; i < n; ++i) s += x[i] * y[i]; _mm_storeu_ps(t, vs1); s += t[0] + t[1] + t[2] + t[3]; _mm_storeu_ps(t, vs2); s += t[0] + t[1] + t[2] + t[3]; return s; } #endif ``` 參閱 [Intel® Intrinsics Guide](https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html) `__m128` 宣告一個 128 位元的 `_mm_setzero_ps()` :Return vector of type `__m128` with all elements set to zero. `__m128 _mm_loadu_ps (float const* mem_addr)`:將 4 個單精度浮點(32位元)共 128 位元,存入指定記憶體。 `_mm_add_ps()`、`mm_mul_ps()` 分別是對 4 個浮點數的加法語乘法。 `_mm_storeu_ps (float* mem_addr, __m128 a)` 將 128 位元存入 `t` 之中。 非 8 的倍數則逐個相乘計算內積,最後再將所有 `__m128` 加進去總和。 ##### `mat_mul0()` Naive 最基礎的矩陣相乘 ##### `mat_mul1()` Transposed 因為記憶體的排列會是循序的,如果取 column 的時候要跳躍較遠的位置才可以從記憶體取到值,無法善用 locality,所以先將 b 矩陣轉置,會需要多一個額外空間,但是可以善用 cache 減少記憶體存取開銷。 ##### `mat_mul4()` sdot w/o hints 單純在轉置矩陣矩陣乘法的操作上使用 `sdot_1()`。 ##### `mat_mul3()` sdot with hints 單純在轉置矩陣矩陣乘法的操作上使用 `sdot_8()`。 ##### `mat_mul2()` SSE sdot 單純在轉置矩陣矩陣乘法的操作上使用 `sdot_sse()`。 ##### `mat_mul7()` SSE+tiling sdot ```c float **mat_mul7(int n_a_rows, int n_a_cols, float *const *a, int n_b_cols, float *const *b) { int i, j, ii, jj, x = 16, n_b_rows = n_a_cols; float **m, **bT; m = mat_init(n_a_rows, n_b_cols); bT = mat_transpose(n_b_rows, n_b_cols, b); for (i = 0; i < n_a_rows; i += x) { for (j = 0; j < n_b_cols; j += x) { int je = n_b_cols < j + x? n_b_cols : j + x; int ie = n_a_rows < i + x? n_a_rows : i + x; for (ii = i; ii < ie; ++ii) for (jj = j; jj < je; ++jj) m[ii][jj] += sdot_sse(n_a_cols, a[ii], bT[jj]); } } mat_destroy(bT); return m; } ``` 這邊程式的意思為拆成一個一個小的 16 * 16 的小塊做計算,以塊狀表示順序為: | Column 1 | Column 2 | Column 3 | | -------- | -------- | -------- | | 1 | 2 | 3 | | 4 | 5 | 6 | | 7 | 8 | 9 | 但是這邊我不懂它怎麼利用 locality 的性質,就直觀來講每一個塊中他的 `bT` 會一直更換,這樣顯的更沒有善用 locality。 :::danger `<cblas.h>` 沒看 ::: #### 實驗與測試 安裝 OpenBLAS ```shell sudo apt-get install libopenblas-dev ``` 編譯 ```shell make CBLAS=/path/to/cblas/prefix ``` 使用命令,`-n` 為矩陣大小,`-a` 為其提供的不同算法: ```shell ./matmul -n INT -a INT ``` 我使用我的設備 `13th Gen Intel(R) Core(TM) i7-13700` 以矩陣 2000 * 2000, 3000 * 3000 去計算,以下為實驗結果: | Implementation | -a | Linux, -n 2000 | Linux, -n 3000 | | --------------- | --- | -------------- | -------------- | | Naive | 0 | 10.093 sec | 45.4654 sec | | Transposed | 1 | 5.46336 sec | 13.1423 sec | | sdot w/o hints | 4 | 5.57785 sec | 13.1282 sec | | sdot with hints | 3 | 3.10309 sec | 6.12113 sec | | SSE sdot | 2 | 3.12613 sec | 6.2497 sec | | SSE+tiling sdot | 7 | 3.05097 sec | 3.9111 sec | | OpenBLAS sdot | 5 | 3.13899 sec | 6.46783 sec | | OpenBLAS sgemm | 6 | 1.50497 sec | 2.82676 sec | 以上實驗使用的是 libopenblas-dev 會遇到的問題是它會預設開啟多個 thread 計算,若是想要固定在單執行緒執行需要使用 [OpenBLAS](https://github.com/OpenMathLib/OpenBLAS) 自行編譯,以下是作者的 OpenBLAS 設定檔。 ```shell TARGET=CORE2 BINARY=64 USE_THREAD=0 NO_SHARED=1 ONLY_CBLAS=1 NO_LAPACK=1 NO_LAPACKE=1 ``` 以下完成 OpenBLAS 設定的編譯: ```shell git clone https://github.com/OpenMathLib/OpenBLAS.git cd OpenBLAS make TARGET=CORE2 BINARY=64 USE_THREAD=0 NO_SHARED=1 ONLY_CBLAS=1 NO_LAPACK=1 NO_LAPACKE=1 FC=gfortran sudo make TARGET=CORE2 BINARY=64 USE_THREAD=0 NO_SHARED=1 ONLY_CBLAS=1 NO_LAPACK=1 NO_LAPACKE=1 FC=gfortran PREFIX=/usr/local install ``` 應該會正常編譯安裝完成 回到 `matmul`,因為剛剛設定位置是 `/usr/local` ```shell make CBLAS=/usr/local ./matmul -n INT -a INT ``` 得到以下結果: | Implementation | -a | Linux, -n 2000 | Linux, -n 3000 | | --------------- | --- | -------------- | -------------- | | Naive | 0 | 7.44295 sec | 45.4654 sec | | Transposed | 1 | 3.05836 sec | 13.1423 sec | | sdot w/o hints | 4 | 3.09377 sec | 13.1282 sec | | sdot with hints | 3 | 0.672151 sec | 6.12113 sec | | SSE sdot | 2 | 0.698778 sec | 6.2497 sec | | SSE+tiling sdot | 7 | 0.540385 sec | 3.9111 sec | | OpenBLAS sdot | 5 | 0.682843 sec | 6.46783 sec | | OpenBLAS sgemm | 6 | 0.281617 sec | 2.82676 sec | 在這邊發現一些編譯設定上面的錯誤,`TARGET=CORE2` 針對的架構優化跟我的設備 `13th Gen Intel(R) Core(TM) i7-13700` 不符,但是 target list 中找不到相符的設備,故先捨棄;然後打開 `LAPACKE` 的設定使 `fortran` 的 `SGEMM` 能正常運行。 故以上 OpenBLAS 的設定改為: ```shell git clone https://github.com/OpenMathLib/OpenBLAS.git cd OpenBLAS make BINARY=64 USE_THREAD=0 NO_SHARED=1 FC=gfortran sudo make BINARY=64 USE_THREAD=0 NO_SHARED=1 FC=gfortran PREFIX=/usr/local install ``` 再到 matmul 重新編譯: ```shell make CBLAS=/usr/local ``` 得以下結果: | Implementation | -a | Linux, -n 2000 | Linux, -n 3000 | | --------------- | --- | -------------- | -------------- | | Naive | 0 | 7.58799 sec | 42.3012 sec | | Transposed | 1 | 3.01531 sec | 10.9528 sec | | sdot w/o hints | 4 | 3.01675 sec | 10.9439 sec | | sdot with hints | 3 | 0.702078 sec | 4.04731 sec | | SSE sdot | 2 | 0.600825 sec | 4.21123 sec | | SSE+tiling sdot | 7 | 0.523011 sec | 1.8697 sec | | OpenBLAS sdot | 5 | 0.597181 sec | 4.17523 sec | | OpenBLAS sgemm | 6 | 0.287184 sec | 0.963865 sec | 在這邊我發現到一開始使用 `libopenblas-dev` 並行計算(以 `htop` 查看),但是我後來將 OpenBLAS 重新編譯只使用一個 thread,有更好的效能表現,所以需要實驗在多少大小矩陣之下平行化的優勢會超過他的成本。 #### Matrix 相乘手法 以下舉例為 a = M * N; b = N * K,比較進存取記憶體次數: ##### Naive 最直觀的手段為 讀取:$2*MNK$ 寫入:$MK$ 總和:$2MNK+MK$ ##### Transposed 專案中的解釋為: >Transposing the second matrix for cache efficiency 不計記憶體生成與銷毀成本。 多新增一個轉置矩陣 `bT` 來改進記憶體存取模式。這種方法可以提高記憶體的連續性,從而改善快取的命中率。 轉置:讀取 $MK$,寫入 $MK$,共 $2MK$ 矩陣相乘:讀取 $2MNK$,寫入 $MK$,共 $2MNK+MK$ 總和:$2MNK+3MK$ :::danger TODO 後面 BLAS 的操作計算要計算記憶體存取次數 ::: ### [Matrix_Multiply_using_Arm_Neon_and_Avx](https://github.com/ruthreshx/Matrix_Multiply_using_Arm_Neon_and_Avx) 專案中提到 neon based matrix multiplication,NEON 是 ARM 處理器中的 SIMD 指令集。NEON 向量寄存器是 128 位,可以包含 2 個 64 位元、4 個 32 位元、8 個 16 位元或 16 個 8 位元。 先附上作者的實驗結果 | Matrix Size | int16(s) | int16neon(s) | int32(s) | int32neon(s) | float32(s) | float32neon(s) | float64(s) | float64neon(s) | | ----------- | --------- | ------------ | -------- | ------------ | ---------- | -------------- | ---------- | -------------- | | 8 x 8 | 0.000042 | 0.000039 | 0.000035 | 0.000022 | 0.000054 | 0.000027 | 0.000030 | 0.000027 | | 32 x 32 | 0.001439 | 0.001305 | 0.001502 | 0.001403 | 0.001559 | 0.000730 | 0.001686 | 0.008837 | | 50 x 50 | 0.005361 | 0.003472 | 0.005773 | 0.002039 | 0.005832 | 0.002198 | 0.005926 | 0.002301 | | 100 x 100 | 0.023119 | 0.016946 | 0.021783 | 0.007522 | 0.022988 | 0.008413 | 0.023099 | 0.011161 | | 200 x 200 | 0.158035 | 0.123986 | 0.149948 | 0.059083 | 0.159806 | 0.067440 | 0.163524 | 0.089471 | | 300 x 300 | 0.532597 | 0.435758 | 0.545451 | 0.199766 | 0.542208 | 0.228595 | 0.582677 | 0.319158 | | 400 x 400 | 1.314986 | 1.004088 | 1.367670 | 0.471891 | 1.397881 | 0.547197 | 1.725430 | 0.787738 | | 1000 x 1000 | 21.098629 | 3.644162 | | | | | | | 實驗環境: ```shell Architecture: aarch64 CPU op-mode(s): 32-bit, 64-bit Byte Order: Little Endian CPU(s): 4 On-line CPU(s) list: 0-3 Vendor ID: ARM Model name: Cortex-A72 Model: 3 Thread(s) per core: 1 Core(s) per cluster: 4 Socket(s): - Cluster(s): 1 Stepping: r0p3 CPU max MHz: 1800.0000 CPU min MHz: 600.0000 BogoMIPS: 108.00 ``` 這邊使用我手上的樹莓派 4 進行測試,發現作者預設便是使用並行計算,所以我將此執行緒限定在 CPU 0 上做比較。 | Matrix Size | int16(s) | int16neon(s) | int32(s) | int32neon(s) | float32(s) | float32neon(s) | float64(s) | float64neon(s) | | ----------- | --------- | ------------ | -------- | ------------ | ---------- | -------------- | ---------- | -------------- | | 1000 x 1000 | 21.098629 | 3.644162 | | | | | | | ### 研讀 [LLaMA Now Goes Faster on CPUs](https://justine.lol/matmul/) #### 介紹 為 llamafile 寫了 84 個新的矩陣乘法核心,使讀取 prompts/images 更快。與 llama.cpp 相比,當在 CPU 上使用 F16 和 Q8_0 加權時,llamafile 的 prompt 評估時間應該會快 30% 至 500%。核心在適合於 L2 快取的矩陣上比 [MKL](https://sites.cs.ucsb.edu/~tyang/class/240a17/slides/intel-mkl-gemm.pdf) 快 2 倍,這使它們成為一項進行中的工作,因為加速效果最適合於 prompts 少於 1000 個 token。 >So far, I've only written optimized kernels for the q8_0, f16, q4_1, q4_0, and f32 data types. I think both q8_0 and f16 are really solid choices. Possibly even f32 if you've got plenty of RAM. 這邊它提到它將浮點數運算轉為定點數運算,已免浮點數運算時間太慢,就有如 jserv 在 [lab0-c ttt](https://hackmd.io/@sysprog/linux2024-ttt) 提到的浮點數轉定點數,以上是它提到的支援的轉換格式。q8_0、f16 會是它推薦的選擇。==怎麼做 quantization 會是新的瓶頸==。 #### 實驗比較 ##### Raspberry Pi 作者有對 Raspberry Pi 做一系列的比較,對不同的 weights, data type 與不同的硬體做比較。使 LLM 的運算有機會下放到邊緣的運算。這邊作者有使用 Raspberry Pi 5,它們還引入了對 ARMv8.2 dotprod 和 fp16 算術 ISA 的支持,光是這兩點在 2023 就對 llama.cpp 使用 fp 16 有了 10 倍的能提昇。作者提到它使用一個針對 AVX512 的撰寫核心模組去做嘗試結果再獲得 2 倍的效能提昇。 補充:[AVX-512](https://en.wikipedia.org/wiki/AVX-512) 是一種面向 x86 架構的 SIMD 技術,允許使用 512 位元的 SIMD 指令。 在 ARMv8.2 fp16 ISA 會產生更多的錯誤,因為在 llamafile 中==並未==採用 [Kahan summation](https://en.wikipedia.org/wiki/Kahan_summation_algorithm) 來計算 fp16 內積,會引起更大的誤差。所以 Q8_0 表現更好的 perplexity > So Q8_0 weights actually end up having slightly better perplexity, because it uses the dotprod ISA which lets us updot signed 8-bit integers into a 32-bit compute type which absorbs errors. 然而更快的 fp16 並不是沒用處,許多開發人員認為 fp16 的誤差是可以無視的。 這邊使用一個垃圾郵件過濾作為範例,架設在 email server 中的 TinyLLaMA 過濾器: ```shell llamafile -m TinyLlama-1.1B-Chat-v1.0.f16.gguf \ --grammar 'root ::= "yes" | "no"' --temp 0 -c 0 \ --no-display-prompt --log-disable -p "<|user|> Can you say for certain that the following email is spam? To: jtunney@gmail.com From: Federal-Tax-DebtHelp <ConfirmationEmail.inzk@janents.com> Subject: Reduce your payments to what you can afford Reduce your payments to what you can afford [IMG] [IMG] [IMG] </s> <|assistant|>" ``` 輸出 yes/no: ```shell jart@pi5:~/scratch$ time ./spam.sh yes ``` 在樹莓派花了 3 秒鐘。 其中有一些設定: - --temp 0 turns off the random number generator (we don't want improvisation for a spam filter) - Here we see why prompt eval time is king. Token generation speed (eval time) doesn't matter for this use case, since we're using the --grammar flag to force the LLM to only print a single "yes\n" or "no\n" token. - I piped the original email through links -codepage utf-8 -force-html -width 400 -dump /dev/stdin which reduced the number of tokens and removed hidden HTML content that was put there by the spammer to make the email look like ham to naive Bayesian filters. - The -c 0 flag configures TinyLLaMA to use the maximum context size, which is 2048 tokens. That's the largest prompt we can give it. To help avoid feeding in too much text, you can pipe the email through sed 's/ */ /g' | dd bs=1 count=7000 to remove superfluous spaces and place an upper limit on its size. ##### Gaming Hardware 使用 Alderlake 在 fp16 的表現相對樹莓派有 5 倍提昇,且 ARMv8.2 有 rounding errors 的產生,x86 核心使用 float32 計算,減少 float 計算偏差的產生。剛剛的過濾郵件在 Alderlake 執行比樹莓派快 7 倍,為 420 ms。這表示在小工作量時,CUDA 還未啟動便由已經由 CPU 完成。 這邊提到 llamafile 避免跑在 E-cores。這是 llamafile 可以比 llama.cpp 更快的原因之一,將 llm 限制在 P-cores 上,剩下的核心執行電腦其他的程式。這邊提到在 llama.cpp 使用 lockstep 來調度執行緒,所以當有一核心需要化更多時間時,其他核心都要等待它。 ##### Apple Hardware 這邊作者提到 Mac Studio 是 llama.cpp 開發團隊最專注的硬體,所以效能提昇的範圍對作者來說比較有限。 Apple 將其微處理器垂直整合,將其 RAM 整合進 CPU 之中,所以在需要記憶體溝通的操作加速,比如說 latency-bound operations 的 token generation,因為 CPU 不用再花時間記憶體存取。 ##### Professional Hardware Llamafile 不僅支援無 GPU 的環境,也支援高端 CPU;其中,AMD Ryzen Threadripper PRO 7995WX 基於 Zen4 架構,擁有 96 個支援 AVX512 的核心。 7995WX 的 x86 指令集架構(ISA)在執行 prompt 時的速度遠超 M2 Ultra,然而在 eval 速度上卻相差無幾。為了全面測試此晶片的性能,作者必須對 llama.cpp 增加對 bf16 和 AVX512 的支援。作者表示,經過他的改進,llama 現在可以在 Zen4 架構上達到 2.8 倍的速度提升。 >One thing I like about AVX512 is that Google's Gemma model can solve math riddles on AVX512 but not on AVX2 because the bigger vectors usually make it easier to reduce rounding errors. Its VDPBF16PS instruction helps us updot bf16 similar to VNNI and ARM dotprod. 這邊 bigger vector 與 reduce rounding errors 的關聯我不是很清楚。但是文中提到 [VDPBF16PS](https://www.felixcloutier.com/x86/vdpbf16ps) 他是一個 SIMD 內積指令,針對 2 個 bf16 做內積並以一個單精度浮點數輸出。因為 Mistral, TinyLLaMA 本身使用 bf16 計算,如果要特別將 bf16 轉為 fb16 只有 13% 的數字會準確表示(因為兩者表示範圍不一樣),在 Mistral 7b 實際操作中 99.71% 的計算在那 13% 之中;但是作者還是希望可以提供準確的計算,所以要在 llamafile 中加入 bf16 會是他的下一步。 #### 技術細節 在 LLM 的 trasformer 之中會有以下的計算: > e.g. rope, transpose, reshape, softmax, rms_norm, etc 全部可以用矩陣相乘來加速計算,使用的函式為 GGML_OP_MUL_MAT,使用 linux perf 效能計算工具,有 95% 的時間在進行矩陣相乘。 以下是最基本 $O(n^3)$ 的矩陣相乘比: python: ```python def matmul(A, B): assert len(B) == len(A[0]) return [[sum(A[i][l] * B[l][j] for l in range(len(B))) for j in range(len(B[0]))] for i in range(len(A))] ``` 使用 python 作者的單位是 0.042 gigaflops,這邊補充 gigaflops 為 一個GFLOPS(gigaFLOPS)等於每秒 $10^9$ 次的浮點運算 fortran: ```fortran SUBROUTINE SGEMM(TRANSA,TRANSB,M,N,K,ALPHA,A,LDA,B,LDB,BETA,C,LDC) * .. Scalar Arguments .. REAL ALPHA,BETA INTEGER K,LDA,LDB,LDC,M,N CHARACTER TRANSA,TRANSB * .. Array Arguments .. REAL A(LDA,*),B(LDB,*),C(LDC,*) [...] * * Form C := alpha*A*B + beta*C. * DO 90 J = 1,N IF (BETA.EQ.ZERO) THEN DO 50 I = 1,M C(I,J) = ZERO 50 CONTINUE ELSE IF (BETA.NE.ONE) THEN DO 60 I = 1,M C(I,J) = BETA*C(I,J) 60 CONTINUE END IF DO 80 L = 1,K IF (B(L,J).NE.ZERO) THEN TEMP = ALPHA*B(L,J) DO 70 I = 1,M C(I,J) = C(I,J) + TEMP*A(I,L) 70 CONTINUE END IF 80 CONTINUE 90 CONTINUE [...] RETURN END ``` 直接將它對接近 numpy,numpy 在這邊的處理是一樣的算法,但是調用 fortran 完成高效運算。表現為 29 gigaflops。 C++: ```cpp // multiplies matrices on cpu // with column major ordering // // m×k * k×n → m×n // k×m * k×n → m×n if aᵀ // m×k * n×k → m×n if bᵀ // k×m * n×k → m×n if aᵀ and bᵀ // template <typename T, typename TA, typename TB, typename TC> void GEMM(bool aᵀ, bool bᵀ, int m, int n, int k, T α, const TA *A, int lda, const TB *B, int ldb, T β, TC *C, int ldc) { assert(m >= 0 && n >= 0 && k >= 0); assert(lda >= std::max(1, aᵀ ? k : m)); assert(ldb >= std::max(1, bᵀ ? n : k)); assert(ldc >= std::max(1, m)); #pragma omp parallel for collapse(2) if (m * n * k > 300000) for (int i = 0; i < m; ++i) for (int j = 0; j < n; ++j) { T d = 0; for (int l = 0; l < k; ++l) { T a = A[aᵀ ? lda * i + l : lda * l + i]; T b = B[bᵀ ? ldb * l + j : ldb * j + l]; d += a * b; } if (β) { T c = C[ldc * j + i]; C[ldc * j + i] = α * d + β * c; } else { C[ldc * j + i] = α * d; } } } ``` :::warning 這段程式還不夠了解 $\alpha, \beta$ flag 是作什麼用的? 參閱 [intel GEMM](https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-dpcpp/2023-2/gemm.html), [知乎 GEMM](https://zhuanlan.zhihu.com/p/66958390) GEMM 跟下面 LLM 在最小的迴圈都採取 vector ,需要理解他的原理 怎麼塞暫存器的原理在 [深入淺出 GEMM](https://zhuanlan.zhihu.com/p/435908830) ::: 使用 C++ 達到 47 gigaflops。在這邊提到他是使用 BLAS (基本線性代數子程序)。 這邊我想補充的是它所使用的 `#pragma omp parallel for collapse(2) if (m * n * k > 300000)` 這是一個 OpenMP 函式。OpenMP 是一個支援多平台共享記憶體平行編程的應用程式介面,常用於提高程式的運行效率,特別是在多核心處理器上。 - `#pragma omp parallel for` 指令用於將 for 迴圈的迭代分配給多個執行緒,從而實現迴圈的平行執行。 - `collapse(2)`: 他的使用是將兩層迴圈的迭代會被視為同一個迴圈 以下來自 [OpenMP](https://www.openmp.org/spec-html/5.2/openmpsu30.html) 的文件: >The collapse clause associates one or more loops with the directive on which it appears for the purpose of identifying the portion of the depth of the canonical loop nest to which to apply the semantics of the directive. The argument n specifies the number of loops of the associated loop nest to which to apply those semantics. - `if (m * n * k > 300000)`: 這個條件判斷確保只有在數據量足夠大時才啟用平行化,這是因為對於小的矩陣,平行化帶來的管理和同步開銷可能會抵消由平行處理帶來的性能提升。 這邊作者根據自己的評估提到 intel 架構很難有人可以超越它,但是僅限於單線程的情況之下,他們的多線程還在實驗階段。它提到若是使用一個 `./configure` 開啟多線程可以將矩陣乘法達到 85 gigaflops。 > Based on my own evaluation, what BLIS says is true. However that is only true for single-threaded performance. Their multithreading mode is still experimental, but if I use a ./configure flag to turn it on, then I'm able to boost performance to 85 gigaflops. ```cpp template <typename T> void LLMM(int m, int n, int k, const T *A, int lda, const T *B, int ldb, T *C, int ldc) { #pragma omp parallel for collapse(2) if (m * n * k > 300000) for (int i = 0; i < m; ++i) for (int j = 0; j < n; ++j) { T d = 0; for (int l = 0; l < k; ++l) d += A[lda * i + l] * B[ldb * j + l]; C[ldc * j + i] = d; } } ``` 使用少即是多戰術,省略 $\alpha, \beta$,因為他們從未參與計算。在 LLM 的設計之下,矩陣 A 幾乎總是轉置的,而矩陣 B 幾乎從不轉置,它可以用連續記憶體向量化。 > which means inner dimension dot product can vectorize over contiguous memory. > The m/k dimensions are usually evenly divisible by 64. 最後加速為 233 gigaflops 在 llama.cpp 上最大的弱點是 prompt, image 消化的速度(轉成 token 的速度)所以加入很多 BLAS libraries 進去專案之中。最強的是 MKL 開發的 384 gigaflops。但是 MKL 是 closed source,這邊提到 intel 開發核心數學模組 50 了,沒有對外開放,保留了所有祕密。所以這邊做這提到 llama.cpp 怎麼使用自己的方法提昇 prompt processing speed。 ```cpp void SLLMM(int m, int n, int k, const float *A, int lda, const float *B, int ldb, float *C, int ldc) { #pragma omp parallel for collapse(2) if (m * n * k > 300000) for (int i = 0; i < m; ++i) for (int j = 0; j < n; ++j) { __m256 c = _mm256_setzero_ps(); for (int l = 0; l < k; l += 8) c = _mm256_fmadd_ps(_mm256_loadu_ps(A + lda * i + l), _mm256_loadu_ps(B + ldb * j + l), c); C[ldc * j + i] = hsum(c); } } ``` CPU 數學核心的技巧是利用更少的記憶體引用來利用指令級並行性,如果使用 `-O3 -ffast-math -march=native` 去編譯,會如以下結果: ```cpp void SLLMM(int m, int n, int k, const float *A, int lda, const float *B, int ldb, float *C, int ldc) { #pragma omp parallel for collapse(2) if (m * n * k > 300000) for (int i = 0; i < m; ++i) for (int j = 0; j < n; ++j) { __m256 c = _mm256_setzero_ps(); for (int l = 0; l < k; l += 8) c = _mm256_fmadd_ps(_mm256_loadu_ps(A + lda * i + l), _mm256_loadu_ps(B + ldb * j + l), c); C[ldc * j + i] = hsum(c); } } ``` 在 llama.cpp 中會將最裡面的迴圈展開 ```cpp void SLLMM2(int m, int n, int k, const float *A, int lda, const float *B, int ldb, float *C, int ldc) { #pragma omp parallel for collapse(2) if (m * n * k > 300000) for (int i = 0; i < m; ++i) for (int j = 0; j < n; ++j) { __m256 c0 = _mm256_setzero_ps(); __m256 c1 = _mm256_setzero_ps(); for (int l = 0; l < k; l += 16) { c0 = _mm256_fmadd_ps(_mm256_loadu_ps(A + lda * i + l + 0), _mm256_loadu_ps(B + ldb * j + l + 0), c0); c1 = _mm256_fmadd_ps(_mm256_loadu_ps(A + lda * i + l + 8), _mm256_loadu_ps(B + ldb * j + l + 8), c1); } C[ldc * j + i] = hsum(c0) + hsum(c1); } } ``` 那其實以上的操作不會有太明顯的效能改進,因為現在編譯器跟 CPU 已經可以很好的自己解決這些問題。所以可以從上一層 for 迴圈進行改動。 在以下的程式: ```cpp void SLLMM4(int m, int n, int k, const float *A, int lda, const float *B, int ldb, float *C, int ldc) { #pragma omp parallel for collapse(2) if (m * n * k > 300000) for (int i = 0; i < m; ++i) for (int j = 0; j < n; j += 4) { __m256 c0 = _mm256_setzero_ps(); __m256 c1 = _mm256_setzero_ps(); __m256 c2 = _mm256_setzero_ps(); __m256 c3 = _mm256_setzero_ps(); for (int l = 0; l < k; l += 8) { __m256 a0 = _mm256_loadu_ps(A + lda * (i + 0) + l); __m256 k0 = _mm256_loadu_ps(B + ldb * (j + 0) + l); __m256 k1 = _mm256_loadu_ps(B + ldb * (j + 1) + l); __m256 k2 = _mm256_loadu_ps(B + ldb * (j + 2) + l); __m256 k3 = _mm256_loadu_ps(B + ldb * (j + 3) + l); c0 = _mm256_fmadd_ps(a0, k0, c0); c1 = _mm256_fmadd_ps(a0, k1, c1); c2 = _mm256_fmadd_ps(a0, k2, c2); c3 = _mm256_fmadd_ps(a0, k3, c3); } C[ldc * (j + 0) + (i + 0)] = hsum(c0); C[ldc * (j + 1) + (i + 0)] = hsum(c1); C[ldc * (j + 2) + (i + 0)] = hsum(c2); C[ldc * (j + 3) + (i + 0)] = hsum(c3); } } ``` 可以到暫存器中的 a0 會在其他矩陣計算直接用到。 我們可以 unroll 兩層,如以下,同樣會有效果提昇: ```cpp void SLLMM3X4(int m, int n, int k, const float *A, int lda, const float *B, int ldb, float *C, int ldc) { #pragma omp parallel for collapse(2) if (m * n * k > 300000) for (int i = 0; i < m; i += 3) for (int j = 0; j < n; j += 4) { __m256 c00 = _mm256_setzero_ps(); __m256 c01 = _mm256_setzero_ps(); __m256 c02 = _mm256_setzero_ps(); __m256 c03 = _mm256_setzero_ps(); __m256 c10 = _mm256_setzero_ps(); __m256 c11 = _mm256_setzero_ps(); __m256 c12 = _mm256_setzero_ps(); __m256 c13 = _mm256_setzero_ps(); __m256 c20 = _mm256_setzero_ps(); __m256 c21 = _mm256_setzero_ps(); __m256 c22 = _mm256_setzero_ps(); __m256 c23 = _mm256_setzero_ps(); for (int l = 0; l < k; l += 8) { __m256 k0 = _mm256_loadu_ps(B + ldb * (j + 0) + l); __m256 k1 = _mm256_loadu_ps(B + ldb * (j + 1) + l); __m256 k2 = _mm256_loadu_ps(B + ldb * (j + 2) + l); __m256 k3 = _mm256_loadu_ps(B + ldb * (j + 3) + l); __m256 a0 = _mm256_loadu_ps(A + lda * (i + 0) + l); c00 = _mm256_fmadd_ps(a0, k0, c00); c01 = _mm256_fmadd_ps(a0, k1, c01); c02 = _mm256_fmadd_ps(a0, k2, c02); c03 = _mm256_fmadd_ps(a0, k3, c03); __m256 a1 = _mm256_loadu_ps(A + lda * (i + 1) + l); c10 = _mm256_fmadd_ps(a1, k0, c10); c11 = _mm256_fmadd_ps(a1, k1, c11); c12 = _mm256_fmadd_ps(a1, k2, c12); c13 = _mm256_fmadd_ps(a1, k3, c13); __m256 a2 = _mm256_loadu_ps(A + lda * (i + 2) + l); c20 = _mm256_fmadd_ps(a2, k0, c20); c21 = _mm256_fmadd_ps(a2, k1, c21); c22 = _mm256_fmadd_ps(a2, k2, c22); c23 = _mm256_fmadd_ps(a2, k3, c23); } C[ldc * (j + 0) + (i + 0)] = hsum(c00); C[ldc * (j + 1) + (i + 0)] = hsum(c01); C[ldc * (j + 2) + (i + 0)] = hsum(c02); C[ldc * (j + 3) + (i + 0)] = hsum(c03); C[ldc * (j + 0) + (i + 1)] = hsum(c10); C[ldc * (j + 1) + (i + 1)] = hsum(c11); C[ldc * (j + 2) + (i + 1)] = hsum(c12); C[ldc * (j + 3) + (i + 1)] = hsum(c13); C[ldc * (j + 0) + (i + 2)] = hsum(c20); C[ldc * (j + 1) + (i + 2)] = hsum(c21); C[ldc * (j + 2) + (i + 2)] = hsum(c22); C[ldc * (j + 3) + (i + 2)] = hsum(c23); } } ``` 經過以上展開的操作在做這的設備上 Vectorized outer product with OpenMP 為 810 gigaflops。計算為 512 * 512 的矩陣。通樣的使用 MLK 計算同樣大小的矩陣為 295 gigaflops。 因為要與 llama.cpp 專案協調,所以不能使用 OpenMp 與 BLAS 函式庫。其每一個核心會生成一個執行緒,執行緒會備 spinlock barrier 所限制,在下次要執行操作的時候會釋放已平行計算。 令執行緒 id 為 `ith`,執行緒數量為 `nth`。沒有使用 futexes 或是 semaphore,因為核心模組的調度會大幅減少 tokens / sec。如果我們讓 ith=0 的線程呼叫一個自己生成線程的 BLAS API,那麼這些線程會立即被所有返回到 spinlock barrier 的 ith>0 線程搶占資源。所以它定義一個新的核心模組框架: ```c #define BEGIN_KERNEL(RM, RN) \ int ytiles = (m - m0) / RM; \ int xtiles = (n - n0) / RN; \ int tiles = ytiles * xtiles; \ int duty = (tiles + nth - 1) / nth; \ int start = duty * ith; \ int end = start + duty; \ if (end > tiles) \ end = tiles; \ for (int job = start; job < end; ++job) { \ int i = m0 + job / xtiles * RM; \ int j = n0 + job % xtiles * RN; #define END_KERNEL() } ```