nx_diffusion_loop.nx source
↩ module page · 245 lines · 8523 B
1// nx_diffusion_loop.nx -- N-step diffusion sampler composer.
2//
3// L4 brick. Closes the image-gen forward path: given noise + a
4// UNet block per layer + a denoising sigma schedule, produces a
5// denoised latent over N sampler steps.
6//
7// Composes:
8// nx_sampler (DPM++ 2M Euler-first + 2nd-order-mid,
9// + Karras schedule via nx_root)
10// nx_unet_block (per-layer ResBlock; caller stacks N layers)
11// NxTensor (L1 4D containers)
12// nx_loop.LoopVerdict (control)
13//
14// ===== High-level shape ==========================================
15//
16// noise -> x_0
17// for i in 0..n_steps:
18// sigma_now = schedule[i]
19// sigma_next = schedule[i+1]
20//
21// // Denoising prediction: caller's UNet (stack of nx_unet_block).
22// // v1 takes a callback function pointer so the composer is
23// // architecture-agnostic.
24// d_now = denoise(x, sigma_now)
25//
26// // DPM++ 2M step.
27// if i == 0:
28// x = nx_sampler_step_euler(x, d_now, sigma_now, sigma_next)
29// else:
30// x = nx_sampler_step_2nd(x, d_now, d_prev, ratio)
31//
32// d_prev = d_now
33//
34// return x
35//
36// ===== v1 scope ==================================================
37//
38// v1 makes the denoiser callback signature explicit:
39//
40// func denoise(x_lanes: *i64, n_lanes: nx_int, sigma_q14: i64,
41// out_d: *i64, ctx: i64) -> nx_int
42//
43// Caller pre-allocates a denoise context (their UNet weights bundle)
44// and passes its pointer cast to i64 as `ctx`. The denoiser is free
45// to interpret ctx however it wants (typically a *NxUnetWeights
46// bundle).
47//
48// For v1 we DON'T validate the denoiser callback's verdict; we
49// trust the caller (per the substrate's "trusting internally"
50// pattern). v2 may add a verdict check.
51//
52// Time-conditioning (classifier-free guidance, prompt embedding
53// injection): caller bakes this INTO their denoise callback;
54// the loop composer is unaware.
55//
56// genealogy_id: ho_2020_ddpm + lu_2022_dpm_solver_pp + karras_2022_edm +
57// rombach_2022_stable_diffusion_sampling_loop
58// lineage_id: substrate_diffusion_loop_v1_dpmpp_2m
59
60// nx_safety_envelope:
61// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
62// sil_target: SIL1
63// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
64// verdict: NOT_YET_EVALUATED
65
66import "nx_syscalls.nx"
67import "nx_tier.nx"
68import "nx_loop.nx"
69import "nx_sampler.nx"
70
71const NX_DL_Q10: nx_int = 1024
72const NX_DL_Q14: nx_int = 16384
73
74// ===== Sealed-enum: DiffusionLoopVerdict ==========================
75
76const NX_DL_OK: nx_int = 0
77const NX_DL_ERR_BAD_DIMS: nx_int = 1
78const NX_DL_ERR_BAD_SCHEDULE: nx_int = 2
79const NX_DL_ERR_OOM: nx_int = 3
80const NX_DL_ERR_DENOISER: nx_int = 4
81const NX_DL_N_VERDICTS: nx_int = 5
82
83func nx_dl_verdict_is_valid(v: nx_int) -> nx_int {
84 if v < 0 { return 0 }
85 if v >= NX_DL_N_VERDICTS { return 0 }
86 return 1
87}
88
89// ===== Compute the per-step ratio q10 =============================
90//
91// ratio_i = sigma[i+1] / sigma[i] (in Q10).
92// For log-uniform schedule this is constant; for Karras it varies
93// per step. Computed here from the Q14 schedule.
94
95func _dl_ratio_q10(sigma_now_q14: i64, sigma_next_q14: i64) -> nx_int {
96 if sigma_now_q14 <= 0 { return 0 }
97 return (sigma_next_q14 * NX_DL_Q10) / sigma_now_q14
98}
99
100// ===== Main loop ==================================================
101//
102// x_io: [n_lanes] i64 Q10 (in: noise; out: denoised latent)
103// n_lanes: number of latent elements (e.g. N*C*H*W flattened)
104// sigmas: [n_steps + 1] i64 Q14 (last entry = 0)
105// n_steps: number of sampler iterations
106// denoise_fn: func(x_lanes, n_lanes, sigma_q14, out_d, ctx) -> nx_int
107// ctx: caller-owned opaque context (e.g. UNet weights bundle ptr)
108//
109// Scratch:
110// d_now, d_prev: [n_lanes] each
111// x_scratch: [n_lanes]
112// All allocated internally via sys_mmap (substrate convention).
113//
114// Returns NX_DL_OK on success or a sealed verdict.
115
116func nx_diffusion_loop(x_io: *i64, n_lanes: nx_int,
117 sigmas: *i64, n_steps: nx_int,
118 denoise_fn: func(*i64, nx_int, i64, *i64, i64) -> nx_int,
119 ctx: i64) -> nx_int {
120 if n_lanes <= 0 { return NX_DL_ERR_BAD_DIMS }
121 if n_steps <= 0 { return NX_DL_ERR_BAD_SCHEDULE }
122
123 let d_now: *i64 = sys_mmap(n_lanes * 8) as *i64
124 let d_prev: *i64 = sys_mmap(n_lanes * 8) as *i64
125 let x_scratch: *i64 = sys_mmap(n_lanes * 8) as *i64
126
127 var step: nx_int = 0
128 var iter: nx_int = 0
129 var verdict: nx_int = NX_LOOP_RUNNING
130 let BUDGET: nx_int = n_steps
131 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
132 let sigma_now: i64 = sigmas[step]
133 let sigma_next: i64 = sigmas[step + 1]
134
135 // Denoise: caller's UNet predicts the clean signal at sigma_now.
136 let v_d: nx_int = denoise_fn(x_io, n_lanes, sigma_now, d_now, ctx)
137 if v_d != 0 { verdict = NX_LOOP_ABORTED }
138
139 if verdict == NX_LOOP_RUNNING {
140 // Compute ratio for this step.
141 let ratio: nx_int = _dl_ratio_q10(sigma_now, sigma_next)
142
143 // Sampler step: Euler on first step, 2nd-order DPM++ on subsequent.
144 if step == 0 {
145 nx_sampler_step_euler(x_io, x_scratch, n_lanes, d_now, ratio)
146 }
147 if step > 0 {
148 nx_sampler_step_2nd(x_io, x_scratch, n_lanes, d_now, d_prev, ratio)
149 }
150
151 // Copy scratch -> x_io for next iteration.
152 var k: nx_int = 0
153 var k_iter: nx_int = 0
154 var k_verdict: nx_int = NX_LOOP_RUNNING
155 let K_BUDGET: nx_int = n_lanes
156 while k_verdict == NX_LOOP_RUNNING && k_iter < K_BUDGET {
157 x_io[k] = x_scratch[k]
158 k = k + 1
159 k_iter = k_iter + 1
160 }
161
162 // Save d_now as d_prev for the next step.
163 var dp: nx_int = 0
164 var dp_iter: nx_int = 0
165 var dp_verdict: nx_int = NX_LOOP_RUNNING
166 while dp_verdict == NX_LOOP_RUNNING && dp_iter < K_BUDGET {
167 d_prev[dp] = d_now[dp]
168 dp = dp + 1
169 dp_iter = dp_iter + 1
170 }
171 }
172
173 step = step + 1
174 iter = iter + 1
175 }
176 if verdict == NX_LOOP_ABORTED { return NX_DL_ERR_DENOISER }
177 return NX_DL_OK
178}
179
180// ===== Denoise: identity (smoke testing only) ====================
181//
182// A trivial denoiser that returns d = x (predicts the input is
183// already clean). Used by the self-test to verify the loop wiring
184// without needing a real UNet.
185
186func _dl_identity_denoiser(x_lanes: *i64, n_lanes: nx_int, sigma_q14: i64,
187 out_d: *i64, ctx: i64) -> nx_int {
188 var i: nx_int = 0
189 var iter: nx_int = 0
190 var verdict: nx_int = NX_LOOP_RUNNING
191 let BUDGET: nx_int = n_lanes
192 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
193 out_d[i] = x_lanes[i]
194 i = i + 1
195 iter = iter + 1
196 }
197 return 0
198}
199
200// ===== Self-test ==================================================
201//
202// Use the log-uniform schedule (already in nx_sampler) + identity
203// denoiser. With D = x at every step, the sampler's update reduces
204// to x_{i+1} = ratio*x + (1-ratio)*x = x. Output should equal input.
205
206func main() -> i64 {
207 let N_STEPS: nx_int = 8
208 let N_LANES: nx_int = 4
209
210 let schedule: *i64 = sys_mmap((N_STEPS + 1) * 8) as *i64
211 let v_sched: nx_int = nx_sampler_fill_schedule_log_uniform(
212 schedule, N_STEPS,
213 NX_SMP_PRESET8_SIGMA0_Q14, NX_SMP_PRESET8_RATIO_Q10)
214 if v_sched != NX_SMP_OK { return 5 }
215
216 let x: *i64 = sys_mmap(N_LANES * 8) as *i64
217 var i: nx_int = 0
218 while i < N_LANES { x[i] = (i + 1) * 1000; i = i + 1 }
219
220 // Save initial.
221 let initial: *i64 = sys_mmap(N_LANES * 8) as *i64
222 var k: nx_int = 0
223 while k < N_LANES { initial[k] = x[k]; k = k + 1 }
224
225 let v: nx_int = nx_diffusion_loop(
226 x, N_LANES, schedule, N_STEPS,
227 _dl_identity_denoiser, 0)
228 if v != NX_DL_OK { return 10 + v }
229
230 // Identity denoiser -> x should be unchanged.
231 var j: nx_int = 0
232 while j < N_LANES {
233 if x[j] != initial[j] { return 20 }
234 j = j + 1
235 }
236
237 // --- Verdict gate ---
238 var vi: nx_int = 0
239 while vi < NX_DL_N_VERDICTS {
240 if nx_dl_verdict_is_valid(vi) != 1 { return 30 + vi }
241 vi = vi + 1
242 }
243
244 return 0
245}