# Linux 核心專題 : 開發用以加速 LLaMA 的 Linux 核心模組
> 執行人: csotaku0926

這份專題是 [第九週測驗二](https://hackmd.io/@sysprog/linux2024-quiz9#%E6%B8%AC%E9%A9%97-2) 的延伸，撰寫加速矩陣運算的核心模組

矩陣運算在大型語言模型佔用大部分的運算時間。
考慮核心模組擁有較高的執行優先權、較貼近底層等優勢，可預期利用其克服矩陣運算、資料存取延遲的效能瓶頸。
我將探討可用於核心模組的加速矩陣運算手段，例如批次計算矩陣、降低處理器切換的成本 ; 另外，針對 SIMD 指令集不被支援於大多數處理器架構的核心模組，我也提出 CPU isolation 的解法

從五月到六月這段期間，嘗試做出有貢獻的東西

參考資料:
- [LLaMA Now Goes Faster on CPUs](https://justine.lol/matmul/)
- [2019q1 Homework2 (kcalc)](https://hackmd.io/@foxhoundsk/BySm2rHuV#%E8%87%AA%E6%88%91%E6%AA%A2%E6%9F%A5%E6%B8%85%E5%96%AE)

相關連結:
- [github repos](https://github.com/csotaku0926/kmatmul)
- [COSCUP 2024 議程投稿](https://pretalx.coscup.org/coscup-2024/talk/review/WMFXLJPNLGNLSP7CCYAKCFSJNWPMNKXW)

## 任務簡述
1. 歸納矩陣乘法加速手段，指出效能瓶頸
2. 提出有效的批次存取模型

## 回答疑問
> 實際打 code 前，先搞清楚做的東西產值在哪裡

### 歸納矩陣乘法加速的手段
> [LLaMA Now Goes Faster on CPUs](https://justine.lol/matmul/)

作者提到他撰寫了一系列的矩陣乘法核心程式碼，提示詞評估 (prompt evaluation) 速度快了 30% ~ 500%
其在 CPU 使用了 F16, Q8_0 的型態於權重

- [ ] 已知瓶頸？
> llama 效能瓶頸，什麼東西阻擋運算...

很大的瓶頸是提示詞處理速度 (prompt processing speed)

- 請問 prompt processing speed 牽扯到哪些操作？

這是因為其中牽扯到矩陣-矩陣乘法，這也是 BLAS library 會使用於 llama 的原因

::: spoiler 所以到底 BLAS 可以幹嘛？優勢在那？
BLAS (Basic Linear Algebra Subprograms) 是個應用程式標準 (API)
其提供向量-向量，向量-矩陣，以及矩陣-矩陣的運算函式庫
被使用於高精度的線性代數軟體，例如 LAPACK
然而，llama.cpp 的矩陣-向量運算，相對其他運算而言是計算上更容易的
延遲 (latency) 才是其主要瓶頸
使用龐大函式庫反而導致資料存取延遲加大
:::

<br></br>

另外，資料存取的延遲 (latency) 是進行矩陣-向量運算上的瓶頸

- [ ] 加速的策略？ 作者的貢獻？

透過 Perf Profiler 分析，作者提出在 `llamafile` 中有許多操作 (轉置， reshape 等)，有 95% 的執行時間都花費在矩陣乘法上 (`GGML_OP_MUL_MAT`)

`#pragma omp parallel for` 是 OpenMP 用以實現 for 迴圈平行化的指令，會將迴圈拆解給不同的執行緒執行
`collapse(n)` 子句可以指定迴圈平行化的層數，例如 `n=2` 則會將雙層迴圈拆解為單一個

template <typename T>
void LLMM(int m, int n, int k,
          const T *A, const T *B,
          int lda, int ldb)
{
#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++) 
             for (int l=0; l<k; k++) 
                 d += A[lda * i + l] * B[ldb * j + l]
        
// ...

與以往認知的矩陣乘法不同，B 矩陣的維度被置換了
這意味著，其內積 (inner dot product) 可以向量化 (vectorize) 的方式儲存於連續記憶體，如此可以使得快取以及硬體資源運用 (如 SIMD) 更好的運用
這使得 token 生成速度得到改善

作者推論 BLAS 優化技巧為指令層級平行 (instruction-level parrellisim, ILP) 並使用更少的記憶體參照 (memory reference)
ILP 的理念為使用一條指令就對多組 (獨立的) 資料進行操作

例如，若將以上案例進行編譯優化，`-O3 -ffast-math -march=native` :

// ...
for (int j = 0; j < n; ++j) {
    c = _mm256_setzero_ps();
    for (int l = 0; l < k; l += 8) {
        __mm256 a0 = _mm256_loadu_ps(A + lda * i + l);
        __mm256 k0 = _mm256_loadu_ps(B + ldb * j + l);
        c = _mm256_fmadd_ps(a0, k0, c);
    }
    C[ldc * j + i] = hsum(c);    
}

根據 [Intel 官方文件](https://www.intel.com/content/www/us/en/docs/cpp-compiler/developer-guide-reference/2021-10/mm256-loadu-ps.html):
- `_mm256_setzero_ps` : 初始化一個數值皆為 0 的 float32 vector
- `_mm256_loadu_ps` : 載入 float32 向量至 256-bit AVX 暫存器
- `_mm256_fmadd_ps` : 相乘前兩個參數 vector 數值，與第三個參數相加

這裡使用了 Intel 的 AVX 指令集，容許一次性的對 256-bit 資料做運算

另一個技巧是 loop unrolling，使運算在編譯時期完成，所以有機會做更多最佳化
如果把外層迴圈拆解，可以與多個浮點運算共享 `a0` 暫存器

for (int i = 0; i < m; i++)
    for (int j = 0; j < n; j += 4) {
        __mm256 c0 = _mm256_setzero_ps();
        // ...
        __mm256 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 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 + 3) + (i + 0)] = hsum(c3);
    }

可以將內外迴圈都拆開
不同於以往 dot-based 矩陣乘法，此為 block-based 乘法，更有效運用 CPU cache

作者在 [sgemm.cpp](https://github.com/ggerganov/llama.cpp/pull/6414/files#diff-b89d92ffb96d84a060ec2a3ce66cd939edf2c8af1d3bc9edc133040bc5947eb1) 收錄他的貢獻程式

- [ ] 多執行緒

作者針對多執行緒的場景，設計了由遞迴 `mnpack` 拆解原矩陣為多個子矩陣乘法
根據剩餘矩陣大小，分為 3x4、1x4、4x1、1x1 四種大小的乘法

- [ ] LKM 的優勢？
> 核心模組比起一般行程，有什麼特別的優勢？
> [Writing a Linux Kernel Module — Part 1: Introduction](https://derekmolloy.ie/writing-a-linux-kernel-module-part-1-introduction/)

電腦實際工作的單位是硬體，而核心用於控制硬體完成各項工作。
換言之，若想要的功能沒有核心的支援，即使硬體如何強大也無用武之地

在不需更動整個核心的前提下，可以將獨立的驅動程式編譯成模組，
於執行時間時載入此模組，便可以使用對應的硬體

這可以使動態設置更容易，更節省核心記憶體

與 user-space program 相比，核心模組擁有較高的執行優先權 (execution privilege)
也就是說，更多的 CPU cycle 會被分配給核心模組 ; 然而，這點可能影響到其他行程的運行

另外，核心模組可以被中斷 (interrupt)，可以同時被多個行程使用

### 批次處理
> 只要能夠批次計算矩陣，並降低處理器之間切換的成本，就能估算出優勢 [name=jserv]

至於 CMWQ，則是用來提供一個靈活的並行處理方案，用來解決原先 worker thread 與 CPU 數量綁定、且 work item 不能轉移的窘境。這使得原先的 work queue 設計中即使消耗大量資源，並行效率也未能得到很大改善的原因

### 資料型態
> [matmul](https://github.com/jart/matmul) # 注意看裡面的 data type

- [ ] 資料型態的使用？和效能有何關係？

參考 `matmul` ，其針對不同資料型態定義了操作

以加減乘法而言，可以看到對於不同指令集，都有對應操作:
- `float`
- `__m128`, `__m256`, `__m512`,
- `float32x4` : 將四個 float32 包裝成為單一 128-bit SIMD 暫存器
- 對應指令集為 `__ARM_NEON`
- vaddq_f32

> [enabling disabling simd operations for kernel](https://forums.freebsd.org/threads/enabling-disabled-simd-operations-for-kernel.65480/)

很遺憾的，若要在核心使用 SIMD 暫存器，有效能的問題
核心常常進行 context switch，每次進行都需要保存 FPU/SIMD state (這些暫存器很大)，這需要許多時間成本
而 [x86-64 處理器](https://lore.kernel.org/lkml/[e-mail]/) 已經不支援 SIMD 的使用了，只有 [amd 架構](https://www.kernel.org/doc/Documentation/arm/kernel_mode_neon.rst) 透過 API 呼叫仍允許部份使用
> 你可以指派一個處理器核並抑制 preempt [name=jserv]
> [StackOverflow](https://stackoverflow.com/questions/49414559/linux-kernel-why-preemption-is-disabled-when-use-per-cpu-variable)

### 測量效能的方式
- `perf` : 分析 user program ，或核心程式碼中各函式呼叫佔多少時間
- [教學](https://wiki.csie.ncku.edu.tw/embedded/perf-tutorial)
- eBPF : python BCC module 或 透過 C code 編譯
- [教學](https://github.com/eunomia-bpf/bpf-developer-tutorial/tree/main)
- 雖然想過以此法測量執行緒成本，但 eBPF 似乎只適用於較大型的函式與系統呼叫
- 以 `kprobe` 測量 `worker_thread` 會導致錯誤 (cannot attach kprobe, Cannot assign requested address)
- `ftrace` : 核心內部呼叫函式追蹤
- [man](https://docs.kernel.org/trace/ftrace.html)
- 可用於測量 `matrix_ioctl` 中各函式執行時間

### 其他

雖然核心模組有較高的執行權限等優勢，但應避免過多使用者/核心層級切換，或是過多系統呼叫 (system call)

注意到 `llama.cpp` 定義新的資料型態，用來有效率的載入權重
[PR #27](https://github.com/ggerganov/ggml/pull/27) 使用提出的 SIMD 4-bit 整數量化 (integer quantization)

一般來說，整數量化是一種將 32-bit 浮點數轉化為鄰近的 8-bit 定點數的優化手法
這麼做可以縮小模型並提升推理速度 (inference speed)

但在大型語言模型中，定義稍有不同。以 `q4_0` 而言，原理為 : 使用一個 f32 scaling factor，以及 `QK/2` 個位元組 (每個 bytes 儲存兩個 4-bit 整數) 表示 `QK` 個浮點數
在 `q4_0` 中，scaling factor 設為 `abs(max(x_i))/7` (其實就是把原先浮點數 scale 到 0~7， 浮點數 block 中最大的數 scale 成為 7 )

## 解釋程式碼運作原理
### 矩陣乘法優化 (第九週測驗2)
測驗題中的手法即是透過分配一部分矩陣給不同的 `worker_thread()` 做乘法運算，達成並行的效果
例如將原先 100x100 的矩陣分為十個 10x10 子矩陣的運算，再把所有結果相加
`work_thread()` 中的 `struct completion` 則是核心提供的 [barrier API](https://www.kernel.org/doc/Documentation/scheduler/completion.txt) ，目的是為了達成同步，避免 race condition

由於計算發生在 kernel space，因此計算之前要先將 user space 的矩陣資料透過 `copy_from_user` (`include/linux/uaccess.h`)

- `copy_to_user` : 複製核心中 stack 的內容到 userspace

使用者資料複製完畢後，為每個子矩陣創建一個 `worker_thread`
透過 `kmalloc` 配置核心空間的記憶體，會回傳指向這個記憶體的指標，代表每個`worker_thread` 使用資料的位址
於是，每個執行緒要執行的函式，使用資料的位址，會由 `kthread_run` 處理

int *thread_arg = kmalloc(sizeof(int), GFP_KERNEL);
*thread_arg = i;
kthread_run(worker_thread, thread_arg, "worker_thread");

可以看到，每個執行緒負責 `SUBMAT_SIZE` 個行的運算
不過因為只是做方陣的計算，所以每個 thread 只須傳入起始行的 index 即可

隨後，於 `matrix_fops` 設定了 `/proc/matmul` 這個設定檔的可用操作
包含讀取結果矩陣以及 ioctl 操作
而 `matrix_ioctl` 操作可以讓核心模組讀入使用者資料，以及計算乘法