在 CFU Playground 上加速 MLPerf™ Tiny 影像分類模型 #1:簡介

Yeecy
10 min readJan 8, 2024

--

1.1 CFU Playground

CFU Playground 提供一套完整的框架讓使用者可以專注於研究如何加速 tflite 模型,而不用煩惱相關的軟體環境配置,如果沒有 FPGA 開發板的話,CFU Playground 也整合了 Validator 以供使用者模擬執行。CFU Playground 中內置一顆 VexRiscv CPU,用以執行編譯完的 C/C++ 程式碼(也就是我們的 tflite 模型),透過 RISC-V 指令集預留的指令編碼空間,CPU 的控制邏輯可將我們自行定義的指令轉發到自己設計的電路進行相關操作,此處的電路即為 CFU(custom function unit)。

在 CFU Playground 的論文 CFU Playground: Full-Stack Open-Source Framework for Tiny Machine Learning (TinyML) Acceleration on FPGAs 中可以看到 CFU 和 ALU 一樣都擁有兩條讀入和一條寫回暫存器堆(register file)的接線,換言之,我們可以將 CFU 看作一個與 ALU 平行的計算單元。

舉例來說,假設 abc 都是 32 位元整數(往後以 int32 表示之)且 ab 儲存於記憶體中而 c 不須寫回記憶體,則 c = a + b 這個表達式在 RISC-V 裡需要兩條 lw(load word)和一條 add 完成,兩條 lw 分別將 ab 從記憶體讀入兩個暫存器中,一條 add 得到這兩個暫存器之和並寫回 c 對應的暫存器。

那要怎麼用 CFU 做到同樣的事呢?我們可以先把 a + b 改成自行定義的指令 cfu_op(func7, func3, a, b),其中 func7func3 為我們定義的值,接著在 cfu.v 這個 Verilog 檔案中實作電路,電路的功能為檢查 func7func3 的值是否為我們定義的值,如果是,將傳入的兩個暫存器的值相加並輸出。cfu_op(func7, func3, a, b)a + b 一樣會被翻譯為三條指令,前面兩條 lw 指令不變,而唯一的差別在於後者對應到的是 add 指令,前者則是是我們自己定義的指令。

CFU 和 CPU 之間的溝通介面如下,讀者可以在 CFU Playground 的教學網頁 https://cfu-playground.readthedocs.io/en/latest/interface.html 上看到更多使用案例。

    >--- cmd_valid ----------------------->
<--- cmd_ready -----------------------<
>--- cmd_payload_function_id[9:0] ---->
>--- cmd_payload_inputs_0[31:0] ------>
>--- cmd_payload_inputs_1[31:0] ------>
CPU CFU
<--- rsp_valid -----------------------<
>--- rsp_ready ----------------------->
<--- rsp_payload_outputs_0[31:0] -----<

這裡筆者稍微提一下最簡單的溝通情況,以前面的例子來說,如果 CPU 在指令解碼階段發現用到我們定義的指令,到了指令執行階段,CPU 的控制邏輯會將 cmd_valid 設為 1,cmd_payload_function_id[9:3] 設為 func7cmd_payload_function_id[2:0] 設為 func3cmd_payload_inputs_0[31:0] 設為 acmd_payload_inputs_1[31:0] 設為 b,因為單純的加法電路很簡單,能在一個周期內算出來,所以我們可以在 cmd_valid 為 1 時,把 a + b 之值寫入 rsp_payload_outputs_0[31:0] 並立即將 rsp_valid 設為 1,讓控制邏輯知道 CFU 已經處理完指令,並且計算結果儲存於 rsp_payload_outputs_0[31:0] 中。

為了讓 CFU 支援異步設計,CFU Playground 的作者們額外加上了 cmd_readyrsp_ready 訊號用來溝通。當 rsp_valid 被設為 1 時,CFU 應同時將 cmd_ready 設為 0,表示 CFU 並未準備好接收下一個新指令,因為此時計算結果尚未被 CPU 控制邏輯寫回暫存器,接著在 CPU 控制邏輯將 CFU 的輸出寫回後,CPU 控制邏輯需要將 rsp_ready 設為 1 以通知 CFU 結果已被寫回,讓 CFU 在下一個週期把 rsp_vaildcmd_ready 分別更改為 0 和 1,告知 CPU 控制邏輯 CFU 已經準備好處理下一條指令。

1.2 MLPerf™ Tiny 影像分類模型

MLPerf™ 為當前業界公認的 AI 測試基準之一,提供數個不同的子任務供各家廠商跑分比較,而 MLPerf™ Tiny 顧名思義是針對嵌入式系統所設計的模型推理測試基準,在本文中我們的加速目標為其中的 int8 量化版影像分類模型。

MLPerf™ Tiny 影像分類模型的架構為修改過的 ResNet,模型推理使用框架為 Tensorflow Lite for Microcontrollers(TFLM),訓練的資料集為 CIFAR10,模型架構如下。

MLPerf™ Tiny 影像分類模型架構

該圖由 Netron 產生,讀者可以將自己的模型上傳至 Netron.app 可視化,方便理解模型架構,另外注意到圖中卷積層的資料格式為 NHWC,跟 PyTorch 預設的 NCHW 不同。

因為模型已經被 int8 量化過(模型的輸入和權重皆為 int8),而且筆者也並不打算在硬體層面加速稀疏向量和矩陣運算,所以沒對模型進行如剪枝或重新量化之類的修改,換句話說,能夠更動的只有模型用到的運算,並且我們對程式碼所做的更動不應該影響到模型的輸出結果。

