nx_f32_conv2d_backward.nx source
↩ module page · 89 lines · 4455 B
1// nx_f32_conv2d_backward.nx -- the GRADIENT of the software-f32 2D convolution: the backward of
2// nx_f32_conv2d_forward. From dL/dOut it computes dL/dInput, dL/dWeight, dL/dBias. This is the ONE training
3// primitive the sovereign ML substrate lacked for CONV nets -- the autograd/Adam/cross-entropy tower
4// (nx_nofloat_autograd, nx_f32_adam, nx_ta_transformer_gradcheck) was matmul/attention-only -- so it is the first
5// brick of TRAINING our own sovereign pose net (operator: "build 2"). Same NCHW layout + stride/pad index math as
6// the forward, exact-mirrored: every (n,co,oh,ow,ci,kh,kw) that contributed input[in_idx]*weight[wt_idx] to out
7// now scatter-adds dout*input into dweight, dout*weight into dinput, and dout into dbias. Composes ONLY nx_f32_mul /
8// nx_f32_add (no libm). Gradient-CHECKED by nx_f32_conv2d_backward_gate (exact integer-f32 gradients, no tolerance).
9// license_tier: ORIGINAL
10import "nx_syscalls.nx"
11import "nx_f32.nx"
12
13const NX_F32CVB_OK: i64 = 0
14const NX_F32CVB_ERR_BAD_ARGS: i64 = 4
15
16// dinput [N,C_in,H,W], dweight [C_out,C_in,KH,KW], dbias [C_out] (nullable=0) are ZEROED here then accumulated.
17// dout is [N,C_out,OH,OW]. Caller allocates all four. In-place NOT supported.
18func nx_f32_conv2d_backward(input: *i64, N: i64, C_in: i64, H: i64, W: i64,
19 weight: *i64, C_out: i64, KH: i64, KW: i64,
20 stride: i64, pad: i64,
21 dout: *i64, dinput: *i64, dweight: *i64, dbias: *i64) -> i64 {
22 if stride <= 0 { return NX_F32CVB_ERR_BAD_ARGS }
23 if KH <= 0 { return NX_F32CVB_ERR_BAD_ARGS }
24 if KW <= 0 { return NX_F32CVB_ERR_BAD_ARGS }
25 if C_in <= 0 { return NX_F32CVB_ERR_BAD_ARGS }
26 if C_out <= 0 { return NX_F32CVB_ERR_BAD_ARGS }
27 let OH: i64 = (H + 2 * pad - KH) / stride + 1
28 let OW: i64 = (W + 2 * pad - KW) / stride + 1
29 if OH <= 0 { return NX_F32CVB_ERR_BAD_ARGS }
30 if OW <= 0 { return NX_F32CVB_ERR_BAD_ARGS }
31 let in_chan_stride: i64 = H * W
32 let in_batch_stride: i64 = C_in * H * W
33 let wt_chan_stride: i64 = KH * KW
34 let wt_outchan_stride: i64 = C_in * KH * KW
35 let out_chan_stride: i64 = OH * OW
36 let out_batch_stride: i64 = C_out * OH * OW
37
38 // zero the grad accumulators (backward computes from scratch; +0.0 == bits 0).
39 var z: i64 = 0
40 let nin: i64 = N * C_in * H * W
41 while z < nin { dinput[z] = 0; z = z + 1 }
42 z = 0
43 let nwt: i64 = C_out * C_in * KH * KW
44 while z < nwt { dweight[z] = 0; z = z + 1 }
45 if (dbias as i64) != 0 { z = 0; while z < C_out { dbias[z] = 0; z = z + 1 } }
46
47 var n: i64 = 0
48 while n < N {
49 var co: i64 = 0
50 while co < C_out {
51 var oh: i64 = 0
52 while oh < OH {
53 var ow: i64 = 0
54 while ow < OW {
55 let out_idx: i64 = n * out_batch_stride + co * out_chan_stride + oh * OW + ow
56 let dov: i64 = dout[out_idx]
57 if (dbias as i64) != 0 { dbias[co] = nx_f32_add(dbias[co], dov) }
58 var ci: i64 = 0
59 while ci < C_in {
60 var kh: i64 = 0
61 while kh < KH {
62 let ih: i64 = oh * stride + kh - pad
63 if ih >= 0 { if ih < H {
64 var kw: i64 = 0
65 while kw < KW {
66 let iw: i64 = ow * stride + kw - pad
67 if iw >= 0 { if iw < W {
68 let in_idx: i64 = n * in_batch_stride + ci * in_chan_stride + ih * W + iw
69 let wt_idx: i64 = co * wt_outchan_stride + ci * wt_chan_stride + kh * KW + kw
70 dweight[wt_idx] = nx_f32_add(dweight[wt_idx], nx_f32_mul(dov, input[in_idx]))
71 dinput[in_idx] = nx_f32_add(dinput[in_idx], nx_f32_mul(dov, weight[wt_idx]))
72 } }
73 kw = kw + 1
74 }
75 } }
76 kh = kh + 1
77 }
78 ci = ci + 1
79 }
80 ow = ow + 1
81 }
82 oh = oh + 1
83 }
84 co = co + 1
85 }
86 n = n + 1
87 }
88 return NX_F32CVB_OK
89}