feat: add simdgroup-optimized Metal kernels

[cluster2600]

Summary

Adds 6 new Metal compute kernels using simdgroup cooperative intrinsics (simd_sum, simd_min, simd_shuffle) for hardware-accelerated reductions across 32 SIMD lanes — no shared memory barriers needed.

Follow-up to #166 ("Future Work: SIMD optimization").

New kernels

Kernel	Description
metal_l2_distance_simdgroup	32-thread cooperative L2 distance
metal_inner_product_simdgroup	32-thread cooperative dot product
metal_cosine_similarity_simdgroup	Normalized inner product with 3 parallel reductions
metal_topk_simdgroup	Per-query top-k selection via simd_min lane voting
metal_matmul_tiled	Tiled matmul with threadgroup shared memory (16×16 tiles)
metal_normalize_simdgroup	In-place L2 normalization

Dispatch model

• Threadgroup size: 32 (one simdgroup)
• Grid: (n_database, n_queries) threadgroups
• Each simdgroup collaborates on one (query, database) pair, splitting dimensions across lanes

Fixes to existing kernels

• Replace simd_make_float4 → float4 constructor (MSL compliance)
• Add device address space qualifiers in metal_l2_distance_batch

Merge order

This PR shares a common base with #173, #175, #176. Recommended merge order: #172 → #173 → #175 → #176. Merging any one brings in the shared base commits; the rest then apply cleanly.

Test plan

• Compiles cleanly with metal -std=metal3.1 -W -Werror on macOS with Xcode Metal toolchain
• Integration test with Metal compute pipeline on Apple Silicon

[cluster2600]

Discussion issue opened: #177 — feedback welcome before review.

[Cuiyus]

@greptile

[greptile-apps]

Greptile Summary

This PR adds hardware-accelerated Metal compute kernels for Apple Silicon, introducing 6 new simdgroup-optimized kernels that leverage cooperative SIMD intrinsics (simd_sum, simd_min, simd_shuffle) for efficient vector operations without shared memory barriers.

Key Changes

• New Metal kernels (distance.metal): L2 distance, inner product, cosine similarity, top-k selection, tiled matmul, and normalization using 32-thread simdgroups
• Backend detection system (detect.py): Prioritizes C++ cuVS > Python cuVS > FAISS GPU > MPS > FAISS CPU
• cuVS integration: Python wrappers for CAGRA (graph ANN) and IVF-PQ with correct RAPIDS API usage
• Apple Silicon support (apple_silicon.py): PyTorch MPS and Accelerate framework integration
• CMake modernization: Added CUDA architecture flags, Metal support, and optional cuVS integration
• Comprehensive testing: New test suite and GPU benchmark notebooks

Issues Found

• metal_topk_simdgroup kernel has O(k²) complexity due to repeated scans of already-selected indices (lines 419-421)

The implementation follows proper Metal Shading Language conventions and correctly uses simdgroup cooperative reductions for hardware acceleration.

Confidence Score: 4/5

• This PR is safe to merge with proper testing on Apple Silicon hardware
• The implementation follows Metal best practices and uses correct API patterns for cuVS integration. The only concern is the O(k²) topk selection algorithm which affects performance but not correctness. Integration testing on actual Apple Silicon hardware is needed to verify Metal kernel compilation and runtime behavior.
• Pay attention to src/ailego/gpu/metal/distance.metal for the topk kernel performance characteristics

Important Files Changed

Filename	Overview
src/ailego/gpu/metal/distance.metal	New Metal kernels with simdgroup intrinsics for L2, dot product, cosine, matmul, topk, and normalize. Topk has O(k²) selection overhead.
python/zvec/backends/detect.py	Hardware detection with backend priority: C++ cuVS > Python cuVS > FAISS GPU > MPS > FAISS CPU > NumPy
python/zvec/backends/apple_silicon.py	Apple Silicon backend using PyTorch MPS or Accelerate framework for matrix ops and L2 distance
python/zvec/backends/cuvs_cagra.py	cuVS CAGRA wrapper using correct RAPIDS API (metric="sqeuclidean"), with CuPy array conversions
python/zvec/backends/cuvs_ivf_pq.py	cuVS IVF-PQ wrapper with proper parameter passing and CuPy device array handling
src/ailego/gpu/cuvs/zvec_cuvs.h	C++ cuVS bindings header defining IVF-PQ, CAGRA, and HNSW index interfaces with parameters
src/CMakeLists.txt	CMake config rewritten to support CUDA, Metal, and optional cuVS with proper architecture flags
python/tests/test_backends.py	New test suite covering hardware detection, GPU index operations, and quantization (PQ/OPQ)

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    Start[Vector Search Request] --> Detect[Hardware Detection]
    Detect --> CheckCppCuvs{C++ cuVS<br/>Available?}
    CheckCppCuvs -->|Yes| UseCppCuvs[Use C++ cuVS<br/>CAGRA/IVF-PQ]
    CheckCppCuvs -->|No| CheckPyCuvs{Python cuVS<br/>Available?}
    CheckPyCuvs -->|Yes| UsePyCuvs[Use Python cuVS<br/>CuPy arrays]
    CheckPyCuvs -->|No| CheckFaissGpu{FAISS GPU +<br/>NVIDIA GPU?}
    CheckFaissGpu -->|Yes| UseFaissGpu[Use FAISS GPU]
    CheckFaissGpu -->|No| CheckMps{Apple Silicon<br/>MPS?}
    CheckMps -->|Yes| UseMps[Use Metal Kernels<br/>simdgroup ops]
    CheckMps -->|No| CheckFaissCpu{FAISS CPU?}
    CheckFaissCpu -->|Yes| UseFaissCpu[Use FAISS CPU<br/>+ Accelerate]
    CheckFaissCpu -->|No| UseNumpy[Fallback to NumPy]
    
    UseCppCuvs --> Execute[Execute Search]
    UsePyCuvs --> Execute
    UseFaissGpu --> Execute
    UseMps --> MetalKernels[Metal Kernels:<br/>L2/cosine/topk]
    MetalKernels --> Execute
    UseFaissCpu --> Execute
    UseNumpy --> Execute
    
    Execute --> Results[Return distances<br/>+ indices]
    
    style UseMps fill:#a8dadc
    style MetalKernels fill:#a8dadc
    style UseCppCuvs fill:#f1faee
    style UsePyCuvs fill:#f1faee

Loading

_{Last reviewed commit: b08a835}

[greptile-apps]

_{41 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}