nx_f32_softmax.nx source
↩ module page · 115 lines · 3411 B
1// nx_f32_softmax.nx -- bits-up f32 numerically-stable softmax.
2//
3// L7 / L8 composition brick. Composes L4 mul/add/sub/div + L6 exp.
4//
5// Algorithm (numerically stable via max-subtraction):
6// m = max_i x[i]
7// shifted_i = x[i] - m
8// exps_i = exp(shifted_i) (each in (0, 1])
9// sum_exp = sum_i exps_i
10// out[i] = exps_i / sum_exp
11//
12// Max-subtraction guards against exp overflow without changing the
13// final probabilities (numerator+denominator scale by the same
14// 1/exp(m) factor, which cancels).
15//
16// References absorbed clean-room:
17// Bridle 1990 (original softmax)
18// Standard logsumexp / max-subtract trick (numerical canon)
19//
20// genealogy_id: bridle_1990_softmax + logsumexp_max_subtract
21// lineage_id: substrate_f32_softmax_v1_stable
22
23import "nx_syscalls.nx"
24import "nx_tier.nx"
25import "nx_f32.nx"
26import "nx_f32_div.nx"
27import "nx_f32_exp.nx"
28
29const NX_F32_SM_OK: nx_int = 0
30const NX_F32_SM_ERR_BAD_DIM: nx_int = 1
31const NX_F32_SM_N_VERDICTS: nx_int = 2
32
33func nx_f32_sm_verdict_is_valid(v: nx_int) -> nx_int {
34 if v < 0 { return 0 }
35 if v >= NX_F32_SM_N_VERDICTS { return 0 }
36 return 1
37}
38
39// Returns 1 iff a > b in f32 (NaN-naive; we don't expect NaN inputs).
40// Compose-only -- subtract and check sign + non-zero.
41
42func _f32_gt(a: i64, b: i64) -> nx_int {
43 let diff: i64 = nx_f32_sub(a, b)
44 let cls: nx_int = nx_f32_classify(diff)
45 if cls == NX_F32_CLS_ZERO { return 0 }
46 if cls == NX_F32_CLS_NAN { return 0 }
47 if nx_f32_sign(diff) == 1 { return 0 }
48 return 1
49}
50
51// ---- HARDWARE TWIN (2026-09-14, search R0 cross-encoder): the same three passes on the __f32 intrinsics; the shift
52// is the sign-flipped add nx_f32_sub computes, exp is the tower's own (already on the intrinsics). DIFFERENTIALLY
53// GATED bit for bit against nx_f32_softmax_sw below by nx_f32_bricks_hw_gate.
54func nx_f32_softmax(x: *i64, n: nx_int, out: *i64) -> nx_int {
55 if n <= 0 { return NX_F32_SM_ERR_BAD_DIM }
56 var m: i64 = x[0]
57 var i: nx_int = 1
58 while i < n {
59 if _f32_gt(x[i], m) == 1 { m = x[i] }
60 i = i + 1
61 }
62 let nm: i64 = m ^ 0x80000000
63 var sum: i64 = 0
64 var j: nx_int = 0
65 while j < n {
66 let e: i64 = nx_f32_exp(__f32_add(x[j], nm))
67 out[j] = e
68 sum = __f32_add(sum, e)
69 j = j + 1
70 }
71 var k: nx_int = 0
72 while k < n {
73 out[k] = __f32_div(out[k], sum)
74 k = k + 1
75 }
76 return NX_F32_SM_OK
77}
78
79// Stable softmax over n f32 values. In-place permitted (out == x). THE ORACLE, kept verbatim.
80//
81// Args:
82// x n f32 input raw-bit values
83// n positive count
84// out n f32 output raw-bit slots
85
86func nx_f32_softmax_sw(x: *i64, n: nx_int, out: *i64) -> nx_int {
87 if n <= 0 { return NX_F32_SM_ERR_BAD_DIM }
88
89 // Pass 1: find max.
90 var m: i64 = x[0]
91 var i: nx_int = 1
92 while i < n {
93 if _f32_gt(x[i], m) == 1 { m = x[i] }
94 i = i + 1
95 }
96
97 // Pass 2: exp(x[i] - m) into out; accumulate sum.
98 var sum: i64 = 0
99 var j: nx_int = 0
100 while j < n {
101 let shifted: i64 = nx_f32_sub(x[j], m)
102 let e: i64 = nx_f32_exp(shifted)
103 out[j] = e
104 sum = nx_f32_add(sum, e)
105 j = j + 1
106 }
107
108 // Pass 3: divide by sum.
109 var k: nx_int = 0
110 while k < n {
111 out[k] = nx_f32_div(out[k], sum)
112 k = k + 1
113 }
114 return NX_F32_SM_OK
115}