Go語言中的SIMD加速：以矩陣加法為例

前些日子，一些資深Gopher，比如fasthttp^[1]的作者Aliaksandr Valialkin^[2]因函數(shù)迭代器^[3]加入Go 1.23版本^[4]而抱怨Go的演進走錯了方向：朝著增加復(fù)雜性和隱式代碼執(zhí)行的方向發(fā)展，而沒有專注于Go語言的基本設(shè)計哲學(xué)——簡單性、生產(chǎn)力和性能。Valialkin希望Go團隊能專注于一些性能打磨和優(yōu)化的環(huán)節(jié)，比如使用SIMD提升一些計算場景下Go代碼的性能，避免Go的某些領(lǐng)地被以性能和安全性著稱的Rust^[5]搶去！

無獨有偶，在Go項目issues中，我們也能看到很多有關(guān)希望Go支持SIMD指令的issue，比如近期的一個proposal^[6]，就期望Go團隊可以在標準庫中添加simd包以支持高性能的SIMD計算，就像Rust std::simd那樣。當然，早期這類issue也有很多，比如：issue 53171^[7]、issue 58610^[8]等。

那么什么是SIMD指令？在Go官方尚未支持simd包或SIMD計算的情況下，如何在Go中使用SIMD指令進行計算加速呢？在這篇文章中，我們就來做個入門版介紹，并以一個最簡單的矩陣加法的示例來展示一下SIMD指令的加速效果。

1. SIMD指令簡介

SIMD是“單指令多數(shù)據(jù)”(Single Instruction Multiple Data)的縮寫。與之對應(yīng)的則是SISD（Single Instruction, Single Data），即“單指令單數(shù)據(jù)”。

在大學(xué)學(xué)習(xí)匯編時，用于舉例的匯編指令通常是SISD指令，比如常見的ADD、MOV、LEA、XCHG等。這些指令每執(zhí)行一次，僅處理一個數(shù)據(jù)項。早期的x86架構(gòu)下，SISD指令處理的數(shù)據(jù)僅限于8字節(jié)（64位）或更小的數(shù)據(jù)。隨著處理器架構(gòu)的發(fā)展，特別是x86-64架構(gòu)的引入，SISD指令也能處理更大的數(shù)據(jù)項，使用更大的寄存器。但SISD指令每次仍然只處理一個數(shù)據(jù)項，即使這個數(shù)據(jù)項可能比較大。

相反，SIMD指令是一種特殊的指令集，它可以讓處理器可以同時處理多個數(shù)據(jù)項，提高計算效率。我們可以用下面這個更為形象生動的比喻來體會SIMD和SISD的差別。

想象你是一個廚師，需要切100個蘋果。普通的方式是一次切一個蘋果，這就像普通的SISD處理器指令。而SIMD指令就像是你突然多了幾雙手，可以同時切4個或8個蘋果。顯然，多手同時工作會大大提高切蘋果的速度。

具體來說，SIMD指令的優(yōu)勢在于以下幾點：

• 并行處理：一條指令可以同時對多個數(shù)據(jù)進行相同的操作。
• 數(shù)據(jù)打包：將多個較小的數(shù)據(jù)(如32位浮點數(shù))打包到一個較大的寄存器(如256位)中。
• 提高數(shù)據(jù)吞吐量：每個時鐘周期可以處理更多的數(shù)據(jù)。

這種并行處理方式特別適合于需要大量重復(fù)計算的任務(wù)，如圖像處理、音頻處理、科學(xué)計算等。通過使用SIMD指令，可以顯著提高這些應(yīng)用的性能。

主流的x86-64(amd64)和arm系列CPU都有對SIMD指令的支持。以x86-64為例，該CPU體系下支持的SIMD指令就包括MMX(MultiMedia eXtensions)、SSE (Streaming SIMD Extensions)、SSE2、SSE3、SSSE3、SSE4、AVX(Advanced Vector Extensions)、AVX2以及AVX-512等。ARM架構(gòu)下也有對應(yīng)的SIMD指令集，包括VFP (Vector Floating Point)、NEON (Advanced SIMD)、SVE (Scalable Vector Extension)、SVE2以及Helium (M-Profile Vector Extension, MVE)等。

