nx_f32_maxpool2d.nx source
↩ module page · 61 lines · 2538 B
1// nx_f32_maxpool2d.nx -- software-f32 2D max-pooling (NCHW), a STANDARD vision op the f32 tower lacked. Almost every
2// CNN stem (ResNet/HRNet pose backbones) does a 3x3 stride-2 maxpool after the first conv; the pose port needs it.
3// Padding positions are EXCLUDED from the max (standard maxpool: pad with -inf, i.e. never selected). Composes only
4// nx_f32_gt (IEEE-754 magnitude compare), so it is correct for any f32 (unlike raw bit compare on negatives).
5// license_tier: ORIGINAL
6import "nx_syscalls.nx"
7import "nx_f32.nx"
8
9const NX_MP_OK: i64 = 0
10const NX_MP_ERR_ARGS: i64 = 4
11
12// input [N,C,H,W] -> out [N,C,OH,OW], OH=(H+2pad-KH)/stride+1. Caller allocates out.
13func nx_f32_maxpool2d(input: *i64, N: i64, C: i64, H: i64, W: i64, KH: i64, KW: i64, stride: i64, pad: i64, out: *i64) -> i64 {
14 if stride <= 0 { return NX_MP_ERR_ARGS }
15 if KH <= 0 { return NX_MP_ERR_ARGS }
16 if KW <= 0 { return NX_MP_ERR_ARGS }
17 let OH: i64 = (H + 2 * pad - KH) / stride + 1
18 let OW: i64 = (W + 2 * pad - KW) / stride + 1
19 if OH <= 0 { return NX_MP_ERR_ARGS }
20 if OW <= 0 { return NX_MP_ERR_ARGS }
21 let cs: i64 = H * W
22 let bs: i64 = C * H * W
23 let ocs: i64 = OH * OW
24 let obs: i64 = C * OH * OW
25 var n: i64 = 0
26 while n < N {
27 var c: i64 = 0
28 while c < C {
29 var oh: i64 = 0
30 while oh < OH {
31 var ow: i64 = 0
32 while ow < OW {
33 var best: i64 = 0
34 var seen: i64 = 0
35 var kh: i64 = 0
36 while kh < KH {
37 let ih: i64 = oh * stride + kh - pad
38 if ih >= 0 { if ih < H {
39 var kw: i64 = 0
40 while kw < KW {
41 let iw: i64 = ow * stride + kw - pad
42 if iw >= 0 { if iw < W {
43 let v: i64 = input[n * bs + c * cs + ih * W + iw]
44 if seen == 0 { best = v; seen = 1 } else { if nx_f32_gt(v, best) == 1 { best = v } }
45 } }
46 kw = kw + 1
47 }
48 } }
49 kh = kh + 1
50 }
51 out[n * obs + c * ocs + oh * OW + ow] = best
52 ow = ow + 1
53 }
54 oh = oh + 1
55 }
56 c = c + 1
57 }
58 n = n + 1
59 }
60 return NX_MP_OK
61}