nx_f32_upsample.nx source
↩ module page · 75 lines · 3057 B
1// nx_f32_upsample.nx -- software-f32 nearest-neighbour Nx upsample (NCHW), VAE-decode brick.
2//
3// sd-server -> Nishi migration: VAE decoders upscale via "resize-conv" = nearest-neighbour upsample
4// followed by nx_f32_conv2d (avoids the checkerboard artefacts of transposed conv). out[n,c,oy,ox] =
5// in[n,c,oy/scale,ox/scale] -- pure index replication of f32 bit-patterns, so it's exact (no arithmetic).
6// The i64 ref `nx_upsample.nx` is the Q10 version; this is the f32-tier one that pairs with the gated
7// f32 conv/resblock so real dequantized weights flow through the decode path.
8// license_tier: ORIGINAL
9import "nx_syscalls.nx"
10import "nx_f32_cvt.nx"
11
12const NX_F32US_OK: i64 = 0
13const NX_F32US_ERR: i64 = 1
14
15// input [N,C,H,W] -> out [N,C,H*scale,W*scale], nearest-neighbour. flat *i64 f32 bits. out != input.
16func nx_f32_upsample_nn(input: *i64, N: i64, C: i64, H: i64, W: i64, scale: i64, out: *i64) -> i64 {
17 if scale <= 0 { return NX_F32US_ERR }
18 let OH: i64 = H * scale
19 let OW: i64 = W * scale
20 let in_chan_stride: i64 = H * W
21 let in_batch_stride: i64 = C * H * W
22 let out_chan_stride: i64 = OH * OW
23 let out_batch_stride: i64 = C * OH * OW
24 var n: i64 = 0
25 while n < N {
26 var c: i64 = 0
27 while c < C {
28 let in_base: i64 = n * in_batch_stride + c * in_chan_stride
29 let out_base: i64 = n * out_batch_stride + c * out_chan_stride
30 var oy: i64 = 0
31 while oy < OH {
32 let iy: i64 = oy / scale
33 var ox: i64 = 0
34 while ox < OW {
35 let ix: i64 = ox / scale
36 out[out_base + oy * OW + ox] = input[in_base + iy * W + ix]
37 ox = ox + 1
38 }
39 oy = oy + 1
40 }
41 c = c + 1
42 }
43 n = n + 1
44 }
45 return NX_F32US_OK
46}
47
48// ===== Self-test (inline gate) ====================================
49// 1x1x2x2 [[10,20],[30,40]] -> 2x -> 4x4 with each source pixel replicated into a 2x2 block (bit-exact).
50func main() -> i64 {
51 let inb: *i64 = sys_mmap(64 * 8) as *i64
52 let out: *i64 = sys_mmap(64 * 8) as *i64
53 inb[0] = nx_i32_to_f32(10); inb[1] = nx_i32_to_f32(20)
54 inb[2] = nx_i32_to_f32(30); inb[3] = nx_i32_to_f32(40)
55 let v: i64 = nx_f32_upsample_nn(inb, 1, 1, 2, 2, 2, out)
56 if v != NX_F32US_OK { return 10 }
57 let t10: i64 = nx_i32_to_f32(10)
58 let t20: i64 = nx_i32_to_f32(20)
59 let t30: i64 = nx_i32_to_f32(30)
60 let t40: i64 = nx_i32_to_f32(40)
61 // expected 4x4: 10 10 20 20 / 10 10 20 20 / 30 30 40 40 / 30 30 40 40
62 if out[0] != t10 { return 20 }
63 if out[1] != t10 { return 21 }
64 if out[2] != t20 { return 22 }
65 if out[3] != t20 { return 23 }
66 if out[4] != t10 { return 24 }
67 if out[7] != t20 { return 25 }
68 if out[8] != t30 { return 26 }
69 if out[10] != t40 { return 27 }
70 if out[15] != t40 { return 28 }
71 // bad scale
72 let vb: i64 = nx_f32_upsample_nn(inb, 1, 1, 2, 2, 0, out)
73 if vb == NX_F32US_OK { return 30 }
74 return 0
75}