Go SIMDパッケージの使い方—Go 1.27で3.9倍、効かないサイズもある

Go SIMDパッケージの使い方—Go 1.27で3.9倍、効かないサイズもある | mohablog

Go 1.27 の実験的な simd パッケージを go1.27rc2 で測りました。float32 65536要素の総和が 35,653 ns/op から 9,018 ns/opif の入った集計では 178,926 ns/op から 9,961 ns/op。ただし要素数が 8 のときはスカラーのループの方が速い。

目次

GOEXPERIMENT=simd を付けないとビルドが通らない

動く最小コードはこれです。simd は標準ライブラリなので外部依存はありません。

package main

import (
	"fmt"
	"runtime"
	"simd"
)

func main() {
	fmt.Println("GOARCH:", runtime.GOARCH)

	var f32 simd.Float32s
	var f64 simd.Float64s
	fmt.Println("Float32s.Len():", f32.Len())
	fmt.Println("Float64s.Len():", f64.Len())

	v := simd.LoadFloat32s([]float32{1, 2, 3, 4})
	w := simd.BroadcastFloat32s(10)

	out := make([]float32, v.Len())
	v.Add(w).Store(out)
	fmt.Println("Add:", out)
}

go run . をそのまま叩くと弾かれます。

$ go run .
package simddemo
	imports simd: build constraints exclude all Go files in /Users/moha/sdk/go1.27rc2/src/simd
package simddemo
	imports simd/internal/bridge: build constraints exclude all Go files in /Users/moha/sdk/go1.27rc2/src/simd/internal/bridge

有効化はビルド時の環境変数

Go 1.27 リリースノートの “New experimental simd package” にこう書かれています。

This package is enabled by setting the environment variable GOEXPERIMENT=simd at build time.

付けて実行すると通ります。

$ GOEXPERIMENT=simd go run .
GOARCH: arm64
Float32s.Len(): 4
Float64s.Len(): 2
Add: [11 12 13 14]

GOEXPERIMENT 自体は Go 1.26 の Green Tea GC でも使われた仕組みで、Green Tea GCとは—Go 1.26で変わるGCの仕組みと効果を実測GOEXPERIMENT=greenteagc を扱っています。ただし simdGo 1 互換性保証の対象外。API はマイナーリリースで変わり得ます。

go.mod の go ディレクティブは 1.26 のままでも通る

go 1.26 と書いた go.mod で simd を import してビルドしたところ、警告なしで通りました。go test も同じ。Go 1.27 から go test が既定で走らせる stdversion vet チェックは、go ディレクティブより新しい標準ライブラリのシンボルを報告しますが、simd はここに引っかかりません。

ツールチェーンだけ 1.27 にすれば、モジュールの宣言バージョンは据え置きで試せます。

Load・Add・Store の3手順で書く

float32 スライスの総和をベクトル化します。パッケージドキュメントの “Obtaining SIMD values” セクションにある通り、値の入り口は Load<Types>Broadcast<Types> の2つ。

func SumSIMD(xs []float32) float32 {
	var acc simd.Float32s // ゼロ値がそのままゼロベクトル
	n := acc.Len()

	i := 0
	for ; i+n <= len(xs); i += n {
		acc = acc.Add(simd.LoadFloat32s(xs[i:]))
	}
	if i < len(xs) {
		v, _ := simd.LoadFloat32sPart(xs[i:])
		acc = acc.Add(v)
	}

	lanes := make([]float32, n)
	acc.Store(lanes)

	var s float32
	for _, x := range lanes {
		s += x
	}
	return s
}

10要素で流すと、4要素ずつ2回まわって端数が2要素。

入力: [1 2 3 4 5 6 7 8 9 10]
  端数 2 要素を Part でロード
合計: 55

端数は LoadFloat32sPart で埋める

LoadFloat32s はベクトル長に満たないスライスを渡すとパニックします。足りないときは LoadFloat32sPart。読めた数を第2戻り値で返し、残りはゼロで埋めます。

v, n := simd.LoadFloat32sPart([]float32{1, 2, 3})
LoadFloat32sPart([1 2 3]) -> loaded=3 vector=[1 2 3 0]

