nx_f32_groupnorm.nx source
↩ module page · 146 lines · 6544 B
1// nx_f32_groupnorm.nx -- software-f32 Group Normalisation (Wu & He 2018), the R3 vision-op re-tier.
2//
3// sd-server -> Nishi migration, rung 3 ("unify the two numeric universes"): the existing diffusion
4// GroupNorm `nx_groupnorm.nx` is i64 Q10 (10 fractional bits) and has only ever run on identity weights;
5// this is the SAME math in the sovereign software-f32 tier (`nx_f32_*`, bit-exact IEEE-754 emulated in
6// i64) so the REAL dequantized GGUF weights (`nx_q4k_to_f32` etc.) flow through it. GroupNorm(32) is the
7// load-bearing norm in SD/Flux/Z-Image UNet + VAE ResBlocks.
8//
9// Tensor representation: a FLAT `*i64` whose every element is an f32 bit-pattern, laid out [N,C,H,W]
10// row-major, with dims passed explicitly. This deliberately avoids `nx_tensor`'s F32 dtype, which is
11// DTYPE_GATED / unusable today (TODO_COMPILER_F32_F64 parked) -- the flat-f32-bits convention is how the
12// whole `nx_f32_*` stack already works.
13//
14// Composes ONLY gated primitives: nx_f32_add/sub/mul/sqrt (nx_f32.nx), nx_f32_div (nx_f32_div.nx),
15// nx_i32_to_f32 (nx_f32_cvt.nx). No libm, no third-party. eps = 1e-5 (Z-Image norm_eps), built from
16// cvt+div so there is no hand-coded float bit-pattern.
17//
18// Math (per group g of a sample n): mean,var over (C/G * H * W) elements; y = (x-mean)/sqrt(var+eps)*gamma+beta.
19// license_tier: ORIGINAL
20import "nx_syscalls.nx"
21import "nx_f32.nx"
22import "nx_f32_div.nx"
23import "nx_f32_cvt.nx"
24const NX_MAGIC_100000: i64 = 100000
25
26const NX_F32GN_OK: i64 = 0
27const NX_F32GN_ERR_BAD_GROUPS: i64 = 5
28
29// eps = 1e-5 = 1.0 / 100000.0, assembled from gated cvt + div (no hand-coded bit pattern).
30func nx_f32gn_eps() -> i64 { return nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(NX_MAGIC_100000)) }
31
32// Forward GroupNorm. x,out: flat *i64 of f32 bits, layout [N,C,H,W]. gamma,beta: *i64 f32 bits, length C
33// (beta nullable: pass 0 for no bias). Returns 0 ok / 5 bad group count. In-place (out==x) supported.
34func nx_f32_groupnorm_forward(x: *i64, N: i64, C: i64, H: i64, W: i64, n_groups: i64,
35 gamma: *i64, beta: *i64, out: *i64) -> i64 {
36 if n_groups <= 0 { return NX_F32GN_ERR_BAD_GROUPS }
37 if C - (C / n_groups) * n_groups != 0 { return NX_F32GN_ERR_BAD_GROUPS }
38 let C_per_group: i64 = C / n_groups
39 let chan_stride: i64 = H * W
40 let batch_stride: i64 = C * H * W
41 let group_size: i64 = C_per_group * chan_stride
42 let gs_f32: i64 = nx_i32_to_f32(group_size)
43 let eps: i64 = nx_f32gn_eps()
44 var n: i64 = 0
45 while n < N {
46 var g: i64 = 0
47 while g < n_groups {
48 let c_start: i64 = g * C_per_group
49 let c_end: i64 = c_start + C_per_group
50 // Pass 1: mean over the group (f32 accumulate; 0 == +0.0).
51 var sum: i64 = 0
52 var c: i64 = c_start
53 while c < c_end {
54 let chan_base: i64 = n * batch_stride + c * chan_stride
55 var p: i64 = 0
56 while p < chan_stride { sum = nx_f32_add(sum, x[chan_base + p]); p = p + 1 }
57 c = c + 1
58 }
59 let mean: i64 = nx_f32_div(sum, gs_f32)
60 // Pass 2: variance (centred sum of squares).
61 var ssq: i64 = 0
62 var c2: i64 = c_start
63 while c2 < c_end {
64 let chan_base: i64 = n * batch_stride + c2 * chan_stride
65 var p: i64 = 0
66 while p < chan_stride {
67 let dx: i64 = nx_f32_sub(x[chan_base + p], mean)
68 ssq = nx_f32_add(ssq, nx_f32_mul(dx, dx))
69 p = p + 1
70 }
71 c2 = c2 + 1
72 }
73 let varf: i64 = nx_f32_div(ssq, gs_f32)
74 let std: i64 = nx_f32_sqrt(nx_f32_add(varf, eps))
75 // Pass 3: normalise + per-channel affine.
76 var c3: i64 = c_start
77 while c3 < c_end {
78 let chan_base: i64 = n * batch_stride + c3 * chan_stride
79 let gv: i64 = gamma[c3]
80 var bv: i64 = 0
81 if (beta as i64) != 0 { bv = beta[c3] }
82 var p: i64 = 0
83 while p < chan_stride {
84 let centred: i64 = nx_f32_sub(x[chan_base + p], mean)
85 let norm: i64 = nx_f32_div(centred, std)
86 out[chan_base + p] = nx_f32_add(nx_f32_mul(norm, gv), bv)
87 p = p + 1
88 }
89 c3 = c3 + 1
90 }
91 g = g + 1
92 }
93 n = n + 1
94 }
95 return NX_F32GN_OK
96}
97
98// ===== Self-test (inline gate) ====================================
99// Mirrors nx_groupnorm.nx's closed-form invariants, in f32:
100// (a) constant input -> output ~0 per group (var=0 + eps)
101// (b) shift invariance: GroupNorm(x + k) == GroupNorm(x)
102// (c) bad group count -> ERR_BAD_GROUPS
103// f32 magnitude compare via (bits & 0x7FFFFFFF) -- sign bit cleared; for finite values the
104// masked pattern is monotonic in magnitude, so |a| < tol <=> (a_bits & 0x7FFFFFFF) < tol_bits.
105func main() -> i64 {
106 let N: i64 = 1
107 let C: i64 = 4 // 2 groups of 2 channels
108 let H: i64 = 2
109 let W: i64 = 2
110 let nel: i64 = N * C * H * W
111 let x: *i64 = sys_mmap(nel * 8) as *i64
112 let y: *i64 = sys_mmap(nel * 8) as *i64
113 let gamma: *i64 = sys_mmap(C * 8) as *i64
114 let beta: *i64 = sys_mmap(C * 8) as *i64
115 let one: i64 = nx_i32_to_f32(1)
116 var i: i64 = 0
117 while i < C { gamma[i] = one; beta[i] = 0; i = i + 1 }
118 let tolb: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(100)) & 0x7FFFFFFF // 1e-2
119
120 // (a) constant input 5.0 -> output ~0
121 let five: i64 = nx_i32_to_f32(5)
122 i = 0
123 while i < nel { x[i] = five; i = i + 1 }
124 let va: i64 = nx_f32_groupnorm_forward(x, N, C, H, W, 2, gamma, beta, y)
125 if va != NX_F32GN_OK { return 10 }
126 i = 0
127 while i < nel { if (y[i] & 0x7FFFFFFF) >= tolb { return 20 } i = i + 1 }
128
129 // (b) shift invariance: group1 input = group0 input + 1000
130 i = 0
131 while i < 8 { x[i] = nx_i32_to_f32(i + 1); x[8 + i] = nx_i32_to_f32(i + 1 + 1000); i = i + 1 }
132 let vb: i64 = nx_f32_groupnorm_forward(x, N, C, H, W, 2, gamma, beta, y)
133 if vb != NX_F32GN_OK { return 25 }
134 i = 0
135 while i < 8 {
136 let d: i64 = nx_f32_sub(y[8 + i], y[i])
137 if (d & 0x7FFFFFFF) >= tolb { return 30 }
138 i = i + 1
139 }
140
141 // (c) bad group count (4 / 3 doesn't divide)
142 let vbad: i64 = nx_f32_groupnorm_forward(x, N, C, H, W, 3, gamma, beta, y)
143 if vbad != NX_F32GN_ERR_BAD_GROUPS { return 40 }
144
145 return 0
146}