Skip to main content

aria_kernel/
lib.rs

1//! Scalar (+ aarch64 NEON / x86 AVX2 + optional CUDA) kernels for Aria engine.
2
3mod compute;
4mod cuda_rt;
5mod error;
6mod ops;
7
8pub use compute::{cpu_simd_label, resolve_compute, ComputeBackend, ComputePref};
9pub use cuda_rt::CudaContext;
10pub use error::EngineError;
11pub use ops::*;
12
13/// Runtime SIMD selection. Tests force [`SimdMode::Scalar`].
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
15pub enum SimdMode {
16    #[default]
17    Scalar,
18    Neon,
19    Avx2,
20}
21
22impl SimdMode {
23    /// Prefer Neon on aarch64, AVX2 on x86_64 when detected, else Scalar.
24    pub fn auto() -> Self {
25        #[cfg(target_arch = "aarch64")]
26        {
27            Self::Neon
28        }
29        #[cfg(target_arch = "x86_64")]
30        {
31            if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
32                Self::Avx2
33            } else {
34                Self::Scalar
35            }
36        }
37        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
38        {
39            Self::Scalar
40        }
41    }
42}