Skip to main content

vector_ta/utilities/
kernel_validation.rs

1use crate::utilities::enums::Kernel;
2#[cfg(feature = "python")]
3use pyo3::exceptions::PyValueError;
4#[cfg(feature = "python")]
5use pyo3::PyResult;
6
7#[cfg(feature = "python")]
8pub fn validate_kernel(kernel: Option<&str>, batch: bool) -> PyResult<Kernel> {
9    match kernel {
10        None | Some("auto") => Ok(if batch { Kernel::Auto } else { Kernel::Auto }),
11        Some("scalar") => Ok(if batch {
12            Kernel::ScalarBatch
13        } else {
14            Kernel::Scalar
15        }),
16        Some("avx2") => {
17            #[cfg(all(feature = "nightly-avx", target_arch = "x86_64"))]
18            {
19                if std::arch::is_x86_feature_detected!("avx2")
20                    && std::arch::is_x86_feature_detected!("fma")
21                {
22                    Ok(if batch {
23                        Kernel::Avx2Batch
24                    } else {
25                        Kernel::Avx2
26                    })
27                } else {
28                    Err(PyValueError::new_err(
29                        "AVX2 kernel requested but not available on this CPU",
30                    ))
31                }
32            }
33            #[cfg(not(all(feature = "nightly-avx", target_arch = "x86_64")))]
34            Err(PyValueError::new_err(
35                "AVX2 kernel not compiled in this build",
36            ))
37        }
38        Some("avx512") => {
39            #[cfg(all(feature = "nightly-avx", target_arch = "x86_64"))]
40            {
41                if std::arch::is_x86_feature_detected!("avx512f")
42                    && std::arch::is_x86_feature_detected!("fma")
43                {
44                    Ok(if batch {
45                        Kernel::Avx512Batch
46                    } else {
47                        Kernel::Avx512
48                    })
49                } else {
50                    Err(PyValueError::new_err(
51                        "AVX512 kernel requested but not available on this CPU",
52                    ))
53                }
54            }
55            #[cfg(not(all(feature = "nightly-avx", target_arch = "x86_64")))]
56            Err(PyValueError::new_err(
57                "AVX512 kernel not compiled in this build",
58            ))
59        }
60        Some(k) => Err(PyValueError::new_err(format!("Unknown kernel: {}", k))),
61    }
62}