SIMD優化之ARM純彙編開發

ARM純彙編開發

注：這篇文章是兩年前寫的，現在更新到CSDN。當時認知不足，其中可能有不少錯誤，敬請行家指正。

爲什麼要用純彙編

開發效率高

這裏可能讓很多人大跌眼鏡了，純彙編開發效率高？
首先，這個是有限定條件的，需要反覆調優的重度運算場景（比如卷積），純彙編開發效率最高。
其次，這裏的純彙編並不是整個代碼用匯編寫，是指的將足夠重的函數提取出來，用純彙編實現。

參數試驗

爲什麼呢，在用C開發時，受到toolchtain制約，我們會花費很多的時間在試驗編譯器上（往往是試驗pragma、-xxxx等等）。但android/ios、arm32/arm64 的編譯器並不一樣，往往在ios調得好好的，一放到android arm32架構上就慢得一踏糊塗，反覆試驗了幾十次，結果到了最後，可能還是要用內聯彙編。

而如果一開始就下定決心用純彙編開發，優化策略相應的簡單很多，基本上循環展開、指令重排之後，就有立竿見影的效果，試驗時間大幅降低，因此總體開發效率反而更快。

方案試驗

在優化過程中，我們會不停地去試各種方案，如果停留在C層面，很容易得到跟理論不一致的結果，比如：

量化權重和特徵之後運算起來跟沒量化的速度差不多。
用了 winograd ，還沒有 im2col + gemm 快。

如果代碼是比較好的純彙編實現，性能表現基本就跟理論一樣，該快的一定會快，不如預期的在寫彙編就會發現設計時沒考慮到的缺陷，就基本沒有試驗這一說了。

代碼調試

一般來說，用匯編處理的是邏輯簡單，運算複雜的場景，該調試的在前一步C/C++過流程時就已經調試好。
運算密集型的場景，C/C++ 的調試也是很無力的。squeezenet 第一層卷積，$227*227*3$ 的輸入， $113*113*64$ 的輸出，你想一個個斷點去檢查正誤根本不可能。只能把結果打出來比對。這種情況下，C/C++和彙編的調試難度差不了多少。

性能穩定

純彙編實現不需要擔心工具鏈對性能影響，無論是哪個工具鏈編譯參數怎麼變，對性能影響都有限。

這裏再提一下內聯彙編，內聯彙編雖然可以繞開有一點麻煩的函數傳參，但在用的寄存器很多時還是會有問題（看編譯器），不如純彙編中按標準靠譜，因此，內聯彙編我們所見到的，用的寄存器都不是很多，無法充分發揮cpu算力。

一個卷積運算，一樣的滑動窗口算法，一樣地使用neon，純彙編重寫後，arm32 速度提升了 100% 以上，arm64架構提升30%-50%。足見純彙編重寫是一種十分有效的優化方法。

彙編開發

開發難點

學習成本

前面已經分析過，用匯編開發其實效率是高的，但之所以人們覺得開發彙編慢，主要是入門很困難，資料很少。arm 指令集沒有個一兩年時間，很難說做到熟練。

忍受重複

習慣看和寫這種代碼，其實也需要一點時間訓練，聰明人可能不屑於寫，但在你真正能自己弄出編譯器之前，還是得忍一下。

vmax.f32 q0, q0, q15
vmax.f32 q1, q1, q15
vmax.f32 q2, q2, q15
vmax.f32 q3, q3, q15
vmax.f32 q4, q4, q15
vmax.f32 q5, q5, q15
vmax.f32 q6, q6, q15
vmax.f32 q7, q7, q15

基本流程

1、梳理代碼，設計實現方案，提煉出核心的運算部分，先用C實現一遍，保證正確
2、32位彙編代碼初步實現，保證功能正確
3、彙編代碼優化：這一步優化只做循環展開和指令重排，如果有更好的計算方案，先退回 C 重新寫
4、64位的也支持一下：替換一下寄存器名和指令名（可以寫腳本完成），然後微調一下函數前後讀參數、入出棧與返回的地方
（可選）64位的進一步優化一下，畢竟寄存器多了一倍

Procedure Call Standard【函數調用標準】

函數調用標準是寫純彙編時一定要掌握的，不然會出現很多莫名奇妙的錯誤
詳細的文檔參見arm官網
http://infocenter.arm.com/help/index.jsp?topic=/com.arm.doc.ihi0042f/index.html

我這裏簡單總結了一下

ARM 32(v7a)

用完恢復是指相應的寄存器在函數返回前必須恢復進入時的值，比如我們要在代碼中用 q4，就必須在函數前寫一句

vpush {q4}

函數返回前寫一句：

vpop {q4}

這裏面寫的不能用的寄存器，如果仔細看了 call standard，會發現其實還是能使用的。但建議沒完全搞懂前，最好別用，一般也不需要用那麼多。

ARM 64(v8)

值得注意的是，arm64 的傳參爲浮點時，會傳到 v0.s[0], v0.s[1] …… 而非通用寄存器，這個很坑，建議不要用浮點傳參

彙編優化實例

C的Relu 代碼

void ReluForward(float* dst, const float* src, size_t sizeDiv4)
{
      for (int i=0; i<4*sizeDiv4; ++i)
      {
           dst[i] = src[i] >0 ? src[i] : 0;
      }
}

有些同學想必會看循環內的 4sizeDiv4 不順眼，心想：這不應該在前面寫個 int size = 4sizeDiv4，免得每次循環時都計算麼? 這裏是特地這麼寫的：
1、4*sizeDiv4 這個表達式由於 4 及 sizeDiv4 在循環時未發生改變，-O2之後編譯器是不會重複生成計算該表達式的語句的，請停止你的優化強迫症。
2、C++的代碼主要就是讓你明白這函數幹嘛用的，別關注性能，寫彙編時再摳。

