原文：blog.fanscore.cn/a/62/
1. 背景
前一段時間公司上線了一套Go實現的推薦系統，上線後發現MMR層雖然只有純計算但耗時十分離譜，透過pprof定位問題所在之後進行了最佳化，雖然降低了非常多但是我們認為其中還有最佳化空間。
可以看到日常平均耗時126ms，P95 360ms。
MMR層主要耗時集中在了餘弦相似度的計算部分，這部分我們使用的gonum庫進行計算，其底層在x86平臺上利用了SSE指令集進行了加速。
SSE指令集已經非常古老了，xmm暫存器只能儲存兩個雙精度浮點數，每次只能並行進行兩個雙精度浮點數的計算，而AVX2指令集可以平行計算四個，理論上可以獲得兩倍的效能提升，因此我們決定自己使用AVX2指令集手寫彙編的方式替代掉gonum庫。
1.1 餘弦相似度演算法
餘弦相似度的計算公式為
對應的程式碼為
import "gonum.org/v1/gonum/floats"
func CosineSimilarity(a, b []float64) float64 {
dotProduct := floats.Dot(a, b) // 計算a和b的點積
normA := floats.Norm(a, 2) // 計算向量a的L2範數
normB := floats.Norm(b, 2) // 計算向量b的L2範數
return dotProduct / (normA * normB)
}
2. Dot點積計算加速
gonum點積計算Dot的部分彙編程式碼如下：
TEXT ·DotUnitary(SB), NOSPLIT, $0
...
loop_uni:
// sum += x[i] * y[i] unrolled 4x.
MOVUPD 0(R8)(SI*8), X0
MOVUPD 0(R9)(SI*8), X1
MOVUPD 16(R8)(SI*8), X2
MOVUPD 16(R9)(SI*8), X3
MULPD X1, X0
MULPD X3, X2
ADDPD X0, X7
ADDPD X2, X8
ADDQ $4, SI // i += 4
SUBQ $4, DI // n -= 4
JGE loop_uni // if n >= 0 goto loop_uni
...
end_uni:
ADDPD X8, X7
MOVSD X7, X0
UNPCKHPD X7, X7
ADDSD X0, X7
MOVSD X7, sum+48(FP) // Return final sum.
RET
可以看到其中使用xmm暫存器平行計算兩個雙精度浮點數，並且還採用了迴圈展開的最佳化手段，一個迴圈中同時進行4個元素的計算。
我們利用AVX2指令集平行計算四個雙精度浮點數進行加速
loop_uni:
// sum += x[i] * y[i] unrolled 8x.
VMOVUPD 0(R8)(SI*8), Y0 // Y0 = x[i:i+4]
VMOVUPD 0(R9)(SI*8), Y1 // Y1 = y[i:i+4]
VMOVUPD 32(R8)(SI*8), Y2 // Y2 = x[i+4:i+8]
VMOVUPD 32(R9)(SI*8), Y3 // Y3 = x[i+4:i+8]
VMOVUPD 64(R8)(SI*8), Y4 // Y4 = x[i+8:i+12]
VMOVUPD 64(R9)(SI*8), Y5 // Y5 = y[i+8:i+12]
VMOVUPD 96(R8)(SI*8), Y6 // Y6 = x[i+12:i+16]
VMOVUPD 96(R9)(SI*8), Y7 // Y7 = x[i+12:i+16]
VFMADD231PD Y0, Y1, Y8 // Y8 = Y0 * Y1 + Y8
VFMADD231PD Y2, Y3, Y9
VFMADD231PD Y4, Y5, Y10
VFMADD231PD Y6, Y7, Y11
ADDQ $16, SI // i += 16
CMPQ DI, SI
JG loop_uni // if len(x) > i goto loop_uni
可以看到我們每個迴圈中同時用到8個ymm暫存器即一次迴圈計算16個數，而且還用到了VFMADD231PD指令同時進行乘法累積的計算。
最終Benchmark結果：
BenchmarkDot 一個迴圈中計算8個數
BenchmarkDot-2 14994770 78.85 ns/op
BenchmarkDot16 一個迴圈中計算16個數
BenchmarkDot16-2 22867993 53.46 ns/op
BenchmarkGonumDot Gonum點積計算
BenchmarkGonumDot-2 8264486 144.4 ns/op
可以看到點積部分我們得到了大約2.7倍的效能提升
3. L2範數計算加速
gonum庫中進行L2範數計算的演算法並不是常規的a1^2 + a2^2 ... + aN^2這種計算，而是採用了Netlib演算法，減少了溢位和下溢，其Go原始碼如下：
func L2NormUnitary(x []float64) (norm float64) {
var scale float64
sumSquares := 1.0
for _, v := range x {
if v == 0 {
continue
}
absxi := math.Abs(v)
if math.IsNaN(absxi) {
return math.NaN()
}
if scale < absxi {
s := scale / absxi
sumSquares = 1 + sumSquares*s*s
scale = absxi
} else {
s := absxi / scale
sumSquares += s * s
}
}
if math.IsInf(scale, 1) {
return math.Inf(1)
}
return scale * math.Sqrt(sumSquares)
}
其彙編程式碼比較晦澀難懂，但管中窺豹再結合Go原始碼可以看出來沒有用到並行能力，每次迴圈只計算一個數
TEXT ·L2NormUnitary(SB), NOSPLIT, $0
...
loop:
MOVSD (X_)(IDX*8), ABSX // absxi = x[i]
...
我們最佳化之後的核心程式碼如下：
loop:
VMOVUPD 0(R8)(SI*8), Y0 // Y0 = x[i:i+4]
VMOVUPD 32(R8)(SI*8), Y1 // Y1 = y[i+4:i+8]
VMOVUPD 64(R8)(SI*8), Y2 // Y2 = x[i+8:i+12]
VMOVUPD 96(R8)(SI*8), Y3 // Y3 = x[i+12:i+16]
VMOVUPD 128(R8)(SI*8), Y4 // Y4 = x[i+16:i+20]
VMOVUPD 160(R8)(SI*8), Y5 // Y5 = y[i+20:i+24]
VMOVUPD 192(R8)(SI*8), Y6 // Y6 = x[i+24:i+28]
VMOVUPD 224(R8)(SI*8), Y7 // Y7 = x[i+28:i+32]
VFMADD231PD Y0, Y0, Y8 // Y8 = Y0 * Y0 + Y8
VFMADD231PD Y1, Y1, Y9
VFMADD231PD Y2, Y2, Y10
VFMADD231PD Y3, Y3, Y11
VFMADD231PD Y4, Y4, Y12
VFMADD231PD Y5, Y5, Y13
VFMADD231PD Y6, Y6, Y14
VFMADD231PD Y7, Y7, Y15
ADDQ $32, SI // i += 32
CMPQ DI, SI
JG loop // if len(x) > i goto loop
我們採用原始的演算法計算以利用到平行計算的能力，並且迴圈展開，一次迴圈中同時計算32個數，最終Benchmark結果：
BenchmarkAVX2L2Norm
BenchmarkAVX2L2Norm-2 29381442 40.99 ns/op
BenchmarkGonumL2Norm
BenchmarkGonumL2Norm-2 1822386 659.4 ns/op
可以看到得到了大約16倍的效能提升
4. 總結
透過這次最佳化我們在餘弦相似度計算部分最終得到了(144.4 + 659.4 * 2) / (53.46 + 40.99 * 2) = 10.8倍的效能提升，效果還是非常顯著的。相較於《記一次SIMD指令最佳化計算的失敗經歷》這次失敗的初次嘗試，本次還是非常成功的，切實感受到了SIMD的威力。
另外在本次最佳化過程中也漲了不少姿勢
AVX-512指令降頻問題
AVX-512指令因為並行度更高理論上效能也更高，但AVX-512指令會造成CPU降頻，因此業界使用非常慎重，這一點可以參考位元組的json解析庫sonic的這個issue: https://github.com/bytedance/sonic/issues/319
迴圈展開最佳化
在一次迴圈中做更多的工作，優點有很多：
- 減少迴圈控制的開銷，迴圈變數的更新和條件判斷次數更少，降低了分支預測失敗的可能性
- 增加指令並行性，更多的指令可以在流水線中並行執行
但一次迴圈使用過多的暫存器從實際Benchmark看效能確實更好，但是否存在隱患我沒有看到相關的資料，希望這方面的專家可以指教一下。