nx_f32_resblock.nx source
↩ module page · 102 lines · 4610 B
1// nx_f32_resblock.nx -- software-f32 ResBlock (GroupNorm->SiLU->Conv 3x3, twice, + residual), R3/R7 rung.
2//
3// sd-server -> Nishi migration: the canonical conv ResBlock is the repeating unit of the VAE decoder and
4// the SD/UNet stages. This composes the now-gated sovereign f32 bricks:
5// nx_f32_groupnorm + nx_f32_silu (nx_f32_activations) + nx_f32_conv2d
6// into: out = x + Conv2( SiLU( GN2( Conv1( SiLU( GN1(x) ) ) ) ) )
7// v1: C_in == C_out, 3x3 convs (pad 1, stride 1), same spatial dims, no time-embedding (the i64 ref
8// nx_unet_block also omits it; time-embed injection is the next addition). Real dequantized weights flow.
9//
10// Buffers: flat *i64 of f32 bits. x [N,C,H,W] (preserved), out [N,C,H,W] (result), scratch [N,C,H,W]
11// (caller-provided intermediate). x, out, scratch MUST be three distinct buffers. gamma/beta [C], conv
12// weights [C,C,3,3], conv bias [C] (nullable=0).
13// license_tier: ORIGINAL
14import "nx_syscalls.nx"
15import "nx_f32.nx"
16import "nx_f32_cvt.nx"
17import "nx_f32_groupnorm.nx"
18import "nx_f32_conv2d.nx"
19import "nx_f32_activations.nx"
20
21const NX_F32RB_OK: i64 = 0
22
23// element-wise SiLU over a flat f32 buffer (in place).
24func nx_f32rb_silu_inplace(buf: *i64, nel: i64) -> i64 {
25 var i: i64 = 0
26 while i < nel { buf[i] = nx_f32_silu(buf[i]); i = i + 1 }
27 return 0
28}
29
30func nx_f32_resblock_forward(x: *i64, N: i64, C: i64, H: i64, W: i64, n_groups: i64,
31 gn1_gamma: *i64, gn1_beta: *i64, conv1_w: *i64, conv1_b: *i64,
32 gn2_gamma: *i64, gn2_beta: *i64, conv2_w: *i64, conv2_b: *i64,
33 out: *i64, scratch: *i64) -> i64 {
34 let nel: i64 = N * C * H * W
35 // scratch = SiLU(GN1(x))
36 let e1: i64 = nx_f32_groupnorm_forward(x, N, C, H, W, n_groups, gn1_gamma, gn1_beta, scratch)
37 if e1 != NX_F32GN_OK { return 100 + e1 }
38 nx_f32rb_silu_inplace(scratch, nel)
39 // out = Conv1(scratch) (3x3, pad1, stride1, C->C)
40 let e2: i64 = nx_f32_conv2d_forward(scratch, N, C, H, W, conv1_w, C, 3, 3, 1, 1, conv1_b, out)
41 if e2 != NX_F32CV_OK { return 200 + e2 }
42 // scratch = SiLU(GN2(out))
43 let e3: i64 = nx_f32_groupnorm_forward(out, N, C, H, W, n_groups, gn2_gamma, gn2_beta, scratch)
44 if e3 != NX_F32GN_OK { return 300 + e3 }
45 nx_f32rb_silu_inplace(scratch, nel)
46 // out = Conv2(scratch)
47 let e4: i64 = nx_f32_conv2d_forward(scratch, N, C, H, W, conv2_w, C, 3, 3, 1, 1, conv2_b, out)
48 if e4 != NX_F32CV_OK { return 400 + e4 }
49 // residual: out = x + out
50 var i: i64 = 0
51 while i < nel { out[i] = nx_f32_add(out[i], x[i]); i = i + 1 }
52 return NX_F32RB_OK
53}
54
55// ===== Self-test (inline gate) ====================================
56// (a) zero conv weights -> the residual branch is 0 -> out == x (bit-exact: proves residual + plumbing)
57// (b) identity convs -> out = x + SiLU(GN2(SiLU(GN1(x)))) != x somewhere (the branch contributes)
58func main() -> i64 {
59 let N: i64 = 1
60 let C: i64 = 2
61 let H: i64 = 2
62 let W: i64 = 2
63 let G: i64 = 2
64 let nel: i64 = N * C * H * W
65 let x: *i64 = sys_mmap(nel * 8) as *i64
66 let out: *i64 = sys_mmap(nel * 8) as *i64
67 let scr: *i64 = sys_mmap(nel * 8) as *i64
68 let g1: *i64 = sys_mmap(C * 8) as *i64
69 let b1: *i64 = sys_mmap(C * 8) as *i64
70 let g2: *i64 = sys_mmap(C * 8) as *i64
71 let b2: *i64 = sys_mmap(C * 8) as *i64
72 let w1: *i64 = sys_mmap(C * C * 9 * 8) as *i64
73 let w2: *i64 = sys_mmap(C * C * 9 * 8) as *i64
74 let one: i64 = nx_i32_to_f32(1)
75
76 var i: i64 = 0
77 while i < C { g1[i] = one; b1[i] = 0; g2[i] = one; b2[i] = 0; i = i + 1 }
78 i = 0
79 while i < nel { x[i] = nx_i32_to_f32(i + 1); i = i + 1 }
80
81 // (a) zero conv weights -> out == x
82 i = 0
83 while i < C * C * 9 { w1[i] = 0; w2[i] = 0; i = i + 1 }
84 let va: i64 = nx_f32_resblock_forward(x, N, C, H, W, G, g1, b1, w1, 0 as *i64, g2, b2, w2, 0 as *i64, out, scr)
85 if va != NX_F32RB_OK { return 10 }
86 i = 0
87 while i < nel { if out[i] != x[i] { return 20 } i = i + 1 }
88
89 // (b) identity convs (centre 1.0 where c_out == c_in) -> branch contributes -> out != x somewhere
90 i = 0
91 while i < C * C * 9 { w1[i] = 0; w2[i] = 0; i = i + 1 }
92 var co: i64 = 0
93 while co < C { w1[(co * C + co) * 9 + 4] = one; w2[(co * C + co) * 9 + 4] = one; co = co + 1 }
94 let vb: i64 = nx_f32_resblock_forward(x, N, C, H, W, G, g1, b1, w1, 0 as *i64, g2, b2, w2, 0 as *i64, out, scr)
95 if vb != NX_F32RB_OK { return 30 }
96 var diff: i64 = 0
97 i = 0
98 while i < nel { if out[i] != x[i] { diff = 1 } i = i + 1 }
99 if diff == 0 { return 40 }
100
101 return 0
102}