nx_unet_block.nx source

1// nx_unet_block.nx -- diffusion-model ResBlock composer. 2// 3// L4 composer for the UNet building block used by Stable Diffusion 4// 1.x/2.x/3, Flux, Z-Image. Pure composition -- no new math. 5// 6// ===== Block shape (canonical SD ResBlock) ======================= 7// 8// x: [N, C, H, W] input 9// 10// 1. h = GroupNorm(x, num_groups=32, gamma_1, beta_1) 11// 2. h = SiLU(h) 12// 3. h = Conv2D(h, W_1) : [N, C, H, W] 13// 4. h = GroupNorm(h, num_groups=32, gamma_2, beta_2) 14// 5. h = SiLU(h) 15// 6. h = Conv2D(h, W_2) : [N, C, H, W] 16// 7. Skip: residual = x (v1: assumes C_in = C_out) 17// 8. Output = h + residual 18// 19// ===== v1 scope ================================================== 20// 21// Assumes C_in == C_out throughout (skip-connection is identity). 22// For "down" blocks in a U-shape (C_in != C_out): caller passes 23// a separate 1x1 skip-conv layer and we'd extend this composer 24// in v2. v1 satisfies the same-channel-count case, which is the 25// majority of ResBlocks within a stage. 26// 27// Time-embedding injection (diffusion-step conditioning): 28// canonical UNets add a per-block timestep-embedding vector to h 29// between steps 3 and 4 (after first conv, before second norm). 30// v1 omits this -- caller is expected to add it post-hoc or use the 31// queued v2 of this primitive that accepts a time_emb tensor. 32// 33// Bits-up composition (pure, no new math): 34// NxTensor (L1) 35// nx_conv2d (L3) 36// nx_groupnorm (L3, shipped 40d6e1c2) 37// nx_silu (L3) 38// nx_loop.LoopVerdict (control) 39// 40// genealogy_id: ho_2020_ddpm_unet + ronneberger_2015_u_net + 41// rombach_2022_stable_diffusion_unet + 42// wu_he_2018_group_norm 43// lineage_id: substrate_unet_block_v1_same_channel 44 45// nx_safety_envelope: 46// intended_use: AUTO_APPLIED -- primitive-specific tuning queued 47// sil_target: SIL1 48// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail] 49// verdict: NOT_YET_EVALUATED 50 51import "nx_syscalls.nx" 52import "nx_tier.nx" 53import "nx_loop.nx" 54import "nx_tensor.nx" 55import "nx_conv2d.nx" 56import "nx_groupnorm.nx" 57import "nx_silu.nx" 58 59const NX_UB_Q10: nx_int = 1024 60const NX_UB_DEFAULT_GROUPS: nx_int = 32 61 62// ===== Sealed-enum: UnetBlockVerdict ============================== 63 64const NX_UB_OK: nx_int = 0 65const NX_UB_ERR_BAD_DIMS: nx_int = 1 66const NX_UB_ERR_SHAPE_MISMATCH: nx_int = 2 67const NX_UB_ERR_INTERNAL: nx_int = 3 68const NX_UB_N_VERDICTS: nx_int = 4 69 70func nx_ub_verdict_is_valid(v: nx_int) -> nx_int { 71 if v < 0 { return 0 } 72 if v >= NX_UB_N_VERDICTS { return 0 } 73 return 1 74} 75 76// ===== Forward pass ============================================== 77// 78// x: *NxTensor [N, C, H, W] Q10 input (preserved; not modified) 79// W_1: *NxTensor [C, C, 3, 3] first conv weights 80// gamma_1: *i64 [C] first GroupNorm scale 81// beta_1: *i64 [C] first GroupNorm bias (nullable) 82// W_2: *NxTensor [C, C, 3, 3] second conv weights 83// gamma_2: *i64 [C] second GroupNorm scale 84// beta_2: *i64 [C] second GroupNorm bias (nullable) 85// n_groups: nx_int GroupNorm group count (typically 32) 86// out: *NxTensor [N, C, H, W] Q10 output 87// scratch: *NxTensor [N, C, H, W] caller-owned scratch (size matches) 88 89func nx_unet_block_forward( 90 x: *NxTensor, 91 W_1: *NxTensor, gamma_1: *i64, beta_1: *i64, 92 W_2: *NxTensor, gamma_2: *i64, beta_2: *i64, 93 n_groups: nx_int, 94 out: *NxTensor, scratch: *NxTensor) -> nx_int { 95 96 if x.dtype != NX_DT_I64 { return NX_UB_ERR_BAD_DIMS } 97 if out.dtype != NX_DT_I64 { return NX_UB_ERR_BAD_DIMS } 98 if scratch.dtype != NX_DT_I64 { return NX_UB_ERR_BAD_DIMS } 99 if x.ndim != 4 { return NX_UB_ERR_BAD_DIMS } 100 if x.shape[0] != out.shape[0] { return NX_UB_ERR_SHAPE_MISMATCH } 101 if x.shape[1] != out.shape[1] { return NX_UB_ERR_SHAPE_MISMATCH } 102 if x.shape[2] != out.shape[2] { return NX_UB_ERR_SHAPE_MISMATCH } 103 if x.shape[3] != out.shape[3] { return NX_UB_ERR_SHAPE_MISMATCH } 104 105 // Step 1+2: scratch = SiLU(GroupNorm(x, gamma_1, beta_1)). 106 let v_gn1: nx_int = nx_groupnorm_forward(x, n_groups, gamma_1, beta_1, scratch) 107 if v_gn1 != NX_GN_OK { return NX_UB_ERR_INTERNAL } 108 let v_silu1: nx_int = nx_silu_forward(scratch, scratch) 109 if v_silu1 != NX_SILU_OK { return NX_UB_ERR_INTERNAL } 110 111 // Step 3: out = Conv2D(scratch, W_1). 112 let v_cv1: nx_int = nx_conv2d_forward(scratch, W_1, 0 as *i64, out) 113 if v_cv1 != NX_CV2_OK { return NX_UB_ERR_INTERNAL } 114 115 // Step 4+5: scratch = SiLU(GroupNorm(out, gamma_2, beta_2)). 116 let v_gn2: nx_int = nx_groupnorm_forward(out, n_groups, gamma_2, beta_2, scratch) 117 if v_gn2 != NX_GN_OK { return NX_UB_ERR_INTERNAL } 118 let v_silu2: nx_int = nx_silu_forward(scratch, scratch) 119 if v_silu2 != NX_SILU_OK { return NX_UB_ERR_INTERNAL } 120 121 // Step 6: out = Conv2D(scratch, W_2). 122 let v_cv2: nx_int = nx_conv2d_forward(scratch, W_2, 0 as *i64, out) 123 if v_cv2 != NX_CV2_OK { return NX_UB_ERR_INTERNAL } 124 125 // Step 7+8: out = out + x (residual skip; v1 assumes C_in = C_out). 126 let n_elem: nx_int = x.shape[0] * x.shape[1] * x.shape[2] * x.shape[3] 127 let pi: *i64 = x.storage as *i64 128 let po: *i64 = out.storage as *i64 129 var i: nx_int = 0 130 var iter: nx_int = 0 131 var verdict: nx_int = NX_LOOP_RUNNING 132 let BUDGET: nx_int = n_elem 133 while verdict == NX_LOOP_RUNNING && iter < BUDGET { 134 po[i] = po[i] + pi[i] 135 i = i + 1 136 iter = iter + 1 137 } 138 return NX_UB_OK 139} 140 141// ===== Self-test ================================================== 142// 143// Builds a tiny UNet block: [N=1, C=4, H=2, W=2] with zero weights + 144// identity gammas + zero betas. Verifies the function COMPLETES OK 145// (structural-acceptance gate). Numerical correctness against a 146// reference model is the next workstream (needs a real 147// stable-diffusion UNet block + golden activation snapshot). 148// 149// Closed-form invariant: zero-weight forward leaves the residual 150// path unchanged. x + 0 = x. We verify this. 151 152func main() -> i64 { 153 let N: nx_int = 1 154 let C: nx_int = 4 155 let H: nx_int = 2 156 let W: nx_int = 2 157 158 let sh: *nx_int = sys_mmap(4 * 8) as *nx_int 159 sh[0]=N; sh[1]=C; sh[2]=H; sh[3]=W 160 161 let wt_sh: *nx_int = sys_mmap(4 * 8) as *nx_int 162 wt_sh[0]=C; wt_sh[1]=C; wt_sh[2]=3; wt_sh[3]=3 163 164 let err: *nx_int = sys_mmap(8) as *nx_int 165 err[0] = 0 166 let x: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 4, err) 167 let out: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 4, err) 168 let scratch: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 4, err) 169 let W_1: *NxTensor = nx_t_alloc(NX_DT_I64, wt_sh, 4, err) 170 let W_2: *NxTensor = nx_t_alloc(NX_DT_I64, wt_sh, 4, err) 171 if err[0] != 0 { return 5 } 172 173 // Fill x with a known pattern. 174 let px: *i64 = x.storage as *i64 175 var i: nx_int = 0 176 while i < N * C * H * W { px[i] = (i + 1) * 100; i = i + 1 } 177 178 // Zero W_1 and W_2 (already zero from sys_mmap). 179 // gamma = Q10 unit, beta = 0. 180 let gamma_1: *i64 = nx_groupnorm_gamma_unit(C) 181 let beta_1: *i64 = nx_groupnorm_beta_zero(C) 182 let gamma_2: *i64 = nx_groupnorm_gamma_unit(C) 183 let beta_2: *i64 = nx_groupnorm_beta_zero(C) 184 185 // 2 groups (C=4 / 2 channels each). 186 let v: nx_int = nx_unet_block_forward( 187 x, W_1, gamma_1, beta_1, 188 W_2, gamma_2, beta_2, 189 2, out, scratch) 190 if v != NX_UB_OK { return 10 + v } 191 192 // With zero weights, conv steps produce zero output. Residual 193 // path adds back x. So out should equal x bit-exact. 194 let po: *i64 = out.storage as *i64 195 var j: nx_int = 0 196 while j < N * C * H * W { 197 if po[j] != px[j] { return 20 } 198 j = j + 1 199 } 200 201 // --- Verdict gate --- 202 var vi: nx_int = 0 203 while vi < NX_UB_N_VERDICTS { 204 if nx_ub_verdict_is_valid(vi) != 1 { return 30 + vi } 205 vi = vi + 1 206 } 207 208 return 0 209}