c-neon

void ReluCNeon(float* dst, const float* src, size_t sizeDiv4)
{
    float32x4_t limit = vdupq_n_f32(0.0f);
    for (int i=0; i<sizeDiv4; ++i)
    {
        float32x4_t value = vld1q_f32(src);
        value = vmaxq_f32(value, limit);
        vst1q_f32(dst, value);

        dst+=4;
        src+=4;
    }
}

基礎彙編

由於ios和android上面函數編譯的符號不一致，這裏引入一個頭文件，定義一個函數聲明宏，去屏蔽這種差異：
ArmAsmGlobal.h

.macro asm_function fname
#ifdef __APPLE__
.globl _\fname
_\fname:
#else
.global \fname
\fname:
#endif

//彙編：ReluBasic
#include "ArmAsmGlobal.h"
asm_function ReluBasic

//按照 arm32 的 函數調用標準，以下變量由調用方傳至寄存器
//r0: dst, r1: src, r2: sizeDiv4

push {lr}
vmov.i32 q15, #0

cmp r2, #0
beq End //跳轉：beq 表示 r2 等於0時跳轉
Loop://標誌，供跳轉用
vld1.32 {q0}, [r1]!
vmax.f32 q0, q0, q15
vst1.32 {q0}, [r0]!
subs r2, r2, #1// 這一句 相當於 sub r2, r2, #1  &&  cmp r2, #0
bne Loop //跳轉：bne 表示 r2 不等於0時跳轉

End:
pop {pc}

彙編優化

我們注意到循環主體，語句前後有較強依賴關係：

vld1.32 {q0}, [r1]!
vmax.f32 q0, q0, q15 //q0 依賴於 前一行的讀
vst1.32 {q0}, [r0]! //q0 依賴於前一行的算

ARM 的CPU一般都有雙通道發射能力（跟多核多線程不是同一個概念），在執行如下類型的語句時，可以併發執行，提升效率：

vld1.32 {q0}, [r1]!
vmax.f32 q1, q1, q15 //不使用 q0，無依賴關係

爲了讓我們的彙編代碼解除語句前後的依賴關係，先進行一次循環展開：

//彙編：ReluUnroll
#include "ArmAsmGlobal.h"
asm_function ReluUnroll

vmov.i32 q15, #0
push {lr}

L4:
cmp r2, #3
ble L1

L4Loop:

vld1.32 {q0, q1}, [r1]!

vld1.32 {q2, q3}, [r1]!
vmax.f32 q0, q0, q15
vmax.f32 q1, q1, q15
vmax.f32 q2, q2, q15
vmax.f32 q3, q3, q15

vst1.32 {q0, q1}, [r0]!
vst1.32 {q2, q3}, [r0]!

sub r2, r2, #4
cmp r2, #4
bge L4Loop


L1:
cmp r2, #0
beq End

L1Loop:
vld1.32 {q0}, [r1]!
vmax.f32 q0, q0, q15
vst1.32 {q0}, [r0]!
subs r2, r2, #1
bne L1Loop


End:
pop {pc}

展開之後，L4Loop 內部的語句已經大部分解除了依賴，但還不完全，爲了完全解除，我們需要用個小技巧【彙編重點技巧】：

//彙編：ReluUnrollReorder
#include "ArmAsmGlobal.h"
asm_function ReluUnrollReorder

push {lr}
vmov.i32 q15, #0

L4:
cmp r2, #3
ble L1

vld1.32 {q0, q1}, [r1]!
vmax.f32 q0, q0, q15
vld1.32 {q2, q3}, [r1]!
vmax.f32 q1, q1, q15

sub r2, r2, #4
cmp r2, #3
ble L4End

L4Loop:

vst1.32 {q0, q1}, [r0]!
vmax.f32 q2, q2, q15
vld1.32 {q0, q1}, [r1]!
vmax.f32 q3, q3, q15
vst1.32 {q2, q3}, [r0]!
vmax.f32 q0, q0, q15
vld1.32 {q2, q3}, [r1]!
sub r2, r2, #4
vmax.f32 q1, q1, q15
cmp r2, #4
bge L4Loop

L4End:
vst1.32 {q0, q1}, [r0]!
vmax.f32 q2, q2, q15
vmax.f32 q3, q3, q15
vst1.32 {q2, q3}, [r0]!


L1:
cmp r2, #0
beq End

L1Loop:
vld1.32 {q0}, [r1]!
vmax.f32 q0, q0, q15
vst1.32 {q0}, [r0]!
subs r2, r2, #1
bne L1Loop


End:
pop {pc}

這個技巧就是將循環主體代碼拆成兩半，原先的 Loop[AB] 就變成了 A->Loop[BA]->B，然後 BA 由於順序顛倒，可以實現錯排併發。

性能對比

魅藍 mental 上測試
sizeDiv4 = 100000，連續跑10000次（由於 relu 是一個十分簡單的op，跑大批量的才能看到效果）
C-neon Cost time : 4856.960449 ms
彙編ReluBasic Cost time : 4716.672363 ms
彙編ReluUnroll Cost time : 2814.848145 ms
彙編ReluUnrollReorder Cost time : 2359.424072 ms

可以看到：
1、最簡單的彙編和用 neon api 的 C差不大多
2、同樣是彙編，ReluUnrollReorder較ReluBasic足足提升了100%