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