総和のようにゼロが単位元になる演算なら、埋まったゼロをそのまま足して構いません。積や最小値では単位元が変わるので、Broadcast で 1 や +Inf を敷いてから Masked で差し替える形になります。

Len() を読んでからループ幅を決める

ループ幅に 4 や 8 を直接書かない。Len() を変数に取ってから回します。レーン数が実行環境で変わるためです。

65536要素の総和で3.95倍

ベンチマークコード

var sink float32

func BenchmarkSumScalar(b *testing.B) {
	xs := data(65536)
	for b.Loop() {
		sink = SumScalar(xs)
	}
}

func BenchmarkSumSIMD(b *testing.B) {
	xs := data(65536)
	for b.Loop() {
		sink = SumSIMD(xs)
	}
}

b.Loop は Go 1.24 で入ったベンチマークループ。b.N の手書きループと違い、中身がコンパイラの最適化で消えません。詳しくはGoのb.Loopで正確なベンチマークを書く—b.Nとの違いと落とし穴に書いています。

計測結果

$ GOEXPERIMENT=simd go test -bench=. -benchmem -count=5 .
goos: darwin
goarch: arm64
pkg: simdbench
cpu: Apple M5 Pro
BenchmarkSumScalar-18    	   33669	     35653 ns/op	       0 B/op	       0 allocs/op
BenchmarkSumScalar-18    	   33648	     35659 ns/op	       0 B/op	       0 allocs/op
BenchmarkSumSIMD-18      	  132811	      9031 ns/op	       0 B/op	       0 allocs/op
BenchmarkSumSIMD-18      	  133098	      9018 ns/op	       0 B/op	       0 allocs/op

35,653 ns から 9,018 ns で 3.95倍。arm64 の Neon は128bit、float32 なら4レーンなので、理論上限の4倍にほぼ張り付いています。lanes := make([]float32, n) を毎回書いていますが 0 allocs/op。エスケープせずスタックに載っています。

アキュムレータを4本に分けると1.75倍

このループは acc = acc.Add(...) が毎回1つ前の結果を待ちます。加算のレイテンシぶんパイプラインが空くので、アキュムレータを分けて依存を切ります。

var a0, a1, a2, a3 simd.Float32s
n := a0.Len()

i := 0
for ; i+4*n <= len(xs); i += 4 * n {
	a0 = a0.Add(simd.LoadFloat32s(xs[i:]))
	a1 = a1.Add(simd.LoadFloat32s(xs[i+n:]))
	a2 = a2.Add(simd.LoadFloat32s(xs[i+2*n:]))
	a3 = a3.Add(simd.LoadFloat32s(xs[i+3*n:]))
}
acc := a0.Add(a1).Add(a2.Add(a3))
BenchmarkSumSIMD-18     	  133411	      8998 ns/op
BenchmarkSumSIMD4-18    	  232448	      5158 ns/op

9,020 ns から 5,158 ns。SIMD化しただけの版から 1.75倍、スカラーからは 6.91倍。float32 の加算に結合則は成り立たないので、分割の仕方で末尾の丸めが変わります。総和の完全一致を要求するテストがあるなら、期待値もこの実装で取り直す必要があります。

要素数が少ないと逆に遅くなる

同じ実装をサイズ別に測ると、8要素でスカラーに負けます。

要素数スカラーSIMD倍率
82.626 ns3.048 ns0.86倍
164.465 ns3.800 ns1.18倍
328.214 ns4.906 ns1.67倍
6415.24 ns7.596 ns2.01倍
25697.15 ns25.71 ns3.78倍
1024491.4 ns123.3 ns3.99倍
6553635,659 ns9,020 ns3.95倍

境界は16要素。ロードと最後の水平加算が固定費として乗るぶん、それより下ではレーン並列の取り分が足りません。

分岐のある集計では18倍になる

4レーンで4倍。ここまでの数字はその範囲です。if の入った集計では、倍率が理論値を超えます。

if を Masked に置き換える

正の値だけ足す処理を、スカラーとSIMDの両方で書きます。スカラー側は素直に if

func SumPositiveScalar(xs []float32) float32 {
	var s float32
	for _, x := range xs {
		if x > 0 {
			s += x
		}
	}
	return s
}

