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;
39
40cfg_if::cfg_if! {
41    if #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] {
42        const CACHE_LINE_SIZE: usize = 64;
43
44        #[inline(always)]
45        unsafe fn prefetch_exactly<const N: usize>(ptr: *const i8) {
46            use std::arch::x86_64::*;
47            for i in 0..N {
48                _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0);
49            }
50        }
51
52        #[inline(always)]
53        unsafe fn prefetch_at_most<const N: usize>(ptr: *const i8, bytes: usize) {
54            use std::arch::x86_64::*;
55            for i in 0..N {
56                if CACHE_LINE_SIZE * i >= bytes {
57                    break;
58                }
59                _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0);
60            }
61        }
62
63        /// Prefetch the given vector in chunks of 64 bytes, which is a cache line size.
64        /// Only the first `MAX_BLOCKS` chunks will be prefetched.
65        #[inline]
66        pub fn prefetch_hint_max<const MAX_CACHE_LINES: usize, T>(vec: &[T]) {
67            let vecsize = std::mem::size_of_val(vec);
68            if vecsize >= MAX_CACHE_LINES * 64 {
69                // SAFETY: Pointer is in-bounds and use of the intrinsic is cfg gated.
70                unsafe { prefetch_exactly::<MAX_CACHE_LINES>(vec.as_ptr().cast()) }
71            } else {
72                // SAFETY: Pointer is in-bounds and use of the intrinsic is cfg gated.
73                unsafe { prefetch_at_most::<MAX_CACHE_LINES>(vec.as_ptr().cast(), vecsize) }
74            }
75        }
76
77        /// Prefetch the given vector in chunks of 64 bytes, which is a cache line size.
78        /// The entire vector will be prefetched.
79        #[inline]
80        pub fn prefetch_hint_all<T>(vec: &[T]) {
81            use std::arch::x86_64::*;
82
83            let vecsize = std::mem::size_of_val(vec);
84            let num_prefetch_blocks = vecsize.div_ceil(64);
85            let vec_ptr = vec.as_ptr() as *const i8;
86            for d in 0..num_prefetch_blocks {
87                // SAFETY: Pointer is in-bounds and use of the intrinsic is gated by the
88                // `cfg`-guard on this function.
89                unsafe {
90                    std::arch::x86_64::_mm_prefetch(vec_ptr.add(d * CACHE_LINE_SIZE), _MM_HINT_T0);
91                }
92            }        }
93    } else {
94        pub fn prefetch_hint_max<const MAX_CACHE_LINES: usize, T>(_vec: &[T]) {}
95        pub fn prefetch_hint_all<T>(_vec: &[T]) {}
96    }
97}
98
99#[cfg(test)]
100mod test_util;
diskann_vector/lib.rs

diskann_vector/
lib.rs
