diskann-quantization 0.61.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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

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

use diskann_utils::views::rowmajor::{self, Matrix};

use crate::matrix_kernels::test_util::TestDistr;

/// Generate a test MaxSim problem `[M x K] . [K x N]` where both matrices are row-major.
pub(super) fn generate_f32(
    m: usize,
    k: usize,
    n: usize,
    rng: &mut impl rand::Rng,
) -> (rowmajor::Owned<f32>, rowmajor::Owned<f32>, Vec<f32>) {
    let ref_a = TestDistr::matrix::<f32>(m, k, rng);
    let ref_b = TestDistr::matrix::<f32>(k, n, rng);

    let ref_c: Vec<f32> = ref_a
        .rows()
        .map(|a_row| {
            let mut max_ip = f32::NEG_INFINITY;
            for b_col in 0..n {
                let mut ip = 0.0;
                for (k, a) in a_row.iter().enumerate() {
                    ip = a.mul_add(*ref_b.element(k, b_col), ip);
                }
                max_ip = max_ip.max(ip);
            }

            max_ip
        })
        .collect();

    (ref_a, ref_b, ref_c)
}

/// Generate a test integer MaxSim problem `[M x K] . [K x N]` where both matrices are
/// row-major.
pub(super) fn generate_i8(
    m: usize,
    k: usize,
    n: usize,
    rng: &mut impl rand::Rng,
) -> (rowmajor::Owned<i8>, rowmajor::Owned<i8>, Vec<i32>) {
    let ref_a = TestDistr::matrix::<i8>(m, k, rng);
    let ref_b = TestDistr::matrix::<i8>(k, n, rng);

    let ref_c: Vec<i32> = ref_a
        .rows()
        .map(|a_row| {
            let mut max_ip = i32::MIN;
            for b_col in 0..n {
                let mut ip = 0;
                for (k, a) in a_row.iter().enumerate() {
                    ip += i32::from(*a) * i32::from(*ref_b.element(k, b_col));
                }
                max_ip = max_ip.max(ip);
            }

            max_ip
        })
        .collect();

    (ref_a, ref_b, ref_c)
}

#[derive(Debug, Clone)]
pub(super) struct TestDims {
    pub(super) a_panels_per_tile: usize,
    pub(super) total_a_rows: usize,
    pub(super) b_cols_per_tile: usize,
    pub(super) total_b_cols: usize,
    pub(super) k: usize,
}

impl TestDims {
    fn from_tuple(
        [
            a_panels_per_tile,
            total_a_rows,
            b_cols_per_tile,
            total_b_cols,
            k,
        ]: [usize; 5],
    ) -> Self {
        Self {
            a_panels_per_tile,
            total_a_rows,
            b_cols_per_tile,
            total_b_cols,
            k,
        }
    }
}

/// Test dimensions for `packed_x_unpacked` kernels.
pub(super) fn packed_x_unpacked_test_dims(mr: usize, nr: usize) -> Vec<TestDims> {
    [
        [1, 1, 1, 1, 1],                        // Smallest logical output
        [1, mr / 2, 1, 1, 1],                   // Partial first A panel
        [1, mr - 1, 2 * nr, nr, 3],             // Nearly full first A panel
        [2, mr + 1, 2 * nr, nr, 3],             // Partial second A panel
        [2, 2 * mr - 1, 2 * nr, 2 * nr + 1, 5], // Partial A panel and split B
        [2, 2 * mr + 1, 2 * nr, 2 * nr + 1, 5], // Split A and B with a partial panel
        [1, mr * 3, 1, 3, 1],                   // Unit advancement, no reuse.
        [2, mr * 2, 2 * nr, 2 * nr, 3],         // Values a direct multiple of the blocking.
        [2, mr * 3, 2 * nr, nr, 3],
        [2, mr, 2 * nr, 2 * nr + 1, 3],
        [2, mr * 3, 2 * nr, 2 * nr + 1, 5],
        [2, mr * 5, 2 * nr, 4 * nr + 1, 1],
    ]
    .map(TestDims::from_tuple)
    .into_iter()
    .collect()
}