nx_f32_conv2d_fast.nx source
↩ module page · 78 lines · 3564 B
1// nx_f32_conv2d_fast.nx -- FAST software-f32 conv2d (NCHW) via im2col + the FORK-PARALLEL nx_f32_matmul_t. The
2// serial nx_f32_conv2d_forward is the training bottleneck (~4-10s/step for a real pose net); this unfolds each conv
3// into a matrix multiply (col[OH*OW, C_in*KH*KW] @ weight[C_out, C_in*KH*KW]^T) so the MAC work runs on the parallel
4// matmul (~7x). Same NCHW layout + weight order as nx_f32_conv2d, so it composes identically. Output is bit-EXACT
5// vs the serial conv when the accumulation has no rounding (integer-valued); for real weights it differs by <=1 ulp
6// only in bias placement (bias added last vs first) -- negligible for train/infer. Composes nx_f32_matmul_t + add.
7// license_tier: ORIGINAL
8import "nx_syscalls.nx"
9import "nx_f32.nx"
10import "nx_f32_matmul_t.nx"
11
12const NX_CVF_OK: i64 = 0
13const NX_CVF_ERR_ARGS: i64 = 4
14
15func nx_f32_conv2d_fast(input: *i64, N: i64, C_in: i64, H: i64, W: i64,
16 weight: *i64, C_out: i64, KH: i64, KW: i64,
17 stride: i64, pad: i64, bias: *i64, out: *i64) -> i64 {
18 if stride <= 0 { return NX_CVF_ERR_ARGS }
19 if KH <= 0 { return NX_CVF_ERR_ARGS }
20 if KW <= 0 { return NX_CVF_ERR_ARGS }
21 let OH: i64 = (H + 2 * pad - KH) / stride + 1
22 let OW: i64 = (W + 2 * pad - KW) / stride + 1
23 if OH <= 0 { return NX_CVF_ERR_ARGS }
24 if OW <= 0 { return NX_CVF_ERR_ARGS }
25 let K: i64 = C_in * KH * KW
26 let M: i64 = OH * OW
27 let in_bs: i64 = C_in * H * W
28 let out_bs: i64 = C_out * OH * OW
29 let col: *i64 = sys_mmap(8 * M * K) as *i64
30 let tmp: *i64 = sys_mmap(8 * M * C_out) as *i64
31 var n: i64 = 0
32 while n < N {
33 let inb: i64 = n * in_bs
34 // im2col: col[oh*OW+ow, ci*KH*KW+kh*KW+kw] = input[ci, oh*stride+kh-pad, ow*stride+kw-pad] (0 if OOB)
35 var oh: i64 = 0
36 while oh < OH {
37 var ow: i64 = 0
38 while ow < OW {
39 let row: i64 = (oh * OW + ow) * K
40 var ci: i64 = 0
41 while ci < C_in {
42 let icb: i64 = inb + ci * H * W
43 var kh: i64 = 0
44 while kh < KH {
45 let ih: i64 = oh * stride + kh - pad
46 var kw: i64 = 0
47 while kw < KW {
48 let iw: i64 = ow * stride + kw - pad
49 var v: i64 = 0
50 if ih >= 0 { if ih < H { if iw >= 0 { if iw < W { v = input[icb + ih * W + iw] } } } }
51 col[row + ci * KH * KW + kh * KW + kw] = v
52 kw = kw + 1
53 }
54 kh = kh + 1
55 }
56 ci = ci + 1
57 }
58 ow = ow + 1
59 }
60 oh = oh + 1
61 }
62 // tmp[M, C_out] = col[M, K] @ weight[C_out, K]^T (weight is [C_out, C_in*KH*KW] = the NCHW conv weight)
63 nx_f32_matmul_t(col, weight, tmp, M, K, C_out)
64 // transpose + bias: out[co, oh, ow] = tmp[oh*OW+ow, co] + bias[co]
65 var co: i64 = 0
66 while co < C_out {
67 var bv: i64 = 0
68 if (bias as i64) != 0 { bv = bias[co] }
69 let ob: i64 = n * out_bs + co * M
70 var p: i64 = 0
71 while p < M { out[ob + p] = nx_f32_add(tmp[p * C_out + co], bv); p = p + 1 }
72 co = co + 1
73 }
74 n = n + 1
75 }
76 sys_munmap(col as *u8, 8 * M * K); sys_munmap(tmp as *u8, 8 * M * C_out) // free scratch (else training OOMs)
77 return NX_CVF_OK
78}