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

diskann_vector/
lib.rs
