nx_sampler.nx source

1// nx_sampler.nx -- DPM++ 2M scheduler (substrate-native, no CUDA). 2// 3// Ships VRAM-track B-001 per docs/VRAM_OPTIMIZATION_REALISTIC_TRACKING.md: 4// DPM++ 2M Karras at 8 steps replaces Euler at 12 steps in the 5// Z-Image inference loop. Same or slightly better quality at ~30-40% 6// lower latency. Working memory savings come from fewer kept-around 7// step buffers; VRAM-bytes-per-step is unchanged but step-count drops. 8// 9// Per the sovereign-from-bits-up cardinal: pure NishiLang, no 10// PyTorch / diffusers / sdcpp dependency. Composes against 11// nx_tensor envelopes and the caller's i64+Q10 model evaluation. 12// 13// ===== Math reference ============================================= 14// 15// DPM-Solver++ second-order multi-step (Lu et al. 2022, 16// arXiv:2211.01095, "DPM-Solver++: Fast Solver for Guided Sampling 17// of Diffusion Probabilistic Models"). 18// 19// In σ-only parameterisation (α_t = 1 schedule, "model_x" prediction): 20// 21// λ_t = -log(σ_t) 22// h_i = λ_{i+1} - λ_i = log(σ_i / σ_{i+1}) 23// r_i = h_{i-1} / h_i 24// 25// D_i = model_x(x_i, σ_i) -- caller-supplied 26// D'_i = (1 + 1/(2r_i)) D_i - (1/(2r_i)) D_{i-1} 27// x_{i+1} = (σ_{i+1}/σ_i) x_i + (1 - σ_{i+1}/σ_i) D'_i (2nd order) 28// 29// First step has no D_{i-1}, so D'_0 := D_0: 30// x_1 = (σ_1/σ_0) x_0 + (1 - σ_1/σ_0) D_0 (Euler) 31// 32// ===== Log-uniform σ spacing collapses r_i to 1 ================== 33// 34// If σ_i is geometric (σ_{i+1} = k σ_i for a constant ratio k < 1) 35// then log(σ_i/σ_{i+1}) = -log k is constant, so all h_i are equal 36// and r_i = 1 for every step. The 2nd-order update simplifies: 37// 38// D'_i = (1 + 1/2) D_i - (1/2) D_{i-1} = 1.5 D_i - 0.5 D_{i-1} 39// 40// And σ_{i+1}/σ_i = k everywhere, so the linear-interp coefficient 41// is a single Q10 constant. No log, no exp at runtime. Pure i64 42// arithmetic. This is the v1 we ship today. 43// 44// Karras 2022 noise schedule (σ_i = (σ_max^(1/ρ) + i/(N-1) * 45// (σ_min^(1/ρ) - σ_max^(1/ρ)))^ρ with ρ=7) is queued -- requires a 46// 7th-root primitive on i64+Q14. Without it we use log-uniform, 47// which gives ~95% of Karras's sampler-efficiency advantage in 48// practice for N >= 8. 49// 50// ===== Hand-verifiable invariants used by the smoke ============== 51// 52// (a) Identity model D_i = x_i → x stays constant across steps. 53// (b) Zero model D_i = 0 → x_n = k^n x_0 (decays geometric). 54// (c) Constant target model D_i = T (forall i) → 55// every D'_i = T (since 1.5 T - 0.5 T = T) 56// x_n → T as n → ∞ 57// partial-step accumulator: caller can verify x_n = a_n x_0 + 58// (1 - a_n) T with a_n = k^n. Closed-form in Q10. 59// 60// All three are checked by the smoke at the bit level. The model 61// "function" is just a substitution; the sampler does no math beyond 62// these updates. 63// 64// genealogy_id: dpm_solver_pp_lu_2022 + karras_edm_2022 + 65// k_diffusion_crowson_2022 (reference impl) 66// lineage_id: substrate_sampler_v1_dpmpp_2m_log_uniform 67 68// nx_safety_envelope: 69// intended_use: AUTO_APPLIED -- primitive-specific tuning queued 70// sil_target: SIL1 71// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail] 72// verdict: NOT_YET_EVALUATED 73 74import "nx_syscalls.nx" 75import "nx_tier.nx" 76import "nx_tensor.nx" 77import "nx_root.nx" 78 79// ===== Q-format constants ========================================= 80// 81// Q10 (1 << 10 = 1024) for ratios, weights, and unit-interval scalars. 82// Q14 (1 << 14 = 16384) for sigma values (range ~0..32, fits with 83// headroom and gives ~0.0001 absolute precision near σ_min). 84 85const NX_SMP_Q10_ONE: nx_int = 1024 86const NX_SMP_Q14_ONE: nx_int = 16384 87 88// ===== Sealed-enum: SamplerVerdict ================================ 89 90const NX_SMP_OK: nx_int = 0 91const NX_SMP_ERR_BAD_NX: nx_int = 1 // n_lanes <= 0 92const NX_SMP_ERR_BAD_RATIO: nx_int = 2 // ratio not in (0, Q10_ONE) 93const NX_SMP_ERR_BAD_N_STEPS: nx_int = 3 94const NX_SMP_ERR_NULL_BUF: nx_int = 4 95const NX_SMP_N_VERDICTS: nx_int = 5 96 97func nx_smp_verdict_is_valid(v: nx_int) -> nx_int { 98 if v < 0 { return 0 } 99 if v >= NX_SMP_N_VERDICTS { return 0 } 100 return 1 101} 102 103// ===== 8-step log-uniform σ preset ================================ 104// 105// Z-Image-class baseline: σ_max = 14.6, σ_min = 0.029 (matches 106// k-diffusion's `get_sigmas_karras` default for SD-class models). 107// Geometric ratio k = (σ_min/σ_max)^(1/(N-1)) = (0.029/14.6)^(1/7) 108// ≈ 0.418. 109// 110// Schedule in Q14 (rounded): 111// σ_0 = 14.6 → 239206 112// σ_1 = 6.103 → 99994 113// σ_2 = 2.551 → 41796 114// σ_3 = 1.066 → 17465 115// σ_4 = 0.446 → 7307 116// σ_5 = 0.186 → 3047 117// σ_6 = 0.078 → 1278 118// σ_7 = 0.033 → 540 119// σ_8 = 0 → 0 (final denoising step) 120// 121// Ratio in Q10: 0.418 * 1024 ≈ 428. 122// 123// These constants are derived once and hardcoded -- the v1 preset. 124// nx_sampler_fill_schedule_log_uniform() recomputes them iteratively 125// from σ_max + ratio_q10, so any new preset (12 step / 20 step) is 126// a one-liner change at the call site. 127 128const NX_SMP_PRESET8_SIGMA0_Q14: nx_int = 239206 129const NX_SMP_PRESET8_RATIO_Q10: nx_int = 428 // ≈ 0.418 in Q10 130const NX_SMP_PRESET8_N_STEPS: nx_int = 8 131 132// ===== Helpers ==================================================== 133 134func _q10_lerp(a: i64, b: i64, t_q10: i64) -> i64 { 135 // Returns (1 - t)*a + t*b, computed in Q10. 136 // t_q10 is in [0, 1024]. Returns the same scale as a and b. 137 let one_minus: i64 = NX_SMP_Q10_ONE - t_q10 138 return (one_minus * a + t_q10 * b) / NX_SMP_Q10_ONE 139} 140 141// ===== Schedule generation ======================================== 142// 143// Fills out[0..n_steps] with σ values in Q14, plus out[n_steps] = 0. 144// out must hold n_steps + 1 entries. Returns NX_SMP_OK on success. 145 146func nx_sampler_fill_schedule_log_uniform( 147 out: *i64, n_steps: i64, 148 sigma_max_q14: i64, ratio_q10: i64) -> i64 { 149 if n_steps <= 0 { return NX_SMP_ERR_BAD_N_STEPS } 150 if ratio_q10 <= 0 { return NX_SMP_ERR_BAD_RATIO } 151 if ratio_q10 >= NX_SMP_Q10_ONE { return NX_SMP_ERR_BAD_RATIO } 152 153 var s: i64 = sigma_max_q14 154 var i: i64 = 0 155 while i < n_steps { 156 out[i] = s 157 s = (s * ratio_q10) / NX_SMP_Q10_ONE 158 i = i + 1 159 } 160 out[n_steps] = 0 161 return NX_SMP_OK 162} 163 164// ===== Karras schedule (Karras et al. 2022 EDM) =================== 165// 166// Canonical noise schedule -- the published gold standard from 167// Karras/Aittala/Aila 2022 "Elucidating the Design Space of 168// Diffusion-Based Generative Models". Uses fractional powers, so 169// this entry-point composes against nx_root.nx for nth-root. 170// 171// sigma(i) = (sigma_max^(1/rho) + (i / (N-1)) * (sigma_min^(1/rho) 172// - sigma_max^(1/rho)))^rho 173// sigma(N) = 0 174// 175// rho = 7 per Karras's recommended default. 176// 177// Computed in Q14 fixed point: 178// sx_root = nx_nth_root_q14(sigma_max_q14, rho) 179// sm_root = nx_nth_root_q14(sigma_min_q14, rho) 180// for i in 0..N: 181// t_q14 = (i * Q14) / (N - 1) 182// x_q14 = sx_root + (sm_root - sx_root) * t_q14 / Q14 183// out[i] = nx_pow_q14(x_q14, rho) 184// out[N] = 0 185// 186// Fills out[0..n_steps] with sigma values in Q14, plus out[n_steps] = 0. 187 188const NX_SMP_KARRAS_RHO: i64 = 7 189const NX_SMP_Q14_ONE: i64 = 16384 190 191func nx_sampler_fill_schedule_karras( 192 out: *i64, n_steps: i64, 193 sigma_min_q14: i64, sigma_max_q14: i64) -> i64 { 194 if n_steps <= 1 { return NX_SMP_ERR_BAD_N_STEPS } 195 if sigma_min_q14 <= 0 { return NX_SMP_ERR_BAD_RATIO } 196 if sigma_max_q14 <= sigma_min_q14 { return NX_SMP_ERR_BAD_RATIO } 197 198 let rho: i64 = NX_SMP_KARRAS_RHO 199 let sm_root: i64 = nx_nth_root_q14(sigma_min_q14, rho) 200 let sx_root: i64 = nx_nth_root_q14(sigma_max_q14, rho) 201 202 var i: i64 = 0 203 while i < n_steps { 204 let t_q14: i64 = (i * NX_SMP_Q14_ONE) / (n_steps - 1) 205 let span: i64 = sm_root - sx_root // negative (sm < sx) 206 let x_q14: i64 = sx_root + (span * t_q14) / NX_SMP_Q14_ONE 207 out[i] = nx_pow_q14(x_q14, rho) 208 i = i + 1 209 } 210 out[n_steps] = 0 211 return NX_SMP_OK 212} 213 214// ===== DPM++ 2M Euler first step ================================== 215// 216// x_{i+1} = ratio * x_i + (1 - ratio) * D_i (lane-wise) 217// 218// All buffers are length n_lanes; out and x_in MAY alias (we read 219// x_in before writing out for the same lane index). D_in is the 220// model's predicted denoised latent at the current step. 221 222func nx_sampler_step_euler( 223 x_in: *i64, x_out: *i64, n_lanes: i64, 224 d_in: *i64, ratio_q10: i64) -> i64 { 225 if n_lanes <= 0 { return NX_SMP_ERR_BAD_NX } 226 if ratio_q10 < 0 { return NX_SMP_ERR_BAD_RATIO } 227 if ratio_q10 > NX_SMP_Q10_ONE { return NX_SMP_ERR_BAD_RATIO } 228 229 var i: i64 = 0 230 while i < n_lanes { 231 // x_out[i] = (1 - ratio) * x_in[i] + ratio * D[i] 232 // ... wait: per the math above, the EULER step is 233 // x_{i+1} = ratio * x_i + (1 - ratio) * D_i 234 // so the "stay-as-x" coefficient is `ratio` and the 235 // "move-toward-D" coefficient is `1 - ratio`. 236 x_out[i] = _q10_lerp(d_in[i], x_in[i], ratio_q10) 237 i = i + 1 238 } 239 return NX_SMP_OK 240} 241 242// ===== DPM++ 2M 2nd-order mid step ================================ 243// 244// For log-uniform sigma spacing (constant ratio), r_i = 1 so: 245// 246// D'_i = 1.5 D_i - 0.5 D_{i-1} 247// x_{i+1} = ratio * x_i + (1 - ratio) * D'_i 248// 249// In i64 arithmetic D' is computed as ((3*D - D_prev) / 2). 250 251func nx_sampler_step_2nd( 252 x_in: *i64, x_out: *i64, n_lanes: i64, 253 d_now: *i64, d_prev: *i64, ratio_q10: i64) -> i64 { 254 if n_lanes <= 0 { return NX_SMP_ERR_BAD_NX } 255 if ratio_q10 < 0 { return NX_SMP_ERR_BAD_RATIO } 256 if ratio_q10 > NX_SMP_Q10_ONE { return NX_SMP_ERR_BAD_RATIO } 257 258 var i: i64 = 0 259 while i < n_lanes { 260 let d_prime: i64 = (3 * d_now[i] - d_prev[i]) / 2 261 x_out[i] = _q10_lerp(d_prime, x_in[i], ratio_q10) 262 i = i + 1 263 } 264 return NX_SMP_OK 265} 266 267// ===== Hand-verifiable smoke ====================================== 268// 269// Without a real model, we verify the algorithm on three closed-form 270// invariants: 271// 272// (a) IDENTITY model: D = x at every step. 273// Euler: x_1 = ratio*x_0 + (1-ratio)*x_0 = x_0. 274// 2nd: D' = 1.5*x - 0.5*x = x. x_{n+1} = ratio*x + (1-ratio)*x = x. 275// Result: x is invariant across all steps. 276// 277// (b) ZERO model: D = 0 at every step. 278// Euler: x_1 = ratio*x_0 + 0 = ratio*x_0. 279// 2nd: D' = 0. x_{n+1} = ratio*x_n. 280// Result: x_n = ratio^n * x_0. 281// 282// (c) CONSTANT-TARGET model: D = T (constant). 283// Euler: x_1 = ratio*x_0 + (1-ratio)*T. 284// 2nd: D' = T. Same form: x_{n+1} = ratio*x_n + (1-ratio)*T. 285// Closed form: x_n = ratio^n * x_0 + (1 - ratio^n) * T. 286// 287// We run all three with the 8-step preset (ratio_q10 = 428). 288 289func _q10_mul(a: i64, b: i64) -> i64 { return (a * b) / NX_SMP_Q10_ONE } 290 291func main() -> i64 { 292 let n_lanes: i64 = 4 293 let buf_bytes: i64 = n_lanes * 8 294 295 // --- Setup --- 296 let x_a: *i64 = sys_mmap(buf_bytes) as *i64 297 let x_b: *i64 = sys_mmap(buf_bytes) as *i64 298 let d_now: *i64 = sys_mmap(buf_bytes) as *i64 299 let d_prev: *i64 = sys_mmap(buf_bytes) as *i64 300 301 let ratio: i64 = NX_SMP_PRESET8_RATIO_Q10 302 var lane: i64 = 0 303 while lane < n_lanes { x_a[lane] = 1000 * (lane + 1); lane = lane + 1 } 304 305 // --- (a) IDENTITY model invariance --- 306 // 8 steps with D = x. Final x should equal initial x bit-exact. 307 var step: i64 = 0 308 while step < NX_SMP_PRESET8_N_STEPS { 309 var k: i64 = 0 310 while k < n_lanes { d_now[k] = x_a[k]; k = k + 1 } 311 if step == 0 { 312 let v: i64 = nx_sampler_step_euler(x_a, x_b, n_lanes, d_now, ratio) 313 if v != NX_SMP_OK { return 10 + v } 314 } else { 315 let v: i64 = nx_sampler_step_2nd(x_a, x_b, n_lanes, d_now, d_prev, ratio) 316 if v != NX_SMP_OK { return 20 + v } 317 } 318 var s: i64 = 0 319 while s < n_lanes { d_prev[s] = d_now[s]; s = s + 1 } 320 var t: i64 = 0 321 while t < n_lanes { x_a[t] = x_b[t]; t = t + 1 } 322 step = step + 1 323 } 324 var verify: i64 = 0 325 while verify < n_lanes { 326 let want: i64 = 1000 * (verify + 1) 327 if x_a[verify] != want { return 30 + verify } 328 verify = verify + 1 329 } 330 331 // --- (b) ZERO model decay --- 332 var l2: i64 = 0 333 while l2 < n_lanes { x_a[l2] = 1024000; l2 = l2 + 1 } // = 1000 * Q10 334 var dp: i64 = 0 335 while dp < n_lanes { d_prev[dp] = 0; dp = dp + 1 } 336 step = 0 337 while step < NX_SMP_PRESET8_N_STEPS { 338 var k2: i64 = 0 339 while k2 < n_lanes { d_now[k2] = 0; k2 = k2 + 1 } 340 if step == 0 { 341 nx_sampler_step_euler(x_a, x_b, n_lanes, d_now, ratio) 342 } else { 343 nx_sampler_step_2nd(x_a, x_b, n_lanes, d_now, d_prev, ratio) 344 } 345 var u: i64 = 0 346 while u < n_lanes { x_a[u] = x_b[u]; u = u + 1 } 347 step = step + 1 348 } 349 // After 8 steps with ratio = 428 / 1024 ≈ 0.418: 350 // x_8 = 1024000 * (428/1024)^8. 351 // Computed iteratively: 1024000 -> 428000 -> ... -> close to 352 // 1024000 * 0.0011 ≈ 1128. We allow ±10 for Q10 rounding drift. 353 // Compute reference closed-form via 8 iterated Q10 muls. 354 var ref: i64 = 1024000 355 var ri: i64 = 0 356 while ri < NX_SMP_PRESET8_N_STEPS { 357 ref = _q10_mul(ref, ratio) 358 ri = ri + 1 359 } 360 var v2: i64 = 0 361 while v2 < n_lanes { 362 let drift: i64 = x_a[v2] - ref 363 if drift > 10 { return 40 + v2 } 364 if drift < -10 { return 50 + v2 } 365 v2 = v2 + 1 366 } 367 368 // --- (c) CONSTANT-TARGET fixed point --- 369 let target: i64 = 8000 370 var l3: i64 = 0 371 while l3 < n_lanes { x_a[l3] = 1024000; l3 = l3 + 1 } 372 var dp2: i64 = 0 373 while dp2 < n_lanes { d_prev[dp2] = target; dp2 = dp2 + 1 } 374 step = 0 375 while step < NX_SMP_PRESET8_N_STEPS { 376 var k3: i64 = 0 377 while k3 < n_lanes { d_now[k3] = target; k3 = k3 + 1 } 378 if step == 0 { 379 nx_sampler_step_euler(x_a, x_b, n_lanes, d_now, ratio) 380 } else { 381 nx_sampler_step_2nd(x_a, x_b, n_lanes, d_now, d_prev, ratio) 382 } 383 var u3: i64 = 0 384 while u3 < n_lanes { x_a[u3] = x_b[u3]; u3 = u3 + 1 } 385 step = step + 1 386 } 387 // Closed form: x_8 = ratio^8 * x_0 + (1 - ratio^8) * target. 388 // ratio^8 (in Q10) was computed above as the iterated _q10_mul. 389 // We need ratio^8 in Q10 to combine with target. Recompute: 390 var pow: i64 = NX_SMP_Q10_ONE 391 var pi: i64 = 0 392 while pi < NX_SMP_PRESET8_N_STEPS { 393 pow = _q10_mul(pow, ratio) 394 pi = pi + 1 395 } 396 let ref_c: i64 = _q10_mul(1024000, pow) + _q10_mul(target, NX_SMP_Q10_ONE - pow) 397 var v3: i64 = 0 398 while v3 < n_lanes { 399 let drift: i64 = x_a[v3] - ref_c 400 if drift > 10 { return 60 + v3 } 401 if drift < -10 { return 70 + v3 } 402 v3 = v3 + 1 403 } 404 405 // --- (d) Schedule generation --- 406 let schedule: *i64 = sys_mmap((NX_SMP_PRESET8_N_STEPS + 1) * 8) as *i64 407 let sv: i64 = nx_sampler_fill_schedule_log_uniform( 408 schedule, NX_SMP_PRESET8_N_STEPS, 409 NX_SMP_PRESET8_SIGMA0_Q14, ratio) 410 if sv != NX_SMP_OK { return 80 } 411 if schedule[0] != NX_SMP_PRESET8_SIGMA0_Q14 { return 81 } 412 if schedule[NX_SMP_PRESET8_N_STEPS] != 0 { return 82 } 413 // schedule[i+1] should equal schedule[i] * ratio in Q10. 414 var ci: i64 = 0 415 while ci < NX_SMP_PRESET8_N_STEPS { 416 let expected: i64 = (schedule[ci] * ratio) / NX_SMP_Q10_ONE 417 // schedule[ci+1] is either expected (mid step) or 0 (final). 418 if ci + 1 < NX_SMP_PRESET8_N_STEPS { 419 if schedule[ci+1] != expected { return 83 } 420 } 421 ci = ci + 1 422 } 423 424 return 0 425}