Skip to main content

diskann_quantization/multi_vector/distance/
isa.rs

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