diskann_vector/
lib.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5//! # vector
6//!
7//! This crate contains SIMD accelerated functions for operating on vector data. Note that the name 'vector'
8//! does not exclusively mean embedding vectors, but any array of data appropriate for SIMD. Therefor, aside
9//! from fast implementations of distance for real vectors, this crate also includes things like SIMD
10//! accelerated contains for slices.
11#![cfg_attr(
12    not(test),
13    warn(
14        clippy::panic,
15        clippy::unwrap_used,
16        clippy::expect_used,
17        clippy::undocumented_unsafe_blocks
18    )
19)]
20
21mod half;
22pub use half::Half;
23
24mod traits;
25pub use traits::{
26    DistanceFunction, DistanceFunctionMut, Norm, PreprocessedDistanceFunction, PureDistanceFunction,
27};
28
29mod value;
30pub use value::{MathematicalValue, SimilarityScore};
31
32mod unaligned;
33pub use unaligned::{AsUnaligned, UnalignedSlice};
34
35pub mod contains;
36pub mod conversion;
37pub mod distance;
38pub mod norm;
39pub mod sparse;
40
41#[cfg(any(
42    all(target_arch = "x86_64", target_feature = "avx2"),
43    target_arch = "aarch64",
44))]
45const CACHE_LINE_SIZE: usize = 64;
46
47/// Prefetch the cache line containing `ptr`.
48///
49/// # Safety
50///
51/// `ptr` must point to a byte within a live allocation.
52#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
53#[inline(always)]
54pub unsafe fn prefetch_cache_line(_ptr: *const i8) {
55    #[cfg(not(miri))]
56    {
57        use std::arch::x86_64::*;
58        // SAFETY: Inherited from caller and intrinsic is cfg-gated to supported targets.
59        unsafe { _mm_prefetch(_ptr, _MM_HINT_T0) };
60    }
61}
62
63/// Prefetch the cache line containing `ptr`.
64///
65/// # Safety
66///
67/// `ptr` must point to a byte within a live allocation.
68#[cfg(target_arch = "aarch64")]
69#[inline(always)]
70pub unsafe fn prefetch_cache_line(_ptr: *const i8) {
71    #[cfg(not(miri))]
72    {
73        // SAFETY: The `prfm` hint does not dereference the pointer in Rust memory model terms.
74        unsafe {
75            std::arch::asm!(
76                "prfm pldl2keep, [{ptr}]",
77                ptr = in(reg) _ptr,
78                options(nostack, preserves_flags, readonly),
79            );
80        }
81    }
82}
83
84/// Prefetch the cache line containing `ptr`.
85///
86/// # Safety
87///
88/// `ptr` must point to a byte within a live allocation.
89#[cfg(not(any(
90    all(target_arch = "x86_64", target_feature = "avx2"),
91    target_arch = "aarch64",
92)))]
93#[inline(always)]
94pub unsafe fn prefetch_cache_line(_ptr: *const i8) {}
95
96#[cfg(any(
97    all(target_arch = "x86_64", target_feature = "avx2"),
98    target_arch = "aarch64",
99))]
100#[inline(always)]
101unsafe fn prefetch_exactly<const N: usize>(ptr: *const i8) {
102    for i in 0..N {
103        // SAFETY: Pointer arithmetic stays within the caller-provided prefetch range.
104        unsafe { prefetch_cache_line(ptr.add(i * CACHE_LINE_SIZE)) };
105    }
106}
107
108#[cfg(any(
109    all(target_arch = "x86_64", target_feature = "avx2"),
110    target_arch = "aarch64",
111))]
112#[inline(always)]
113unsafe fn prefetch_at_most<const N: usize>(ptr: *const i8, bytes: usize) {
114    for i in 0..N {
115        if CACHE_LINE_SIZE * i >= bytes {
116            break;
117        }
118
119        // SAFETY: `i` is bounded by `bytes`, so computed addresses remain in range.
120        unsafe { prefetch_cache_line(ptr.add(i * CACHE_LINE_SIZE)) };
121    }
122}
123
124/// Prefetch the given vector in chunks of 64 bytes, which is a cache line size.
125/// Only the first `MAX_CACHE_LINES` chunks will be prefetched.
126#[cfg(any(
127    all(target_arch = "x86_64", target_feature = "avx2"),
128    target_arch = "aarch64",
129))]
130#[inline]
131pub fn prefetch_hint_max<const MAX_CACHE_LINES: usize, T>(vec: &[T]) {
132    let vecsize = std::mem::size_of_val(vec);
133    if vecsize >= MAX_CACHE_LINES * CACHE_LINE_SIZE {
134        // SAFETY: Pointer originates from `vec` and prefetch count is bounded.
135        unsafe { prefetch_exactly::<MAX_CACHE_LINES>(vec.as_ptr().cast()) }
136    } else {
137        // SAFETY: Pointer originates from `vec`; `bytes` bound caps traversal.
138        unsafe { prefetch_at_most::<MAX_CACHE_LINES>(vec.as_ptr().cast(), vecsize) }
139    }
140}
141
142/// Prefetch the given vector in chunks of 64 bytes, which is a cache line size.
143/// The entire vector will be prefetched.
144#[cfg(any(
145    all(target_arch = "x86_64", target_feature = "avx2"),
146    target_arch = "aarch64",
147))]
148#[inline]
149pub fn prefetch_hint_all<T>(vec: &[T]) {
150    let vecsize = std::mem::size_of_val(vec);
151    let num_prefetch_blocks = vecsize.div_ceil(CACHE_LINE_SIZE);
152    let vec_ptr = vec.as_ptr() as *const i8;
153    for d in 0..num_prefetch_blocks {
154        // SAFETY: Loop bounds are derived from `vec` size and this is a prefetch hint.
155        unsafe { prefetch_cache_line(vec_ptr.add(d * CACHE_LINE_SIZE)) };
156    }
157}
158
159#[cfg(not(any(
160    all(target_arch = "x86_64", target_feature = "avx2"),
161    target_arch = "aarch64",
162)))]
163pub fn prefetch_hint_max<const MAX_CACHE_LINES: usize, T>(_vec: &[T]) {}
164
165#[cfg(not(any(
166    all(target_arch = "x86_64", target_feature = "avx2"),
167    target_arch = "aarch64",
168)))]
169pub fn prefetch_hint_all<T>(_vec: &[T]) {}
170
171#[cfg(test)]
172mod test_util;
diskann_vector/lib.rs

diskann_vector/
lib.rs
