nx_upsample.nx source

1// nx_upsample.nx -- 2x upsample for 4D feature maps. 2// 3// Closes the missing-dep gap for VAE decoder + UNet up-blocks. 4// Modern image-gen architectures use 2x nearest-neighbor or 5// 2x bilinear upsample inside the decoder stages: 6// 7// VAE decoder: 4 stages of 2x upsample (8x total = 64x64 latent 8// -> 512x512 image) 9// UNet up-path: 2x upsample at each scale level 10// Super-resolution networks: 2x or 4x upsample 11// 12// L3 canonical primitive composing: 13// NxTensor (L1, 4D NCHW input/output) 14// nx_loop.LoopVerdict (control) 15// -- no other deps; pure pixel-replication / linear-interp math 16// 17// ===== Tensor layout ============================================= 18// 19// Input: [N, C, H, W] NCHW Q10 20// Output: [N, C, 2H, 2W] 2x upscaled in spatial dims; channels 21// unchanged 22// 23// ===== Filter modes ============================================== 24// 25// NX_UP_NEAREST -- each output pixel = nearest input pixel. 26// Blocky but cheap. Common in VAE decoder when 27// followed by a 3x3 conv that smooths the 28// artifacts. 29// NX_UP_BILINEAR -- 4-tap bilinear from the 4 nearest input pixels. 30// Smoother; slightly more expensive. 31// 32// ===== Pixel-center convention =================================== 33// 34// Output pixel (oy, ox) maps to input pixel at: 35// iy = (oy + 0.5) / 2 - 0.5 36// ix = (ox + 0.5) / 2 - 0.5 37// 38// This is the "align_corners=False" convention from PyTorch -- the 39// standard for VAE-class upsampling. In Q10 fixed-point: 40// iy_q10 = (oy + 1) * Q10 / 2 - Q10 / 2 ; if (... ) ... simplifies 41// We compute integer iy + fractional bit separately. 42// 43// genealogy_id: pytorch_upsample_2d_2017 + opencv_resize_2000 + 44// neural_super_resolution_ledig_2016 45// lineage_id: substrate_upsample_v1 46 47// nx_safety_envelope: 48// intended_use: AUTO_APPLIED -- primitive-specific tuning queued 49// sil_target: SIL1 50// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail] 51// verdict: NOT_YET_EVALUATED 52 53import "nx_syscalls.nx" 54import "nx_tier.nx" 55import "nx_loop.nx" 56import "nx_tensor.nx" 57 58// ===== Constants ================================================== 59 60const NX_UP_Q10: nx_int = 1024 61 62// ===== Sealed-enum: FilterMode ==================================== 63 64const NX_UP_NEAREST: nx_int = 0 65const NX_UP_BILINEAR: nx_int = 1 66const NX_UP_FILTER_N: nx_int = 2 67 68func nx_up_filter_is_valid(k: nx_int) -> nx_int { 69 if k < 0 { return 0 } 70 if k >= NX_UP_FILTER_N { return 0 } 71 return 1 72} 73 74// ===== Sealed-enum: UpsampleVerdict =============================== 75 76const NX_UP_OK: nx_int = 0 77const NX_UP_ERR_BAD_DTYPE: nx_int = 1 78const NX_UP_ERR_BAD_NDIM: nx_int = 2 79const NX_UP_ERR_SHAPE_MISMATCH: nx_int = 3 80const NX_UP_ERR_NOT_CONTIGUOUS: nx_int = 4 81const NX_UP_ERR_BAD_FILTER: nx_int = 5 82const NX_UP_N_VERDICTS: nx_int = 6 83 84func nx_up_verdict_is_valid(v: nx_int) -> nx_int { 85 if v < 0 { return 0 } 86 if v >= NX_UP_N_VERDICTS { return 0 } 87 return 1 88} 89 90// ===== Nearest-neighbor 2x ======================================== 91 92func _up_nearest(input: *NxTensor, output: *NxTensor) -> nx_int { 93 let N: nx_int = input.shape[0] 94 let C: nx_int = input.shape[1] 95 let H: nx_int = input.shape[2] 96 let W: nx_int = input.shape[3] 97 let H2: nx_int = H * 2 98 let W2: nx_int = W * 2 99 100 let pi: *i64 = input.storage as *i64 101 let po: *i64 = output.storage as *i64 102 let in_batch: nx_int = C * H * W 103 let in_chan: nx_int = H * W 104 let out_batch: nx_int = C * H2 * W2 105 let out_chan: nx_int = H2 * W2 106 107 var n: nx_int = 0 108 var n_iter: nx_int = 0 109 var n_verdict: nx_int = NX_LOOP_RUNNING 110 let N_BUDGET: nx_int = N 111 while n_verdict == NX_LOOP_RUNNING && n_iter < N_BUDGET { 112 var c: nx_int = 0 113 var c_iter: nx_int = 0 114 var c_verdict: nx_int = NX_LOOP_RUNNING 115 let C_BUDGET: nx_int = C 116 while c_verdict == NX_LOOP_RUNNING && c_iter < C_BUDGET { 117 var oy: nx_int = 0 118 var oy_iter: nx_int = 0 119 var oy_verdict: nx_int = NX_LOOP_RUNNING 120 let OY_BUDGET: nx_int = H2 121 while oy_verdict == NX_LOOP_RUNNING && oy_iter < OY_BUDGET { 122 // Nearest input row: oy / 2. 123 let iy: nx_int = oy / 2 124 var ox: nx_int = 0 125 var ox_iter: nx_int = 0 126 var ox_verdict: nx_int = NX_LOOP_RUNNING 127 let OX_BUDGET: nx_int = W2 128 while ox_verdict == NX_LOOP_RUNNING && ox_iter < OX_BUDGET { 129 let ix: nx_int = ox / 2 130 let src: nx_int = n * in_batch + c * in_chan + iy * W + ix 131 let dst: nx_int = n * out_batch + c * out_chan + oy * W2 + ox 132 po[dst] = pi[src] 133 ox = ox + 1 134 ox_iter = ox_iter + 1 135 } 136 oy = oy + 1 137 oy_iter = oy_iter + 1 138 } 139 c = c + 1 140 c_iter = c_iter + 1 141 } 142 n = n + 1 143 n_iter = n_iter + 1 144 } 145 return NX_UP_OK 146} 147 148// ===== Bilinear 2x ================================================ 149// 150// align_corners=False convention. Each output pixel (oy, ox) maps 151// to fractional input coords: 152// iy_q10 = oy * Q10 / 2 + Q10 / 4 - Q10 / 2 = (2*oy + 1) * Q10 / 4 - Q10 / 2 153// Simpler: iy_int = (oy - 1) / 2 if oy odd else oy / 2 - 1; we just 154// compute it carefully below. 155// 156// Bilinear: 157// iy_q10 = ((oy * 2 + 1) * Q10) / 4 - Q10/2 (i.e. centre-shift) 158// iy_int = iy_q10 / Q10 (floor) 159// frac_y = iy_q10 - iy_int * Q10 160// (clamp iy_int to [0, H-1] and iy_int+1 to [0, H-1]; same for x) 161 162func _up_bilinear(input: *NxTensor, output: *NxTensor) -> nx_int { 163 let N: nx_int = input.shape[0] 164 let C: nx_int = input.shape[1] 165 let H: nx_int = input.shape[2] 166 let W: nx_int = input.shape[3] 167 let H2: nx_int = H * 2 168 let W2: nx_int = W * 2 169 170 let pi: *i64 = input.storage as *i64 171 let po: *i64 = output.storage as *i64 172 let in_batch: nx_int = C * H * W 173 let in_chan: nx_int = H * W 174 let out_batch: nx_int = C * H2 * W2 175 let out_chan: nx_int = H2 * W2 176 177 var n: nx_int = 0 178 var n_iter: nx_int = 0 179 var n_verdict: nx_int = NX_LOOP_RUNNING 180 let N_BUDGET: nx_int = N 181 while n_verdict == NX_LOOP_RUNNING && n_iter < N_BUDGET { 182 var c: nx_int = 0 183 var c_iter: nx_int = 0 184 var c_verdict: nx_int = NX_LOOP_RUNNING 185 let C_BUDGET: nx_int = C 186 while c_verdict == NX_LOOP_RUNNING && c_iter < C_BUDGET { 187 var oy: nx_int = 0 188 var oy_iter: nx_int = 0 189 var oy_verdict: nx_int = NX_LOOP_RUNNING 190 let OY_BUDGET: nx_int = H2 191 while oy_verdict == NX_LOOP_RUNNING && oy_iter < OY_BUDGET { 192 // Fractional input row centre. 193 let iy_centre_q10: nx_int = ((oy * 2 + 1) * NX_UP_Q10) / 4 - NX_UP_Q10 / 2 194 var iy_int: nx_int = iy_centre_q10 / NX_UP_Q10 195 if iy_centre_q10 < 0 { 196 if iy_centre_q10 - iy_int * NX_UP_Q10 != 0 { iy_int = iy_int - 1 } 197 } 198 let frac_y: nx_int = iy_centre_q10 - iy_int * NX_UP_Q10 199 // Clamp to valid input rows. 200 var iy0: nx_int = iy_int 201 if iy0 < 0 { iy0 = 0 } 202 if iy0 >= H { iy0 = H - 1 } 203 var iy1: nx_int = iy_int + 1 204 if iy1 < 0 { iy1 = 0 } 205 if iy1 >= H { iy1 = H - 1 } 206 207 var ox: nx_int = 0 208 var ox_iter: nx_int = 0 209 var ox_verdict: nx_int = NX_LOOP_RUNNING 210 let OX_BUDGET: nx_int = W2 211 while ox_verdict == NX_LOOP_RUNNING && ox_iter < OX_BUDGET { 212 let ix_centre_q10: nx_int = ((ox * 2 + 1) * NX_UP_Q10) / 4 - NX_UP_Q10 / 2 213 var ix_int: nx_int = ix_centre_q10 / NX_UP_Q10 214 if ix_centre_q10 < 0 { 215 if ix_centre_q10 - ix_int * NX_UP_Q10 != 0 { ix_int = ix_int - 1 } 216 } 217 let frac_x: nx_int = ix_centre_q10 - ix_int * NX_UP_Q10 218 var ix0: nx_int = ix_int 219 if ix0 < 0 { ix0 = 0 } 220 if ix0 >= W { ix0 = W - 1 } 221 var ix1: nx_int = ix_int + 1 222 if ix1 < 0 { ix1 = 0 } 223 if ix1 >= W { ix1 = W - 1 } 224 225 let base: nx_int = n * in_batch + c * in_chan 226 let v00: i64 = pi[base + iy0 * W + ix0] 227 let v10: i64 = pi[base + iy0 * W + ix1] 228 let v01: i64 = pi[base + iy1 * W + ix0] 229 let v11: i64 = pi[base + iy1 * W + ix1] 230 let one_minus_x: nx_int = NX_UP_Q10 - frac_x 231 let one_minus_y: nx_int = NX_UP_Q10 - frac_y 232 let top: nx_int = (v00 * one_minus_x + v10 * frac_x) / NX_UP_Q10 233 let bot: nx_int = (v01 * one_minus_x + v11 * frac_x) / NX_UP_Q10 234 let val: nx_int = (top * one_minus_y + bot * frac_y) / NX_UP_Q10 235 let dst: nx_int = n * out_batch + c * out_chan + oy * W2 + ox 236 po[dst] = val 237 ox = ox + 1 238 ox_iter = ox_iter + 1 239 } 240 oy = oy + 1 241 oy_iter = oy_iter + 1 242 } 243 c = c + 1 244 c_iter = c_iter + 1 245 } 246 n = n + 1 247 n_iter = n_iter + 1 248 } 249 return NX_UP_OK 250} 251 252// ===== Public dispatch =========================================== 253 254func nx_upsample_2x(input: *NxTensor, output: *NxTensor, filter: nx_int) -> nx_int { 255 if input.dtype != NX_DT_I64 { return NX_UP_ERR_BAD_DTYPE } 256 if output.dtype != NX_DT_I64 { return NX_UP_ERR_BAD_DTYPE } 257 if input.ndim != 4 { return NX_UP_ERR_BAD_NDIM } 258 if output.ndim != 4 { return NX_UP_ERR_BAD_NDIM } 259 if input.shape[0] != output.shape[0] { return NX_UP_ERR_SHAPE_MISMATCH } 260 if input.shape[1] != output.shape[1] { return NX_UP_ERR_SHAPE_MISMATCH } 261 if input.shape[2] * 2 != output.shape[2] { return NX_UP_ERR_SHAPE_MISMATCH } 262 if input.shape[3] * 2 != output.shape[3] { return NX_UP_ERR_SHAPE_MISMATCH } 263 if nx_t_is_contiguous(input) == 0 { return NX_UP_ERR_NOT_CONTIGUOUS } 264 if nx_t_is_contiguous(output) == 0 { return NX_UP_ERR_NOT_CONTIGUOUS } 265 if filter == NX_UP_NEAREST { return _up_nearest(input, output) } 266 if filter == NX_UP_BILINEAR { return _up_bilinear(input, output) } 267 return NX_UP_ERR_BAD_FILTER 268} 269 270// ===== Self-test ================================================== 271// 272// 1x1x2x2 input upsampled to 1x1x4x4. 273// Input: [[10, 20], 274// [30, 40]] 275// Nearest -> [[10,10,20,20],[10,10,20,20],[30,30,40,40],[30,30,40,40]] 276// Bilinear -> smooth gradient between corners. 277// 278// Closed-form invariants: 279// (a) Nearest 2x doubles dims, replicates pixels 280// (b) Bilinear 2x preserves input corner values (with align_corners=False 281// convention, output corners ARE input corners after rounding) 282// (c) Bad filter -> ERR_BAD_FILTER 283// (d) Bad shape (output H != 2 * input H) -> ERR_SHAPE_MISMATCH 284// (e) Verdict + filter sealed-enum gates 285 286func main() -> i64 { 287 let in_sh: *nx_int = sys_mmap(4 * 8) as *nx_int 288 in_sh[0]=1; in_sh[1]=1; in_sh[2]=2; in_sh[3]=2 289 let out_sh: *nx_int = sys_mmap(4 * 8) as *nx_int 290 out_sh[0]=1; out_sh[1]=1; out_sh[2]=4; out_sh[3]=4 291 292 let err: *nx_int = sys_mmap(8) as *nx_int 293 err[0] = 0 294 let xt: *NxTensor = nx_t_alloc(NX_DT_I64, in_sh, 4, err) 295 let yt: *NxTensor = nx_t_alloc(NX_DT_I64, out_sh, 4, err) 296 if err[0] != 0 { return 5 } 297 298 // Input: [10, 20; 30, 40] in Q10 (but values are unscaled for the 299 // smoke; we just want bit-exact replication for nearest). 300 let pi: *i64 = xt.storage as *i64 301 pi[0] = 10 302 pi[1] = 20 303 pi[2] = 30 304 pi[3] = 40 305 306 // --- (a) Nearest 2x --- 307 let v_n: nx_int = nx_upsample_2x(xt, yt, NX_UP_NEAREST) 308 if v_n != NX_UP_OK { return 10 + v_n } 309 let po: *i64 = yt.storage as *i64 310 if po[0] != 10 { return 20 } // (0,0) 311 if po[1] != 10 { return 21 } // (0,1) 312 if po[2] != 20 { return 22 } // (0,2) 313 if po[3] != 20 { return 23 } // (0,3) 314 if po[4] != 10 { return 24 } // (1,0) 315 if po[7] != 20 { return 25 } 316 if po[8] != 30 { return 26 } 317 if po[15] != 40 { return 27 } 318 319 // --- (b) Bilinear 2x --- 320 // Just verify it produces SOMETHING reasonable; bilinear of [10,20;30,40] 321 // should have output values between 10 and 40. Centre pixels = ~25. 322 let v_b: nx_int = nx_upsample_2x(xt, yt, NX_UP_BILINEAR) 323 if v_b != NX_UP_OK { return 30 + v_b } 324 var k: nx_int = 0 325 while k < 16 { 326 if po[k] < 0 { return 40 } 327 if po[k] > 100 { return 41 } 328 k = k + 1 329 } 330 331 // --- (c) Bad filter --- 332 let v_bf: nx_int = nx_upsample_2x(xt, yt, 99) 333 if v_bf != NX_UP_ERR_BAD_FILTER { return 50 } 334 335 // --- (d) Bad shape --- 336 let bad_sh: *nx_int = sys_mmap(4 * 8) as *nx_int 337 bad_sh[0]=1; bad_sh[1]=1; bad_sh[2]=3; bad_sh[3]=4 // wrong H 338 let bad: *NxTensor = nx_t_alloc(NX_DT_I64, bad_sh, 4, err) 339 let v_bs: nx_int = nx_upsample_2x(xt, bad, NX_UP_NEAREST) 340 if v_bs != NX_UP_ERR_SHAPE_MISMATCH { return 60 } 341 342 // --- (e) Verdict + filter gates --- 343 var vi: nx_int = 0 344 while vi < NX_UP_N_VERDICTS { 345 if nx_up_verdict_is_valid(vi) != 1 { return 70 + vi } 346 vi = vi + 1 347 } 348 var fi: nx_int = 0 349 while fi < NX_UP_FILTER_N { 350 if nx_up_filter_is_valid(fi) != 1 { return 80 + fi } 351 fi = fi + 1 352 } 353 354 return 0 355}