thermite 0.2.0-beta.0

High-performance, generic, ISA-portable SIMD library with a policy-configurable transcendental math library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use core::cell::UnsafeCell;
use core::sync::atomic::{AtomicU8, Ordering};

use super::InstructionSet;

const UNINITIALIZED: u8 = 0;
const INITIALIZING: u8 = 1;
const INITIALIZED: u8 = 2;

pub struct DetectInstructionSet {
    state: AtomicU8,
    isa: UnsafeCell<Option<InstructionSet>>,
}

// SAFETY: We ensure proper synchronization using atomic operations
unsafe impl Sync for DetectInstructionSet {}

impl DetectInstructionSet {
    pub const fn new() -> Self {
        Self {
            state: AtomicU8::new(UNINITIALIZED),
            isa: UnsafeCell::new(None),
        }
    }

    #[inline(never)] #[rustfmt::skip]
    fn initialize(&self) -> InstructionSet {
        let previous_state = self.state.compare_exchange(UNINITIALIZED, INITIALIZING, Ordering::AcqRel, Ordering::Acquire);

        match previous_state.unwrap_or_else(|s| s) {
            INITIALIZING => while self.state.load(Ordering::Acquire) != INITIALIZED {
                core::hint::spin_loop(); // Wait for initialization to complete
            }
            UNINITIALIZED => {
                // SAFETY: We are the only thread initializing at this point
                unsafe { *self.isa.get() = Some(Self::detect_internal()) };
                self.state.store(INITIALIZED, Ordering::Release);
            }
            _ => {} // Already initialized
        }

        // SAFETY: At this point, the state is guaranteed to be INITIALIZED
        unsafe { (*self.isa.get()).unwrap_unchecked() }
    }

    #[inline(always)]
    pub fn get_or_init(&self) -> InstructionSet {
        if self.state.load(Ordering::Acquire) == INITIALIZED {
            // SAFETY: The state is INITIALIZED, so isa is guaranteed to be Some
            return unsafe { (*self.isa.get()).unwrap_unchecked() };
        }

        self.initialize()
    }

    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    fn detect_internal() -> InstructionSet {
        let mut best = InstructionSet::Scalar;

        if core_detect::is_x86_feature_detected!("avx512f") {
            best = InstructionSet::X86V4; // TODO: Check if more AVX512 features are needed
        } else if core_detect::is_x86_feature_detected!("avx2") && core_detect::is_x86_feature_detected!("fma") {
            best = InstructionSet::X86V3;
        } else if core_detect::is_x86_feature_detected!("sse4.2") && core_detect::is_x86_feature_detected!("popcnt") {
            best = InstructionSet::X86V2;
        } else if core_detect::is_x86_feature_detected!("sse2") {
            best = InstructionSet::X86V1;
        }

        best
    }

    #[cfg(all(feature = "neon", any(target_arch = "arm", target_arch = "aarch64")))]
    fn detect_internal() -> InstructionSet {
        InstructionSet::NEON // Assume Neon is always available on ARM with the feature enabled
    }

    #[cfg(all(feature = "wasm", target_arch = "wasm32"))]
    fn detect_internal() -> InstructionSet {
        InstructionSet::WASM32 // Assume Wasm SIMD is always available on wasm32 with the feature enabled
    }

    #[cfg(all(feature = "wasm", target_arch = "wasm64"))]
    fn detect_internal() -> InstructionSet {
        InstructionSet::WASM64 // Assume Wasm SIMD is always available on wasm64 with the feature enabled
    }

    #[cfg(not(any(
        all(feature = "neon", any(target_arch = "arm", target_arch = "aarch64")),
        all(feature = "wasm", any(target_arch = "wasm32", target_arch = "wasm64")),
        any(target_arch = "x86", target_arch = "x86_64")
    )))]
    fn detect_internal() -> InstructionSet {
        InstructionSet::Scalar // Fallback for unsupported architectures
    }
}