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 {}