注：在Linux上，你可以通過lscpu或cat /proc/cpuinfo來查看當前主機cpu支持的SIMD指令集的種類。注：Go在Go 1.11版本才開始支持AVX-512指令^[9]。

每類SIMD指令集都有其特定的優(yōu)勢和應(yīng)用場景，以x86-64下的SIMD指令集為例：

• MMX主要用于早期的多媒體處理；
• SSE系列逐步改進了浮點運算和整數(shù)運算能力，廣泛應(yīng)用于圖形處理和音視頻編碼；
• AVX系列大幅提高了并行處理能力，特別適合科學(xué)計算和高性能計算場景。

這些指令集的演進反映了處理器技術(shù)的發(fā)展和應(yīng)用需求的變化。從支持64位計算的MMX到支持512位計算的AVX-512，SIMD指令的并行處理能力不斷提升，更多更大的寄存器加入進來，為各種復(fù)雜的計算任務(wù)提供了強大的硬件支持。

注：SSE和AVX各自有16個寄存器，SSE的16個寄存器為XMM0-XMM15，XMM是128位寄存器，而YMM是256位寄存器。支持AVX的x86-64處理器包含16個256位大小的寄存器，從YMM0到Y(jié)MM15。每個YMM寄存器的低128位是相對應(yīng)的XMM寄存器。大多數(shù)AVX指令可以使用任何一個XMM或者YMM寄存器作為SIMD操作數(shù)。AVX512將每個AVXSIMD寄存器的大小從256位擴展到512位，稱為ZMM寄存器；符合AVX512標準的處理器包含32個ZMM寄存器，從ZMM0~ZMM31。YMM和XMM寄存器分別對應(yīng)于每個ZMM寄存器的低256位和低128位。

既然SIMD指令這么好，那么在Go中應(yīng)該如何使用SIMD指令呢？接下來我們就來看看。

2. 在Go中如何使用SIMD指令

Go主要面向的是云計算領(lǐng)域、微服務(wù)領(lǐng)域，這些領(lǐng)域中對計算性能的要求相對沒那么極致。以至于在一些對性能要求較高的場景，比如高性能計算、 圖形學(xué)、數(shù)字信號處理等領(lǐng)域，很多gopher會遇到對Go計算性能進行優(yōu)化的需求。

純計算領(lǐng)域，怎么優(yōu)化呢？此時此刻，Go官方并沒有提供對SIMD提供支持的simd包。

一種想法是使用cgo機制在Go中調(diào)用更快的C或C++，但cgo的負擔又不能不考慮，cgo不是go^[10]，很多人不愿意引入cgo。

另外一種想法就是再向下一層，直接上匯編，在匯編中直接利用SIMD指令實現(xiàn)并行計算。但手寫匯編難度是很高的，手寫Plan9風格、資料甚少的Go匯編難度則更高。那么有什么方法避免直接手搓匯編呢？目前看大致有這么幾種(如果有更好的方法，歡迎在評論區(qū)提出你的建議)：

