nx_f32_upsample_bilinear.nx source
↩ module page · 80 lines · 3657 B
1// nx_f32_upsample_bilinear.nx -- software-f32 BILINEAR upsample (NCHW, integer scale, align_corners=FALSE to match
2// PyTorch/HF nn.functional.interpolate). The LAST vision op the pose port needed: ViTPose's SIMPLE decoder is
3// ReLU -> bilinear-upsample(scale_factor) -> Conv2d, so a FAITHFUL port (bit-comparable to the reference) needs true
4// bilinear, not nearest. Trick: with an INTEGER scale the sample coordinate src=(dst+0.5)/scale-0.5 is rational with
5// denominator 2*scale, so ALL coordinate + weight math is exact INTEGER (floor, neighbor indices, the 4 bilinear
6// weights that sum to (2*scale)^2); only the pixel combine + final /den^2 use f32. No f32 floor needed. Border:
7// out-of-range neighbor indices clamp to [0,H-1]/[0,W-1] (PyTorch edge behavior). license_tier: ORIGINAL
8import "nx_syscalls.nx"
9import "nx_f32.nx"
10import "nx_f32_cvt.nx"
11import "nx_f32_div.nx"
12
13const NX_UB_OK: i64 = 0
14const NX_UB_ERR_ARGS: i64 = 4
15
16func nx_f32_upsample_bilinear(input: *i64, N: i64, C: i64, H: i64, W: i64, scale: i64, out: *i64) -> i64 {
17 if scale <= 0 { return NX_UB_ERR_ARGS }
18 let OH: i64 = H * scale
19 let OW: i64 = W * scale
20 let den: i64 = 2 * scale
21 let den2f: i64 = nx_i32_to_f32(den * den)
22 let cs: i64 = H * W
23 let bs: i64 = C * H * W
24 let ocs: i64 = OH * OW
25 let obs: i64 = C * OH * OW
26 var n: i64 = 0
27 while n < N {
28 var c: i64 = 0
29 while c < C {
30 let base: i64 = n * bs + c * cs
31 var oy: i64 = 0
32 while oy < OH {
33 // src_y * den = 2*oy + 1 - scale ; floor + positive remainder
34 let SY: i64 = 2 * oy + 1 - scale
35 var qy: i64 = SY / den
36 var ry: i64 = SY - qy * den
37 if ry < 0 { qy = qy - 1; ry = ry + den }
38 var y0: i64 = qy
39 var y1: i64 = qy + 1
40 if y0 < 0 { y0 = 0 }
41 if y0 > H - 1 { y0 = H - 1 }
42 if y1 < 0 { y1 = 0 }
43 if y1 > H - 1 { y1 = H - 1 }
44 let iwy: i64 = den - ry // weight toward y0 (numerator)
45 var ox: i64 = 0
46 while ox < OW {
47 let SX: i64 = 2 * ox + 1 - scale
48 var qx: i64 = SX / den
49 var rx: i64 = SX - qx * den
50 if rx < 0 { qx = qx - 1; rx = rx + den }
51 var x0: i64 = qx
52 var x1: i64 = qx + 1
53 if x0 < 0 { x0 = 0 }
54 if x0 > W - 1 { x0 = W - 1 }
55 if x1 < 0 { x1 = 0 }
56 if x1 > W - 1 { x1 = W - 1 }
57 let iwx: i64 = den - rx
58 let w00: i64 = nx_i32_to_f32(iwy * iwx)
59 let w01: i64 = nx_i32_to_f32(iwy * rx)
60 let w10: i64 = nx_i32_to_f32(ry * iwx)
61 let w11: i64 = nx_i32_to_f32(ry * rx)
62 let v00: i64 = input[base + y0 * W + x0]
63 let v01: i64 = input[base + y0 * W + x1]
64 let v10: i64 = input[base + y1 * W + x0]
65 let v11: i64 = input[base + y1 * W + x1]
66 var acc: i64 = nx_f32_mul(w00, v00)
67 acc = nx_f32_add(acc, nx_f32_mul(w01, v01))
68 acc = nx_f32_add(acc, nx_f32_mul(w10, v10))
69 acc = nx_f32_add(acc, nx_f32_mul(w11, v11))
70 out[n * obs + c * ocs + oy * OW + ox] = nx_f32_div(acc, den2f)
71 ox = ox + 1
72 }
73 oy = oy + 1
74 }
75 c = c + 1
76 }
77 n = n + 1
78 }
79 return NX_UB_OK
80}