Gunosy Tech Blog

Gunosy Tech Blogは株式会社Gunosyのエンジニアが知見を共有する技術ブログです。

GoでSIMDを駆使して高速な内積演算を行う

Go

本記事は、Gunosy Advent Calendar 2020 14日目の記事です。

昨日はeastさんの「RedisでEVALを使うとこんなにお得！GunosyでのEVAL活用例 - Gunosy Tech Blog」でした。

最近、確率統計のことを考えながらパチンコを打つと面白いことに気づきました。 @gumigumi4f です。 この記事はグノシーでのパーソナライズにおいて必要な内積演算の高速化をGo上でSIMDを駆使して行うという内容になっています。

はじめに

ニュースアプリであるグノシーでは、ユーザーからのリクエストに対してリアルタイムに最適なリスト面を生成しています (いわゆるパーソナライズ)。 しかしながら、ユーザーからのリクエストは膨大で、時には数万req/sの負荷にも耐えうる必要があります。

また、レスポンスタイムが遅いとユーザー満足度を下げてしまう原因になります。 弊社では 50ms or die という某社の標語を基準に、ユーザーへのレスポンスタイムが50msになるように日々最適化を行っております。

そんな制約の中でパーソナライズを行おうと思ったときに、ちょっとでも負荷を軽減したい、ちょっとでも処理を高速化したい、そんなワガママを実現してくれるのが SIMD です。

SIMD とは

SIMD とは Single Instruction Multiple Data の略で、そのまま捉えれば「一回の命令で複数のデータを処理できる」となります。

wikipedia:SIMD

単純な例としてIntel CPUに搭載されている AVX (Advanced Vector Extension) で考えてみます。 AVXにおいては、256bitのデータを一度に処理することが可能となっています。これはfloat32で考えると約8個の数字を一度の命令で処理することが可能ということになります。

つまり以下の内積演算を考えたときに

similarity := 0
for i := 0; i < size; i++ {
    similarity += vx[i] * vy[i]
}

AVXを駆使すると以下のように書けるようになるということです (実際はこんな書き方できません)。

similarity := 0
for i := 0; i < size - 8; i += 8 {
    similarity += AVXdot(vx[i:i+8], vy[i:i+8])
}

これによって従来必要だった演算の数を 1/8 に減らすことができ、スループットの向上が望めます。

amd64系のCPUではAVXがSIMDとして提供されており*1、arm64系ではNEONがSIMD拡張として提供されています*2。 前者は演算幅が256bit、後者は128bitとなっています。

なぜ、SIMDがパーソナライズの高速化につながるのか

グノシーでのパーソナライズは、ユーザーとアイテムをそれぞれ事前に分散表現として表しておき、リクエスト毎にアイテムの候補群とユーザーのベクトルの類似度を計算、適切なリストを生成するといった流れになっています。

ここでのベクトルの類似度計算としてWord2Vec等でも用いられているコサイン類似度について考えてみます。 コサイン類似度は以下の式で表されます。

\cos (\vec{q}, \vec{d}) = \frac{\vec{q} \cdot \vec{d}}{|\vec{q}||\vec{d}|}

各ベクトルの大きさで割るという演算が分母に含まれていますが、これらは各ベクトルの大きさを事前に1にしておくことで無視できます。 つまり、リアルタイムにコサイン類似度を計算する場合、必要な処理はは分子の \vec{q} \cdot \vec{d} のみになります。 SIMD拡張はここの内積演算に対して使用することが可能であるため、全体の結果としてパーソナライズの高速化が期待できます。

ここでは実際の例としてグノシーでのパーソナライズを挙げSIMDの優位性を説明しましたが、これは通常のディープラーニングなどを考えた場合でも同様です。 行列積が計算量の大半を占めるようなディープラーニングでは、AVX等のSIMD命令が高速化に非常に役に立ちます。 TensorflowではAVX命令を搭載していないamd64系CPUではそもそも実行すらできないことからも、その重要性がわかると思います *3。

Go で SIMD

グノシーにおいてユーザーからのリクエストに対してパーソナライズを行うAPIはGo言語によって書かれています。

C++などを用いる場合、gccのコンパイルオプションとして -march=native -O2 などをつければ、頭のいいコンパイラが自動的にSIMD命令を使用できるところで使用してくれたりするのですが、Go言語の場合はそうではありません。 また、今後もGo言語がSIMD拡張に対応することはなさそうです*4。

そのため、Go言語でSIMD拡張を用いる場合CGOを使用してC言語にて拡張をする必要があります。

というわけで amd64 (avx) と arm64 (neon) の両方のアーキテクチャに対応したベクトル演算高速化ライブラリを作成しました。

github.com

元のライブラリである https://github.com/monochromegane/go-avx からforkし、NEON拡張への対応と一部修正を加えたものになります。

使い方

Before

