nx_f32_activations.nx source
↩ module page · 67 lines · 2487 B
1// nx_f32_activations.nx -- bits-up f32 ML activation functions.
2//
3// L7 composition brick. Composes L4 mul/add/sub/div + L6 exp.
4// Provides sigmoid, SiLU/Swish, tanh -- the standard activation
5// set used in Llama / Mistral / Qwen / GPT-class transformers.
6//
7// All functions: scalar f32 in, scalar f32 out. Caller broadcasts
8// across tensor elements. No libm.
9//
10// Sigmoid: sigma(x) = 1 / (1 + exp(-x))
11// Endpoints: x -> +inf : 1.0
12// x -> -inf : +0
13// Our exp clamps very-negative-arg to 0 -> sigma -> 1,
14// and very-positive-arg to +inf -> sigma -> 0. Both
15// extreme cases handled correctly without underflow path.
16//
17// SiLU / Swish: silu(x) = x * sigma(x)
18// Used as the gate activation in Llama FFN SwiGLU.
19// Hendrycks+Gimpel 2016 (GELU); Ramachandran+ 2017 (Swish);
20// Elfwing+ 2017 (SiLU).
21//
22// Tanh: tanh(x) = 2 * sigma(2x) - 1
23// Cheaper than the canonical (exp(x)-exp(-x))/(exp(x)+exp(-x))
24// form via the sigmoid identity.
25
26import "nx_syscalls.nx"
27import "nx_tier.nx"
28import "nx_f32.nx"
29import "nx_f32_div.nx"
30import "nx_f32_exp.nx"
31
32const NX_F32_ACT_ONE: i64 = 0x3F800000 // 1.0
33const NX_F32_ACT_TWO: i64 = 0x40000000 // 2.0
34
35// ===== Sigmoid =====================================================
36
37func nx_f32_sigmoid(x: i64) -> i64 {
38 let neg_x: i64 = nx_f32_neg(x)
39 let e: i64 = nx_f32_exp(neg_x)
40 // HARDWARE SSE add/div (2026-07-10): IEEE-754, bit-identical to the
41 // software path (gated by nx_f32_activations_test) -- the silu hot loop.
42 let denom: i64 = __f32_add(NX_F32_ACT_ONE, e)
43 return __f32_div(NX_F32_ACT_ONE, denom)
44}
45
46// ===== SiLU / Swish ================================================
47
48func nx_f32_silu(x: i64) -> i64 {
49 let s: i64 = nx_f32_sigmoid(x)
50 return __f32_mul(x, s)
51}
52
53// ===== Tanh =======================================================
54//
55// tanh(x) = 2 * sigma(2x) - 1.
56// Verified identity: 2*(1/(1+exp(-2x))) - 1
57// = (2 - (1 + exp(-2x))) / (1 + exp(-2x))
58// = (1 - exp(-2x)) / (1 + exp(-2x))
59// = (e^x - e^{-x}) / (e^x + e^{-x}) (multiplying both by e^x)
60// = tanh(x)
61
62func nx_f32_tanh(x: i64) -> i64 {
63 let two_x: i64 = nx_f32_mul(NX_F32_ACT_TWO, x)
64 let s: i64 = nx_f32_sigmoid(two_x)
65 let two_s: i64 = nx_f32_mul(NX_F32_ACT_TWO, s)
66 return nx_f32_sub(two_s, NX_F32_ACT_ONE)
67}