nx_f32_gelu.nx source
↩ module page · 57 lines · 2985 B
1// nx_f32_gelu.nx -- software-f32 EXACT GELU (erf-based, matching PyTorch nn.GELU default / HF "gelu"), the one
2// transformer activation the f32 tower lacked (the LLM path uses SiLU/SwiGLU; ViT/ViTPose uses GELU). gelu(x) =
3// 0.5*x*(1 + erf(x/sqrt2)). erf is the Abramowitz-Stegun 7.1.26 rational-times-exp approximation (|err| <= 1.5e-7,
4// effectively f32-faithful), composing nx_f32_exp + the gated f32 arithmetic. Reusable ecosystem-wide for any
5// ViT/BERT/GPT-2-class model port, not just the pose benchmark. license_tier: ORIGINAL
6import "nx_syscalls.nx"
7import "nx_f32.nx"
8import "nx_f32_cvt.nx"
9import "nx_f32_div.nx"
10import "nx_f32_exp.nx"
11const K_MAGIC_3275911: i64 = 3275911
12const K_MAGIC_10000000: i64 = 10000000
13const K_MAGIC_254829592: i64 = 254829592
14const K_MAGIC_1000000000: i64 = 1000000000
15const K_MAGIC_284496736: i64 = 284496736
16const K_MAGIC_1421413741: i64 = 1421413741
17const K_MAGIC_1453152027: i64 = 1453152027
18const K_MAGIC_1061405429: i64 = 1061405429
19
20// erf(z) via A&S 7.1.26: t=1/(1+p|z|); erf=sign(z)*(1 - (a1 t + a2 t^2 + ... + a5 t^5) e^{-z^2}).
21func nx_f32_erf(z: i64) -> i64 {
22 let one: i64 = nx_i32_to_f32(1)
23 let p: i64 = nx_f32_div(nx_i32_to_f32(K_MAGIC_3275911), nx_i32_to_f32(K_MAGIC_10000000)) // 0.K_MAGIC_3275911
24 let a1: i64 = nx_f32_div(nx_i32_to_f32(K_MAGIC_254829592), nx_i32_to_f32(K_MAGIC_1000000000)) // 0.K_MAGIC_254829592
25 let a2: i64 = nx_f32_neg(nx_f32_div(nx_i32_to_f32(K_MAGIC_284496736), nx_i32_to_f32(K_MAGIC_1000000000)))
26 let a3: i64 = nx_f32_div(nx_i32_to_f32(K_MAGIC_1421413741), nx_i32_to_f32(K_MAGIC_1000000000))
27 let a4: i64 = nx_f32_neg(nx_f32_div(nx_i32_to_f32(K_MAGIC_1453152027), nx_i32_to_f32(K_MAGIC_1000000000)))
28 let a5: i64 = nx_f32_div(nx_i32_to_f32(K_MAGIC_1061405429), nx_i32_to_f32(K_MAGIC_1000000000))
29 var sgn: i64 = one
30 var az: i64 = z
31 if nx_f32_lt(z, 0) == 1 { sgn = nx_f32_neg(one); az = nx_f32_neg(z) }
32 let t: i64 = nx_f32_div(one, nx_f32_add(one, nx_f32_mul(p, az)))
33 // Horner: ((((a5 t + a4) t + a3) t + a2) t + a1) t
34 var poly: i64 = nx_f32_add(nx_f32_mul(a5, t), a4)
35 poly = nx_f32_add(nx_f32_mul(poly, t), a3)
36 poly = nx_f32_add(nx_f32_mul(poly, t), a2)
37 poly = nx_f32_add(nx_f32_mul(poly, t), a1)
38 poly = nx_f32_mul(poly, t)
39 let ez: i64 = nx_f32_exp(nx_f32_neg(nx_f32_mul(az, az)))
40 let e: i64 = nx_f32_sub(one, nx_f32_mul(poly, ez))
41 return nx_f32_mul(sgn, e)
42}
43
44func nx_f32_gelu(x: i64) -> i64 {
45 let one: i64 = nx_i32_to_f32(1)
46 let half: i64 = nx_f32_div(one, nx_i32_to_f32(2))
47 let inv_sqrt2: i64 = nx_f32_div(one, nx_f32_sqrt(nx_i32_to_f32(2)))
48 let e: i64 = nx_f32_erf(nx_f32_mul(x, inv_sqrt2))
49 return nx_f32_mul(nx_f32_mul(half, x), nx_f32_add(one, e))
50}
51
52// in-place GELU over n f32 values (the MLP activation between fc1 and fc2).
53func nx_f32_gelu_vec(buf: *i64, n: i64) -> i64 {
54 var i: i64 = 0
55 while i < n { buf[i] = nx_f32_gelu(buf[i]); i = i + 1 }
56 return 0
57}