size := 512
vecx := make([]float32, size)
vecy := make([]float32, size)
for j := 0; j < size; j++ {
    vecx[j] = rand.Float32()
    vecy[j] = rand.Float32()
}

similarity := float32(0)
for k := 0; k < size; k++ {
    similarity += vecx[k] * vecy[k]
}

After

size := 512
vecx := vx.AlignedAlloc(size)
vecy := vx.AlignedAlloc(size)
defer vx.Free(vecx)
defer vx.Free(vecy)
for j := 0; j < size; j++ {
    vecx[j] = rand.Float32()
    vecy[j] = rand.Float32()
}

similarity := vx.Dot(size, vecx, vecy)

注意点として[]float32のsliceをAlignedAllocを通して確保する必要があります。 これはAVXを使用する際に適切なアライメントに配置されたメモリ領域を必要とするからです*5*6。

また、確保したメモリは自動では開放されないため、必ずFreeを用いて開放する必要があります。

ベンチマーク

以下のテストコードにてベンチマークを実施しました。 https://github.com/gumigumi4f/go-vx/blob/master/vx_test.go#L169-L230

想定されるシチュエーションとしては、「記事の候補が16384記事あり、とあるユーザーからのリクエストが飛んできた際に、そのユーザーと各記事のベクトルの類似度を計算する」といったものになります。これは実際にグノシーでのパーソナライズで考えられるシチュエーションの一つとなります。 また、各ベクトルの次元は512次元としています。

func BenchmarkDotVx(b *testing.B) {
    num := 16384
    size := 512

    vx := AlignedAlloc(size)
    for j := 0; j < size; j++ {
        vx[j] = rand.Float32()
    }

    vys := make([][]float32, num)
    for i := range vys {
        vys[i] = AlignedAlloc(size)
        for j := 0; j < size; j++ {
            vys[i][j] = rand.Float32()
        }
    }

    similarities := make([]float32, num)

    b.ResetTimer()
    for i := 0; i < b.N; i++ {
        for j, vy := range vys {
            similarities[j] = Dot(size, vx, vy)
        }
    }

    Free(vx)
    for i := range vys {
        Free(vys[i])
    }
}

func BenchmarkDotNative(b *testing.B) {
    num := 16384
    size := 512

    vx := make([]float32, size)
    for j := 0; j < size; j++ {
        vx[j] = rand.Float32()
    }

    vys := make([][]float32, num)
    for i := range vys {
        vys[i] = make([]float32, size)
        for j := 0; j < size; j++ {
            vys[i][j] = rand.Float32()
        }
    }

    similarities := make([]float32, num)

    b.ResetTimer()
    for i := 0; i < b.N; i++ {
        for j, vy := range vys {
            similarity := float32(0)
            for k := 0; k < size; k++ {
                similarity += vx[k] * vy[k]
            }
            similarities[j] = similarity
        }
    }
}

amd64ではAWSのc5.largeインスタンス、arm64ではAWSのc6g.mediumインスタンスを用いてベンチマークを行っています。

amd64

goos: linux
goarch: amd64
pkg: github.com/gumigumi4f/go-vx
BenchmarkDotVx-2             364           3201061 ns/op
BenchmarkDotNative-2         100          10833781 ns/op
PASS
ok      github.com/gumigumi4f/go-vx     3.953s

arm64

goos: linux
goarch: arm64
pkg: github.com/gumigumi4f/go-vx
BenchmarkDotVx               274           4275091 ns/op
BenchmarkDotNative           100          10599838 ns/op
PASS
ok      github.com/gumigumi4f/go-vx     4.154s

amd64ではおよそ3.38倍、arm64でも2.48倍の高速化を実現することができました。

純粋に書いたコードではどちらの場合もほぼ同じ速度でしたが、256bit幅で演算するc5.largeのほうがSIMD拡張を用いた場合早くなるという結果になりました。

まとめ

Go言語の思想から言ったらおそらくアーキテクチャコテコテのコードは微妙なのかなとは思いつつも、高いパフォーマンスが要求されるような場所においてはこのような最適化も有効なのかな、と思いブログにしたためました。 M1 Macの発売やAWSのGraviton 2等、armの勢いが凄まじい中でこの記事が何かしらのシステムの高速化に役立つことがあれば幸いです。

明日はakinkさんです、よろしくおねがいします！

*1:AVX-512もあるがここでは割愛, https://mag.osdn.jp/20/07/16/103700

*2:https://www.arm.com/ja/why-arm/technologies/neon

*3:https://www.tensorflow.org/install/source?hl=ja

*4:https://groups.google.com/g/golang-nuts/c/I2mTRxIwyQ4

*5:アライメントが取れていない場合でも使用できる命令もあるが、パフォーマンスは劣る

*6:arm64の場合必ずしもアライメントが取れている必要はないらしいが、一部のCPUでスループットが低下するらしい, https://developer.arm.com/documentation/uan0015/b/