nx_f32_conv_transpose2d.nx source
↩ module page · 86 lines · 3679 B
1// nx_f32_conv_transpose2d.nx -- software-f32 2D TRANSPOSED convolution / deconv (NCHW), a STANDARD vision op the
2// f32 tower lacked (only a 1-D conv_transpose in nx_vocops). It is the learned-upsampling head of a SimpleBaseline/
3// ViTPose pose net (3 stride-2 deconvs 8x-upsample the low-res features into full heatmaps), and is broadly used by
4// segmentation/generative decoders. Mechanically it is the SCATTER dual of a conv: each input pixel scatter-adds
5// its value * kernel into a strided output window. Weight layout is PyTorch ConvTranspose2d: [C_in, C_out, KH, KW].
6// OH = (H-1)*stride - 2*pad + KH + out_pad. Composes ONLY nx_f32_mul / nx_f32_add. license_tier: ORIGINAL
7import "nx_syscalls.nx"
8import "nx_f32.nx"
9
10const NX_CT_OK: i64 = 0
11const NX_CT_ERR_ARGS: i64 = 4
12
13func nx_f32_conv_transpose2d(input: *i64, N: i64, C_in: i64, H: i64, W: i64,
14 weight: *i64, C_out: i64, KH: i64, KW: i64,
15 stride: i64, pad: i64, out_pad: i64,
16 bias: *i64, out: *i64) -> i64 {
17 if stride <= 0 { return NX_CT_ERR_ARGS }
18 if KH <= 0 { return NX_CT_ERR_ARGS }
19 if KW <= 0 { return NX_CT_ERR_ARGS }
20 if C_in <= 0 { return NX_CT_ERR_ARGS }
21 if C_out <= 0 { return NX_CT_ERR_ARGS }
22 let OH: i64 = (H - 1) * stride - 2 * pad + KH + out_pad
23 let OW: i64 = (W - 1) * stride - 2 * pad + KW + out_pad
24 if OH <= 0 { return NX_CT_ERR_ARGS }
25 if OW <= 0 { return NX_CT_ERR_ARGS }
26 let in_cs: i64 = H * W
27 let in_bs: i64 = C_in * H * W
28 let wt_os: i64 = KH * KW // per (ci,co) kernel
29 let wt_cis: i64 = C_out * KH * KW // per input channel
30 let out_cs: i64 = OH * OW
31 let out_bs: i64 = C_out * OH * OW
32
33 // initialize output to bias (or +0.0) -- scatter accumulates on top.
34 var n: i64 = 0
35 while n < N {
36 var co: i64 = 0
37 while co < C_out {
38 var bv: i64 = 0
39 if (bias as i64) != 0 { bv = bias[co] }
40 var p: i64 = 0
41 while p < out_cs { out[n * out_bs + co * out_cs + p] = bv; p = p + 1 }
42 co = co + 1
43 }
44 n = n + 1
45 }
46
47 n = 0
48 while n < N {
49 var ci: i64 = 0
50 while ci < C_in {
51 var h: i64 = 0
52 while h < H {
53 var w: i64 = 0
54 while w < W {
55 let iv: i64 = input[n * in_bs + ci * in_cs + h * W + w]
56 var co: i64 = 0
57 while co < C_out {
58 var kh: i64 = 0
59 while kh < KH {
60 let oh: i64 = h * stride + kh - pad
61 if oh >= 0 { if oh < OH {
62 var kw: i64 = 0
63 while kw < KW {
64 let ow: i64 = w * stride + kw - pad
65 if ow >= 0 { if ow < OW {
66 let wt_idx: i64 = ci * wt_cis + co * wt_os + kh * KW + kw
67 let o_idx: i64 = n * out_bs + co * out_cs + oh * OW + ow
68 out[o_idx] = nx_f32_add(out[o_idx], nx_f32_mul(iv, weight[wt_idx]))
69 } }
70 kw = kw + 1
71 }
72 } }
73 kh = kh + 1
74 }
75 co = co + 1
76 }
77 w = w + 1
78 }
79 h = h + 1
80 }
81 ci = ci + 1
82 }
83 n = n + 1
84 }
85 return NX_CT_OK
86}