nx_f32_gelu.nx source
↩ module page · 112 lines · 5528 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_sw(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_sw(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_sw(nx_f32_mul(x, inv_sqrt2))
49 return nx_f32_mul(nx_f32_mul(half, x), nx_f32_add(one, e))
50}
51
52// ---- HARDWARE TWINS (2026-09-14, search R0 cross-encoder): the SAME arithmetic on the __f32 intrinsics (IEEE
53// round-to-nearest on both sides, so the bits are identical -- DIFFERENTIALLY GATED against the _sw oracles above by
54// nx_f32_bricks_hw_gate, never asserted), with the A-S constants derived ONCE instead of eight software divisions and a
55// square root per element. Measured need: a BERT cross-encoder evaluates GELU three million times per query-document
56// pair, and the software chain made the activation the whole cost.
57static g_gelu_ready: i64
58static g_gelu_one: i64
59static g_gelu_p: i64
60static g_gelu_a1: i64
61static g_gelu_a2: i64
62static g_gelu_a3: i64
63static g_gelu_a4: i64
64static g_gelu_a5: i64
65static g_gelu_half: i64
66static g_gelu_inv_sqrt2: i64
67
68func nx_f32_gelu_init() -> i64 {
69 if g_gelu_ready == 1 { return 0 }
70 let one: i64 = nx_i32_to_f32(1)
71 g_gelu_one = one
72 g_gelu_p = nx_f32_div(nx_i32_to_f32(K_MAGIC_3275911), nx_i32_to_f32(K_MAGIC_10000000))
73 g_gelu_a1 = nx_f32_div(nx_i32_to_f32(K_MAGIC_254829592), nx_i32_to_f32(K_MAGIC_1000000000))
74 g_gelu_a2 = nx_f32_neg(nx_f32_div(nx_i32_to_f32(K_MAGIC_284496736), nx_i32_to_f32(K_MAGIC_1000000000)))
75 g_gelu_a3 = nx_f32_div(nx_i32_to_f32(K_MAGIC_1421413741), nx_i32_to_f32(K_MAGIC_1000000000))
76 g_gelu_a4 = nx_f32_neg(nx_f32_div(nx_i32_to_f32(K_MAGIC_1453152027), nx_i32_to_f32(K_MAGIC_1000000000)))
77 g_gelu_a5 = nx_f32_div(nx_i32_to_f32(K_MAGIC_1061405429), nx_i32_to_f32(K_MAGIC_1000000000))
78 g_gelu_half = nx_f32_div(one, nx_i32_to_f32(2))
79 g_gelu_inv_sqrt2 = nx_f32_div(one, nx_f32_sqrt(nx_i32_to_f32(2)))
80 g_gelu_ready = 1
81 return 0
82}
83
84func nx_f32_erf(z: i64) -> i64 {
85 nx_f32_gelu_init()
86 let one: i64 = g_gelu_one
87 var sgn: i64 = one
88 var az: i64 = z
89 if nx_f32_lt(z, 0) == 1 { sgn = nx_f32_neg(one); az = nx_f32_neg(z) }
90 let t: i64 = __f32_div(one, __f32_add(one, __f32_mul(g_gelu_p, az)))
91 var poly: i64 = __f32_add(__f32_mul(g_gelu_a5, t), g_gelu_a4)
92 poly = __f32_add(__f32_mul(poly, t), g_gelu_a3)
93 poly = __f32_add(__f32_mul(poly, t), g_gelu_a2)
94 poly = __f32_add(__f32_mul(poly, t), g_gelu_a1)
95 poly = __f32_mul(poly, t)
96 let ez: i64 = nx_f32_exp(nx_f32_neg(__f32_mul(az, az)))
97 let e: i64 = __f32_add(one, nx_f32_neg(__f32_mul(poly, ez)))
98 return __f32_mul(sgn, e)
99}
100
101func nx_f32_gelu(x: i64) -> i64 {
102 nx_f32_gelu_init()
103 let e: i64 = nx_f32_erf(__f32_mul(x, g_gelu_inv_sqrt2))
104 return __f32_mul(__f32_mul(g_gelu_half, x), __f32_add(g_gelu_one, e))
105}
106
107// in-place GELU over n f32 values (the MLP activation between fc1 and fc2).
108func nx_f32_gelu_vec(buf: *i64, n: i64) -> i64 {
109 var i: i64 = 0
110 while i < n { buf[i] = nx_f32_gelu(buf[i]); i = i + 1 }
111 return 0
112}