• 使用c2goasm^[11](https://github.com/minio/c2goasm/)轉(zhuǎn)換

我們可以先用c/c++實現(xiàn)對應(yīng)的函數(shù)功能(可以利用類似intel提供的面向simd的intrisic functions^[12])，然后生成匯編代碼(基于clang)，再用c2goasm轉(zhuǎn)換為go語言匯編。不過目前c2goasm已經(jīng)public archive了，并且該方法應(yīng)用受很多因素限制，比如clang版本和特定的編譯選項啥的。親測這種方法上手難度較高。

• 使用uber工程師Michael McLoughlin開源的avo^[13]來生成go匯編

avo(https://github.com/mmcloughlin/avo)是一個go包，它支持以一種相對高級一些的Go語法來編寫匯編，至少你可以不必直面那些晦澀難懂的匯編代碼。但使用avo編寫匯編也不是很容易的事情，你仍然需要大致知道匯編的運作原理和基本的編寫規(guī)則。此外avo與匯編的能力并非完全等價，其作者聲明：avo也還處于實驗階段。

• 使用goplus/llgo集成c/c++生態(tài)

在go中調(diào)用c的cgo機制不受待見，llgo^[14]反其道而行之，將go、python、c/c++等代碼統(tǒng)統(tǒng)轉(zhuǎn)換為llvm中間代碼進而通過clang編譯和優(yōu)化為可執(zhí)行文件。這樣就可以直接利用python、c/c++的生態(tài)，進而利用高性能的c/c++實現(xiàn)（比如支持SIMD指令）。目前l(fā)lgo還不成熟，七牛云老板許式偉正在全力開發(fā)llgo，等llgo成熟后，這后續(xù)可能也是一種選擇。

考慮到Go目前不直接支持intel intrisic functions for SIMD^[15]，要在Go中使用SIMD只能直接使用匯編。而在手搓匯編難度太高的情況下，通過avo生成匯編便是一條可以嘗試的路徑，我們可以將一些計算的核心部分用avo生成的匯編來進行加速。

接下來，我們就來通過一個矩陣加法的示例看看SIMD指令的加速效果。基于SIMD指令的矩陣加法的匯編邏輯，我們采用avo實現(xiàn)。

3. 第一版SIMD優(yōu)化(基于SSE)

我們使用avo先來實現(xiàn)一版基于SSE指令集的矩陣加法。前面說過avo是一個Go庫，我們無需安裝任何二進制程序，直接使用avo庫中的類型和函數(shù)編寫矩陣加法的實現(xiàn)即可：

// simd-in-go/matadd-sse/pkg/asm.go

//go:build ignore
// +build ignore

package main

import (
 "github.com/mmcloughlin/avo/attr"
 . "github.com/mmcloughlin/avo/build"
 . "github.com/mmcloughlin/avo/operand"
)

func main() {
 TEXT("MatrixAddSIMD", attr.NOSPLIT, "func(a, b, c []float32)")
 a := Mem{Base: Load(Param("a").Base(), GP64())}
 b := Mem{Base: Load(Param("b").Base(), GP64())}
 c := Mem{Base: Load(Param("c").Base(), GP64())}
 n := Load(Param("a").Len(), GP64())

 X0 := XMM()
 X1 := XMM()

 Label("loop")
 CMPQ(n, U32(4))
 JL(LabelRef("done"))

 MOVUPS(a.Offset(0), X0)
 MOVUPS(b.Offset(0), X1)
 ADDPS(X1, X0)
 MOVUPS(X0, c.Offset(0))

 ADDQ(U32(16), a.Base)
 ADDQ(U32(16), b.Base)
 ADDQ(U32(16), c.Base)
 SUBQ(U32(4), n)
 JMP(LabelRef("loop"))

 Label("done")
 RET()

 Generate()
}

第一次看上面這段代碼，你是不是覺得即便使用avo來生成矩陣加法的代碼，如果你不了解匯編的編寫和運行模式，你也是無從下手的。簡單說一下這段代碼。

首先，該文件是用于生成矩陣加法的匯編代碼的，因此該asm.go并不會編譯到最終的可執(zhí)行文件中或測試代碼中，這里利用go編譯器構(gòu)建約束將該文件排除在外。

main函數(shù)的第一行的TEXT函數(shù)定義了一個名為MatrixAddSIMD的函數(shù)，使用attr.NOSPLIT屬性表示不需要棧分割，函數(shù)簽名是：

func(a, b, c []float32)

變量a, b, c分別表示輸入矩陣a, b和輸出矩陣c的內(nèi)存地址，使用Load函數(shù)從參數(shù)中加載基地址到GP64返回的通用寄存器。n表示矩陣的長度，使用 Load函數(shù)從參數(shù)中加載長度到GP64返回的通用寄存器。

X0和X1定義了兩個XMM寄存器，用于SIMD操作。

接下來定義了一個循環(huán)，在這個循環(huán)的循環(huán)體中，將通過SSE指令處理輸入的矩陣數(shù)據(jù)：

• MOVUPS(a.Offset(0), X0)：將矩陣a的前16字節(jié)（4 個float32）加載到XMM寄存器X0。
• MOVUPS(b.Offset(0), X1)：將矩陣b的前16字節(jié)（4個float32）加載到XMM寄存器X1。
• ADDPS(X1, X0)：將X1和X0中的數(shù)據(jù)相加，結(jié)果存入X0。
• MOVUPS(X0, c.Offset(0))：將結(jié)果從X0存入矩陣c的前16字節(jié)。
• ADDQ(U32(16), a.Base)：將矩陣a的基地址增加16字節(jié)（4個float32）。
• ADDQ(U32(16), b.Base)：將矩陣b的基地址增加16字節(jié)（4個float32）。
• ADDQ(U32(16), c.Base)：將矩陣c的基地址增加16字節(jié)（4個float32）。
• SUBQ(U32(4), n)：將矩陣長度n減少4。
• JMP(LabelRef("loop"))：無條件跳轉(zhuǎn)到標簽loop，繼續(xù)循環(huán)。

最后調(diào)用Generate函數(shù)生成匯編代碼。

下面我們就來運行該代碼，生成相應(yīng)的匯編代碼以及stub函數(shù)：

$cd matadd-sse/pkg
$make
go run asm.go -out add.s -stubs stub.go

下面是生產(chǎn)的add.s的全部匯編代碼：

// simd-in-go/matadd-sse/pkg/add.s

// Code generated by command: go run asm.go -out add.s -stubs stub.go. DO NOT EDIT.

#include "textflag.h"

// func MatrixAddSIMD(a []float32, b []float32, c []float32)
// Requires: SSE
TEXT ·MatrixAddSIMD(SB), NOSPLIT, $0-72
 MOVQ a_base+0(FP), AX
 MOVQ b_base+24(FP), CX
 MOVQ c_base+48(FP), DX
 MOVQ a_len+8(FP), BX

loop:
 CMPQ   BX, $0x00000004
 JL     done
 MOVUPS (AX), X0
 MOVUPS (CX), X1
 ADDPS  X1, X0
 MOVUPS X0, (DX)
 ADDQ   $0x00000010, AX
 ADDQ   $0x00000010, CX
 ADDQ   $0x00000010, DX
 SUBQ   $0x00000004, BX
 JMP    loop

done:
 RET

這里使用的ADDPS、MOVUPS和ADDQ都是SSE指令：

• ADDPS (Add Packed Single-Precision Floating-Point Values)：這是一個SSE指令，用于對兩個128位的XMM寄存器中的4個單精度浮點數(shù)進行并行加法運算。
• MOVUPS (Move Unaligned Packed Single-Precision Floating-Point Values): 這也是一個SSE指令，用于在內(nèi)存和XMM寄存器之間移動128位的單精度浮點數(shù)數(shù)據(jù)。與MOVAPS(Move Aligned Packed Single-Precision Floating-Point Values) 指令不同，MOVUPS不要求地址對齊，可以處理非對齊的數(shù)據(jù)。

除了生成匯編代碼外，asm.go還生成了一個stub函數(shù)：MatrixAddSIMD，即上面匯編實現(xiàn)的那個函數(shù)。

// simd-in-go/matadd-sse/pkg/stub.go

// Code generated by command: go run asm.go -out add.s -stubs stub.go. DO NOT EDIT.

package pkg

func MatrixAddSIMD(a []float32, b []float32, c []float32)

在matadd-sse/pkg/add-no-simd.go中，我們放置了常規(guī)的矩陣加法的實現(xiàn)：

package pkg

func MatrixAddNonSIMD(a, b, c []float32) {
 n := len(a)
 for i := 0; i < n; i++ {
  c[i] = a[i] + b[i]
 }
}

接下來，我們編寫一些單測代碼，確保一下MatrixAddSIMD和MatrixAddNonSIMD的功能是正確的：

// simd-in-go/matadd-sse/matrix_add_test.go 
package main

import (
 "demo/pkg"
 "testing"
)

func TestMatrixAddNonSIMD(t *testing.T) {
 size := 1024
 a := make([]float32, size)
 b := make([]float32, size)
 c := make([]float32, size)
 expected := make([]float32, size)

 for i := 0; i < size; i++ {
  a[i] = float32(i)
  b[i] = float32(i)
  expected[i] = a[i] + b[i]
 }

 pkg.MatrixAddNonSIMD(a, b, c)

 for i := 0; i < size; i++ {
  if c[i] != expected[i] {
   t.Errorf("MatrixAddNonSIMD: expected %f, got %f at index %d", expected[i], c[i], i)
  }
 }
}

func TestMatrixAddSIMD(t *testing.T) {
 size := 1024
 a := make([]float32, size)
 b := make([]float32, size)
 c := make([]float32, size)
 expected := make([]float32, size)

 for i := 0; i < size; i++ {
  a[i] = float32(i)
  b[i] = float32(i)
  expected[i] = a[i] + b[i]
 }

 pkg.MatrixAddSIMD(a, b, c)

 for i := 0; i < size; i++ {
  if c[i] != expected[i] {
   t.Errorf("MatrixAddSIMD: expected %f, got %f at index %d", expected[i], c[i], i)
  }
 }
}

如我們預(yù)期的那樣，上述單測代碼可以順利通過。接下來，我們再來做一下benchmark，看看使用SSE實現(xiàn)的矩陣加法性能到底提升了多少：

// simd-in-go/matadd-sse/benchmark_test.go 
package main

import (
 "demo/pkg"
 "testing"
)

func BenchmarkMatrixAddNonSIMD(tb *testing.B) {
 size := 1024
 a := make([]float32, size)
 b := make([]float32, size)
 c := make([]float32, size)

 for i := 0; i < size; i++ {
  a[i] = float32(i)
  b[i] = float32(i)
 }

 tb.ResetTimer()
 for i := 0; i < tb.N; i++ {
  pkg.MatrixAddNonSIMD(a, b, c)
 }
}

func BenchmarkMatrixAddSIMD(tb *testing.B) {
 size := 1024
 a := make([]float32, size)
 b := make([]float32, size)
 c := make([]float32, size)

 for i := 0; i < size; i++ {
  a[i] = float32(i)
  b[i] = float32(i)
 }

 tb.ResetTimer()
 for i := 0; i < tb.N; i++ {
  pkg.MatrixAddSIMD(a, b, c)
 }
}

運行這個benchmark，我們得到下面結(jié)果：

$go test -bench .
goos: darwin
goarch: amd64
pkg: demo
... ...
BenchmarkMatrixAddNonSIMD-8     2129426        554.4 ns/op
BenchmarkMatrixAddSIMD-8        3481318        357.4 ns/op
PASS
ok   demo 3.350s

我們看到SIMD實現(xiàn)的確性能優(yōu)秀，幾乎在非SIMD實現(xiàn)的基礎(chǔ)上提升了一倍。但這似乎還并不足以說明SIMD的優(yōu)秀。我們再來擴展一下并行處理的數(shù)據(jù)的數(shù)量和寬度，使用AVX指令再來實現(xiàn)一版矩陣加法，看是否還會有進一步的性能提升。

4. 第二版SIMD優(yōu)化(基于AVX)

下面是基于avo使用AVX指令實現(xiàn)的Go代碼：

// simd-in-go/matadd-avx/pkg/asm.go

//go:build ignore
// +build ignore

package main

import (
 "github.com/mmcloughlin/avo/attr"
 . "github.com/mmcloughlin/avo/build"
 . "github.com/mmcloughlin/avo/operand"
)

func main() {
 TEXT("MatrixAddSIMD", attr.NOSPLIT, "func(a, b, c []float32)")
 a := Mem{Base: Load(Param("a").Base(), GP64())}
 b := Mem{Base: Load(Param("b").Base(), GP64())}
 c := Mem{Base: Load(Param("c").Base(), GP64())}
 n := Load(Param("a").Len(), GP64())

 Y0 := YMM()
 Y1 := YMM()

 Label("loop")
 CMPQ(n, U32(8))
 JL(LabelRef("done"))

 VMOVUPS(a.Offset(0), Y0)
 VMOVUPS(b.Offset(0), Y1)
 VADDPS(Y1, Y0, Y0)
 VMOVUPS(Y0, c.Offset(0))

 ADDQ(U32(32), a.Base)
 ADDQ(U32(32), b.Base)
 ADDQ(U32(32), c.Base)
 SUBQ(U32(8), n)
 JMP(LabelRef("loop"))

 Label("done")
 RET()

 Generate()
}

這里的代碼與上面sse實現(xiàn)的代碼邏輯類似，只是指令換成了avx的指令，包括VMOVUPS、VADDPS等：

• VADDPS (Vectorized Add Packed Single-Precision Floating-Point Values): 是AVX (Advanced Vector Extensions) 指令集中的一個指令，用于對兩個256位的YMM寄存器中的8個單精度浮點數(shù)進行并行加法運算。
• VMOVUPS (Vectorized Move Unaligned Packed Single-Precision Floating-Point Values): 這也是一個AVX指令，用于在內(nèi)存和YMM寄存器之間移動256位的單精度浮點數(shù)數(shù)據(jù)。與MOVUPS指令相比，VMOVUPS可以處理更寬的256位SIMD數(shù)據(jù)。

由于在SSE實現(xiàn)的版本中做了詳細說明，這里就不再贅述代碼邏輯，其他單元測試與benchmark測試的代碼也都完全相同，我們直接看benchmark的結(jié)果：

$go test -bench .
goos: darwin
goarch: amd64
pkg: demo
... ...
BenchmarkMatrixAddNonSIMD-8     2115284        566.6 ns/op
BenchmarkMatrixAddSIMD-8       10703102        111.5 ns/op
PASS
ok   demo 3.088s

我們看到AVX版的矩陣加法的性能是常規(guī)實現(xiàn)的5倍多，是SSE實現(xiàn)的性能的近3倍，在實際生產(chǎn)中，這將大大提升代碼的執(zhí)行效率。

也許還有更優(yōu)化的實現(xiàn)，但我們已經(jīng)達到了基于SIMD加速矩陣加法的目的，這里就不再做繼續(xù)優(yōu)化了，大家如果有什么新的想法和驗證的結(jié)果，可以在評論區(qū)留言告訴我哦！

5. 小結(jié)

在這篇文章中，我們探討了在Go語言中使用SIMD指令進行計算加速的方法。盡管Go官方目前還沒有直接支持SIMD的包，但我們通過使用avo庫生成匯編代碼的方式，成功實現(xiàn)了基于SSE和AVX指令集的矩陣加法優(yōu)化。

我們首先介紹了SIMD指令的基本概念和優(yōu)勢，然后討論了在Go中使用SIMD指令的幾種可能方法。接著，我們通過一個具體的矩陣加法示例，展示了如何使用avo庫生成基于SSE和AVX指令集的匯編代碼。

通過benchmark測試，我們看到基于SSE指令的實現(xiàn)相比常規(guī)實現(xiàn)提升了約1.5倍的性能，而基于AVX指令的實現(xiàn)則帶來了約5倍的性能提升。這充分說明了SIMD指令在并行計算密集型任務(wù)中的強大優(yōu)勢。

雖然直接使用SIMD指令需要一定的匯編知識，增加了代碼的復(fù)雜性，但在一些對性能要求極高的場景下，這種優(yōu)化方法仍然是非常有價值的。我希望這篇文章能為Go開發(fā)者在進行性能優(yōu)化時提供一些新的思路和參考。

當然，這里展示的只是SIMD優(yōu)化的一個簡單示例。在實際應(yīng)用中，可能還需要考慮更多因素，如數(shù)據(jù)對齊、邊界條件處理等。大家可以在此基礎(chǔ)上進行更深入的探索和實踐。

本文涉及的源碼可以在這里^[16]下載 - https://github.com/bigwhite/experiments/blob/master/simd-in-go

本文部分源代碼由deepseek coder v2^[17]實現(xiàn)。

6. 參考資料

• Intel Intrinsics Guide^[18] - https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html
• Go Wiki: AVX512^[19] - https://go.dev/wiki/AVX512
• A Manual for the Plan 9 assembler^[20] - http://doc.cat-v.org/plan_9/4th_edition/papers/asm
• From slow to SIMD: A Go optimization story^[21] - https://sourcegraph.com/blog/slow-to-simd
• Efficient and performance-portable vector software^[22] - https://github.com/google/highway
• 并行處理-SIMD^[23] - https://www.slidestalk.com/u231/simd_computer
• 玩轉(zhuǎn)SIMD指令編程^[24] - https://zhuanlan.zhihu.com/p/591900754

參考資料

想要了解Go更多內(nèi)容，歡迎掃描下方??關(guān)注公眾號，掃描 [實戰(zhàn)群]二維碼  ,即可進群和我們交流~