From 388a1342899e013e4665909cedf5bedba8d67c9e Mon Sep 17 00:00:00 2001 From: mattthew Date: Fri, 21 Aug 2026 16:00:26 -0400 Subject: [PATCH] perf(vector): replace vek with in-tree SIMD distance kernels Drop github.com/viterin/vek in favour of distance kernels maintained here, for the three metrics HNSW actually uses. vek only ships SIMD assembly for amd64, gated on `HasAVX2 && HasFMA && GOOS != "darwin"`. Everything else falls through to scalar Go, which means linux/arm64 -- an officially supported platform -- and every macOS dev box ran vector search with no acceleration at all. Upstream has ruled out changing this: viterin/vek#12 was closed with "there are no plans to add SIMD acceleration", the last release was 2025-08-14, and vek's codegen pipeline targets amd64 by construction. Kernels live in two build-tagged files behind one contract, documented in kernels.go. kernels_simd.go uses the Go 1.27 simd package and needs GOEXPERIMENT=simd at build time; kernels_generic.go is an unrolled scalar fallback so a plain `go build ./...` keeps working without the experiment. The vector width is read at run time, so the same code covers 128-bit Neon and 256/512-bit AVX. Measured on darwin/arm64 at 768 float32 dimensions, ns/op: vek fallback simd dot 575 157 108 euclidean 645 166 111 cosine 646 339 129 Two details in the kernels carry their own weight. Reslicing b to len(a) eliminates the bounds checks on b and is worth ~11%; it is not a length guard, since a short subslice of a longer array reslices back within capacity, so applyDistanceFunction remains the only enforcement point. And the dot kernels hoist their loads into locals instead of indexing inline, because under go1.27.0 on arm64 the inline form generates a 1.6x slower loop (253ns vs 158ns) while the hoisted form matches go1.26.5. Results are no longer bit-identical to the previous implementation: multiple accumulators reassociate the partial sums, and the SIMD path adds fused multiply-add. Relative error against a float64 reference stays within ~4e-7, well below the resolution at which ranking changes. Tests compare with a tolerance and cover the tail cases, degenerate inputs, and allocation behaviour, which matters because the horizontal reduction has to stay off the heap in the hottest loop in vector search. euclideanDistanceSq is renamed to euclideanDistance: it always returned the square-rooted value, and distance_threshold compares against it in the metric domain. Behaviour is unchanged. Also removes the chewxy/math32 and viterin/partial indirect dependencies. Co-Authored-By: Claude Opus 5 (1M context) --- go.mod | 3 - go.sum | 6 - tok/hnsw/helper.go | 19 +-- tok/hnsw/kernels.go | 42 ++++++ tok/hnsw/kernels_generic.go | 165 ++++++++++++++++++++++ tok/hnsw/kernels_simd.go | 240 +++++++++++++++++++++++++++++++ tok/hnsw/kernels_test.go | 271 ++++++++++++++++++++++++++++++++++++ tok/hnsw/persistent_hnsw.go | 2 +- tok/index/helper_test.go | 10 -- 9 files changed, 730 insertions(+), 28 deletions(-) create mode 100644 tok/hnsw/kernels.go create mode 100644 tok/hnsw/kernels_generic.go create mode 100644 tok/hnsw/kernels_simd.go create mode 100644 tok/hnsw/kernels_test.go diff --git a/go.mod b/go.mod index 69fc69494a4..5d0e1c4dcce 100644 --- a/go.mod +++ b/go.mod @@ -47,7 +47,6 @@ require ( github.com/spf13/viper v1.21.0 github.com/stretchr/testify v1.11.1 github.com/twpayne/go-geom v1.6.1 - github.com/viterin/vek v0.4.3 github.com/xdg/scram v1.0.5 go.etcd.io/etcd/raft/v3 v3.5.29 go.opencensus.io v0.24.0 @@ -88,7 +87,6 @@ require ( github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/cenkalti/backoff/v5 v5.0.3 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect - github.com/chewxy/math32 v1.11.1 // indirect github.com/containerd/errdefs/pkg v0.3.0 // indirect github.com/containerd/log v0.1.0 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect @@ -157,7 +155,6 @@ require ( github.com/spf13/afero v1.15.0 // indirect github.com/subosito/gotenv v1.6.0 // indirect github.com/tinylib/msgp v1.6.4 // indirect - github.com/viterin/partial v1.1.0 // indirect github.com/xdg/stringprep v1.0.3 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect github.com/zeebo/xxh3 v1.1.0 // indirect diff --git a/go.sum b/go.sum index 93e7359ca19..b5250e00961 100644 --- a/go.sum +++ b/go.sum @@ -95,8 +95,6 @@ github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XL github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/chewxy/math32 v1.11.1 h1:b7PGHlp8KjylDoU8RrcEsRuGZhJuz8haxnKfuMMRqy8= -github.com/chewxy/math32 v1.11.1/go.mod h1:dOB2rcuFrCn6UHrze36WSLVPKtzPMRAQvBvUwkSsLqs= github.com/chromedp/cdproto v0.0.0-20230802225258-3cf4e6d46a89/go.mod h1:GKljq0VrfU4D5yc+2qA6OVr8pmO/MBbPEWqWQ/oqGEs= github.com/chromedp/chromedp v0.9.2/go.mod h1:LkSXJKONWTCHAfQasKFUZI+mxqS4tZqhmtGzzhLsnLs= github.com/chromedp/sysutil v1.0.0/go.mod h1:kgWmDdq8fTzXYcKIBqIYvRRTnYb9aNS9moAV0xufSww= @@ -524,10 +522,6 @@ github.com/tinylib/msgp v1.6.4 h1:mOwYbyYDLPj35mkA2BjjYejgJk9BuHxDdvRnb6v2ZcQ= github.com/tinylib/msgp v1.6.4/go.mod h1:RSp0LW9oSxFut3KzESt5Voq4GVWyS+PSulT77roAqEA= github.com/twpayne/go-geom v1.6.1 h1:iLE+Opv0Ihm/ABIcvQFGIiFBXd76oBIar9drAwHFhR4= github.com/twpayne/go-geom v1.6.1/go.mod h1:Kr+Nly6BswFsKM5sd31YaoWS5PeDDH2NftJTK7Gd028= -github.com/viterin/partial v1.1.0 h1:iH1l1xqBlapXsYzADS1dcbizg3iQUKTU1rbwkHv/80E= -github.com/viterin/partial v1.1.0/go.mod h1:oKGAo7/wylWkJTLrWX8n+f4aDPtQMQ6VG4dd2qur5QA= -github.com/viterin/vek v0.4.3 h1:cogdlNjd6EJYtNbmTN0lJCey2htrfSo1AHWpc6DVncQ= -github.com/viterin/vek v0.4.3/go.mod h1:A4JRAe8OvbhdzBL5ofzjBS0J29FyUrf95tQogvtHHUc= github.com/xdg/scram v1.0.5 h1:TuS0RFmt5Is5qm9Tm2SoD89OPqe4IRiFtyFY4iwWXsw= github.com/xdg/scram v1.0.5/go.mod h1:lB8K/P019DLNhemzwFU4jHLhdvlE6uDZjXFejJXr49I= github.com/xdg/stringprep v1.0.3 h1:cmL5Enob4W83ti/ZHuZLuKD/xqJfus4fVPwE+/BDm+4= diff --git a/tok/hnsw/helper.go b/tok/hnsw/helper.go index 7b157b20292..32a170ed02d 100644 --- a/tok/hnsw/helper.go +++ b/tok/hnsw/helper.go @@ -20,8 +20,6 @@ import ( c "github.com/dgraph-io/dgraph/v25/tok/constraints" "github.com/dgraph-io/dgraph/v25/tok/index" "github.com/pkg/errors" - "github.com/viterin/vek" - "github.com/viterin/vek/vek32" ) const ( @@ -96,20 +94,25 @@ func applyDistanceFunction[T c.Float](a, b []T, floatBits int, funcName string, // function, hence it takes in a floatBits parameter, // but doesn't actually use it. func dotProduct[T c.Float](a, b []T, floatBits int) (T, error) { - return applyDistanceFunction(a, b, floatBits, "dot product", vek32.Dot, vek.Dot) + return applyDistanceFunction(a, b, floatBits, "dot product", dotF32, dotF64) } // This needs to implement signature of SimilarityType[T].distanceScore // function, hence it takes in a floatBits parameter. func cosineSimilarity[T c.Float](a, b []T, floatBits int) (T, error) { - return applyDistanceFunction(a, b, floatBits, "cosine distance", vek32.CosineSimilarity, vek.CosineSimilarity) + return applyDistanceFunction(a, b, floatBits, "cosine distance", cosineSimF32, cosineSimF64) } +// euclideanDistance returns the metric-domain euclidean distance, i.e. square rooted. +// Callers rely on that: DistanceThreshold in SearchWithOptions compares against it +// directly. The square root is monotonic, so dropping it would leave ranking unchanged +// but would require squaring the threshold at those comparison sites. +// // This needs to implement signature of SimilarityType[T].distanceScore // function, hence it takes in a floatBits parameter, // but doesn't actually use it. -func euclideanDistanceSq[T c.Float](a, b []T, floatBits int) (T, error) { - return applyDistanceFunction(a, b, floatBits, "euclidean distance", vek32.Distance, vek.Distance) +func euclideanDistance[T c.Float](a, b []T, floatBits int) (T, error) { + return applyDistanceFunction(a, b, floatBits, "euclidean distance", euclideanF32, euclideanF64) } // Used for distance, since shorter distance is better @@ -217,7 +220,7 @@ type SimilarityType[T c.Float] struct { func GetSimType[T c.Float](indexType string, floatBits int) SimilarityType[T] { switch { case indexType == Euclidean: - return SimilarityType[T]{indexType: Euclidean, distanceScore: euclideanDistanceSq[T], + return SimilarityType[T]{indexType: Euclidean, distanceScore: euclideanDistance[T], insortHeap: insortPersistentHeapAscending[T], isBetterScore: isBetterScoreForDistance[T], isSimilarityMetric: false} case indexType == Cosine: @@ -229,7 +232,7 @@ func GetSimType[T c.Float](indexType string, floatBits int) SimilarityType[T] { insortHeap: insortPersistentHeapDescending[T], isBetterScore: isBetterScoreForSimilarity[T], isSimilarityMetric: true} default: - return SimilarityType[T]{indexType: Euclidean, distanceScore: euclideanDistanceSq[T], + return SimilarityType[T]{indexType: Euclidean, distanceScore: euclideanDistance[T], insortHeap: insortPersistentHeapAscending[T], isBetterScore: isBetterScoreForDistance[T], isSimilarityMetric: false} } diff --git a/tok/hnsw/kernels.go b/tok/hnsw/kernels.go new file mode 100644 index 00000000000..a08debec458 --- /dev/null +++ b/tok/hnsw/kernels.go @@ -0,0 +1,42 @@ +/* + * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package hnsw + +// This file documents the distance-kernel contract shared by the two build-tagged +// implementations in kernels_simd.go (//go:build goexperiment.simd) and +// kernels_generic.go (//go:build !goexperiment.simd). +// +// Each implementation provides, for both float32 and float64: +// +// dotF32/dotF64 - sum(a[i]*b[i]) +// euclideanSqF32/euclideanSqF64 - sum((a[i]-b[i])^2), NOT square-rooted +// euclideanF32/euclideanF64 - sqrt of the above, the metric-domain distance +// cosineSimF32/cosineSimF64 - dot(a,b) / sqrt(dot(a,a)*dot(b,b)) +// +// Contract for every kernel: +// +// - Callers guarantee len(a) == len(b), and applyDistanceFunction enforces it by +// returning an error before any kernel is reached. Each kernel then reslices b to +// len(a) so the compiler can eliminate bounds checks on b inside the loop. Note +// that this reslice is not itself a length check: b[:len(a)] succeeds whenever +// cap(b) >= len(a), so passing a short subslice of a longer array reads the +// elements beyond its length rather than panicking. The guard is the wrapper, not +// the kernel. +// - Zero-length input yields 0 for dot and euclidean, and NaN for cosine (0/0). +// This matches the behaviour of a zero vector and is strictly safer than the +// vek implementation this replaced, which panicked on empty input. +// - Results are not bit-identical to a naive left-to-right summation. Both +// implementations use multiple independent accumulators, and the SIMD path +// additionally uses fused multiply-add, so partial sums are reassociated and +// rounded differently. Relative error against a float64 reference stays within +// a few ULP of float32 (~4e-7 measured at 768 dimensions), which is far below +// the resolution at which ranking decisions differ. Tests must compare with a +// tolerance rather than for exact equality. +// +// Why multiple accumulators: the obvious `sum += a[i]*b[i]` loop is bound by the +// latency of the floating-point add dependency chain, not by throughput. Splitting +// into independent partial sums lets the CPU keep several adds in flight, which is +// worth roughly 2x on its own before any vectorisation. diff --git a/tok/hnsw/kernels_generic.go b/tok/hnsw/kernels_generic.go new file mode 100644 index 00000000000..9a6835e9034 --- /dev/null +++ b/tok/hnsw/kernels_generic.go @@ -0,0 +1,165 @@ +//go:build !goexperiment.simd + +/* + * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package hnsw + +import "math" + +// Portable distance kernels used when the build does not enable GOEXPERIMENT=simd. +// See kernels.go for the contract these must satisfy. +// +// These are deliberately unrolled into independent accumulators rather than written +// as the obvious single-accumulator loop. Dot product and squared euclidean use a +// 4-way unroll; cosine already has three independent chains (dot, |a|, |b|) so it +// uses a 2-way unroll, giving six, which is enough to saturate the FP pipeline +// without risking register spills. + +func dotF32(a, b []float32) float32 { + b = b[:len(a)] + var s0, s1, s2, s3 float32 + i := 0 + for ; i+4 <= len(a); i += 4 { + // The loads are hoisted into locals rather than indexed inline in the + // accumulation. That is not cosmetic: under go1.27.0 on arm64 the inline form + // generates a 1.6x slower loop (253ns vs 158ns at 768 dimensions), while the + // hoisted form matches go1.26.5. The other two kernels here happen to hoist + // already, via their difference and product temporaries. + a0, b0 := a[i], b[i] + a1, b1 := a[i+1], b[i+1] + a2, b2 := a[i+2], b[i+2] + a3, b3 := a[i+3], b[i+3] + s0 += a0 * b0 + s1 += a1 * b1 + s2 += a2 * b2 + s3 += a3 * b3 + } + for ; i < len(a); i++ { + s0 += a[i] * b[i] + } + return (s0 + s1) + (s2 + s3) +} + +func dotF64(a, b []float64) float64 { + b = b[:len(a)] + var s0, s1, s2, s3 float64 + i := 0 + for ; i+4 <= len(a); i += 4 { + // The loads are hoisted into locals rather than indexed inline in the + // accumulation. That is not cosmetic: under go1.27.0 on arm64 the inline form + // generates a 1.6x slower loop (253ns vs 158ns at 768 dimensions), while the + // hoisted form matches go1.26.5. The other two kernels here happen to hoist + // already, via their difference and product temporaries. + a0, b0 := a[i], b[i] + a1, b1 := a[i+1], b[i+1] + a2, b2 := a[i+2], b[i+2] + a3, b3 := a[i+3], b[i+3] + s0 += a0 * b0 + s1 += a1 * b1 + s2 += a2 * b2 + s3 += a3 * b3 + } + for ; i < len(a); i++ { + s0 += a[i] * b[i] + } + return (s0 + s1) + (s2 + s3) +} + +func euclideanSqF32(a, b []float32) float32 { + b = b[:len(a)] + var s0, s1, s2, s3 float32 + i := 0 + for ; i+4 <= len(a); i += 4 { + d0 := a[i] - b[i] + d1 := a[i+1] - b[i+1] + d2 := a[i+2] - b[i+2] + d3 := a[i+3] - b[i+3] + s0 += d0 * d0 + s1 += d1 * d1 + s2 += d2 * d2 + s3 += d3 * d3 + } + for ; i < len(a); i++ { + d := a[i] - b[i] + s0 += d * d + } + return (s0 + s1) + (s2 + s3) +} + +func euclideanSqF64(a, b []float64) float64 { + b = b[:len(a)] + var s0, s1, s2, s3 float64 + i := 0 + for ; i+4 <= len(a); i += 4 { + d0 := a[i] - b[i] + d1 := a[i+1] - b[i+1] + d2 := a[i+2] - b[i+2] + d3 := a[i+3] - b[i+3] + s0 += d0 * d0 + s1 += d1 * d1 + s2 += d2 * d2 + s3 += d3 * d3 + } + for ; i < len(a); i++ { + d := a[i] - b[i] + s0 += d * d + } + return (s0 + s1) + (s2 + s3) +} + +func euclideanF32(a, b []float32) float32 { + return float32(math.Sqrt(float64(euclideanSqF32(a, b)))) +} + +func euclideanF64(a, b []float64) float64 { + return math.Sqrt(euclideanSqF64(a, b)) +} + +func cosineSimF32(a, b []float32) float32 { + b = b[:len(a)] + var d0, d1, x0, x1, y0, y1 float32 + i := 0 + for ; i+2 <= len(a); i += 2 { + av0, bv0 := a[i], b[i] + av1, bv1 := a[i+1], b[i+1] + d0 += av0 * bv0 + d1 += av1 * bv1 + x0 += av0 * av0 + x1 += av1 * av1 + y0 += bv0 * bv0 + y1 += bv1 * bv1 + } + for ; i < len(a); i++ { + d0 += a[i] * b[i] + x0 += a[i] * a[i] + y0 += b[i] * b[i] + } + dot, na, nb := d0+d1, x0+x1, y0+y1 + return dot / float32(math.Sqrt(float64(na)*float64(nb))) +} + +func cosineSimF64(a, b []float64) float64 { + b = b[:len(a)] + var d0, d1, x0, x1, y0, y1 float64 + i := 0 + for ; i+2 <= len(a); i += 2 { + av0, bv0 := a[i], b[i] + av1, bv1 := a[i+1], b[i+1] + d0 += av0 * bv0 + d1 += av1 * bv1 + x0 += av0 * av0 + x1 += av1 * av1 + y0 += bv0 * bv0 + y1 += bv1 * bv1 + } + for ; i < len(a); i++ { + d0 += a[i] * b[i] + x0 += a[i] * a[i] + y0 += b[i] * b[i] + } + dot, na, nb := d0+d1, x0+x1, y0+y1 + return dot / math.Sqrt(na*nb) +} diff --git a/tok/hnsw/kernels_simd.go b/tok/hnsw/kernels_simd.go new file mode 100644 index 00000000000..becc4a7d1d2 --- /dev/null +++ b/tok/hnsw/kernels_simd.go @@ -0,0 +1,240 @@ +//go:build goexperiment.simd + +/* + * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package hnsw + +import ( + "math" + "simd" +) + +// Vectorised distance kernels, used when the build enables GOEXPERIMENT=simd. +// See kernels.go for the contract these must satisfy. +// +// The vector width is chosen once per program execution by the simd package: 128 bits +// on arm64 (Neon) and wasm, and 256 or 512 bits on amd64 depending on AVX2/AVX-512 +// support. Nothing here assumes a width; Len() is queried at run time and the tail +// is handled with the zero-filling Part loads. +// +// Dot product and squared euclidean use four independent vector accumulators, so the +// effective instruction-level parallelism is four times the lane count. Cosine uses two, +// across its three chains, for six vector accumulators total. +// +// maxF32Lanes and maxF64Lanes bound the horizontal-reduction scratch buffers so they +// stack-allocate. They cover the widest vector the simd package can select, 512 bits. + +const ( + maxF32Lanes = 16 + maxF64Lanes = 8 +) + +// hsumF32 reduces a vector accumulator to a scalar. The scratch array is fixed size so +// it does not escape; a make() sized from Len() would heap-allocate on every distance +// computation, which is the hottest loop in vector search. +func hsumF32(v simd.Float32s) float32 { + var buf [maxF32Lanes]float32 + n := v.Len() + v.Store(buf[:n]) + var s float32 + for _, f := range buf[:n] { + s += f + } + return s +} + +func hsumF64(v simd.Float64s) float64 { + var buf [maxF64Lanes]float64 + n := v.Len() + v.Store(buf[:n]) + var s float64 + for _, f := range buf[:n] { + s += f + } + return s +} + +func dotF32(a, b []float32) float32 { + b = b[:len(a)] + var acc0, acc1, acc2, acc3 simd.Float32s + w := acc0.Len() + i := 0 + for ; i+4*w <= len(a); i += 4 * w { + acc0 = simd.LoadFloat32s(a[i:]).MulAdd(simd.LoadFloat32s(b[i:]), acc0) + acc1 = simd.LoadFloat32s(a[i+w:]).MulAdd(simd.LoadFloat32s(b[i+w:]), acc1) + acc2 = simd.LoadFloat32s(a[i+2*w:]).MulAdd(simd.LoadFloat32s(b[i+2*w:]), acc2) + acc3 = simd.LoadFloat32s(a[i+3*w:]).MulAdd(simd.LoadFloat32s(b[i+3*w:]), acc3) + } + for ; i+w <= len(a); i += w { + acc0 = simd.LoadFloat32s(a[i:]).MulAdd(simd.LoadFloat32s(b[i:]), acc0) + } + if i < len(a) { + av, _ := simd.LoadFloat32sPart(a[i:]) + bv, _ := simd.LoadFloat32sPart(b[i:]) + acc0 = av.MulAdd(bv, acc0) + } + return hsumF32(acc0.Add(acc1).Add(acc2.Add(acc3))) +} + +func dotF64(a, b []float64) float64 { + b = b[:len(a)] + var acc0, acc1, acc2, acc3 simd.Float64s + w := acc0.Len() + i := 0 + for ; i+4*w <= len(a); i += 4 * w { + acc0 = simd.LoadFloat64s(a[i:]).MulAdd(simd.LoadFloat64s(b[i:]), acc0) + acc1 = simd.LoadFloat64s(a[i+w:]).MulAdd(simd.LoadFloat64s(b[i+w:]), acc1) + acc2 = simd.LoadFloat64s(a[i+2*w:]).MulAdd(simd.LoadFloat64s(b[i+2*w:]), acc2) + acc3 = simd.LoadFloat64s(a[i+3*w:]).MulAdd(simd.LoadFloat64s(b[i+3*w:]), acc3) + } + for ; i+w <= len(a); i += w { + acc0 = simd.LoadFloat64s(a[i:]).MulAdd(simd.LoadFloat64s(b[i:]), acc0) + } + if i < len(a) { + av, _ := simd.LoadFloat64sPart(a[i:]) + bv, _ := simd.LoadFloat64sPart(b[i:]) + acc0 = av.MulAdd(bv, acc0) + } + return hsumF64(acc0.Add(acc1).Add(acc2.Add(acc3))) +} + +func euclideanSqF32(a, b []float32) float32 { + b = b[:len(a)] + var acc0, acc1, acc2, acc3 simd.Float32s + w := acc0.Len() + i := 0 + for ; i+4*w <= len(a); i += 4 * w { + d0 := simd.LoadFloat32s(a[i:]).Sub(simd.LoadFloat32s(b[i:])) + d1 := simd.LoadFloat32s(a[i+w:]).Sub(simd.LoadFloat32s(b[i+w:])) + d2 := simd.LoadFloat32s(a[i+2*w:]).Sub(simd.LoadFloat32s(b[i+2*w:])) + d3 := simd.LoadFloat32s(a[i+3*w:]).Sub(simd.LoadFloat32s(b[i+3*w:])) + acc0 = d0.MulAdd(d0, acc0) + acc1 = d1.MulAdd(d1, acc1) + acc2 = d2.MulAdd(d2, acc2) + acc3 = d3.MulAdd(d3, acc3) + } + for ; i+w <= len(a); i += w { + d := simd.LoadFloat32s(a[i:]).Sub(simd.LoadFloat32s(b[i:])) + acc0 = d.MulAdd(d, acc0) + } + if i < len(a) { + av, _ := simd.LoadFloat32sPart(a[i:]) + bv, _ := simd.LoadFloat32sPart(b[i:]) + d := av.Sub(bv) + acc0 = d.MulAdd(d, acc0) + } + return hsumF32(acc0.Add(acc1).Add(acc2.Add(acc3))) +} + +func euclideanSqF64(a, b []float64) float64 { + b = b[:len(a)] + var acc0, acc1, acc2, acc3 simd.Float64s + w := acc0.Len() + i := 0 + for ; i+4*w <= len(a); i += 4 * w { + d0 := simd.LoadFloat64s(a[i:]).Sub(simd.LoadFloat64s(b[i:])) + d1 := simd.LoadFloat64s(a[i+w:]).Sub(simd.LoadFloat64s(b[i+w:])) + d2 := simd.LoadFloat64s(a[i+2*w:]).Sub(simd.LoadFloat64s(b[i+2*w:])) + d3 := simd.LoadFloat64s(a[i+3*w:]).Sub(simd.LoadFloat64s(b[i+3*w:])) + acc0 = d0.MulAdd(d0, acc0) + acc1 = d1.MulAdd(d1, acc1) + acc2 = d2.MulAdd(d2, acc2) + acc3 = d3.MulAdd(d3, acc3) + } + for ; i+w <= len(a); i += w { + d := simd.LoadFloat64s(a[i:]).Sub(simd.LoadFloat64s(b[i:])) + acc0 = d.MulAdd(d, acc0) + } + if i < len(a) { + av, _ := simd.LoadFloat64sPart(a[i:]) + bv, _ := simd.LoadFloat64sPart(b[i:]) + d := av.Sub(bv) + acc0 = d.MulAdd(d, acc0) + } + return hsumF64(acc0.Add(acc1).Add(acc2.Add(acc3))) +} + +func euclideanF32(a, b []float32) float32 { + return float32(math.Sqrt(float64(euclideanSqF32(a, b)))) +} + +func euclideanF64(a, b []float64) float64 { + return math.Sqrt(euclideanSqF64(a, b)) +} + +func cosineSimF32(a, b []float32) float32 { + b = b[:len(a)] + var dot0, dot1, na0, na1, nb0, nb1 simd.Float32s + w := dot0.Len() + i := 0 + for ; i+2*w <= len(a); i += 2 * w { + av0 := simd.LoadFloat32s(a[i:]) + bv0 := simd.LoadFloat32s(b[i:]) + av1 := simd.LoadFloat32s(a[i+w:]) + bv1 := simd.LoadFloat32s(b[i+w:]) + dot0 = av0.MulAdd(bv0, dot0) + dot1 = av1.MulAdd(bv1, dot1) + na0 = av0.MulAdd(av0, na0) + na1 = av1.MulAdd(av1, na1) + nb0 = bv0.MulAdd(bv0, nb0) + nb1 = bv1.MulAdd(bv1, nb1) + } + for ; i+w <= len(a); i += w { + av := simd.LoadFloat32s(a[i:]) + bv := simd.LoadFloat32s(b[i:]) + dot0 = av.MulAdd(bv, dot0) + na0 = av.MulAdd(av, na0) + nb0 = bv.MulAdd(bv, nb0) + } + if i < len(a) { + av, _ := simd.LoadFloat32sPart(a[i:]) + bv, _ := simd.LoadFloat32sPart(b[i:]) + dot0 = av.MulAdd(bv, dot0) + na0 = av.MulAdd(av, na0) + nb0 = bv.MulAdd(bv, nb0) + } + dot := hsumF32(dot0.Add(dot1)) + na := hsumF32(na0.Add(na1)) + nb := hsumF32(nb0.Add(nb1)) + return dot / float32(math.Sqrt(float64(na)*float64(nb))) +} + +func cosineSimF64(a, b []float64) float64 { + b = b[:len(a)] + var dot0, dot1, na0, na1, nb0, nb1 simd.Float64s + w := dot0.Len() + i := 0 + for ; i+2*w <= len(a); i += 2 * w { + av0 := simd.LoadFloat64s(a[i:]) + bv0 := simd.LoadFloat64s(b[i:]) + av1 := simd.LoadFloat64s(a[i+w:]) + bv1 := simd.LoadFloat64s(b[i+w:]) + dot0 = av0.MulAdd(bv0, dot0) + dot1 = av1.MulAdd(bv1, dot1) + na0 = av0.MulAdd(av0, na0) + na1 = av1.MulAdd(av1, na1) + nb0 = bv0.MulAdd(bv0, nb0) + nb1 = bv1.MulAdd(bv1, nb1) + } + for ; i+w <= len(a); i += w { + av := simd.LoadFloat64s(a[i:]) + bv := simd.LoadFloat64s(b[i:]) + dot0 = av.MulAdd(bv, dot0) + na0 = av.MulAdd(av, na0) + nb0 = bv.MulAdd(bv, nb0) + } + if i < len(a) { + av, _ := simd.LoadFloat64sPart(a[i:]) + bv, _ := simd.LoadFloat64sPart(b[i:]) + dot0 = av.MulAdd(bv, dot0) + na0 = av.MulAdd(av, na0) + nb0 = bv.MulAdd(bv, nb0) + } + dot := hsumF64(dot0.Add(dot1)) + na := hsumF64(na0.Add(na1)) + nb := hsumF64(nb0.Add(nb1)) + return dot / math.Sqrt(na*nb) +} diff --git a/tok/hnsw/kernels_test.go b/tok/hnsw/kernels_test.go new file mode 100644 index 00000000000..03f11cca456 --- /dev/null +++ b/tok/hnsw/kernels_test.go @@ -0,0 +1,271 @@ +/* + * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +package hnsw + +import ( + "fmt" + "math" + "math/rand" + "testing" + + "github.com/stretchr/testify/require" +) + +// These tests exercise whichever kernel implementation the build selected, so they run +// identically with and without GOEXPERIMENT=simd. Every assertion is a tolerance +// comparison against a float64 reference: the kernels reassociate partial sums across +// independent accumulators and, on the SIMD path, use fused multiply-add, so results are +// deliberately not bit-identical to a naive summation. See kernels.go. + +// relTol is a generous bound on float32 accumulation error over the dimensions tested. +// Measured error at 768 dimensions is ~4e-7; ranking decisions are unaffected well +// before this threshold. +const relTol = 1e-5 + +// kernelDims covers the tail cases that the unrolls and the partial vector loads have to +// handle: shorter than one vector, not a multiple of the unroll factor, and exactly on +// the boundary for widths up to 512 bits. +var kernelDims = []int{0, 1, 2, 3, 4, 5, 7, 8, 15, 16, 17, 31, 32, 33, 63, 64, 65, 128, 384, 768, 1000, 1536} + +func randVec32(n int, seed int64) []float32 { + r := rand.New(rand.NewSource(seed)) + v := make([]float32, n) + for i := range v { + v[i] = r.Float32()*2 - 1 + } + return v +} + +func randVec64(n int, seed int64) []float64 { + r := rand.New(rand.NewSource(seed)) + v := make([]float64, n) + for i := range v { + v[i] = r.Float64()*2 - 1 + } + return v +} + +// refDot, refEuclideanSq and refCosine accumulate in float64 regardless of input width, +// giving a reference the kernels can be measured against. +func refDot(a, b []float64) float64 { + var s float64 + for i := range a { + s += a[i] * b[i] + } + return s +} + +func refEuclideanSq(a, b []float64) float64 { + var s float64 + for i := range a { + d := a[i] - b[i] + s += d * d + } + return s +} + +func refCosine(a, b []float64) float64 { + return refDot(a, b) / math.Sqrt(refDot(a, a)*refDot(b, b)) +} + +func widen(v []float32) []float64 { + out := make([]float64, len(v)) + for i, f := range v { + out[i] = float64(f) + } + return out +} + +// requireClose compares against the reference with a relative tolerance, falling back to +// an absolute bound when the reference is near zero. +func requireClose(t *testing.T, got, want float64, what string) { + t.Helper() + if math.Abs(want) < 1e-6 { + require.InDelta(t, want, got, 1e-6, what) + return + } + require.InEpsilon(t, want, got, relTol, what) +} + +func TestKernelsFloat32(t *testing.T) { + for _, n := range kernelDims { + t.Run(fmt.Sprintf("d=%d", n), func(t *testing.T) { + a, b := randVec32(n, 1), randVec32(n, 2) + a64, b64 := widen(a), widen(b) + + requireClose(t, float64(dotF32(a, b)), refDot(a64, b64), "dot") + requireClose(t, float64(euclideanSqF32(a, b)), refEuclideanSq(a64, b64), "euclideanSq") + requireClose(t, float64(euclideanF32(a, b)), math.Sqrt(refEuclideanSq(a64, b64)), "euclidean") + + // Cosine of a zero-length or zero-magnitude vector is 0/0. Skip the + // comparison there; TestKernelsDegenerate covers it explicitly. + if n > 0 { + requireClose(t, float64(cosineSimF32(a, b)), refCosine(a64, b64), "cosine") + } + }) + } +} + +func TestKernelsFloat64(t *testing.T) { + for _, n := range kernelDims { + t.Run(fmt.Sprintf("d=%d", n), func(t *testing.T) { + a, b := randVec64(n, 3), randVec64(n, 4) + + requireClose(t, dotF64(a, b), refDot(a, b), "dot") + requireClose(t, euclideanSqF64(a, b), refEuclideanSq(a, b), "euclideanSq") + requireClose(t, euclideanF64(a, b), math.Sqrt(refEuclideanSq(a, b)), "euclidean") + + if n > 0 { + requireClose(t, cosineSimF64(a, b), refCosine(a, b), "cosine") + } + }) + } +} + +// TestKernelsSelfDistance pins the identities the HNSW search relies on: a vector is at +// distance zero from itself and has cosine similarity 1. +func TestKernelsSelfDistance(t *testing.T) { + for _, n := range []int{1, 8, 17, 768} { + a := randVec32(n, 5) + require.Zero(t, euclideanSqF32(a, a), "euclideanSq(a,a) must be exactly 0") + require.Zero(t, euclideanF32(a, a), "euclidean(a,a) must be exactly 0") + require.InDelta(t, 1.0, float64(cosineSimF32(a, a)), relTol, "cosine(a,a)") + + d := randVec64(n, 6) + require.Zero(t, euclideanSqF64(d, d), "euclideanSq(d,d) must be exactly 0") + require.InDelta(t, 1.0, cosineSimF64(d, d), relTol, "cosine(d,d)") + } +} + +// TestKernelsDegenerate documents the empty and zero-vector behaviour promised in +// kernels.go. The vek implementation this replaced panicked on empty input. +func TestKernelsDegenerate(t *testing.T) { + var empty32 []float32 + require.Zero(t, dotF32(empty32, empty32)) + require.Zero(t, euclideanSqF32(empty32, empty32)) + require.Zero(t, euclideanF32(empty32, empty32)) + require.True(t, math.IsNaN(float64(cosineSimF32(empty32, empty32))), "cosine of empty is 0/0") + + zeros := make([]float32, 16) + require.Zero(t, dotF32(zeros, zeros)) + require.Zero(t, euclideanSqF32(zeros, zeros)) + require.True(t, math.IsNaN(float64(cosineSimF32(zeros, zeros))), "cosine of zero vector is 0/0") + + var empty64 []float64 + require.Zero(t, dotF64(empty64, empty64)) + require.Zero(t, euclideanSqF64(empty64, empty64)) + require.True(t, math.IsNaN(cosineSimF64(empty64, empty64)), "cosine of empty is 0/0") +} + +// TestKernelsMismatchedLength pins the real boundary behaviour, which is subtler than it +// looks. Kernels reslice b to len(a) for bounds-check elimination, so a genuinely short b +// panics, but a short *subslice of a longer array* does not: the reslice stays within +// capacity and silently reads the elements past b's length. That is why the length guard +// lives in applyDistanceFunction rather than in the kernels, and why nothing should call +// a kernel directly. See TestDistanceScoreLengthMismatch for the enforced path. +func TestKernelsMismatchedLength(t *testing.T) { + a := randVec32(16, 7) + + // Insufficient capacity: the reslice panics. + short := randVec32(8, 13) + require.Panics(t, func() { dotF32(a, short) }) + require.Panics(t, func() { euclideanSqF32(a, short) }) + require.Panics(t, func() { cosineSimF32(a, short) }) + + // Sufficient capacity: the reslice succeeds and reads past len(b). Documented here + // so the behaviour is deliberate rather than a latent surprise. + sub := a[:8] + require.NotPanics(t, func() { dotF32(a, sub) }) + require.Equal(t, dotF32(a, a), dotF32(a, sub), + "a short subslice is silently widened back to the full array") +} + +// TestDistanceScoreLengthMismatch covers the wrapper that guards the kernels. +func TestDistanceScoreLengthMismatch(t *testing.T) { + a := randVec32(16, 8) + for name, fn := range map[string]func(a, b []float32, floatBits int) (float32, error){ + "dot": dotProduct[float32], + "cosine": cosineSimilarity[float32], + "euclidean": euclideanDistance[float32], + } { + _, err := fn(a, a[:8], 32) + require.Error(t, err, name) + require.Contains(t, err.Error(), "different lengths", name) + } +} + +// TestKernelsNoAllocs guards against a regression that would matter: the horizontal +// reduction on the SIMD path uses a fixed-size stack array precisely so the hottest loop +// in vector search stays allocation free. +func TestKernelsNoAllocs(t *testing.T) { + a, b := randVec32(768, 9), randVec32(768, 10) + c, d := randVec64(768, 11), randVec64(768, 12) + allocs := testing.AllocsPerRun(100, func() { + _ = dotF32(a, b) + _ = euclideanSqF32(a, b) + _ = cosineSimF32(a, b) + _ = dotF64(c, d) + _ = euclideanSqF64(c, d) + _ = cosineSimF64(c, d) + }) + require.Zero(t, allocs, "distance kernels must not allocate") +} + +var benchDims = []int{384, 768, 1536} + +func BenchmarkKernelsFloat32(b *testing.B) { + for _, n := range benchDims { + x, y := randVec32(n, 1), randVec32(n, 2) + b.Run(fmt.Sprintf("dot/d=%d", n), func(b *testing.B) { + b.SetBytes(int64(n * 4 * 2)) + for b.Loop() { + _ = dotF32(x, y) + } + }) + b.Run(fmt.Sprintf("euclideanSq/d=%d", n), func(b *testing.B) { + b.SetBytes(int64(n * 4 * 2)) + for b.Loop() { + _ = euclideanSqF32(x, y) + } + }) + b.Run(fmt.Sprintf("euclidean/d=%d", n), func(b *testing.B) { + b.SetBytes(int64(n * 4 * 2)) + for b.Loop() { + _ = euclideanF32(x, y) + } + }) + b.Run(fmt.Sprintf("cosine/d=%d", n), func(b *testing.B) { + b.SetBytes(int64(n * 4 * 2)) + for b.Loop() { + _ = cosineSimF32(x, y) + } + }) + } +} + +func BenchmarkKernelsFloat64(b *testing.B) { + for _, n := range benchDims { + x, y := randVec64(n, 1), randVec64(n, 2) + b.Run(fmt.Sprintf("dot/d=%d", n), func(b *testing.B) { + b.SetBytes(int64(n * 8 * 2)) + for b.Loop() { + _ = dotF64(x, y) + } + }) + b.Run(fmt.Sprintf("euclideanSq/d=%d", n), func(b *testing.B) { + b.SetBytes(int64(n * 8 * 2)) + for b.Loop() { + _ = euclideanSqF64(x, y) + } + }) + b.Run(fmt.Sprintf("cosine/d=%d", n), func(b *testing.B) { + b.SetBytes(int64(n * 8 * 2)) + for b.Loop() { + _ = cosineSimF64(x, y) + } + }) + } +} diff --git a/tok/hnsw/persistent_hnsw.go b/tok/hnsw/persistent_hnsw.go index 864e7e98637..71632d9c088 100644 --- a/tok/hnsw/persistent_hnsw.go +++ b/tok/hnsw/persistent_hnsw.go @@ -105,7 +105,7 @@ func (ph *persistentHNSW[T]) applyOptions(o opt.Options) error { } ph.simType = okSimType } else { - ph.simType = SimilarityType[T]{indexType: Euclidean, distanceScore: euclideanDistanceSq[T], + ph.simType = SimilarityType[T]{indexType: Euclidean, distanceScore: euclideanDistance[T], insortHeap: insortPersistentHeapAscending[T], isBetterScore: isBetterScoreForDistance[T], isSimilarityMetric: false} } diff --git a/tok/index/helper_test.go b/tok/index/helper_test.go index bfff6c0405c..f11c79a5d38 100644 --- a/tok/index/helper_test.go +++ b/tok/index/helper_test.go @@ -17,7 +17,6 @@ import ( "github.com/dgraph-io/dgraph/v25/protos/pb" c "github.com/dgraph-io/dgraph/v25/tok/constraints" - "github.com/viterin/vek/vek32" "google.golang.org/protobuf/proto" ) @@ -292,15 +291,6 @@ func BenchmarkDotProduct(b *testing.B) { b.Skip() } - b.Run(fmt.Sprintf("vek:size=%d", len(data)), - func(b *testing.B) { - temp := make([]float32, num) - BytesAsFloatArray[float32](data, &temp, 32) - for k := 0; k < b.N; k++ { - vek32.Dot(temp, temp) - } - }) - b.Run(fmt.Sprintf("dotProduct:size=%d", len(data)), func(b *testing.B) {