在圖中可以看到該模型用到最多的是卷積層,而且實際跑過模型之後也會發現卷積層確實花了最多時間,所以主要的加速重點在於卷積層,也就是說接下來大多數的時間都會討論如何加速卷積計算!

1.3 向量乘積累加單元

CPU 中有一類指令被稱為 SIMD(single instruction multiple data)指令,這類指令的特點為可以用一條指令完成多條同構(isomorphic)指令所做的運算,我們來看下面的程式碼。

void vadd4(int* a, int* b, int* __restrict c) {
(void)__builtin_assume_aligned(a, 128);
(void)__builtin_assume_aligned(b, 128);
(void)__builtin_assume_aligned(c, 128);

c[0] = a[0] + b[0];
c[1] = a[1] + b[1];
c[2] = a[2] + b[2];
c[3] = a[3] + b[3];
}

/* x86-64 clang 17.0.1 -O3, generated with godbolt.org
vadd(int*, int*, int*):
movdqa xmm0, xmmword ptr [rsi]
paddd xmm0, xmmword ptr [rdi]
movdqa xmmword ptr [rdx], xmm0
ret
*/

在編譯器不進行向量化(vectorization)的情況下,我們可以預期 c[i] = a[i] + b[i] 會被翻譯為三條指令,第一條指令將 a[i] 讀入暫存器,第二條指令將該暫存器的值加上 b[i](x86 允許另一個運算元儲存在記憶體裡),第三條指令將該暫存器的值寫回 c[i],也就是說共需要 12 條指令。

然而從上面註解裡我們發現編譯器進行向量化後只需要三條指令,而這正是因為編譯器使用了 x86–64 裡的 SSE2 指令集提供的 SIMD 指令,第一條 movdqa 將從 a 指向的記憶體位址裡連續儲存的四個 int32 讀入 xmm0 暫存器,第二條 paddd 指令將 xmm0 裡的四個 int32 跟從 b 指向的記憶體位址裡連續儲存的四個 int32 各別相加,第三條 movdqa 將 xmm0 裡儲存的四個 int32 連續寫回 c 指向的記憶體位址,完全跟前面的純量指令等價。

那為什麼筆者要用 x86–64 而不是 RISC-V 舉例呢?原因很單純,因為標準的 RISC-V 裡只有純量指令,沒有 SIMD 指令(除非 CPU 實作了 P 指令集擴充),但有了 CFU,我們能夠自己實作 SIMD 指令。

現在我們來看哪裡會需要 SIMD 指令。深度學習模型裡最常見的運算模式為乘積累加,用全連接層(fully connected layer)舉例,為求簡潔,設該層的輸入維度為 m,輸出維度為 1,令 xᵢ 表示第 i 個輸入,wᵢ 表示第 i 個權重,b 表示偏差值,則唯一的輸出值 y 可以表示為 y = x₁×w₁ + x₂×w₂ + … + xₘ×wₘ + b,這個式子的計算可以被改寫如下。

y ← b
y ← x₁×w₁ + y
y ← x₂×w₂ + y
...
y ← xₘ×wₘ + y*

m 次的乘積累加運算,不難想像當全連接層的輸出維度為 n 時,需要 mn 次乘積累加。

在 CFU Playground 的教學 https://cfu-playground.readthedocs.io/en/latest/step-by-step.html 裡便有簡單的乘積累加硬體設計,其設計精神簡單講是因為模型經過 int8 量化,所以模型的權重和輸入都是 int8,又因為一個 int32 可以存放四個 int8,所以我們可以將四個 int8 輸入和四個 int8 權重各別組合成兩個 int32 作為 CFU 的兩個運算元,接著在 CFU 裡,同時將兩個 int32 中相同位置的 int8 相乘,得到四個 int16 的乘積(int8 乘上 int8 的結果為 int16),再將四個乘積加起來輸出,用一條指令完成四次乘積累加計算,另外從數學定義來說,這個操作就是向量內積,所以也能稱為向量內積指令。

需要注意的是在嚴格的 SIMD 定義中,前面的向量乘積累加運算並非為一種 SIMD 指令,因為這樣的操作只輸出一個元素,而非四個 int8,不過因為這不是什麼專業的學術文章,筆者懶得也不需要分那麼清楚,所以後續提到 SIMD 可能指代的是廣義的向量指令。

最後腦筋急轉彎一下,如果我們沒有 CFU,有沒有辦法模擬出針對 int8 的 SIMD 指令呢?實際上可以,假設要做的是 int8 加法,那麼可以用 int32 的從左至右數來第二個和第四個位元組(如果不需偵測算數溢位的話使用第一和第三個位元組也行)存放 int8 的值,接著將兩個 int32 相加,再取出第二和第四個位元組即為兩次 int8 加法所得之值。舉個例子,假如我們需要做 0x54 + 0x87 和 0x10 + 0x24,那麼可以構造兩個 int32 為 0x00540087 和 0x00100024,把兩個值加起來會得到 0x006400ab,那麼 0x64 和 0xab 分別是 0x54 + 0x10 和 0x10 + 0x24 的結果,也就是用一次 int32 加法完成兩次 int8 加法,當然這裡沒有考慮構造和取出 int8 的代價就是了。

--

--

Yeecy
Yeecy

Written by Yeecy

A Ph.D. student at NYCU CS and a compiler engineer at ICEshell Co., Ltd. You can find more information about me on my GitHub page github.com/ADNRs.

No responses yet