nx_f32_conv2d.nx source
↩ module page · 129 lines · 5888 B
1// nx_f32_conv2d.nx -- software-f32 multi-channel 2D convolution (NCHW), the R3 vision-op re-tier.
2//
3// sd-server -> Nishi migration, rung 3: the existing `nx_conv2d.nx` is i64 Q10, HARDCODED 3x3, stride-1,
4// pad-1, and has only ever run on identity weights. This is the same direct convolution in the sovereign
5// software-f32 tier (`nx_f32_*`, IEEE-754-in-i64, gated) AND generalized to arbitrary KH/KW + stride + pad,
6// so it covers everything the DiT/VAE use: 1x1 projections, 3x3 ResBlock convs, and stride-2 downsamples.
7// Real dequantized GGUF weights flow through it.
8//
9// Tensors: flat `*i64` of f32 bit-patterns. input [N,C_in,H,W], weight [C_out,C_in,KH,KW], bias [C_out]
10// (nullable=0), out [N,C_out,OH,OW] with OH=(H+2*pad-KH)/stride+1, OW=(W+2*pad-KW)/stride+1. Caller
11// allocates out for those dims. Zero-padding. In-place NOT supported (in != out). Composes only gated
12// primitives nx_f32_add / nx_f32_mul. No libm, no Q-scaling (f32 multiply needs no rescale).
13// license_tier: ORIGINAL
14import "nx_syscalls.nx"
15import "nx_f32.nx"
16import "nx_f32_cvt.nx"
17
18const NX_F32CV_OK: i64 = 0
19const NX_F32CV_ERR_BAD_ARGS: i64 = 4
20
21func nx_f32_conv2d_forward(input: *i64, N: i64, C_in: i64, H: i64, W: i64,
22 weight: *i64, C_out: i64, KH: i64, KW: i64,
23 stride: i64, pad: i64,
24 bias: *i64, out: *i64) -> i64 {
25 if stride <= 0 { return NX_F32CV_ERR_BAD_ARGS }
26 if KH <= 0 { return NX_F32CV_ERR_BAD_ARGS }
27 if KW <= 0 { return NX_F32CV_ERR_BAD_ARGS }
28 if C_in <= 0 { return NX_F32CV_ERR_BAD_ARGS }
29 if C_out <= 0 { return NX_F32CV_ERR_BAD_ARGS }
30 let OH: i64 = (H + 2 * pad - KH) / stride + 1
31 let OW: i64 = (W + 2 * pad - KW) / stride + 1
32 if OH <= 0 { return NX_F32CV_ERR_BAD_ARGS }
33 if OW <= 0 { return NX_F32CV_ERR_BAD_ARGS }
34 let in_chan_stride: i64 = H * W
35 let in_batch_stride: i64 = C_in * H * W
36 let wt_chan_stride: i64 = KH * KW
37 let wt_outchan_stride: i64 = C_in * KH * KW
38 let out_chan_stride: i64 = OH * OW
39 let out_batch_stride: i64 = C_out * OH * OW
40 var n: i64 = 0
41 while n < N {
42 var co: i64 = 0
43 while co < C_out {
44 var bv: i64 = 0 // f32 +0.0 if no bias
45 if (bias as i64) != 0 { bv = bias[co] }
46 var oh: i64 = 0
47 while oh < OH {
48 var ow: i64 = 0
49 while ow < OW {
50 var acc: i64 = bv // accumulate in f32, seeded with the bias
51 var ci: i64 = 0
52 while ci < C_in {
53 var kh: i64 = 0
54 while kh < KH {
55 let ih: i64 = oh * stride + kh - pad
56 if ih >= 0 { if ih < H {
57 var kw: i64 = 0
58 while kw < KW {
59 let iw: i64 = ow * stride + kw - pad
60 if iw >= 0 { if iw < W {
61 let in_idx: i64 = n * in_batch_stride + ci * in_chan_stride + ih * W + iw
62 let wt_idx: i64 = co * wt_outchan_stride + ci * wt_chan_stride + kh * KW + kw
63 acc = nx_f32_add(acc, nx_f32_mul(input[in_idx], weight[wt_idx]))
64 } }
65 kw = kw + 1
66 }
67 } }
68 kh = kh + 1
69 }
70 ci = ci + 1
71 }
72 let out_idx: i64 = n * out_batch_stride + co * out_chan_stride + oh * OW + ow
73 out[out_idx] = acc
74 ow = ow + 1
75 }
76 oh = oh + 1
77 }
78 co = co + 1
79 }
80 n = n + 1
81 }
82 return NX_F32CV_OK
83}
84
85// ===== Self-test (inline gate) ====================================
86// (a) 3x3 identity kernel (centre 1.0) -> output == input (bit-exact)
87// (b) 1x1 identity kernel -> output == input (bit-exact, proves 1x1)
88// (c) 3x3 all-ones on constant 2.0 -> interior pixel == 18.0 (9*2, exact in f32)
89// (d) bad stride 0 -> ERR_BAD_ARGS
90func main() -> i64 {
91 let inb: *i64 = sys_mmap(64 * 8) as *i64
92 let wt: *i64 = sys_mmap(64 * 8) as *i64
93 let out: *i64 = sys_mmap(64 * 8) as *i64
94 let one: i64 = nx_i32_to_f32(1)
95 let zero: i64 = 0
96
97 // (a) 3x3 identity, 1ch 4x4. input[h,w] = h*10+w.
98 var h: i64 = 0
99 while h < 4 { var w: i64 = 0; while w < 4 { inb[h * 4 + w] = nx_i32_to_f32(h * 10 + w); w = w + 1 } h = h + 1 }
100 var i: i64 = 0
101 while i < 9 { wt[i] = zero; i = i + 1 }
102 wt[4] = one // centre of the 3x3
103 let va: i64 = nx_f32_conv2d_forward(inb, 1, 1, 4, 4, wt, 1, 3, 3, 1, 1, 0 as *i64, out)
104 if va != NX_F32CV_OK { return 10 }
105 i = 0
106 while i < 16 { if out[i] != inb[i] { return 20 } i = i + 1 }
107
108 // (b) 1x1 identity (pad 0)
109 wt[0] = one
110 let vb: i64 = nx_f32_conv2d_forward(inb, 1, 1, 4, 4, wt, 1, 1, 1, 1, 0, 0 as *i64, out)
111 if vb != NX_F32CV_OK { return 25 }
112 i = 0
113 while i < 16 { if out[i] != inb[i] { return 26 } i = i + 1 }
114
115 // (c) 3x3 all-ones on constant 2.0 -> interior(1,1) == 18.0
116 var j: i64 = 0
117 while j < 16 { inb[j] = nx_i32_to_f32(2); j = j + 1 }
118 i = 0
119 while i < 9 { wt[i] = one; i = i + 1 }
120 let vc: i64 = nx_f32_conv2d_forward(inb, 1, 1, 4, 4, wt, 1, 3, 3, 1, 1, 0 as *i64, out)
121 if vc != NX_F32CV_OK { return 30 }
122 if out[1 * 4 + 1] != nx_i32_to_f32(18) { return 31 }
123
124 // (d) bad stride
125 let vd: i64 = nx_f32_conv2d_forward(inb, 1, 1, 4, 4, wt, 1, 3, 3, 0, 1, 0 as *i64, out)
126 if vd == NX_F32CV_OK { return 40 }
127
128 return 0
129}