Add avx

Author: hupe1980

internal/math32/floats.go

@@ -1,8 +1,8 @@
 package math32
 
 var (
-	useAVX512 bool // nolint unused
-	useNEON   bool // nolint unused
+	useAVX  bool // nolint unused
+	useNEON bool // nolint unused
 )
 
 // Dot two vectors.

internal/math32/floats_amd64.go

@@ -3,26 +3,46 @@
 package math32
 
 import (
+	"unsafe"
+
 	"golang.org/x/sys/cpu"
 )
 
 func init() {
-	useAVX512 = cpu.X86.HasAVX512
+	useAVX = cpu.X86.HasAVX
 }
 
+//go:noescape
+func _dot_product_avx(a, b unsafe.Pointer, n uintptr, result unsafe.Pointer)
+
+//go:noescape
+func _squared_l2_avx(a, b unsafe.Pointer, n uintptr, result unsafe.Pointer)
+
 func dot(a, b []float32) float32 {
 	switch {
-	case useAVX512:
-		return dotGeneric(a, b) // TODO
+	case useAVX:
+		var ret float32
+
+		if len(a) > 0 {
+			_dot_product_avx(unsafe.Pointer(&a[0]), unsafe.Pointer(&b[0]), uintptr(len(a)), unsafe.Pointer(&ret))
+		}
+
+		return ret
 	default:
 		return dotGeneric(a, b)
 	}
 }
 
 func squaredL2(a, b []float32) float32 {
 	switch {
-	case useAVX512:
-		return squaredL2Generic(a, b) // TODO
+	case useAVX:
+		var ret float32
+
+		if len(a) > 0 {
+			_squared_l2_avx(unsafe.Pointer(&a[0]), unsafe.Pointer(&b[0]), uintptr(len(a)), unsafe.Pointer(&ret))
+		}
+
+		return ret
 	default:
 		return squaredL2Generic(a, b)
 	}

internal/math32/floats_arm64.go

@@ -13,7 +13,7 @@ func init() {
 }
 
 //go:noescape
-func _vdot_neon(a unsafe.Pointer, b unsafe.Pointer, n uintptr, ret unsafe.Pointer)
+func _dot_product_neon(a unsafe.Pointer, b unsafe.Pointer, n uintptr, ret unsafe.Pointer)
 
 //go:noescape
 func _squared_l2_neon(a, b unsafe.Pointer, n uintptr, result unsafe.Pointer)
@@ -24,7 +24,7 @@ func dot(a, b []float32) float32 {
 		var ret float32
 
 		if len(a) > 0 {
-			_vdot_neon(unsafe.Pointer(&a[0]), unsafe.Pointer(&b[0]), uintptr(len(a)), unsafe.Pointer(&ret))
+			_dot_product_neon(unsafe.Pointer(&a[0]), unsafe.Pointer(&b[0]), uintptr(len(a)), unsafe.Pointer(&ret))
 		}
 
 		return ret