diskann-vector 0.50.1

DiskANN is a fast approximate nearest neighbor search library for high dimensional data
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */
//! # vector
//!
//! This crate contains SIMD accelerated functions for operating on vector data. Note that the name 'vector'
//! does not exclusively mean embedding vectors, but any array of data appropriate for SIMD. Therefor, aside
//! from fast implementations of distance for real vectors, this crate also includes things like SIMD
//! accelerated contains for slices.
#![cfg_attr(
    not(test),
    warn(
        clippy::panic,
        clippy::unwrap_used,
        clippy::expect_used,
        clippy::undocumented_unsafe_blocks
    )
)]

mod half;
pub use half::Half;

mod traits;
pub use traits::{
    DistanceFunction, DistanceFunctionMut, Norm, PreprocessedDistanceFunction, PureDistanceFunction,
};

mod value;
pub use value::{MathematicalValue, SimilarityScore};

pub mod contains;
pub mod conversion;
pub mod distance;
pub mod norm;

cfg_if::cfg_if! {
    if #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] {
        const CACHE_LINE_SIZE: usize = 64;

        #[inline(always)]
        unsafe fn prefetch_exactly<const N: usize>(ptr: *const i8) {
            use std::arch::x86_64::*;
            for i in 0..N {
                _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0);
            }
        }

        #[inline(always)]
        unsafe fn prefetch_at_most<const N: usize>(ptr: *const i8, bytes: usize) {
            use std::arch::x86_64::*;
            for i in 0..N {
                if CACHE_LINE_SIZE * i >= bytes {
                    break;
                }
                _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0);
            }
        }

        /// Prefetch the given vector in chunks of 64 bytes, which is a cache line size.
        /// Only the first `MAX_BLOCKS` chunks will be prefetched.
        #[inline]
        pub fn prefetch_hint_max<const MAX_CACHE_LINES: usize, T>(vec: &[T]) {
            let vecsize = std::mem::size_of_val(vec);
            if vecsize >= MAX_CACHE_LINES * 64 {
                // SAFETY: Pointer is in-bounds and use of the intrinsic is cfg gated.
                unsafe { prefetch_exactly::<MAX_CACHE_LINES>(vec.as_ptr().cast()) }
            } else {
                // SAFETY: Pointer is in-bounds and use of the intrinsic is cfg gated.
                unsafe { prefetch_at_most::<MAX_CACHE_LINES>(vec.as_ptr().cast(), vecsize) }
            }
        }

        /// Prefetch the given vector in chunks of 64 bytes, which is a cache line size.
        /// The entire vector will be prefetched.
        #[inline]
        pub fn prefetch_hint_all<T>(vec: &[T]) {
            use std::arch::x86_64::*;

            let vecsize = std::mem::size_of_val(vec);
            let num_prefetch_blocks = vecsize.div_ceil(64);
            let vec_ptr = vec.as_ptr() as *const i8;
            for d in 0..num_prefetch_blocks {
                // SAFETY: Pointer is in-bounds and use of the intrinsic is gated by the
                // `cfg`-guard on this function.
                unsafe {
                    std::arch::x86_64::_mm_prefetch(vec_ptr.add(d * CACHE_LINE_SIZE), _MM_HINT_T0);
                }
            }        }
    } else {
        pub fn prefetch_hint_max<const MAX_CACHE_LINES: usize, T>(_vec: &[T]) {}
        pub fn prefetch_hint_all<T>(_vec: &[T]) {}
    }
}

#[cfg(test)]
mod test_util;