nexar 0.1.2-beta.1

Distributed runtime with QUIC transport, stream-multiplexed messaging, and built-in collectives
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

//! SIMD-accelerated reduction dispatch for multiple architectures and data types.

#[cfg(target_arch = "x86_64")]
mod x86_64;

#[cfg(target_arch = "aarch64")]
mod aarch64;

use crate::types::ReduceOp;

/// Try SIMD-accelerated f32 reduction. Returns `true` if handled.
///
/// # Safety
/// `dst` and `src` must both have at least `count * 4` bytes.
pub(crate) unsafe fn reduce_f32_simd(
    dst: &mut [u8],
    src: &[u8],
    count: usize,
    op: ReduceOp,
) -> bool {
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx512f") {
            unsafe { x86_64::reduce_f32_op_avx512(dst, src, count, op) };
            return true;
        }
        if is_x86_feature_detected!("avx2") {
            unsafe { x86_64::reduce_f32_op_avx2(dst, src, count, op) };
            return true;
        }
    }
    #[cfg(target_arch = "aarch64")]
    {
        unsafe { aarch64::reduce_f32_op_neon(dst, src, count, op) };
        return true;
    }
    #[allow(unreachable_code)]
    false
}

/// Try SIMD-accelerated f64 reduction. Returns `true` if handled.
///
/// # Safety
/// `dst` and `src` must both have at least `count * 8` bytes.
pub(crate) unsafe fn reduce_f64_simd(
    dst: &mut [u8],
    src: &[u8],
    count: usize,
    op: ReduceOp,
) -> bool {
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx512f") {
            unsafe { x86_64::reduce_f64_op_avx512(dst, src, count, op) };
            return true;
        }
        if is_x86_feature_detected!("avx2") {
            unsafe { x86_64::reduce_f64_op_avx2(dst, src, count, op) };
            return true;
        }
    }
    #[cfg(target_arch = "aarch64")]
    {
        unsafe { aarch64::reduce_f64_op_neon(dst, src, count, op) };
        return true;
    }
    #[allow(unreachable_code)]
    false
}

/// Try SIMD-accelerated bf16 reduction. Returns `true` if handled.
///
/// # Safety
/// `dst` and `src` must both have at least `count * 2` bytes.
pub(crate) unsafe fn reduce_bf16_simd(
    dst: &mut [u8],
    src: &[u8],
    count: usize,
    op: ReduceOp,
) -> bool {
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx2") {
            unsafe { x86_64::reduce_bf16_op_avx2(dst, src, count, op) };
            return true;
        }
    }
    // aarch64: scalar fallback (NEON fp16 support inconsistent across chips)
    let _ = (dst, src, count, op);
    false
}

/// Try SIMD-accelerated f16 reduction. Returns `true` if handled.
///
/// # Safety
/// `dst` and `src` must both have at least `count * 2` bytes.
pub(crate) unsafe fn reduce_f16_simd(
    dst: &mut [u8],
    src: &[u8],
    count: usize,
    op: ReduceOp,
) -> bool {
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("f16c") {
            unsafe { x86_64::reduce_f16_op_f16c(dst, src, count, op) };
            return true;
        }
    }
    // aarch64: scalar fallback
    let _ = (dst, src, count, op);
    false
}