Go 1.26のリリースの足音が聞こえてきたので1.26の新機能のうち、気になる機能を有志で紹介していく連載を行います。

1.26の新機能のダイジェスト

久々に結構大きめな変更が多くなっていますね

• new(プリミティブ)
• 1.25で導入されたGCが有効化
• cgo, memory allocate高速化
• goroutine leak profiler
• compilerがスライスをよりスタックに割り当てられるように
• SIMDを扱う実験的パッケージ（AMD64のみサポート）
• runtime/secretという機密情報を扱う実験的パッケージ
• image/jpegが高速に
• io.ReadAllの消費メモリが少なく
• log/slogで複数の出力先に一気に出力できるMultiHandlerが追加
• testing.ArtifactDir
• 暗号アルゴリズムの更新たくさん

全体的に高速化が行われています。なのでアップデートするだけで恩恵がいろいろあります。

テストのArtifactDir()は-artifactsフラグを付けてテストを実行すると、T.ArtifactDir()メソッドが中間生成物の置き場を用意してくれて、そこに格納することであとから検死できるようになるというものです。以前、TempDir()が使いにくいというのを技術ブログで書きましたが、それのソリューションが提供された感じですね。

SIMD

それでは、今回追加されたsimd/archsimdパッケージについて紹介します。

SIMDというのは1つの命令で複数のデータをまとめて処理する命令セットです。4つ、8つ・・・のデータをまとめて足し算したり、掛け算したり。行列計算などを高速に行うのに使います。

インテル系では主に以下のようにSIMDの命令セットの世代があります。それぞれとレジスタサイズの対応は以下の通りです。

• MMX: 64ビット
• SSE: 128ビット
• AVX: 256ビット
• AVX512: 512ビット

なお、AVX512はそれほど大量のデータを処理する必要とするアプリケーションが一部に限られるせいか、10th/11th Gen Core CPUでは対応していましたが、その後のチップからは取り除かれてしまいました。それだけ半導体を使うのに使用されないのであればその分通常のコア数を増やそうという感じですね。Eコア・Pコアに分かれた時代からは使えなくなりました。あとはCPUだけが早くてもメモリ転送がボトルネックになりそうですしね。

など、上記の命令セットの名前はインテルですが、AMD、ARM(NEON)なども備えています。

例えば、AVXは256ビットですが、これを8ビットx32本、と使ったり、16ビットx16本として使ったり色々できます。simd/archsimdパッケージを見ると、いろんな型がありますが、よく見ると、Float32x8といった型になっております。たくさんありますが、データ型と要素数の組み合わせになっております。合計のビット数は128ビット、256ビット、512ビットがあることがわかります。だいたい、和だったり、どのような演算をするかはメソッドで提供されていますが、int8/uint8などは和はあるが、積などはないなど、型によって使えるメソッドに違いがあります。

なお、Goがバイナリを最後に生成するのに使うアセンブラは、アーキテクチャ限定でSIMD命令にはいくつか対応していたので、コードを.goではなく、.sで作成すればSIMDは今までも活用できました。Go 1.24で導入されたmapの新アルゴリズムのSwiss Tableは、SIMDを活用して高速検索を行うというものでした。サードパーティのライブラリでSIMDを使った行列演算ライブラリなどもいくつかありましたが、それが今回公式でできるようになりました。

ベクトルの内積計算

データ作成は後で説明するのですが、embedding化で作ったドキュメントのベクトルデータと、検索用語のベクトルデータの内積を取ることで、どれだけそれぞれの内容が近いかを数値化できます。ベクトル化すると、例えば「猫」について説明した文章があったとして、猫に近い「動物」や「トラ」といった言葉では距離が近いと判定されます。それにより、類似文書検索ができます。転置インデックスを使った全文検索だと「猫」で検索しないとヒットしませんが、ベクトル検索では近い内容も探せます。その分、ベクトルを計算するコーパスの質が大事で、今まではなかなか自由に簡単に使えるものがなかったのですが、生成AIのおかげでモデルを簡単にインストールしたり利用できるようになったので、簡単にベクトル検索ができる時代になりました。

まず、ollamaをインストールして、今回のembeddingで使うモデルを使えるようにしてローカルサーバーを起動しておきます。

ollama pull nomic-embed-text
$ ollama serve

今回実験で使ったnomic-embed-textという軽量なモデルを使うと、768次元のベクトルとしてベクトルが計算できます。SIMDがなかったら、ループを回してそれぞれの項の積をしてから和を計算します。SIMDを使う場合、128ビットだと16ビットの要素の積を8つ同時に行えますし、内積用のメソッドを使うと和も求められます。

ベンチマークを取ってみると、6.6倍ほど高速になっていることがわかります。効果バツグンですね。なお、AMD64しか対応していないので、M3のMacBook AirでAMD64ビルドしたものをRosetta2でテストしています。そのうち早くネイティブ対応して欲しいですね。

なお、Rosettaのエミュレーションが公式には128ビットまでということでその範囲のテストになっていますが、非公式に256ビットとかも行けます。ただ、実行してみると128ビットの2倍程度の時間がかかっているので128ビット演算器で256ビットのエミュレーションをしている感じでしょうか？

