refactor(math): merge float kernels into ailego math

.github/workflows/ci.yml

@@ -98,6 +98,6 @@ jobs:
           sudo apt-get update
           sudo apt-get install -y qemu-user
       - name: Test RISC-V RVV
-        run: make -C internal/floats test-rvv
+        run: make -C internal/ailego/math test-rvv
       - name: Test LoongArch LASX
-        run: make -C internal/floats test-lasx
+        run: make -C internal/ailego/math test-lasx

go.mod

@@ -1,6 +1,6 @@
 module github.com/gorse-io/xvec
 
-go 1.27
+go 1.27.0
 
 require (
 	github.com/RoaringBitmap/roaring/v2 v2.24.0
@@ -36,7 +36,7 @@ require (
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/golang/snappy v0.0.5-0.20231225225746-43d5d4cd4e0e // indirect
 	github.com/google/uuid v1.6.0 // indirect
-	github.com/gorse-io/goat v0.2.1-0.20260618151728-201cbcf325ad // indirect
+	github.com/gorse-io/goat v0.2.2-0.20260901023055-3025b7527c27 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/klauspost/asmfmt v1.3.2 // indirect
 	github.com/klauspost/compress v1.17.11 // indirect

go.sum

@@ -72,8 +72,8 @@ github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/gorse-io/goat v0.2.1-0.20260618151728-201cbcf325ad h1:HESEie1ivHg8yI9Qnej4Cu+AP8w8gQ9gqpYcnV5zC7Q=
-github.com/gorse-io/goat v0.2.1-0.20260618151728-201cbcf325ad/go.mod h1:NjwKvGIyhFveMjnUdYrngJYvYLqK0aK7Qdp7QJWkBLo=
+github.com/gorse-io/goat v0.2.2-0.20260901023055-3025b7527c27 h1:7Ui/Xhc9+zMN6cR8A5GVFffssiCdDCVGB0WBrlLeMBw=
+github.com/gorse-io/goat v0.2.2-0.20260901023055-3025b7527c27/go.mod h1:gRNfHg2VIBM1rLZXr6LtS2BX+qVI8DKckQC9LBESUMk=
 github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
 github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=

internal/ailego/math/Makefile

@@ -2,28 +2,97 @@ GO ?= go
 GOAT ?= $(GO) tool goat
 GOAT_TARGET_AMD64 ?= amd64
 GOAT_TARGET_ARM64 ?= arm64
+GOAT_TARGET_RVV ?= riscv64
+GOAT_TARGET_LASX ?= loong64
 GOAT_OUTPUT_DIR ?= .goat/mathutil
 MV ?= mv
+CLANG_RESOURCE_DIR ?= $(shell clang -print-resource-dir)
+GOAT_RVV_INCLUDE ?= $(CLANG_RESOURCE_DIR)/include
+LASX_OBJDUMP ?= loongarch64-linux-gnu-objdump
+RISCV64_QEMU ?= qemu-riscv64
+RISCV64_QEMU_FLAGS ?= -cpu max,v=true
+LOONG64_QEMU ?= qemu-loongarch64
+LOONG64_QEMU_FLAGS ?= -cpu max
+TEST_RVV_BINARY ?= .test-rvv.riscv64.test
+TEST_LASX_BINARY ?= .test-lasx.loong64.test
 
-.PHONY: generate avx2 avx512 neon clean
+.PHONY: generate avx avx2 avx512 neon rvv lasx fht-avx2 fht-avx512 fht-neon distance-avx distance-avx512 distance-neon distance-rvv distance-lasx clean test-rvv test-lasx
 
-generate: avx2 avx512 neon clean
+generate: fht-avx2 fht-avx512 fht-neon distance-avx distance-avx512 distance-neon distance-rvv distance-lasx clean
 
-avx2:
+avx: distance-avx
+
+avx2: fht-avx2
+
+avx512: fht-avx512 distance-avx512
+
+neon: fht-neon distance-neon
+
+rvv: distance-rvv
+
+lasx: distance-lasx
+
+fht-avx2:
 	mkdir -p $(GOAT_OUTPUT_DIR)
 	$(GOAT) src/fht_avx2.c --target $(GOAT_TARGET_AMD64) -O3 -mavx2 -o $(GOAT_OUTPUT_DIR)
 	$(MV) $(GOAT_OUTPUT_DIR)/fht_avx2.go $(GOAT_OUTPUT_DIR)/fht_avx2.s .
 
-avx512:
+fht-avx512:
 	mkdir -p $(GOAT_OUTPUT_DIR)
 	$(GOAT) src/fht_avx512.c --target $(GOAT_TARGET_AMD64) -O3 -mavx512f -mavx512dq -o $(GOAT_OUTPUT_DIR)
 	$(MV) $(GOAT_OUTPUT_DIR)/fht_avx512.go $(GOAT_OUTPUT_DIR)/fht_avx512.s .
 
-neon:
+fht-neon:
 	mkdir -p $(GOAT_OUTPUT_DIR)
 	$(GOAT) src/fht_neon.c --target $(GOAT_TARGET_ARM64) -O3 -o $(GOAT_OUTPUT_DIR)
 	$(MV) $(GOAT_OUTPUT_DIR)/fht_neon.go $(GOAT_OUTPUT_DIR)/fht_neon.s .
 
+distance-avx:
+	$(RM) -r $(GOAT_OUTPUT_DIR)
+	mkdir -p $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/euclidean_distance_matrix_fp32_sse.c --target $(GOAT_TARGET_AMD64) -O3 -msse2 -o $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/inner_product_matrix_fp32_sse.c --target $(GOAT_TARGET_AMD64) -O3 -msse2 -o $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/mips_euclidean_distance_matrix_fp32_sse.c --target $(GOAT_TARGET_AMD64) -O3 -msse2 -o $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/euclidean_distance_matrix_fp32_avx.c --target $(GOAT_TARGET_AMD64) -O3 -mavx -o $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/inner_product_matrix_fp32_avx.c --target $(GOAT_TARGET_AMD64) -O3 -mavx -o $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/mips_euclidean_distance_matrix_fp32_avx.c --target $(GOAT_TARGET_AMD64) -O3 -mavx -o $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/distance_utility_batch_avx.c --target $(GOAT_TARGET_AMD64) -O3 -mavx -o $(GOAT_OUTPUT_DIR)
+	$(MV) $(GOAT_OUTPUT_DIR)/*.go $(GOAT_OUTPUT_DIR)/*.s .
+
+distance-avx512:
+	$(RM) -r $(GOAT_OUTPUT_DIR)
+	mkdir -p $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/euclidean_distance_matrix_fp32_avx512.c --target $(GOAT_TARGET_AMD64) -O3 -mavx -mfma -mavx512f -o $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/inner_product_matrix_fp32_avx512.c --target $(GOAT_TARGET_AMD64) -O3 -mavx -mfma -mavx512f -o $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/mips_euclidean_distance_matrix_fp32_avx512.c --target $(GOAT_TARGET_AMD64) -O3 -mavx -mfma -mavx512f -o $(GOAT_OUTPUT_DIR)
+	$(MV) $(GOAT_OUTPUT_DIR)/*.go $(GOAT_OUTPUT_DIR)/*.s .
+
+distance-neon:
+	$(RM) -r $(GOAT_OUTPUT_DIR)
+	mkdir -p $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/euclidean_distance_matrix_fp32_neon.c --target $(GOAT_TARGET_ARM64) -O3 -o $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/inner_product_matrix_fp32_neon.c --target $(GOAT_TARGET_ARM64) -O3 -o $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/mips_euclidean_distance_matrix_fp32_neon.c --target $(GOAT_TARGET_ARM64) -O3 -o $(GOAT_OUTPUT_DIR)
+	$(MV) $(GOAT_OUTPUT_DIR)/*.go $(GOAT_OUTPUT_DIR)/*.s .
+
+distance-rvv:
+	$(GOAT) src/distance_utility_rvv.c --target $(GOAT_TARGET_RVV) -O3 -march=rv64imafdv -e=-isystem -e=$(GOAT_RVV_INCLUDE)
+
+distance-lasx:
+	OBJDUMP=$(LASX_OBJDUMP) $(GOAT) src/distance_utility_lasx.c --target $(GOAT_TARGET_LASX) -O3 -mlasx
+
 clean:
 	$(RM) src/*.o src/*.s
 	$(RM) -r .goat
+
+test-rvv:
+	@set -e; \
+	trap '$(RM) ./$(TEST_RVV_BINARY)' EXIT; \
+	GOOS=linux GOARCH=riscv64 CGO_ENABLED=0 $(GO) test -c -o $(TEST_RVV_BINARY) .; \
+	$(RISCV64_QEMU) $(RISCV64_QEMU_FLAGS) ./$(TEST_RVV_BINARY)
+
+test-lasx:
+	@set -e; \
+	trap '$(RM) ./$(TEST_LASX_BINARY)' EXIT; \
+	GOOS=linux GOARCH=loong64 CGO_ENABLED=0 $(GO) test -c -o $(TEST_LASX_BINARY) .; \
+	$(LOONG64_QEMU) $(LOONG64_QEMU_FLAGS) ./$(TEST_LASX_BINARY)

internal/ailego/math/distance.go

@@ -17,8 +17,6 @@ package mathutil
 import (
 	"errors"
 	"math"
-
-	"github.com/gorse-io/xvec/internal/floats"
 )
 
 var (
@@ -36,7 +34,7 @@ func L2Squared(left, right []float32) (float32, error) {
 	if err := validateDenseFinite(left, right); err != nil {
 		return 0, err
 	}
-	return finiteScore(float64(floats.L2Squared(left, right)))
+	return finiteScore(float64(squaredEuclidean(left, right)))
 }
 
 // InnerProduct computes the dot-product similarity. Higher scores are better.
@@ -47,7 +45,7 @@ func InnerProduct(left, right []float32) (float32, error) {
 	if err := validateDenseFinite(left, right); err != nil {
 		return 0, err
 	}
-	return finiteScore(float64(floats.InnerProduct(left, right)))
+	return finiteScore(float64(innerProduct(left, right)))
 }
 
 // CosineDistance computes 1-cos(left,right). Lower scores are better. Two zero
@@ -63,7 +61,7 @@ func CosineDistance(left, right []float32) (float32, error) {
 }
 
 func cosineDistance(left, right []float32) float32 {
-	inner, leftNorm, rightNorm := floats.DotNorms(left, right)
+	inner, leftNorm, rightNorm := dotNorms(left, right)
 	if leftNorm == 0 && rightNorm == 0 {
 		return 0
 	}
@@ -91,7 +89,7 @@ func MIPSL2Squared(left, right []float32) (float32, error) {
 }
 
 func mipsL2Squared(left, right []float32) float32 {
-	inner, leftNorm, rightNorm := floats.DotNorms(left, right)
+	inner, leftNorm, rightNorm := dotNorms(left, right)
 	denominator := max(leftNorm, rightNorm)
 	if denominator == 0 {
 		return 0

internal/ailego/math/distance_prevalidated.go

@@ -16,8 +16,6 @@ package mathutil
 
 import (
 	"math"
-
-	"github.com/gorse-io/xvec/internal/floats"
 )
 
 // DenseDistance computes a score for two already validated dense vectors.
@@ -28,12 +26,12 @@ type DenseDistance func(left, right []float32) (float32, error)
 // its inputs. It is intended for index hot paths whose storage boundary has
 // already validated every vector.
 func L2SquaredPrevalidated(left, right []float32) (float32, error) {
-	return finiteScore(float64(floats.L2Squared(left, right)))
+	return finiteScore(float64(squaredEuclidean(left, right)))
 }
 
 // InnerProductPrevalidated computes inner product without validating inputs.
 func InnerProductPrevalidated(left, right []float32) (float32, error) {
-	return finiteScore(float64(floats.InnerProduct(left, right)))
+	return finiteScore(float64(innerProduct(left, right)))
 }
 
 // CosineDistancePrevalidated computes cosine distance without validating inputs.
@@ -44,7 +42,7 @@ func CosineDistancePrevalidated(left, right []float32) (float32, error) {
 // L2MagnitudePrevalidated computes a vector magnitude without validating its
 // components. It is intended for indexes that cache norms at ingestion time.
 func L2MagnitudePrevalidated(vector []float32) (float32, error) {
-	norm := floats.InnerProduct(vector, vector)
+	norm := innerProduct(vector, vector)
 	if norm < 0 {
 		norm = 0
 	}
@@ -61,7 +59,7 @@ func CosineDistanceWithMagnitudesPrevalidated(left, right []float32, leftMagnitu
 	if leftMagnitude == 0 || rightMagnitude == 0 {
 		return 1, nil
 	}
-	cosine := floats.InnerProduct(left, right) / (leftMagnitude * rightMagnitude)
+	cosine := innerProduct(left, right) / (leftMagnitude * rightMagnitude)
 	cosine = min(1, max(-1, cosine))
 	return finiteScore(float64(1 - cosine))
 }
@@ -73,7 +71,7 @@ func CosineDistances2WithMagnitudesPrevalidated(
 	query, first, second []float32,
 	queryMagnitude, firstMagnitude, secondMagnitude float32,
 ) (firstDistance, secondDistance float32, err error) {
-	firstProduct, secondProduct := floats.InnerProducts2(query, first, second)
+	firstProduct, secondProduct := innerProducts2(query, first, second)
 	firstDistance, err = cosineDistanceFromProduct(firstProduct, queryMagnitude, firstMagnitude)
 	if err != nil {
 		return 0, 0, err
@@ -92,7 +90,7 @@ func CosineDistances4WithMagnitudesPrevalidated(
 	query, first, second, third, fourth []float32,
 	queryMagnitude, firstMagnitude, secondMagnitude, thirdMagnitude, fourthMagnitude float32,
 ) (firstDistance, secondDistance, thirdDistance, fourthDistance float32, err error) {
-	firstProduct, secondProduct, thirdProduct, fourthProduct := floats.InnerProducts4(query, first, second, third, fourth)
+	firstProduct, secondProduct, thirdProduct, fourthProduct := innerProducts4(query, first, second, third, fourth)
 	firstDistance, err = cosineDistanceFromProduct(firstProduct, queryMagnitude, firstMagnitude)
 	if err != nil {
 		return 0, 0, 0, 0, err

internal/floats/floats.go → internal/ailego/math/distance_utility.go

@@ -12,9 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package floats provides allocation-free float32 vector kernels. Callers must
-// pass equal, non-empty slices; validation belongs at the API boundary.
-package floats
+// The distance utility functions provide allocation-free float32 vector kernels.
+// Callers must pass equal, non-empty slices; validation belongs at the API boundary.
+package mathutil
 
 type binaryKernel func(left, right []float32) float32
 type batch2Kernel func(query, first, second []float32) (firstProduct, secondProduct float32)
@@ -28,42 +28,42 @@ var kernels = struct {
 	dot4     batch4Kernel
 	products productsKernel
 }{
-	l2:       l2SquaredScalar,
+	l2:       squaredEuclideanScalar,
 	dot:      innerProductScalar,
 	dot2:     innerProducts2Scalar,
 	dot4:     innerProducts4Scalar,
 	products: dotNormsScalar,
 }
 
-// L2Squared returns the squared Euclidean distance between left and right.
-func L2Squared(left, right []float32) float32 {
+// squaredEuclidean returns the squared Euclidean distance between left and right.
+func squaredEuclidean(left, right []float32) float32 {
 	return kernels.l2(left, right)
 }
 
-// InnerProduct returns the dot product of left and right.
-func InnerProduct(left, right []float32) float32 {
+// innerProduct returns the dot product of left and right.
+func innerProduct(left, right []float32) float32 {
 	return kernels.dot(left, right)
 }
 
-// InnerProducts2 computes the dot product of one query with two candidates in
+// innerProducts2 computes the dot product of one query with two candidates in
 // one pass. SIMD implementations reuse each loaded query block for both
 // candidates, which is the dominant HNSW one-to-many scoring pattern.
-func InnerProducts2(query, first, second []float32) (firstProduct, secondProduct float32) {
+func innerProducts2(query, first, second []float32) (firstProduct, secondProduct float32) {
 	return kernels.dot2(query, first, second)
 }
 
-// InnerProducts4 computes the dot product of one query with four candidates
+// innerProducts4 computes the dot product of one query with four candidates
 // in one pass, amortizing each query load across four independent products.
-func InnerProducts4(query, first, second, third, fourth []float32) (firstProduct, secondProduct, thirdProduct, fourthProduct float32) {
+func innerProducts4(query, first, second, third, fourth []float32) (firstProduct, secondProduct, thirdProduct, fourthProduct float32) {
 	return kernels.dot4(query, first, second, third, fourth)
 }
 
-// DotNorms computes the dot product and both squared norms in one pass.
-func DotNorms(left, right []float32) (dot, leftNorm, rightNorm float32) {
+// dotNorms computes the dot product and both squared norms in one pass.
+func dotNorms(left, right []float32) (dot, leftNorm, rightNorm float32) {
 	return kernels.products(left, right)
 }
 
-func l2SquaredScalar(left, right []float32) (sum float32) {
+func squaredEuclideanScalar(left, right []float32) (sum float32) {
 	for index, leftValue := range left {
 		difference := leftValue - right[index]
 		sum += difference * difference