SIMD側に if は書けません。レーンごとに真偽が違うためです。代わりに比較でマスクを作り、Masked で不要なレーンをゼロに潰してから足します。

func SumPositiveSIMD(xs []float32) float32 {
	var acc, zero simd.Float32s
	n := acc.Len()

	i := 0
	for ; i+n <= len(xs); i += n {
		v := simd.LoadFloat32s(xs[i:])
		acc = acc.Add(v.Masked(v.Greater(zero)))
	}
	if i < len(xs) {
		v, _ := simd.LoadFloat32sPart(xs[i:])
		acc = acc.Add(v.Masked(v.Greater(zero)))
	}

	lanes := make([]float32, n)
	acc.Store(lanes)
	var s float32
	for _, x := range lanes {
		s += x
	}
	return s
}

マスクの中身は String() で覗けます。

入力: [1 -2 3 -4 5 -6 7 -8]
正の値のみ合計: 16
mask: {1,0,1,0}
Masked: [1 0 3 0]

65536要素、符号をランダムに振ったデータで測ります。

BenchmarkSumPositiveScalar-18    	    6285	    179314 ns/op
BenchmarkSumPositiveScalar-18    	    6249	    178926 ns/op
BenchmarkSumPositiveSIMD-18      	  120778	      9961 ns/op
BenchmarkSumPositiveSIMD-18      	  119757	     10061 ns/op

178,926 ns から 9,961 ns。18.0倍。4レーンしかないのに18倍出ています。

分岐予測が当たると差は2.6倍に縮む

超過分はレーン並列ではなく、スカラー側が分岐予測を外している損です。裏を取るため、同じデータを符号順にソートしてから測り直しました。前半が全部負、後半が全部正になるので、分岐はほぼ予測どおりに進みます。

func sortedData(n int) []float32 {
	xs := signedData(n)
	sort.Slice(xs, func(i, j int) bool { return xs[i] < xs[j] })
	return xs
}
BenchmarkSortedScalar-18    	   46632	     25738 ns/op
BenchmarkSortedSIMD-18      	  119748	     10013 ns/op
データの並びスカラーSIMD倍率
符号がランダム178,926 ns9,961 ns18.0倍
符号でソート済み25,738 ns10,010 ns2.6倍

スカラー版は並べ替えただけで 178,926 ns から 25,738 ns、7.0倍速くなりました。同じ命令数、同じメモリアクセス量です。対してSIMD版は 9,961 ns と 10,010 ns で、差は 0.5%。マスクに分岐は無いので、データの並びに影響されません。

SIMD化の効果をレーン数から見積もると、この差を取りこぼします。予測しづらい条件分岐が内側にあるループほど、置き換えの取り分は大きい。分岐が一定方向に流れているループなら、期待値はレーン数どまりです。

ベクトル長は書いた時点では決まらない

パッケージドキュメントの “SIMD Types” セクションにこうあります。

In all cases, the vector length is at least 128 bits, and within a given program execution, all vectors have the same length.

「その実行の中では」同じ長さ。コンパイル時ではありません。長さを決めているのは標準ライブラリ内の archMaxVectorSize() で、アーキごとにファイルが分かれています。

arm64 は128bit固定、SVE は未対応

// src/simd/midway_arm64.go
func archMaxVectorSize() (size, allFeatureSize int) {
	// This describes Neon, SVE is still TBD.
	size = 128
	if cpu.ARM64.HasPMULL {
		allFeatureSize = 128
	}
	return
}

コメントの通り SVE は “still TBD”。CPU の対応状況を見ずに Neon の128bitで決め打ちます。arm64 なら Float32s.Len() は 4、Float64s.Len() は 2 で固定です。

amd64 は AVX512 があれば512bit

// src/simd/midway_amd64.go
func archMaxVectorSize() (size, allFeatureSize int) {
	if archsimd.X86.AVX() {
		size = 128
		allFeatureSize = 128
	}
	if archsimd.X86.AVX2() {
		size = 256
		// ...
	}
	if archsimd.X86.AVX512() {
		size = 512
		// ...
	}
	return
}

