chromem-go
chromem-go copied to clipboard
Use SIMD for dot product calculation on supported CPUs
: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:
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
can you rebase this branch please, just would like to test/push the changes from 4341c0e
Oh nice!
Rebased ✔️
thanks https://github.com/philippgille/chromem-go/pull/75 is ready for review
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
}