Summary

• Replaces the external numkong dependency with in-tree C kernels for bf16 distance computation (dot product, L2, cosine, norm_l2)
• Follows the existing f16 kernel pattern: C source compiled via build.rs with per-architecture flags, runtime CPU dispatch via SIMD_SUPPORT
• Kernels are only enabled when the CPU supports the required instructions (NEON on aarch64, AVX2/AVX-512 on x86_64, LSX/LASX on loongarch64), with scalar fallback otherwise
• Gated behind the existing fp16kernels feature flag

Benchmark Results

Tested on two platforms with 1M x 1024-dim vectors:

Apple Silicon (M-series, NEON)

AMD Ryzen 5 4500 (Zen 2, AVX2)

Why the approach works

BF16-to-f32 conversion is a simple left-shift by 16 bits. The C kernels compiled with architecture-specific flags (-march=haswell, -mtune=apple-m1, etc.) plus -ffast-math and vectorization pragmas give the compiler more freedom to emit tight SIMD code than LLVM gets from the Rust scalar loops. ARM benefits more because the baseline Rust auto-vectorization was weaker there.

Files Changed

• New: rust/lance-linalg/src/simd/bf16.c — C kernels for dot, L2, cosine, norm_l2
• rust/lance-linalg/build.rs — compile bf16.c for each architecture
• rust/lance-linalg/src/distance/{dot,l2,cosine,norm_l2}.rs — runtime SIMD dispatch for bf16
• rust/lance-linalg/Cargo.toml — removed numkong dependency and feature
• rust/lance-linalg/benches/{dot,l2,cosine}.rs — removed numkong benchmark sections
• Deleted: scripts/bench_numkong.sh

Test plan

• cargo test -p lance-linalg --features fp16kernels — all bf16 tests pass (kernel path)
• cargo test -p lance-linalg — all bf16 tests pass (scalar fallback)
• cargo clippy -p lance-linalg --features fp16kernels --tests --benches -- -D warnings — clean
• Benchmarked on Apple Silicon (ARM NEON)
• Benchmarked on AMD Ryzen 5 4500 (x86_64 AVX2)
• To reproduce:

  git checkout HEAD~1
  TARGET_TIME=3 cargo bench -p lance-linalg --features fp16kernels --bench dot -- --save-baseline before "bf16"
  TARGET_TIME=3 cargo bench -p lance-linalg --features fp16kernels --bench norm_l2 -- --save-baseline before "bf16"
  git checkout -
  TARGET_TIME=3 cargo bench -p lance-linalg --features fp16kernels --bench dot -- --baseline before "bf16"
  TARGET_TIME=3 cargo bench -p lance-linalg --features fp16kernels --bench norm_l2 -- --baseline before "bf16"

🤖 Generated with Claude Code