numr 0.5.2

High-performance numerical computing with multi-backend GPU acceleration (CPU/CUDA/WebGPU)
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
114
115
116
117
118
119
120
121

//! NEON softmax backward kernels for ARM64.
//!
//! Fused 2-pass: SIMD dot product, then SIMD elementwise output * (grad - dot).

#[cfg(target_arch = "aarch64")]
use std::arch::aarch64::*;

use super::super::super::math::aarch64::neon::{hsum_f32, hsum_f64};

const F32_LANES: usize = 4;
const F64_LANES: usize = 2;

/// NEON softmax backward for f32.
///
/// # Safety
/// - CPU must support NEON (always true on AArch64)
/// - All pointers must point to `outer_size * dim_size` valid f32 elements
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn softmax_bwd_f32(
    grad: *const f32,
    output: *const f32,
    d_input: *mut f32,
    outer_size: usize,
    dim_size: usize,
) {
    let chunks = dim_size / F32_LANES;
    let remainder = dim_size % F32_LANES;

    for o in 0..outer_size {
        let g_base = grad.add(o * dim_size);
        let o_base = output.add(o * dim_size);
        let d_base = d_input.add(o * dim_size);

        // Pass 1: SIMD dot product
        let mut dot_acc = vdupq_n_f32(0.0);
        for i in 0..chunks {
            let offset = i * F32_LANES;
            let g = vld1q_f32(g_base.add(offset));
            let out = vld1q_f32(o_base.add(offset));
            dot_acc = vfmaq_f32(dot_acc, g, out);
        }
        let mut dot = hsum_f32(dot_acc);

        for i in 0..remainder {
            let offset = chunks * F32_LANES + i;
            dot += *g_base.add(offset) * *o_base.add(offset);
        }

        // Pass 2: d_input = output * (grad - dot)
        let v_dot = vdupq_n_f32(dot);
        for i in 0..chunks {
            let offset = i * F32_LANES;
            let g = vld1q_f32(g_base.add(offset));
            let out = vld1q_f32(o_base.add(offset));
            let shifted = vsubq_f32(g, v_dot);
            let result = vmulq_f32(out, shifted);
            vst1q_f32(d_base.add(offset), result);
        }

        for i in 0..remainder {
            let offset = chunks * F32_LANES + i;
            *d_base.add(offset) = *o_base.add(offset) * (*g_base.add(offset) - dot);
        }
    }
}

/// NEON softmax backward for f64.
///
/// # Safety
/// - CPU must support NEON (always true on AArch64)
/// - All pointers must point to `outer_size * dim_size` valid f64 elements
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn softmax_bwd_f64(
    grad: *const f64,
    output: *const f64,
    d_input: *mut f64,
    outer_size: usize,
    dim_size: usize,
) {
    let chunks = dim_size / F64_LANES;
    let remainder = dim_size % F64_LANES;

    for o in 0..outer_size {
        let g_base = grad.add(o * dim_size);
        let o_base = output.add(o * dim_size);
        let d_base = d_input.add(o * dim_size);

        // Pass 1: SIMD dot product
        let mut dot_acc = vdupq_n_f64(0.0);
        for i in 0..chunks {
            let offset = i * F64_LANES;
            let g = vld1q_f64(g_base.add(offset));
            let out = vld1q_f64(o_base.add(offset));
            dot_acc = vfmaq_f64(dot_acc, g, out);
        }
        let mut dot = hsum_f64(dot_acc);

        for i in 0..remainder {
            let offset = chunks * F64_LANES + i;
            dot += *g_base.add(offset) * *o_base.add(offset);
        }

        // Pass 2: d_input = output * (grad - dot)
        let v_dot = vdupq_n_f64(dot);
        for i in 0..chunks {
            let offset = i * F64_LANES;
            let g = vld1q_f64(g_base.add(offset));
            let out = vld1q_f64(o_base.add(offset));
            let shifted = vsubq_f64(g, v_dot);
            let result = vmulq_f64(out, shifted);
            vst1q_f64(d_base.add(offset), result);
        }

        for i in 0..remainder {
            let offset = chunks * F64_LANES + i;
            *d_base.add(offset) = *o_base.add(offset) * (*g_base.add(offset) - dot);
        }
    }
}