Skip to main content

diskann_quantization/multi_vector/distance/
isa.rs

1// Copyright (c) Microsoft Corporation. All rights reserved.
2// Licensed under the MIT license.
3
4//! Instruction Set Architecture (ISA) selector for the multi-vector MaxSim
5//! factory.
6
7/// Instruction Set Architecture (ISA) selector for which multi-vector MaxSim
8/// kernel to build.
9///
10/// `#[non_exhaustive]` so adding a variant (e.g. for a new in-tree kernel) is
11/// not a breaking change. Deliberately **not** `Serialize`/`Deserialize` —
12/// callers wanting JSON support maintain their own shadow enum and convert
13/// via `From` / `TryFrom`, so the library is not pinned to a particular
14/// serialization format.
15#[non_exhaustive]
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
17#[allow(non_camel_case_types)]
18pub enum MaxSimIsa {
19    /// Pick the highest ISA the host CPU supports.
20    Auto,
21    /// Pure-scalar (emulated SIMD) kernel — always available.
22    Scalar,
23    /// x86_64 AVX2 + FMA.
24    X86_64_V3,
25    /// x86_64 AVX-512.
26    X86_64_V4,
27    /// AArch64 Neon.
28    Neon,
29    /// Non-SIMD reference fallback. Slow; serves as a correctness baseline.
30    Reference,
31}
32
33impl MaxSimIsa {
34    /// Whether a kernel for this ISA can be built on the current host.
35    /// Feature-gated variants may return `false` even when compiled for
36    /// the matching target architecture.
37    pub fn is_available(self) -> bool {
38        match self {
39            Self::Auto | Self::Scalar | Self::Reference => true,
40            #[cfg(target_arch = "x86_64")]
41            Self::X86_64_V3 => diskann_wide::arch::x86_64::V3::new_checked().is_some(),
42            #[cfg(target_arch = "x86_64")]
43            Self::X86_64_V4 => diskann_wide::arch::x86_64::V4::new_checked().is_some(),
44            #[cfg(not(target_arch = "x86_64"))]
45            Self::X86_64_V3 | Self::X86_64_V4 => false,
46            #[cfg(target_arch = "aarch64")]
47            Self::Neon => diskann_wide::arch::aarch64::Neon::new_checked().is_some(),
48            #[cfg(not(target_arch = "aarch64"))]
49            Self::Neon => false,
50        }
51    }
52}
53
54impl std::fmt::Display for MaxSimIsa {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        let s = match self {
57            Self::Auto => "auto",
58            Self::Scalar => "scalar",
59            Self::X86_64_V3 => "x86-64-v3",
60            Self::X86_64_V4 => "x86-64-v4",
61            Self::Neon => "neon",
62            Self::Reference => "reference",
63        };
64        f.write_str(s)
65    }
66}
67
68/// Returned by [`build_max_sim`](super::build_max_sim) when the requested
69/// ISA cannot be produced on the current host (e.g. x86_64 V4 requested on
70/// a non-AVX512 CPU, or Neon requested on x86_64).
71#[derive(Debug, Clone, Copy)]
72pub struct NotSupported {
73    pub isa: MaxSimIsa,
74    pub reason: &'static str,
75}
76
77impl std::fmt::Display for NotSupported {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        write!(f, "{} not supported: {}", self.isa, self.reason)
80    }
81}
82
83impl std::error::Error for NotSupported {}