diskann-quantization 0.62.0

DiskANN3 is a composable library for bringing scalable, accurate and cost-effective vector indexing to multiple databases.
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
98

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use diskann_wide::arch::Scalar;
#[cfg(target_arch = "aarch64")]
use diskann_wide::arch::aarch64::Neon;
#[cfg(target_arch = "x86_64")]
use diskann_wide::arch::x86_64::{V3, V4};

use super::kernel::{MinMaxMaxSimKernel, validate};
use crate::{
    bits::{Representation, Unsigned},
    matrix_kernels::{
        self as mk,
        maxsim::packed_minmax_x_unpacked_minmax::{Driver, PackedA},
    },
    minmax::MinMaxMeta,
    multi_vector::{MatRef, MaxSimError, Overflow},
};

// Fix ISA and packing layout; query bit widths are needed only during packing.
macro_rules! prepared_kernel {
    ($name:ident, $arch:ty, $element:ty, $mr:literal, $nr:literal, $pack:literal) => {
        #[derive(Debug)]
        pub(super) struct $name<const B_BITS: usize> {
            arch: $arch,
            query: PackedA<$element, $mr, $pack>,
            rows: usize,
            dim: usize,
        }

        impl<const B_BITS: usize> $name<B_BITS> {
            pub(super) fn new<const A_BITS: usize>(
                arch: $arch,
                query: MatRef<'_, MinMaxMeta<A_BITS>>,
            ) -> Result<Self, Overflow>
            where
                Unsigned: Representation<A_BITS>,
            {
                Ok(Self {
                    arch,
                    query: PackedA::new(query)?,
                    rows: query.num_vectors(),
                    dim: query.repr().intrinsic_dim(),
                })
            }
        }

        impl<const B_BITS: usize> MinMaxMaxSimKernel<B_BITS> for $name<B_BITS>
        where
            Unsigned: Representation<B_BITS>,
        {
            fn nrows(&self) -> usize {
                self.rows
            }

            fn compute_max_sim(
                &self,
                document: MatRef<'_, MinMaxMeta<B_BITS>>,
                scores: &mut [f32],
            ) -> Result<(), MaxSimError> {
                if !validate(self.rows, self.dim, document, scores)? {
                    return Ok(());
                }
                // SAFETY: Validation establishes nonempty A, positive K, matching dimensions
                // and exactly one output per query row. Packing and ISA selection are fixed.
                let mut driver = unsafe {
                    Driver::<$arch, $mr, $nr, $pack>::new(
                        self.arch,
                        &self.query,
                        document,
                        scores,
                        mk::Cache::detect(),
                    )
                };
                mk::Drive::drive(&mut driver);
                for score in scores {
                    *score = (-*score).min(f32::MAX);
                }
                Ok(())
            }
        }
    };
}
prepared_kernel!(ScalarBytes, Scalar, u8, 16, 2, 4);
prepared_kernel!(ScalarWords, Scalar, i16, 8, 2, 2);
#[cfg(target_arch = "x86_64")]
prepared_kernel!(V3Bytes, V3, u8, 16, 4, 4);
#[cfg(target_arch = "x86_64")]
prepared_kernel!(V3Words, V3, i16, 16, 4, 2);
#[cfg(target_arch = "x86_64")]
prepared_kernel!(V4Bytes, V4, u8, 16, 8, 4);
#[cfg(target_arch = "x86_64")]
prepared_kernel!(V4Words, V4, i16, 16, 8, 2);
#[cfg(target_arch = "aarch64")]
prepared_kernel!(NeonBytes, Neon, u8, 8, 4, 4);