% GOARCH=amd64 GOEXPERIMENT=simd go test -bench . -benchmem
goos: darwin
goarch: amd64
pkg: simdtest
cpu: VirtualApple @ 2.50GHz
BenchmarkPrimitive_RealData-8            5915974               203.0 ns/op             0 B/op          0 allocs/op
BenchmarkSIMD_RealData-8                39017252                29.33 ns/op            0 B/op          0 allocs/op

以下のコードがテストです。データの読み込みは省略しています。SIMDの方は[8]int16などの固定長配列を前提としているので、[][8]int16という形で、ベクトルを計算単位でまとめた二重配列的な構造にしてから与えています。なお、テストデータはReal World HTTPの原稿で、「クッキーのセキュリティ」という検索用語をそれぞれベクトル化してあらかじめファイルにしてあり、それを利用しています。

package main

import (
 "encoding/json"
 "log"
 "os"
 "simd/archsimd"
 "testing"
)

// プリミティブな内積計算ロジック
func dotProductPrimitive(a, b []int16) int32 {
 var sum int32 = 0
 for i := 0; i < len(a); i++ {
  sum += int32(a[i]) * int32(b[i])
 }
 return sum
}

// SIMDの内積計算ロジック
func dotProductSIMD(a, b [][8]int16) int32 {
 acc := archsimd.Int32x4{}
 for i := 0; i < len(a); i++ {
  vq := archsimd.LoadInt16x8(&a[i])
  vt := archsimd.LoadInt16x8(&b[i])
  dot := vq.DotProductPairs(vt)
  acc = acc.Add(dot)
 }
 return acc.GetElem(0) + acc.GetElem(1) + acc.GetElem(2) + acc.GetElem(3)
}

// ロード済みのデータ(初期化は省略)
var (
 testQuery8     []int8
 testQuery16    []int16
 testQuerySIMD  [][8]int16
 testTarget16   []int16
 testTargetSIMD [][8]int16
)

// --- ベンチマーク関数 ---

func BenchmarkPrimitive_RealData(b *testing.B) {
 b.ResetTimer()
 for i := 0; i < b.N; i++ {
  _ = dotProductPrimitive(testQuery16, testTarget16)
 }
}

func BenchmarkSIMD_RealData(b *testing.B) {
 b.ResetTimer()
 for i := 0; i < b.N; i++ {
  _ = dotProductSIMD(testQuerySIMD, testTargetSIMD)
 }
}

ベクトルの作成

ローカルでollamaを動かして、それにファイルを投げてベクトル化をします。なお、ベクトル化する場合はあまり長すぎるデータを与えてもコンテキストから溢れてエラーになってしまうため、あらかじめ原稿データを小さい単位の.rstファイルに分割しておきました。だいたい節タイトルごとに10行とか20行以下ぐらいに区切り、節タイトルがファイル名となるように事前処理しました。

積をSIMDでやるために計算はint16でやっていますが、データ自体は8ビットのベクトルとしていました。このあたり、どのあたりが良いのかはチューニングしていきたいですね。

package main

import (
 "context"
 "encoding/json"
 "fmt"
 "log"
 "os"
 "path/filepath"

 "github.com/ollama/ollama/api"
)

// 保存用のデータ構造
// key: ファイル名, value: int8のベクトル
type VectorStore map[string][]int8

func main() {
 if len(os.Args) < 2 {
  fmt.Println("使用法: go run main.go <directory_or_file>")
  return
 }

 inputPath := os.Args[1]
 client, _ := api.ClientFromEnvironment()
 ctx := context.Background()
 model := "nomic-embed-text"
 store := make(VectorStore)
 files, _ := filepath.Glob(filepath.Join(inputPath, "*.rst"))

 for _, file := range files {
  content, err := os.ReadFile(file)
  if err != nil {
   continue
  }

  // 1. Ollamaから float32 ベクトルを取得
  req := &api.EmbedRequest{Model: model, Input: string(content)}
  resp, err := client.Embed(ctx, req)
  if err != nil || len(resp.Embeddings) == 0 {
   log.Printf("Skip %s: %v", file, err)
   continue
  }

  // 2. float32 -> int8 量子化
  floatVec := resp.Embeddings[0]
  int8Vec := make([]int8, len(floatVec))
  for i, v := range floatVec {
   val := v * 127.0
   if val > 127 {
    val = 127
   }
   if val < -128 {
    val = -128
   }
   int8Vec[i] = int8(val)
  }

  shortName := filepath.Base(file)
  store[shortName] = int8Vec
  fmt.Printf("Vectorized: %s\n", shortName)
 }

 // 3. JSONファイルとして書き出し
 outputFile := "vectors.json"
 jsonData, _ := json.MarshalIndent(store, "", "  ")
 os.WriteFile(outputFile, jsonData, 0644)

 fmt.Printf("\nSuccess! Saved to %s\n", outputFile)
}

検索の方のプログラムは全文は載せませんが、すでに載せているロジックの組み合わせでいきます。検索用語の方も上記のベクトル化と同じようにベクトル化を行います。そして、[][8]int16にすべてベクトルを変換しておき、あとはドット積の大きい順序でソートして返すだけです。

まとめ

AI界隈でよくでてくる計算がGoでも簡単に実現できるようになりました。類似検索とかベクトルデータベースの実装が今後増えることを期待しています。対応アーキテクチャも今後増えていくでしょう。WASMも対応になると嬉しいですね。

明日は武田さんになります。