chromem-go icon indicating copy to clipboard operation
chromem-go copied to clipboard

Use SIMD for dot product calculation on supported CPUs

Open philippgille opened this issue 11 months ago • 4 comments

:warning: WIP: I'd like to get rid of the golang.org/x/sys/cpu dependency

Using generated Assembler code from https://github.com/viterin/vek. That code is MIT licensed, which means we can use and copy it as long as we adhere to its terms, like preserving a copy of the license and copyright in the licensed material. chromem-go stays AGPL licensed.

Shoutout to @viterin :bowing_man:

philippgille avatar Mar 16 '24 18:03 philippgille

can you rebase this branch please, just would like to test/push the changes from https://github.com/philippgille/chromem-go/pull/75/commits/4341c0eb6a76110c3df09dac4b0b1f5568c74d0d

rinor avatar May 13 '24 16:05 rinor

can you rebase this branch please, just would like to test/push the changes from 4341c0e

Oh nice!

Rebased ✔️

philippgille avatar May 17 '24 19:05 philippgille

thanks https://github.com/philippgille/chromem-go/pull/75 is ready for review

rinor avatar May 17 '24 20:05 rinor

also do you think it's worth trying some non-simd stuff from https://sourcegraph.com/blog/slow-to-simd

Basically dotProduct function might look smth like:

// dotProduct calculates the dot product between two vectors.
// It's the same as cosine similarity for normalized vectors.
// The resulting value represents the similarity, so a higher value means the
// vectors are more similar.
func dotProduct(a, b []float32) (float32, error) {
	// The vectors must have the same length
	if len(a) != len(b) {
		return 0, errors.New("vectors must have the same length")
	}

	dotProduct := float32(0)
	vLen := len(a)

	switch {
	case simd.UseAVX2:
		dotProduct = simd.Dot_AVX2_F32(a, b)
	case vLen%4 == 0:
		for i := 0; i < vLen; i += 4 {
			aTmp := a[i : i+4 : i+4]
			bTmp := b[i : i+4 : i+4]
			s0 := aTmp[0] * bTmp[0]
			s1 := aTmp[1] * bTmp[1]
			s2 := aTmp[2] * bTmp[2]
			s3 := aTmp[3] * bTmp[3]
			dotProduct += s0 + s1 + s2 + s3
		}
	default:
		for i := range a {
			dotProduct += a[i] * b[i]
		}
	}

	return dotProduct, nil
}

rinor avatar May 18 '24 11:05 rinor