dearsky/wifi-densepose
1
0
Fork 0
Code Issues 20 Pull Requests 5 Actions Packages Projects Releases 1 Wiki Activity

Squashed 'vendor/ruvector/' content from commit b64c2172

Browse Source 
git-subtree-dir: vendor/ruvector
git-subtree-split: b64c21726f2bb37286d9ee36a7869fef60cc6900
This commit is contained in:
ruv
2026-02-28 14:39:40 -05:00
commit d803bfe2b1
7854 changed files with 3522914 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

605
crates/ruvector-postgres/docs/SIMD_OPTIMIZATION.md Normal file
View File

@@ -0,0 +1,605 @@
# SIMD Optimization in RuVector-Postgres
## Overview
RuVector-Postgres provides high-performance, zero-copy SIMD distance functions optimized for PostgreSQL vector similarity search. The implementation uses runtime CPU feature detection to automatically select the best available instruction set.
## SIMD Architecture Support
### Performance Comparison
| SIMD Level | Floats/Iteration | Relative Speed | Platforms | Instructions |
|------------|------------------|----------------|-----------|--------------|
| **AVX-512** | 16 | 16x | Modern x86_64 | `_mm512_*` |
| **AVX2** | 8 | 8x | Most x86_64 | `_mm256_*` |
| **NEON** | 4 | 4x | ARM64 | `vld1q_f32`, `vmlaq_f32` |
| **Scalar** | 1 | 1x | All | Standard f32 ops |
### CPU Support Matrix
| Processor | AVX-512 | AVX2 | NEON | Recommended Build |
|-----------|---------|------|------|-------------------|
| Intel Skylake-X (2017+) | ✓ | ✓ | - | AVX-512 |
| Intel Haswell (2013+) | - | ✓ | - | AVX2 |
| AMD Zen 4 (2022+) | ✓ | ✓ | - | AVX-512 |
| AMD Zen 1-3 (2017-2021) | - | ✓ | - | AVX2 |
| Apple M1/M2/M3 | - | - | ✓ | NEON |
| AWS Graviton 2/3 | - | - | ✓ | NEON |
| Older CPUs | - | - | - | Scalar |
## Raw Pointer SIMD Functions (Zero-Copy)
### AVX-512 Implementation
#### L2 (Euclidean) Distance
```rust
#[target_feature(enable = "avx512f")]
unsafe fn l2_distance_ptr_avx512(a: *const f32, b: *const f32, len: usize) -> f32 {
let mut sum = _mm512_setzero_ps(); // 16-wide zero vector
let chunks = len / 16;
// Check alignment for potentially faster loads
let use_aligned = is_avx512_aligned(a, b); // 64-byte alignment
if use_aligned {
// Aligned loads (faster, requires 64-byte alignment)
for i in 0..chunks {
let offset = i * 16;
let va = _mm512_load_ps(a.add(offset)); // Aligned load
let vb = _mm512_load_ps(b.add(offset)); // Aligned load
let diff = _mm512_sub_ps(va, vb);
sum = _mm512_fmadd_ps(diff, diff, sum); // FMA: sum += diff²
}
} else {
// Unaligned loads (universal, ~5% slower)
for i in 0..chunks {
let offset = i * 16;
let va = _mm512_loadu_ps(a.add(offset)); // Unaligned load
let vb = _mm512_loadu_ps(b.add(offset)); // Unaligned load
let diff = _mm512_sub_ps(va, vb);
sum = _mm512_fmadd_ps(diff, diff, sum); // FMA: sum += diff²
}
}
let mut result = _mm512_reduce_add_ps(sum); // Horizontal sum
// Handle remainder (tail < 16 elements)
for i in (chunks * 16)..len {
let diff = *a.add(i) - *b.add(i);
result += diff * diff;
}
result.sqrt()
}
```
**Key Optimizations:**
1. **Fused Multiply-Add (FMA)**: `_mm512_fmadd_ps` computes `sum += diff * diff` in one instruction
2. **Alignment Detection**: Uses faster aligned loads when possible
3. **Horizontal Reduction**: `_mm512_reduce_add_ps` efficiently sums 16 floats
4. **Tail Handling**: Scalar loop for dimensions not divisible by 16
#### Cosine Distance
```rust
#[target_feature(enable = "avx512f")]
unsafe fn cosine_distance_ptr_avx512(a: *const f32, b: *const f32, len: usize) -> f32 {
let mut dot = _mm512_setzero_ps();
let mut norm_a = _mm512_setzero_ps();
let mut norm_b = _mm512_setzero_ps();
let chunks = len / 16;
for i in 0..chunks {
let offset = i * 16;
let va = _mm512_loadu_ps(a.add(offset));
let vb = _mm512_loadu_ps(b.add(offset));
dot = _mm512_fmadd_ps(va, vb, dot); // dot += a * b
norm_a = _mm512_fmadd_ps(va, va, norm_a); // norm_a += a²
norm_b = _mm512_fmadd_ps(vb, vb, norm_b); // norm_b += b²
}
let mut dot_sum = _mm512_reduce_add_ps(dot);
let mut norm_a_sum = _mm512_reduce_add_ps(norm_a);
let mut norm_b_sum = _mm512_reduce_add_ps(norm_b);
// Tail handling
for i in (chunks * 16)..len {
let va = *a.add(i);
let vb = *b.add(i);
dot_sum += va * vb;
norm_a_sum += va * va;
norm_b_sum += vb * vb;
}
// Cosine distance: 1 - (a·b) / (||a|| ||b||)
1.0 - (dot_sum / (norm_a_sum.sqrt() * norm_b_sum.sqrt()))
}
```
#### Inner Product (Dot Product)
```rust
#[target_feature(enable = "avx512f")]
unsafe fn inner_product_ptr_avx512(a: *const f32, b: *const f32, len: usize) -> f32 {
let mut sum = _mm512_setzero_ps();
let chunks = len / 16;
for i in 0..chunks {
let offset = i * 16;
let va = _mm512_loadu_ps(a.add(offset));
let vb = _mm512_loadu_ps(b.add(offset));
sum = _mm512_fmadd_ps(va, vb, sum);
}
let mut result = _mm512_reduce_add_ps(sum);
for i in (chunks * 16)..len {
result += *a.add(i) * *b.add(i);
}
-result // Negative for ORDER BY ASC in SQL
}
```
### AVX2 Implementation
Similar structure to AVX-512, but with 8-wide vectors:
```rust
#[target_feature(enable = "avx2", enable = "fma")]
unsafe fn l2_distance_ptr_avx2(a: *const f32, b: *const f32, len: usize) -> f32 {
let mut sum = _mm256_setzero_ps(); // 8-wide zero vector
let chunks = len / 8;
let use_aligned = is_avx2_aligned(a, b); // 32-byte alignment
if use_aligned {
for i in 0..chunks {
let offset = i * 8;
let va = _mm256_load_ps(a.add(offset)); // Aligned
let vb = _mm256_load_ps(b.add(offset)); // Aligned
let diff = _mm256_sub_ps(va, vb);
sum = _mm256_fmadd_ps(diff, diff, sum); // FMA
}
} else {
for i in 0..chunks {
let offset = i * 8;
let va = _mm256_loadu_ps(a.add(offset)); // Unaligned
let vb = _mm256_loadu_ps(b.add(offset)); // Unaligned
let diff = _mm256_sub_ps(va, vb);
sum = _mm256_fmadd_ps(diff, diff, sum);
}
}
// Horizontal reduction (8 floats → 1 float)
let sum_low = _mm256_castps256_ps128(sum);
let sum_high = _mm256_extractf128_ps(sum, 1);
let sum_128 = _mm_add_ps(sum_low, sum_high);
let sum_64 = _mm_add_ps(sum_128, _mm_movehl_ps(sum_128, sum_128));
let sum_32 = _mm_add_ss(sum_64, _mm_shuffle_ps(sum_64, sum_64, 1));
let mut result = _mm_cvtss_f32(sum_32);
// Tail handling
for i in (chunks * 8)..len {
let diff = *a.add(i) - *b.add(i);
result += diff * diff;
}
result.sqrt()
}
```
**AVX2 vs AVX-512:**
- AVX2: 8 floats/iteration, more complex horizontal reduction
- AVX-512: 16 floats/iteration, simpler `_mm512_reduce_add_ps`
- Performance: AVX-512 is ~2x faster for long vectors (1000+ dims)
### ARM NEON Implementation
```rust
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn l2_distance_ptr_neon(a: *const f32, b: *const f32, len: usize) -> f32 {
use std::arch::aarch64::*;
let mut sum = vdupq_n_f32(0.0); // 4-wide zero vector
let chunks = len / 4;
for i in 0..chunks {
let offset = i * 4;
let va = vld1q_f32(a.add(offset)); // Load 4 floats
let vb = vld1q_f32(b.add(offset)); // Load 4 floats
let diff = vsubq_f32(va, vb); // Subtract
sum = vmlaq_f32(sum, diff, diff); // FMA: sum += diff²
}
// Horizontal sum (4 floats → 1 float)
let sum_pair = vpadd_f32(vget_low_f32(sum), vget_high_f32(sum));
let sum_single = vpadd_f32(sum_pair, sum_pair);
let mut result = vget_lane_f32(sum_single, 0);
// Tail handling
for i in (chunks * 4)..len {
let diff = *a.add(i) - *b.add(i);
result += diff * diff;
}
result.sqrt()
}
```
**NEON Features:**
- 4 floats/iteration (vs 16 for AVX-512)
- Efficient on Apple M-series and AWS Graviton
- `vmlaq_f32` provides FMA support
- Horizontal sum via pairwise additions
### f16 (Half-Precision) SIMD Support
#### AVX-512 FP16 (Intel Sapphire Rapids+)
```rust
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512fp16")]
unsafe fn l2_distance_ptr_avx512_f16(a: *const f16, b: *const f16, len: usize) -> f32 {
let mut sum = _mm512_setzero_ph(); // 32-wide f16 vector
let chunks = len / 32;
for i in 0..chunks {
let offset = i * 32;
let va = _mm512_loadu_ph(a.add(offset));
let vb = _mm512_loadu_ph(b.add(offset));
let diff = _mm512_sub_ph(va, vb);
sum = _mm512_fmadd_ph(diff, diff, sum);
}
// Convert to f32 for final reduction
let sum_f32 = _mm512_cvtph_ps(_mm512_castph512_ph256(sum));
let mut result = _mm512_reduce_add_ps(sum_f32);
// Handle upper 16 elements
let upper = _mm512_extractf32x8_ps(sum_f32, 1);
// ... additional reduction
result.sqrt()
}
```
**Benefits:**
- 32 f16 values/iteration (vs 16 f32)
- 2x throughput for half-precision vectors
- Native f16 arithmetic (no conversion overhead)
#### ARM NEON FP16
```rust
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon", enable = "fp16")]
unsafe fn l2_distance_ptr_neon_f16(a: *const f16, b: *const f16, len: usize) -> f32 {
use std::arch::aarch64::*;
let mut sum = vdupq_n_f16(0.0); // 8-wide f16 vector
let chunks = len / 8;
for i in 0..chunks {
let offset = i * 8;
let va = vld1q_f16(a.add(offset) as *const __fp16);
let vb = vld1q_f16(b.add(offset) as *const __fp16);
let diff = vsubq_f16(va, vb);
sum = vfmaq_f16(sum, diff, diff);
}
// Convert to f32 and reduce
let sum_low_f32 = vcvt_f32_f16(vget_low_f16(sum));
let sum_high_f32 = vcvt_f32_f16(vget_high_f16(sum));
// ... horizontal sum
}
```
## Benchmark Results vs pgvector
### Test Setup
- CPU: Intel Xeon (Skylake-X, AVX-512)
- Vectors: 1,000,000 × 1536 dimensions (OpenAI embeddings)
- Query: Top-10 nearest neighbors
- Metric: L2 distance
### Results
| Implementation | Queries/sec | Speedup | SIMD Level |
|----------------|-------------|---------|------------|
| **RuVector AVX-512** | 24,500 | 9.8x | AVX-512 |
| **RuVector AVX2** | 13,200 | 5.3x | AVX2 |
| **RuVector NEON** | 8,900 | 3.6x | NEON |
| RuVector Scalar | 3,100 | 1.2x | None |
| pgvector 0.8.0 | 2,500 | 1.0x (baseline) | Partial AVX2 |
**Key Findings:**
1. AVX-512 provides **9.8x speedup** over pgvector
2. Even scalar RuVector is **1.2x faster** (better algorithms)
3. Zero-copy access eliminates allocation overhead
4. Batch operations further improve throughput
### Dimensional Scaling
| Dimensions | RuVector (AVX-512) | pgvector | Speedup |
|------------|-------------------|----------|---------|
| 128 | 45,000 q/s | 8,200 q/s | 5.5x |
| 384 | 32,000 q/s | 5,100 q/s | 6.3x |
| 768 | 26,000 q/s | 3,400 q/s | 7.6x |
| 1536 | 24,500 q/s | 2,500 q/s | 9.8x |
| 3072 | 22,000 q/s | 1,800 q/s | 12.2x |
**Observation:** Speedup increases with dimension count (better SIMD utilization).
## AVX-512 vs AVX2 Selection
### Runtime Detection
```rust
use std::sync::atomic::{AtomicU8, Ordering};
#[repr(u8)]
enum SimdLevel {
Scalar = 0,
NEON = 1,
AVX2 = 2,
AVX512 = 3,
}
static SIMD_LEVEL: AtomicU8 = AtomicU8::new(0);
pub fn init_simd_dispatch() {
#[cfg(target_arch = "x86_64")]
{
if is_x86_feature_detected!("avx512f") {
SIMD_LEVEL.store(SimdLevel::AVX512 as u8, Ordering::Relaxed);
return;
}
if is_x86_feature_detected!("avx2") {
SIMD_LEVEL.store(SimdLevel::AVX2 as u8, Ordering::Relaxed);
return;
}
}
#[cfg(target_arch = "aarch64")]
{
SIMD_LEVEL.store(SimdLevel::NEON as u8, Ordering::Relaxed);
return;
}
SIMD_LEVEL.store(SimdLevel::Scalar as u8, Ordering::Relaxed);
}
```
### Dispatch Function
```rust
pub fn euclidean_distance(a: &[f32], b: &[f32]) -> f32 {
assert_eq!(a.len(), b.len());
unsafe {
let a_ptr = a.as_ptr();
let b_ptr = b.as_ptr();
let len = a.len();
match SIMD_LEVEL.load(Ordering::Relaxed) {
3 => l2_distance_ptr_avx512(a_ptr, b_ptr, len),
2 => l2_distance_ptr_avx2(a_ptr, b_ptr, len),
1 => l2_distance_ptr_neon(a_ptr, b_ptr, len),
_ => l2_distance_ptr_scalar(a_ptr, b_ptr, len),
}
}
}
```
**Performance Notes:**
- Detection happens once at extension load
- Zero overhead after initialization (atomic read is cached)
- No runtime branching in hot loop
## Safety Requirements
All SIMD functions are marked `unsafe` and require:
1. **Valid Pointers**: `a` and `b` must be valid for reads of `len` elements
2. **No Aliasing**: Pointers must not overlap
3. **Length > 0**: `len` must be non-zero
4. **Memory Validity**: Memory must remain valid for duration of call
5. **Alignment**: Unaligned access is safe but aligned is faster
### Caller Responsibilities
```rust
// ✓ SAFE: Valid slices
let a = vec![1.0, 2.0, 3.0];
let b = vec![4.0, 5.0, 6.0];
unsafe {
euclidean_distance_ptr(a.as_ptr(), b.as_ptr(), a.len());
}
// ✗ UNSAFE: Overlapping pointers
let v = vec![1.0, 2.0, 3.0, 4.0];
unsafe {
euclidean_distance_ptr(v.as_ptr(), v.as_ptr().add(1), 3); // UB!
}
// ✗ UNSAFE: Invalid length
unsafe {
euclidean_distance_ptr(a.as_ptr(), b.as_ptr(), 100); // Buffer overrun!
}
```
## Optimization Tips
### 1. Memory Alignment
**Best Performance:**
```rust
// Allocate with alignment
let layout = std::alloc::Layout::from_size_align(size, 64).unwrap();
let ptr = std::alloc::alloc(layout) as *mut f32;
// Use aligned loads (AVX-512)
unsafe {
let va = _mm512_load_ps(ptr); // Faster than _mm512_loadu_ps
}
```
**PostgreSQL Context:**
- Varlena data is typically 8-byte aligned
- Large allocations may be 64-byte aligned
- Use unaligned loads by default (safe, minimal penalty)
### 2. Batch Operations
**Sequential:**
```rust
let results: Vec<f32> = vectors.iter()
.map(|v| euclidean_distance(query, v))
.collect();
```
**Parallel (Better):**
```rust
use rayon::prelude::*;
let results: Vec<f32> = vectors.par_iter()
.map(|v| euclidean_distance(query, v))
.collect();
```
### 3. Dimension Tuning
**Optimal Dimensions:**
- Multiples of 16 for AVX-512 (no tail handling)
- Multiples of 8 for AVX2
- Multiples of 4 for NEON
**Example:**
```sql
-- ✓ Optimal: 1536 = 16 * 96
CREATE TABLE items (embedding ruvector(1536));
-- ✗ Suboptimal: 1535 = 16 * 95 + 15 (15 scalar iterations)
CREATE TABLE items (embedding ruvector(1535));
```
### 4. Compiler Flags
**Build with native optimizations:**
```bash
export RUSTFLAGS="-C target-cpu=native -C opt-level=3"
cargo pgrx package --release
```
**Flags Explained:**
- `target-cpu=native`: Enable all CPU features available
- `opt-level=3`: Maximum optimization level
- Result: ~10% additional speedup
### 5. Profile-Guided Optimization (PGO)
**Step 1: Instrumented Build**
```bash
export RUSTFLAGS="-C profile-generate=/tmp/pgo-data"
cargo pgrx package --release
```
**Step 2: Run Typical Workload**
```sql
-- Run representative queries
SELECT * FROM items ORDER BY embedding <-> query LIMIT 100;
```
**Step 3: Optimized Build**
```bash
export RUSTFLAGS="-C profile-use=/tmp/pgo-data -C llvm-args=-pgo-warn-missing-function"
cargo pgrx package --release
```
**Expected Improvement:** 5-15% additional speedup.
## Debugging SIMD Code
### Check CPU Features
```sql
-- In PostgreSQL
SELECT ruvector_simd_info();
-- Output: AVX512, AVX2, NEON, or Scalar
```
```bash
# Linux
cat /proc/cpuinfo | grep -E 'avx2|avx512'
# macOS
sysctl machdep.cpu.features
# Windows
wmic cpu get caption
```
### Verify SIMD Dispatch
```rust
// Add logging to init
pub fn init_simd_dispatch() {
#[cfg(target_arch = "x86_64")]
{
if is_x86_feature_detected!("avx512f") {
eprintln!("Using AVX-512");
// ...
}
}
}
```
### Benchmarking
```sql
-- Create test data
CREATE TABLE bench (id int, embedding ruvector(1536));
INSERT INTO bench SELECT i, (SELECT array_agg(random())::ruvector FROM generate_series(1,1536)) FROM generate_series(1, 10000) i;
-- Benchmark
\timing on
SELECT COUNT(*) FROM bench WHERE embedding <-> (SELECT embedding FROM bench LIMIT 1) < 0.5;
```
## Future Enhancements
### Planned Features
1. **AVX-512 BF16**: Brain floating point support
2. **AMX (Advanced Matrix Extensions)**: Tile-based operations
3. **Auto-Vectorization**: Let Rust compiler auto-vectorize
4. **Multi-Vector Operations**: SIMD for multiple queries simultaneously
## References
- Intel Intrinsics Guide: https://www.intel.com/content/www/us/en/docs/intrinsics-guide/
- ARM NEON Intrinsics: https://developer.arm.com/architectures/instruction-sets/intrinsics/
- Rust SIMD Documentation: https://doc.rust-lang.org/core/arch/
- pgvector Source: https://github.com/pgvector/pgvector