remove pool

bruteforce.go

@@ -48,9 +48,10 @@ func (b *Index[T]) Search(query Vector, k int) []Result[T] {
 		return nil
 	}
 
+	var relevance float64
 	dst := make(minheap[T], 0, k)
 	for _, v := range b.arr {
-		relevance := simd.Cosine(v.Vector, query)
+		simd.Cosine(&relevance, v.Vector, query)
 		result := Result[T]{
 			entry:     v,
 			Relevance: relevance,
@@ -61,7 +62,7 @@ func (b *Index[T]) Search(query Vector, k int) []Result[T] {
 		switch {
 		case dst.Len() < k:
 			dst.Push(result)
-		case relevance > dst[0].Relevance:
+		case result.Relevance > dst[0].Relevance:
 			dst.Pop()
 			dst.Push(result)
 		}

bruteforce_test.go

@@ -13,7 +13,7 @@ import (
 
 /*
 cpu: 13th Gen Intel(R) Core(TM) i7-13700K
-BenchmarkIndex/search-24         	    3807	    316587 ns/op	     264 B/op	       2 allocs/op
+BenchmarkIndex/search-24         	    4029	    298055 ns/op	     272 B/op	       3 allocs/op
 */
 func BenchmarkIndex(b *testing.B) {
 	data, err := loadDataset()

internal/cosine/simd/simd.go

@@ -3,7 +3,6 @@ package simd
 import (
 	"math"
 	"runtime"
-	"sync"
 	"unsafe"
 
 	"github.com/klauspost/cpuid/v2"
@@ -16,28 +15,17 @@ var (
 	hardware = avx2 || apple || neon
 )
 
-var pool = sync.Pool{
-	New: func() any {
-		var x float64
-		return &x
-	},
-}
-
-// Cosine calculates the cosine similarity between two vectors
-func Cosine(a, b []float32) float64 {
+// Cosine calculates the cosine similarity between two vectors and stores the result in the destination
+func Cosine(dst *float64, a, b []float32) {
 	if len(a) != len(b) {
 		panic("vectors must be of same length")
 	}
 
 	switch {
 	case hardware:
-		out := pool.Get().(*float64)
-		f32_cosine_distance(unsafe.Pointer(&a[0]), unsafe.Pointer(&b[0]), unsafe.Pointer(out), uint64(len(a)))
-		result := *out // copy out
-		pool.Put(out)
-		return result
+		f32_cosine_distance(unsafe.Pointer(&a[0]), unsafe.Pointer(&b[0]), unsafe.Pointer(dst), uint64(len(a)))
 	default:
-		return cosine(a, b)
+		*dst = cosine(a, b)
 	}
 }
 

internal/cosine/simd/simd_test.go

@@ -9,8 +9,8 @@ import (
 
 /*
 cpu: 13th Gen Intel(R) Core(TM) i7-13700K
-BenchmarkCosine/std-24         	15074694	        80.61 ns/op	       0 B/op	       0 allocs/op
-BenchmarkCosine/our-24         	45370162	        25.92 ns/op	       0 B/op	       0 allocs/op
+BenchmarkCosine/std-24         	15045380	        80.61 ns/op	       0 B/op	       0 allocs/op
+BenchmarkCosine/our-24         	55741100	        20.85 ns/op	       0 B/op	       0 allocs/op
 */
 func BenchmarkCosine(b *testing.B) {
 	x := randVec()
@@ -24,9 +24,10 @@ func BenchmarkCosine(b *testing.B) {
 	})
 
 	b.Run("our", func(b *testing.B) {
+		var out float64
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			Cosine(x, y)
+			Cosine(&out, x, y)
 		}
 	})
 }
@@ -36,9 +37,10 @@ func TestCosine(t *testing.T) {
 		x := randVec()
 		y := randVec()
 
+		var actual float64
+		Cosine(&actual, x, y)
 		expect := cosine(x, y)
-		actual := Cosine(x, y)
-		assert.InDelta(t, expect, actual, 1e-4, "expected %v, got %v", expect, actual)
+		assert.InDelta(t, expect, actual, 1e-4, "expected %v, got %v", cosine(x, y), actual)
 	}
 }