nx_gsplat.nx source

1// nx_gsplat.nx -- ★SOVEREIGN 3D GAUSSIAN SPLATTING (the frontier rendering rep for photoreal avatars, confirmed 2// #1 by the 2024-2025 SOTA scan: GaussianAvatars/VRGaussianAvatar/HumanSplat all use it). A scene = a cloud of 3// 3D Gaussians {position, scale(σ), colour, opacity}; render = PROJECT each to a 2D splat, DEPTH-SORT, and 4// FRONT-TO-BACK alpha-composite. ALL INTEGER (fx1024 positions, fx256 opacity/transmittance, an integer exp-LUT 5// for the Gaussian falloff), deterministic, VM-vettable. ★The RENDERER needs NO trained weights (only a 6// generator would) -- so it is fully sovereign-buildable TODAY. v0 = ISOTROPIC Gaussians (spherical σ -> a 7// circular screen splat); anisotropic 3D covariance (the Jacobian projection) + SH view-dependent colour are the 8// next rungs. Camera convention MATCHES nx_sdfrender (yaw orbit, FOCAL 586) so splats align with our SDF/mesh. 9// license_tier: ORIGINAL 10import "nx_syscalls.nx" 11import "nx_itrig.nx" 12const K_MAGIC_4096: i64 = 4096 13const K_MAGIC_3970: i64 = 3970 14const K_MAGIC_2000000000: i64 = 2000000000 15const K_MAGIC_65536: i64 = 65536 16const K_MAGIC_999999: i64 = 999999 17 18const GW: i64 = 512 19const GH: i64 = 384 20const GHW: i64 = 256 21const GHH: i64 = 192 22const GFOCAL: i64 = 586 23const GFX: i64 = 1024 // model-unit fixed point (matches sdf) 24const GEXPN: i64 = 176 // exp-LUT entries: k=0..175 -> u=k/16 in [0,11) (11σ² = past 3σ) 25const GNB: i64 = 2048 // depth-sort buckets 26 27func gs_isqrt(v: i64) -> i64 { if v <= 0 { return 0 } var x: i64 = v; var y: i64 = (x + 1) / 2; while y < x { x = y; y = (x + v / x) / 2 } return x } 28 29// exp-LUT: explut[k] = round(256 * exp(-0.5 * k/16)), built by an integer recurrence (factor exp(-0.5/16) 30// ~= 3970/4096 in fx4096, then scaled to fx256). No float, deterministic. 31func gs_build_explut(explut: *i64) -> i64 { 32 var v: i64 = K_MAGIC_4096 // fx4096, exp(0)=1 33 var k: i64 = 0 34 while k < GEXPN { 35 explut[k] = v * 256 / K_MAGIC_4096 // -> fx256 36 v = v * K_MAGIC_3970 / K_MAGIC_4096 // *= exp(-0.5/16) 37 k = k + 1 38 } 39 return 0 40} 41 42// clear the per-pixel accumulators: acc = 0 (fx256 colour sum), trans = 256 (fx256 transmittance = 1.0) 43func gs_clear(acc: *i64, trans: *i64) -> i64 { 44 var i: i64 = 0 45 while i < GW * GH { acc[i*3] = 0; acc[i*3+1] = 0; acc[i*3+2] = 0; trans[i] = 256; i = i + 1 } 46 return 0 47} 48 49// render the Gaussian cloud. gauss = ng*8 i64 {x,y,z,scale,r,g,b,opacity(fx256)}; camera (yaw it4096, camz units). 50// scratch: acc(GW*GH*3), trans(GW*GH), depth(ng), sx(ng), sy(ng), sig(ng), order(ng), count(GNB+1), explut(GEXPN). 51// writes fb (GW*GH packed rgb). Returns the number of splatted (visible) gaussians. 52func gs_render(gauss: *i64, ng: i64, yaw: i64, camz: i64, fb: *i64, acc: *i64, trans: *i64, depth: *i64, sxb: *i64, syb: *i64, sigb: *i64, order: *i64, count: *i64, explut: *i64, bgr: i64, bgg: i64, bgb: i64) -> i64 { 53 let sy4: i64 = it_sin4096(yaw) 54 let cy4: i64 = it_cos4096(yaw) 55 let R: i64 = camz * GFX 56 // --- project all gaussians; mark culled with depth = -1 --- 57 var dmin: i64 = K_MAGIC_2000000000 58 var dmax: i64 = 0 - K_MAGIC_2000000000 59 var vis: i64 = 0 60 var i: i64 = 0 61 while i < ng { 62 let x: i64 = gauss[i*8] 63 let y: i64 = gauss[i*8+1] 64 let z: i64 = gauss[i*8+2] 65 let sc: i64 = gauss[i*8+3] 66 let cz: i64 = (x * sy4 + z * cy4) / K_MAGIC_4096 + R // camera-forward depth 67 if cz > 64 { 68 let cx: i64 = (x * cy4 - z * sy4) / K_MAGIC_4096 69 let sx: i64 = GHW + GFOCAL * cx / cz 70 let sy: i64 = GHH - GFOCAL * y / cz 71 var sig: i64 = GFOCAL * sc / cz // screen-space sigma (px) 72 if sig < 1 { sig = 1 } 73 depth[i] = cz; sxb[i] = sx; syb[i] = sy; sigb[i] = sig 74 if cz < dmin { dmin = cz } 75 if cz > dmax { dmax = cz } 76 vis = vis + 1 77 } else { depth[i] = 0 - 1 } 78 i = i + 1 79 } 80 if vis == 0 { return 0 } 81 // --- counting sort visible gaussians by depth (near->far = ascending cz) --- 82 let span: i64 = dmax - dmin + 1 83 var b: i64 = 0 84 while b <= GNB { count[b] = 0; b = b + 1 } 85 i = 0 86 while i < ng { if depth[i] >= 0 { var bk: i64 = (depth[i] - dmin) * GNB / span; if bk < 0 { bk = 0 } if bk >= GNB { bk = GNB - 1 } count[bk] = count[bk] + 1 } i = i + 1 } 87 var acc2: i64 = 0 // prefix sum -> bucket start offsets 88 b = 0 89 while b < GNB { let c: i64 = count[b]; count[b] = acc2; acc2 = acc2 + c; b = b + 1 } 90 i = 0 91 while i < ng { if depth[i] >= 0 { var bk: i64 = (depth[i] - dmin) * GNB / span; if bk < 0 { bk = 0 } if bk >= GNB { bk = GNB - 1 } order[count[bk]] = i; count[bk] = count[bk] + 1 } i = i + 1 } 92 // --- clear + composite front-to-back --- 93 gs_clear(acc, trans) 94 var oi: i64 = 0 95 while oi < vis { 96 let g: i64 = order[oi] 97 let sx: i64 = sxb[g] 98 let sy: i64 = syb[g] 99 let sig: i64 = sigb[g] 100 let rad: i64 = 3 * sig 101 let s2: i64 = sig * sig 102 let gr: i64 = gauss[g*8+4] 103 let gg: i64 = gauss[g*8+5] 104 let gb: i64 = gauss[g*8+6] 105 let op: i64 = gauss[g*8+7] 106 var py: i64 = sy - rad 107 if py < 0 { py = 0 } 108 var pye: i64 = sy + rad 109 if pye >= GH { pye = GH - 1 } 110 while py <= pye { 111 let dy: i64 = py - sy 112 var px: i64 = sx - rad 113 if px < 0 { px = 0 } 114 var pxe: i64 = sx + rad 115 if pxe >= GW { pxe = GW - 1 } 116 while px <= pxe { 117 let dx: i64 = px - sx 118 let d2: i64 = dx * dx + dy * dy 119 var k: i64 = d2 * 16 / s2 // u=d2/sig2 in 1/16ths 120 if k < GEXPN { 121 let pix: i64 = py * GW + px 122 let tr: i64 = trans[pix] 123 if tr > 1 { 124 let alpha: i64 = op * explut[k] / 256 // fx256 splat opacity at this pixel 125 let contrib: i64 = tr * alpha / 256 // transmittance-weighted (fx256) 126 acc[pix*3] = acc[pix*3] + gr * contrib 127 acc[pix*3+1] = acc[pix*3+1] + gg * contrib 128 acc[pix*3+2] = acc[pix*3+2] + gb * contrib 129 trans[pix] = tr - tr * alpha / 256 // *= (1-alpha) 130 } 131 } 132 px = px + 1 133 } 134 py = py + 1 135 } 136 oi = oi + 1 137 } 138 // --- resolve: colour = accum/256 + background * leftover transmittance --- 139 i = 0 140 while i < GW * GH { 141 let tr: i64 = trans[i] 142 var r: i64 = acc[i*3] / 256 + bgr * tr / 256 143 var gg2: i64 = acc[i*3+1] / 256 + bgg * tr / 256 144 var bb: i64 = acc[i*3+2] / 256 + bgb * tr / 256 145 if r > 255 { r = 255 } 146 if gg2 > 255 { gg2 = 255 } 147 if bb > 255 { bb = 255 } 148 fb[i] = r + gg2 * 256 + bb * K_MAGIC_65536 149 i = i + 1 150 } 151 return vis 152} 153 154// set one isotropic gaussian 155func gs_set(gauss: *i64, i: i64, x: i64, y: i64, z: i64, sc: i64, r: i64, g: i64, b: i64, op: i64) -> i64 { 156 gauss[i*8]=x; gauss[i*8+1]=y; gauss[i*8+2]=z; gauss[i*8+3]=sc; gauss[i*8+4]=r; gauss[i*8+5]=g; gauss[i*8+6]=b; gauss[i*8+7]=op 157 return 0 158} 159 160// ===== ★ANISOTROPIC SURFACE SPLATTING (the defining 3DGS feature: oriented ELLIPSE splats, not blobs) ===== 161// each gaussian = a DISK in the tangent plane (a "surfel") with normal n + in-plane radius rtan -> projects to a 162// 2D ELLIPSE (EWA-style). This aligns splats to the surface -> smooth coverage between mesh verts (kills the 163// isotropic-blob grid/dotty look). 12 i64/gaussian: {x,y,z, nx,ny,nz(fx256 unit), rtan, r,g,b, op, _spare}. 164func gs_set_aniso(gauss: *i64, i: i64, x: i64, y: i64, z: i64, nx: i64, ny: i64, nz: i64, rtan: i64, r: i64, g: i64, b: i64, op: i64) -> i64 { 165 gauss[i*12]=x; gauss[i*12+1]=y; gauss[i*12+2]=z; gauss[i*12+3]=nx; gauss[i*12+4]=ny; gauss[i*12+5]=nz 166 gauss[i*12+6]=rtan; gauss[i*12+7]=r; gauss[i*12+8]=g; gauss[i*12+9]=b; gauss[i*12+10]=op; gauss[i*12+11]=0 167 return 0 168} 169func gs_proj_x(x: i64, y: i64, z: i64, sy4: i64, cy4: i64, R: i64) -> i64 { 170 let cz: i64 = (x*sy4 + z*cy4)/K_MAGIC_4096 + R 171 if cz <= 64 { return 0 - K_MAGIC_999999 } 172 let cx: i64 = (x*cy4 - z*sy4)/K_MAGIC_4096 173 return GHW + GFOCAL*cx/cz 174} 175func gs_proj_y(x: i64, y: i64, z: i64, sy4: i64, cy4: i64, R: i64) -> i64 { 176 let cz: i64 = (x*sy4 + z*cy4)/K_MAGIC_4096 + R 177 if cz <= 64 { return 0 - K_MAGIC_999999 } 178 return GHH - GFOCAL*y/cz 179} 180// anisotropic render: pa/pb/pc/pdet = per-gaussian 2D-covariance scratch (ng each). Same sort + front-to-back 181// composite as the isotropic path; only the per-splat kernel is an oriented ellipse. 182func gs_render_aniso(gauss: *i64, ng: i64, yaw: i64, camz: i64, fb: *i64, acc: *i64, trans: *i64, depth: *i64, sxb: *i64, syb: *i64, pa: *i64, pb: *i64, pc: *i64, pdet: *i64, order: *i64, count: *i64, explut: *i64, bgr: i64, bgg: i64, bgb: i64) -> i64 { 183 let sy4: i64 = it_sin4096(yaw) 184 let cy4: i64 = it_cos4096(yaw) 185 let R: i64 = camz * GFX 186 var dmin: i64 = K_MAGIC_2000000000 187 var dmax: i64 = 0 - K_MAGIC_2000000000 188 var vis: i64 = 0 189 var i: i64 = 0 190 while i < ng { 191 let x: i64 = gauss[i*12] 192 let y: i64 = gauss[i*12+1] 193 let z: i64 = gauss[i*12+2] 194 let cz: i64 = (x*sy4 + z*cy4)/K_MAGIC_4096 + R 195 if cz > 64 { 196 let nx: i64 = gauss[i*12+3] 197 let ny: i64 = gauss[i*12+4] 198 let nz: i64 = gauss[i*12+5] 199 let rt: i64 = gauss[i*12+6] 200 // tangent basis from n (fx256): up = (0,256,0) unless n is near-vertical 201 var ux: i64 = 0 202 var uy: i64 = 256 203 var uz: i64 = 0 204 if ny > 210 { ux = 256; uy = 0; uz = 0 } 205 if ny < 0 - 210 { ux = 256; uy = 0; uz = 0 } 206 var t1x: i64 = (ny*uz - nz*uy)/256 207 var t1y: i64 = (nz*ux - nx*uz)/256 208 var t1z: i64 = (nx*uy - ny*ux)/256 209 let l1: i64 = gs_isqrt(t1x*t1x + t1y*t1y + t1z*t1z) 210 if l1 > 0 { t1x = t1x*256/l1; t1y = t1y*256/l1; t1z = t1z*256/l1 } 211 var t2x: i64 = (ny*t1z - nz*t1y)/256 212 var t2y: i64 = (nz*t1x - nx*t1z)/256 213 var t2z: i64 = (nx*t1y - ny*t1x)/256 214 let l2: i64 = gs_isqrt(t2x*t2x + t2y*t2y + t2z*t2z) 215 if l2 > 0 { t2x = t2x*256/l2; t2y = t2y*256/l2; t2z = t2z*256/l2 } 216 // world in-plane axes a1=t1*rt, a2=t2*rt (t* are fx256 -> /256) 217 let a1x: i64 = t1x*rt/256 218 let a1y: i64 = t1y*rt/256 219 let a1z: i64 = t1z*rt/256 220 let a2x: i64 = t2x*rt/256 221 let a2y: i64 = t2y*rt/256 222 let a2z: i64 = t2z*rt/256 223 let sx: i64 = gs_proj_x(x, y, z, sy4, cy4, R) 224 let sy: i64 = gs_proj_y(x, y, z, sy4, cy4, R) 225 // project the axis endpoints -> screen deltas e1,e2 (the Jacobian applied, computed exactly) 226 let e1x: i64 = gs_proj_x(x+a1x, y+a1y, z+a1z, sy4, cy4, R) - sx 227 let e1y: i64 = gs_proj_y(x+a1x, y+a1y, z+a1z, sy4, cy4, R) - sy 228 let e2x: i64 = gs_proj_x(x+a2x, y+a2y, z+a2z, sy4, cy4, R) - sx 229 let e2y: i64 = gs_proj_y(x+a2x, y+a2y, z+a2z, sy4, cy4, R) - sy 230 // 2D covariance Sigma2d = e1 e1^T + e2 e2^T + isotropic floor (min 1px, antialias) 231 var ca: i64 = e1x*e1x + e2x*e2x + 1 232 let cb: i64 = e1x*e1y + e2x*e2y 233 var cc: i64 = e1y*e1y + e2y*e2y + 1 234 var det: i64 = ca*cc - cb*cb 235 if det < 1 { det = 1 } 236 depth[i] = cz; sxb[i] = sx; syb[i] = sy; pa[i] = ca; pb[i] = cb; pc[i] = cc; pdet[i] = det 237 if cz < dmin { dmin = cz } 238 if cz > dmax { dmax = cz } 239 vis = vis + 1 240 } else { depth[i] = 0 - 1 } 241 i = i + 1 242 } 243 if vis == 0 { return 0 } 244 // counting sort by depth (ascending) 245 let span: i64 = dmax - dmin + 1 246 var b: i64 = 0 247 while b <= GNB { count[b] = 0; b = b + 1 } 248 i = 0 249 while i < ng { if depth[i] >= 0 { var bk: i64 = (depth[i]-dmin)*GNB/span; if bk<0{bk=0} if bk>=GNB{bk=GNB-1} count[bk]=count[bk]+1 } i=i+1 } 250 var pre: i64 = 0 251 b = 0 252 while b < GNB { let c: i64 = count[b]; count[b] = pre; pre = pre + c; b = b + 1 } 253 i = 0 254 while i < ng { if depth[i] >= 0 { var bk: i64 = (depth[i]-dmin)*GNB/span; if bk<0{bk=0} if bk>=GNB{bk=GNB-1} order[count[bk]]=i; count[bk]=count[bk]+1 } i=i+1 } 255 gs_clear(acc, trans) 256 var oi: i64 = 0 257 while oi < vis { 258 let g: i64 = order[oi] 259 let sx: i64 = sxb[g] 260 let sy: i64 = syb[g] 261 let ca: i64 = pa[g] 262 let cb: i64 = pb[g] 263 let cc: i64 = pc[g] 264 let det: i64 = pdet[g] 265 let rx: i64 = 3 * gs_isqrt(ca) // loose ellipse AABB 266 let ry: i64 = 3 * gs_isqrt(cc) 267 let gr: i64 = gauss[g*12+7] 268 let gg: i64 = gauss[g*12+8] 269 let gb: i64 = gauss[g*12+9] 270 let op: i64 = gauss[g*12+10] 271 var py: i64 = sy - ry 272 if py < 0 { py = 0 } 273 var pye: i64 = sy + ry 274 if pye >= GH { pye = GH - 1 } 275 while py <= pye { 276 let dy: i64 = py - sy 277 var px: i64 = sx - rx 278 if px < 0 { px = 0 } 279 var pxe: i64 = sx + rx 280 if pxe >= GW { pxe = GW - 1 } 281 while px <= pxe { 282 let dx: i64 = px - sx 283 // power = d^T Sigma2d^-1 d = (cc dx^2 - 2 cb dx dy + ca dy^2)/det ; LUT idx k = power*16 284 let pnum: i64 = cc*dx*dx - 2*cb*dx*dy + ca*dy*dy 285 let k: i64 = pnum * 16 / det 286 if k < GEXPN { if k >= 0 { 287 let pix: i64 = py*GW + px 288 let tr: i64 = trans[pix] 289 if tr > 1 { 290 let alpha: i64 = op * explut[k] / 256 291 let contrib: i64 = tr * alpha / 256 292 acc[pix*3] = acc[pix*3] + gr*contrib 293 acc[pix*3+1] = acc[pix*3+1] + gg*contrib 294 acc[pix*3+2] = acc[pix*3+2] + gb*contrib 295 trans[pix] = tr - tr*alpha/256 296 } 297 } } 298 px = px + 1 299 } 300 py = py + 1 301 } 302 oi = oi + 1 303 } 304 i = 0 305 while i < GW*GH { 306 let tr: i64 = trans[i] 307 var r: i64 = acc[i*3]/256 + bgr*tr/256 308 var g2: i64 = acc[i*3+1]/256 + bgg*tr/256 309 var bb: i64 = acc[i*3+2]/256 + bgb*tr/256 310 if r>255{r=255} if g2>255{g2=255} if bb>255{bb=255} 311 fb[i] = r + g2*256 + bb*K_MAGIC_65536 312 i = i + 1 313 } 314 return vis 315} 316// ===== ★GAUSSIAN-PARAMETER OPTIMIZER (the generator/"training" side = the real photoreal lift) ===== 317// v0 = DIFFERENTIABLE COLOUR fit: a splatted pixel is LINEAR in the Gaussian colours (pixel = Σ w_i·c_i), so the 318// image-loss gradient wrt each colour is ANALYTIC (dL/dc_i = Σ_px (render-target)·w_i). One gradient-descent 319// step: forward-render, accumulate the per-Gaussian colour gradient over its splat footprint, take a 320// least-squares-preconditioned step (÷ Σw²). Integer, deterministic. This is the 3DGS training mechanism (colour 321// channel); position/covariance optimization + adaptive density = the next rungs. Returns the L1 image loss. 322func gs_color_grad_step(gauss: *i64, ng: i64, target: *i64, yaw: i64, camz: i64, fb: *i64, acc: *i64, trans: *i64, depth: *i64, sxb: *i64, syb: *i64, sigb: *i64, order: *i64, count: *i64, explut: *i64, gradbuf: *i64, wnorm: *i64) -> i64 { 323 gs_render(gauss, ng, yaw, camz, fb, acc, trans, depth, sxb, syb, sigb, order, count, explut, 26, 28, 44) 324 // L1 loss + zero the grad accumulators 325 var loss: i64 = 0 326 var i: i64 = 0 327 while i < GW*GH { 328 var dr: i64 = (fb[i]&255) - (target[i]&255); if dr<0 {dr=0-dr} 329 var dg: i64 = ((fb[i]>>8)&255) - ((target[i]>>8)&255); if dg<0 {dg=0-dg} 330 var db: i64 = ((fb[i]>>16)&255) - ((target[i]>>16)&255); if db<0 {db=0-db} 331 loss = loss + dr + dg + db 332 i = i + 1 333 } 334 i = 0 335 while i < ng { gradbuf[i*3]=0; gradbuf[i*3+1]=0; gradbuf[i*3+2]=0; wnorm[i]=0; i=i+1 } 336 // accumulate per-gaussian colour gradient over its splat footprint (uses the projections gs_render just filled) 337 i = 0 338 while i < ng { 339 if depth[i] >= 0 { 340 let sx: i64 = sxb[i] 341 let sy: i64 = syb[i] 342 let sig: i64 = sigb[i] 343 let rad: i64 = 3*sig 344 let s2: i64 = sig*sig 345 let op: i64 = gauss[i*8+7] 346 var py: i64 = sy-rad 347 if py<0 {py=0} 348 var pye: i64 = sy+rad 349 if pye>=GH {pye=GH-1} 350 while py <= pye { 351 let dy: i64 = py-sy 352 var px: i64 = sx-rad 353 if px<0 {px=0} 354 var pxe: i64 = sx+rad 355 if pxe>=GW {pxe=GW-1} 356 while px <= pxe { 357 let dx: i64 = px-sx 358 let d2: i64 = dx*dx+dy*dy 359 let k: i64 = d2*16/s2 360 if k < GEXPN { 361 let pix: i64 = py*GW+px 362 let w: i64 = op*explut[k]/256 // this gaussian's splat weight (fx256) 363 gradbuf[i*3] = gradbuf[i*3] + ((fb[pix]&255) - (target[pix]&255)) * w 364 gradbuf[i*3+1] = gradbuf[i*3+1] + (((fb[pix]>>8)&255) - ((target[pix]>>8)&255)) * w 365 gradbuf[i*3+2] = gradbuf[i*3+2] + (((fb[pix]>>16)&255) - ((target[pix]>>16)&255)) * w 366 wnorm[i] = wnorm[i] + w*w/256 367 } 368 px = px + 1 369 } 370 py = py + 1 371 } 372 // least-squares-preconditioned step (lr 0.75): c -= 3*grad/(4*wnorm) 373 let nd: i64 = wnorm[i]*4 + 1 374 var nr: i64 = gauss[i*8+4] - gradbuf[i*3]*3/nd 375 var nge: i64 = gauss[i*8+5] - gradbuf[i*3+1]*3/nd 376 var nb: i64 = gauss[i*8+6] - gradbuf[i*3+2]*3/nd 377 if nr<0 {nr=0} if nr>255 {nr=255} 378 if nge<0 {nge=0} if nge>255 {nge=255} 379 if nb<0 {nb=0} if nb>255 {nb=255} 380 gauss[i*8+4]=nr; gauss[i*8+5]=nge; gauss[i*8+6]=nb 381 } 382 i = i + 1 383 } 384 return loss 385} 386 387func gs_w() -> i64 { return GW } 388func gs_h() -> i64 { return GH } 389func gs_expn() -> i64 { return GEXPN } 390func gs_nb() -> i64 { return GNB }