perf(vector): replace vek with in-tree SIMD distance kernels

[matthewmcneely]

Description

This PR drops the github.com/viterin/vek dependency and replaces it with distance
kernels maintained in-tree, covering the three metrics the HNSW index actually uses: dot
product, cosine similarity, and euclidean distance. The vectorized path uses Go 1.27's new
experimental simd package, with a pure-Go fallback so nothing depends on the experiment
being enabled.

Why

vek only ships SIMD assembly for amd64, and it gates that on:

var UseAVX2 bool = cpu.X86.HasAVX2 && cpu.X86.HasFMA && runtime.GOOS != "darwin"

Everything outside that gate falls through to scalar Go loops. That includes
linux/arm64, which we officially support (Graviton, Ampere), plus every macOS dev box.
vek.Info() confirms it at run time, reporting Acceleration:false on arm64. So vector
search on ARM has been doing res += x[i]*y[i] one float at a time.

Upstream is not going to change that. viterin/vek#12
("Support for ARM64?") was closed on 2025-11-03 by the maintainer with "vek runs on ARM64
as pure Go code, but there are no plans to add SIMD acceleration". The last release was
v0.4.3 on 2025-08-14, the last commit was 2025-09-06 (docs only), and there are no open
issues or PRs. Their codegen pipeline (asm/_cpp → asm2avo.py → avo → asm/_avx2/*.s)
is amd64-only by construction, so NEON would need a whole new path. One fork
(orneryd/vek, January 2026) did add NEON, but through cgo and C++, never offered it
upstream, has no release tags, and its own README reports only 2x on vek32.Dot because
of the per-call cgo transition cost.

What's here

File	Purpose
tok/hnsw/kernels.go	The contract both implementations satisfy, and the reasoning behind it
tok/hnsw/kernels_simd.go	//go:build goexperiment.simd, uses the simd package
tok/hnsw/kernels_generic.go	//go:build !goexperiment.simd, unrolled scalar fallback
tok/hnsw/kernels_test.go	Parity, degenerate cases, allocation guard, benchmarks

The fallback is what keeps a plain go build ./... working without GOEXPERIMENT=simd,
which also matters for anyone consuming this as a library. Vector width is read at run
time rather than assumed, so one implementation covers 128-bit Neon and 256/512-bit AVX.

Also removes the chewxy/math32 and viterin/partial indirect dependencies, which only
came in through vek.

Numbers

darwin/arm64, 768 float32 dimensions, ns/op:

Metric	vek (before)	Fallback	GOEXPERIMENT=simd
dot product	575	157	108
euclidean	645	166	111
cosine	646	339	129

Worth noting the middle column: even without the experiment enabled, the fallback beats
vek by 1.9x on cosine and ~3.9x on dot and euclidean. The scalar loop vek falls back to is latency-bound on the
floating-point accumulator dependency chain, so independent partial sums buy most of the
win before any vectorization.

Go 1.27 codegen finding (would appreciate a second pair of eyes)

While measuring the fallback I hit something that is probably worth reporting upstream.
Go 1.27.0 appears to regress the s += a[i] * b[i] accumulation pattern on arm64 by
about 1.6x. At 768 dimensions the inline-indexed form compiles to a 253ns loop under
go1.27.0, where go1.26.5 compiled the identical source at ~155ns.

I tried to break the result and could not: interleaved runs to rule out thermal drift,
same GOARM64=v8.0 baseline on both toolchains, identical FP instruction mix (5x FMADDS
either way, so FMA fusion is happening), identical bounds-check counts under
-d=ssa/check_bce, and no dependence on GOEXPERIMENT=simd. Same instructions, different
schedule.

Hoisting the loads into locals first recovers it completely (158ns, matching go1.26.5):

a0, b0 := a[i], b[i]
s0 += a0 * b0

So the dot kernels here hoist deliberately, with a comment saying why. The euclidean and
cosine kernels were never affected, because they already hoist via their difference and
product temporaries. That asymmetry is what made the regression visible in the first
place: fallback dot was measuring slower than euclideanSq despite doing strictly less
work per element.

Two caveats before anyone files this: I only measured on Apple Silicon, so it needs
confirming on linux/arm64 and amd64. And it is worth grepping for other hot FP
accumulation loops in the tree that hit the same shape, since I only fixed the ones this
PR touches.

Behaviour change to flag

Results are no longer bit-identical to the previous implementation. Multiple
accumulators reassociate the partial sums, and the SIMD path additionally uses fused
multiply-add. Relative error against a float64 reference stays within ~4e-7, which is far
below the resolution at which ranking decisions change, but any test asserting exact float
equality on scores would need a tolerance. The tests here use InEpsilon/InDelta
throughout.

One other detail that looks incidental but is load-bearing: each kernel reslices b to
len(a). That halves the bounds checks (it eliminates the ones on b) and is worth ~11%.
It is not a length guard, though. A short subslice of a longer array reslices back
within capacity and silently reads past its length, so applyDistanceFunction remains the
only place length equality is actually enforced. The contract doc and a test both spell
this out so it does not become a latent surprise.

euclideanDistanceSq is renamed to euclideanDistance, since it always returned the
square-rooted value and distance_threshold compares against it in the metric domain.
Behaviour is unchanged; the name just now matches what it does.

Testing

• Parity against a float64 reference across 22 dimensions, weighted toward the tail and
  boundary cases the unrolls and partial vector loads have to handle (0, 1, 3, 5, 7, 15,
  17, 31, 33, 63, 65, and up).
• Degenerate inputs: empty and zero-magnitude vectors. Note vek used to panic on empty
  input via its checkNotEmpty; these return 0 (or NaN for cosine's 0/0) instead.
• Self-distance identities, since HNSW relies on them.
• A zero-allocation guard. The horizontal reduction uses a fixed-size stack array on
  purpose; sizing it from Len() would heap-allocate on every distance computation, in the
  hottest loop in vector search.
• Ran on go1.26.5, go1.27.0, and go1.27.0 with GOEXPERIMENT=simd, plus posting,
  schema, and types. build, vet, and gofmt clean.

Open question for reviewers

Should CI set GOEXPERIMENT=simd for release builds? Without it we get the fallback,
which is still a solid win over vek but leaves roughly 1.5x on the table. Against that,
the simd API is explicitly documented as not yet stable, so pinning release artifacts to
an experiment has its own cost. I do not have a strong opinion and would rather it be a
deliberate call than a default.

Follow-ups, deliberately not bundled here

• Drop the square root from euclidean. Ranking is unaffected, but it needs the two
  DistanceThreshold comparison sites to square the threshold instead.
• Pre-normalize vectors at insert time, which collapses cosine into a plain dot product
  (129ns → 108ns, using this PR's numbers). Changes stored-data semantics, so it deserves its own
  discussion.
• Profile a real similar_to workload. Every number above is kernel-level; the share of
  query latency that is actually distance math, as opposed to posting-list fetch, is still
  unmeasured, and that is what decides how much of this shows up end to end.

Checklist

• The PR title follows the
  Conventional Commits syntax, leading
  with fix:, feat:, chore:, ci:, etc.
• Code compiles correctly and linting (via trunk) passes locally
• Tests added for new functionality, or regression tests for bug fixes added as applicable

🤖 Generated with Claude Code

---

^{Need help on this PR? Tag @codesmith-bot with what you need. Autofix is disabled.}

[matthewmcneely]

@shiva-istari Took a pass at converting our vector distance calcs with the new SIMD features in go 1.27. Dropping viterin/vek in favor of our own distance functions. If GOEXPERIMENT=simd is set at build time, and Dgraph's running on amd64 we get SIMD processing! Additionally, the fallback scalar algos for non-simd are faster thanks to Claude.

Would be good to run this against your larger vector benchmarks for comparison.