nx_vae_decode_stage.nx source

1// nx_vae_decode_stage.nx -- one stage of a VAE decoder. 2// 3// L4 composer. Closes the latent -> pixel path: VAE decoder = N 4// stacked stages of (ResBlock + optional 2x upsample), ending with 5// final GroupNorm + SiLU + 3x3 conv to RGB. This primitive 6// is one such stage; caller composes the full decoder. 7// 8// ===== Canonical SD VAE decoder topology ========================= 9// 10// Input: latent [N, 4, H, W] (typical: 64x64 for 512x512 image) 11// 12// stage_in: conv2d 4 -> 512 ch 13// stage_mid: ResBlock 512 -> 512 (no upsample) 14// stage_up_0: ResBlock 512 -> 512, upsample 2x 15// stage_up_1: ResBlock 512 -> 512, upsample 2x 16// stage_up_2: ResBlock 512 -> 256, upsample 2x 17// stage_up_3: ResBlock 256 -> 128, no upsample 18// stage_out: GroupNorm + SiLU + conv2d 128 -> 3 19// 20// Output: image [N, 3, 8H, 8W] 21// 22// v1 substrate ships THIS PRIMITIVE for one (ResBlock + optional 23// upsample) stage. The CHANNEL CHANGE (e.g. 512 -> 256) is queued 24// for v2 (needs 1x1 skip conv in the ResBlock). v1 covers the 25// same-channel-count stages which are the majority. 26// 27// Bits-up composition (pure): 28// nx_unet_block_forward (L4 ResBlock, shipped 6ede5571) 29// nx_upsample_2x (L3 upsample, shipped dbd26d43) 30// NxTensor (L1) 31// nx_loop (control) 32// 33// genealogy_id: kingma_2014_vae + rombach_2022_stable_diffusion_autoencoder 34// lineage_id: substrate_vae_decode_stage_v1 35 36// nx_safety_envelope: 37// intended_use: AUTO_APPLIED -- primitive-specific tuning queued 38// sil_target: SIL1 39// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail] 40// verdict: NOT_YET_EVALUATED 41 42import "nx_syscalls.nx" 43import "nx_tier.nx" 44import "nx_loop.nx" 45import "nx_tensor.nx" 46import "nx_unet_block.nx" 47import "nx_upsample.nx" 48 49// ===== Sealed-enum: VaeDecodeStageVerdict ========================= 50 51const NX_VDS_OK: nx_int = 0 52const NX_VDS_ERR_BAD_DIMS: nx_int = 1 53const NX_VDS_ERR_SHAPE_MISMATCH: nx_int = 2 54const NX_VDS_ERR_INTERNAL: nx_int = 3 55const NX_VDS_N_VERDICTS: nx_int = 4 56 57func nx_vds_verdict_is_valid(v: nx_int) -> nx_int { 58 if v < 0 { return 0 } 59 if v >= NX_VDS_N_VERDICTS { return 0 } 60 return 1 61} 62 63// ===== Stage forward ============================================= 64// 65// input: [N, C, H, W] Q10 66// W_1, W_2: ResBlock conv weights (each [C, C, 3, 3]) 67// gamma_1, beta_1, gamma_2, beta_2: GroupNorm scale + bias [C] 68// n_groups: GroupNorm group count (typically 32) 69// do_upsample: 1 = upsample 2x at the end; 0 = no upsample 70// upsample_filter: NX_UP_NEAREST or NX_UP_BILINEAR 71// 72// output: [N, C, H, W] Q10 if do_upsample=0 73// [N, C, 2H, 2W] Q10 if do_upsample=1 74// resblock_scratch: scratch tensor for ResBlock intermediate 75// [N, C, H, W] 76// post_resblock: output of the ResBlock before upsample 77// [N, C, H, W] 78// 79// When do_upsample=1 we put the ResBlock output in `post_resblock` 80// and upsample into `output`. When do_upsample=0 we put the 81// ResBlock output directly in `output` and `post_resblock` is 82// unused (caller can pass it as a throwaway). 83 84func nx_vae_decode_stage_forward( 85 input: *NxTensor, 86 W_1: *NxTensor, gamma_1: *i64, beta_1: *i64, 87 W_2: *NxTensor, gamma_2: *i64, beta_2: *i64, 88 n_groups: nx_int, 89 do_upsample: nx_int, 90 upsample_filter: nx_int, 91 post_resblock: *NxTensor, 92 resblock_scratch: *NxTensor, 93 output: *NxTensor) -> nx_int { 94 95 // Step 1: ResBlock. Output target depends on do_upsample. 96 if do_upsample == 1 { 97 let v_rb: nx_int = nx_unet_block_forward( 98 input, W_1, gamma_1, beta_1, 99 W_2, gamma_2, beta_2, 100 n_groups, post_resblock, resblock_scratch) 101 if v_rb != NX_UB_OK { return NX_VDS_ERR_INTERNAL } 102 103 // Step 2: 2x upsample into the final output buffer. 104 let v_up: nx_int = nx_upsample_2x(post_resblock, output, upsample_filter) 105 if v_up != NX_UP_OK { return NX_VDS_ERR_INTERNAL } 106 } 107 if do_upsample == 0 { 108 // No upsample -- ResBlock writes directly to output. 109 let v_rb: nx_int = nx_unet_block_forward( 110 input, W_1, gamma_1, beta_1, 111 W_2, gamma_2, beta_2, 112 n_groups, output, resblock_scratch) 113 if v_rb != NX_UB_OK { return NX_VDS_ERR_INTERNAL } 114 } 115 return NX_VDS_OK 116} 117 118// ===== Self-test ================================================== 119// 120// Tiny stage smoke: [N=1, C=4, H=2, W=2] -> ResBlock (zero weights) 121// + upsample 2x = [1, 4, 4, 4]. 122// With zero ResBlock weights, ResBlock output == input (residual 123// pass-through). Upsample 2x then doubles spatial dims. 124// 125// Closed-form invariants: 126// (a) do_upsample=1: output dims = [N, C, 2H, 2W] 127// (b) do_upsample=0: output dims = [N, C, H, W]; zero-conv leaves 128// input intact via residual 129// (c) Verdict gate 130 131func main() -> i64 { 132 let N: nx_int = 1 133 let C: nx_int = 4 134 let H: nx_int = 2 135 let W: nx_int = 2 136 137 let in_sh: *nx_int = sys_mmap(4 * 8) as *nx_int 138 in_sh[0]=N; in_sh[1]=C; in_sh[2]=H; in_sh[3]=W 139 let mid_sh: *nx_int = sys_mmap(4 * 8) as *nx_int 140 mid_sh[0]=N; mid_sh[1]=C; mid_sh[2]=H; mid_sh[3]=W 141 let up_sh: *nx_int = sys_mmap(4 * 8) as *nx_int 142 up_sh[0]=N; up_sh[1]=C; up_sh[2]=H*2; up_sh[3]=W*2 143 let wt_sh: *nx_int = sys_mmap(4 * 8) as *nx_int 144 wt_sh[0]=C; wt_sh[1]=C; wt_sh[2]=3; wt_sh[3]=3 145 146 let err: *nx_int = sys_mmap(8) as *nx_int 147 err[0] = 0 148 let input: *NxTensor = nx_t_alloc(NX_DT_I64, in_sh, 4, err) 149 let post_rb: *NxTensor = nx_t_alloc(NX_DT_I64, mid_sh, 4, err) 150 let scratch: *NxTensor = nx_t_alloc(NX_DT_I64, mid_sh, 4, err) 151 let upscaled: *NxTensor = nx_t_alloc(NX_DT_I64, up_sh, 4, err) 152 let W_1: *NxTensor = nx_t_alloc(NX_DT_I64, wt_sh, 4, err) 153 let W_2: *NxTensor = nx_t_alloc(NX_DT_I64, wt_sh, 4, err) 154 if err[0] != 0 { return 5 } 155 156 // Fill input with a pattern. 157 let pi: *i64 = input.storage as *i64 158 var i: nx_int = 0 159 while i < N * C * H * W { pi[i] = (i + 1) * 50; i = i + 1 } 160 161 // Unit gammas, zero betas. 162 let gamma_1: *i64 = sys_mmap(C * 8) as *i64 163 let gamma_2: *i64 = sys_mmap(C * 8) as *i64 164 let beta_1: *i64 = sys_mmap(C * 8) as *i64 165 let beta_2: *i64 = sys_mmap(C * 8) as *i64 166 var k: nx_int = 0 167 while k < C { gamma_1[k] = 1024; gamma_2[k] = 1024; beta_1[k] = 0; beta_2[k] = 0; k = k + 1 } 168 169 // --- (a) Stage with upsample --- 170 let v_a: nx_int = nx_vae_decode_stage_forward( 171 input, W_1, gamma_1, beta_1, W_2, gamma_2, beta_2, 172 2, 1, NX_UP_NEAREST, post_rb, scratch, upscaled) 173 if v_a != NX_VDS_OK { return 10 + v_a } 174 // Upscaled output dims: [1, 4, 4, 4] = 64 elements. 175 // Zero conv -> residual -> post_rb == input. Then nearest upsample 176 // 2x: each input pixel replicated 2x2. The post_rb[0] = input[0] = 50. 177 // After nearest 2x: upscaled[0..3] (= upscaled row 0 cols 0..3 of 178 // channel 0) should be [50, 50, ?, ?] (where ? comes from input 179 // column 1). Verify simple pattern. 180 let pu: *i64 = upscaled.storage as *i64 181 if pu[0] != 50 { return 20 } 182 if pu[1] != 50 { return 21 } 183 // Pixel (row 0, col 2) in upsampled = input row 0 col 1. 184 // input col 1 of channel 0 (N=1, C=4, H=2, W=2 layout, contiguous): 185 // input flat index for (n=0, c=0, h=0, w=1) = 0 + 0 + 0 + 1 = 1. 186 // pi[1] = 100. 187 if pu[2] != 100 { return 22 } 188 189 // --- (b) Stage without upsample --- 190 let out_no_up: *NxTensor = nx_t_alloc(NX_DT_I64, mid_sh, 4, err) 191 if err[0] != 0 { return 30 } 192 let v_b: nx_int = nx_vae_decode_stage_forward( 193 input, W_1, gamma_1, beta_1, W_2, gamma_2, beta_2, 194 2, 0, NX_UP_NEAREST, post_rb, scratch, out_no_up) 195 if v_b != NX_VDS_OK { return 40 + v_b } 196 // Zero conv -> residual -> output == input. 197 let po: *i64 = out_no_up.storage as *i64 198 var j: nx_int = 0 199 while j < N * C * H * W { 200 if po[j] != pi[j] { return 50 } 201 j = j + 1 202 } 203 204 // --- (c) Verdict gate --- 205 var vi: nx_int = 0 206 while vi < NX_VDS_N_VERDICTS { 207 if nx_vds_verdict_is_valid(vi) != 1 { return 60 + vi } 208 vi = vi + 1 209 } 210 211 return 0 212}