AVX512 が使えれば Float32s.Len() は 16。同じソースでレーン数が4倍変わります。手元の M5 Pro で GOARCH=amd64 にクロスビルドして Rosetta 2 上で走らせたところ、結果は Float32s.Len(): 4 でした。AVX が見えない環境では下限の128bitに落ちます。

archsimd はレーン数を型名に持つ

もう一方の simd/archsimdFloat32x4 のように長さを型名へ埋め込んだ低レベルAPI。arm64 で両方書いて測ると差はほとんど出ません。

BenchmarkSumArch-18    	  130195	      9213 ns/op
BenchmarkSumSIMD-18    	  133290	      9008 ns/op

Neon が128bit固定で、portable 側の Float32s も結局 Float32x4 に落ちるためです。ToArch() の戻り値を %T で出すと archsimd.Float32x4 と表示されます。amd64 で AVX512 を名指しで叩きたいとき、シャッフル系の命令を直接使いたいときが archsimd の出番。ただし ToArch()any を返すので、アーキ別の build tag と型アサーションが要ります。

GOEXPERIMENT なしでもビルドできる形にする

実験フラグ前提のコードをそのまま置くと、フラグ無しのCIとローカルビルドが両方 build constraints で落ちます。//go:build goexperiment.simd でファイルを割り、フォールバックを用意します。

ファイルを build tag で割る

//go:build goexperiment.simd

package fallback

import "simd"

const SIMDEnabled = true

func Sum(xs []float32) float32 {
	var acc simd.Float32s
	n := acc.Len()
	i := 0
	for ; i+n <= len(xs); i += n {
		acc = acc.Add(simd.LoadFloat32s(xs[i:]))
	}
	if i < len(xs) {
		v, _ := simd.LoadFloat32sPart(xs[i:])
		acc = acc.Add(v)
	}
	lanes := make([]float32, n)
	acc.Store(lanes)
	return SumScalar(lanes)
}
//go:build !goexperiment.simd

package fallback

const SIMDEnabled = false

func Sum(xs []float32) float32 { return SumScalar(xs) }

公開するのは Sum という1つの名前だけ。呼び出し側は分岐を知りません。SIMDEnabled を定数で公開しておくと、テストとログでどちらが選ばれたか確認できます。

両方の条件でテストを通す

func TestSum(t *testing.T) {
	xs := []float32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
	if got, want := Sum(xs), float32(55); got != want {
		t.Fatalf("Sum() = %v, want %v", got, want)
	}
	t.Logf("SIMDEnabled=%v", SIMDEnabled)
}
$ go test -v -run TestSum .
=== RUN   TestSum
    sum_test.go:10: SIMDEnabled=false
--- PASS: TestSum (0.00s)

$ GOEXPERIMENT=simd go test -v -run TestSum .
=== RUN   TestSum
    sum_test.go:10: SIMDEnabled=true
--- PASS: TestSum (0.00s)

CIでは GOEXPERIMENT=simd 有無の2ジョブを回す。実験フラグが将来 GA になっても外れても、片方を消すだけで済みます。

まとめ

  • GOEXPERIMENT=simd をビルド時に付けないと simd は import できない。go.mod の go ディレクティブは go 1.26 のままでも通る
  • float32 65536要素の総和は 35,653 ns から 9,018 ns へ。アキュムレータを4本に割ると 5,158 ns まで落ち、スカラー比 6.91倍
  • 8要素ではスカラーの方が速い。損益分岐は16要素、1024要素で 3.99倍に頭打ち
  • if 入りの集計を Masked に置き換えると18.0倍。レーン並列ぶんは4倍で、残りはスカラー側の分岐予測ミス。符号でソートすると2.6倍まで縮む
  • レーン数は実行時に決まる。arm64 は Neon の128bit固定で SVE は “still TBD”、amd64 は AVX512 があれば512bit
  • ループ幅は Len() から取り、端数は LoadFloat32sPart に任せる
  • Go 1 互換性保証の対象外。//go:build goexperiment.simd でフォールバックを置き、CIは2ジョブで回す

検証環境は go1.27rc2 darwin/arm64、Apple M5 Pro。amd64 の AVX512 実機では倍率が変わります。

よかったらシェアしてね!
  • URLをコピーしました!
  • URLをコピーしました!
目次