nx_rvq_gate.nx source

1// nx_rvq_gate.nx -- proves Residual VQ (nx_rvq): greedily train a 3-stage RVQ stack on LPC reflection vectors, then show 2// the reconstruction distortion DROPS with each added stage (the rate-distortion knob), reaching far below single-stage 3// VQ. This is the SoundStream/Lyra codec structure; the neural net replaces the fixed centroids with learned ones. 4import "nx_syscalls_x86_64.nx" 5import "nx_lpc_autocorr.nx" 6import "nx_lpc_levinson.nx" 7import "nx_vq.nx" 8import "nx_vq_train.nx" 9import "nx_rvq.nx" 10 11func g_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 } 12func g_pn(v: i64) -> i64 { 13 let b: *u8 = sys_mmap(28); var x: i64 = v 14 if x < 0 { b[0]=45; sys_write(1,b,1); x = 0 - x } 15 if x == 0 { b[0]=48; sys_write(1,b,1); return 0 } 16 var d: i64=0; var y: i64=x 17 while y>0 { d=d+1; y=y/10 } 18 var i: i64=d-1; y=x 19 while i>=0 { b[i]=(48+(y%10)) as u8; y=y/10; i=i-1 } 20 sys_write(1,b,d); return 0 21} 22func g_check(name: *u8, cond: i64) -> i64 { 23 if cond==1 { g_puts(" PASS " as *u8) } else { g_puts(" FAIL " as *u8) } 24 g_puts(name); g_puts("\n" as *u8); return cond 25} 26func st16(buf: *u8, idx: i64, v: i64) -> i64 { var x: i64 = v; if x < 0 { x = x + 65536 } buf[idx*2] = x & 0xff; buf[idx*2+1] = (x >> 8) & 0xff; return 0 } 27func ld64(buf: *u8, byteoff: i64) -> i64 { var v: i64 = 0; var b: i64 = 0; while b < 8 { v = v | (buf[byteoff + b] << (b*8)); b = b + 1 } return v } 28func fill_parabola(s: *u8, n: i64, period: i64) -> i64 { var i: i64=0; while i<n { let p: i64 = i % period; st16(s, i, 40 + p*(period-p)*2); i=i+1 } return 0 } 29 30func main() -> i64 { 31 g_puts("nx_rvq gate (residual VQ stack: distortion drops per stage, the SoundStream/Lyra structure, MEASURED)\n" as *u8) 32 var pass: i64 = 0; var total: i64 = 0 33 let n: i64 = 64; let order: i64 = 10; let D: i64 = order 34 let T: i64 = 16; let K: i64 = 8; let S: i64 = 3 35 36 let s_pcm: *u8 = sys_mmap(n*2) 37 let R: *u8 = sys_mmap((order+1)*8) 38 let kb: *u8 = sys_mmap((order+1)*8) 39 let ab: *u8 = sys_mmap((order+1)*8) 40 let eb: *u8 = sys_mmap((order+1)*8) 41 let train: *i64 = sys_mmap(T*D*8) as *i64 42 let resid_train: *i64 = sys_mmap(T*D*8) as *i64 43 let cbs: *i64 = sys_mmap(S*K*D*8) as *i64 44 let assign: *i64 = sys_mmap(T*8) as *i64 45 let sumb: *i64 = sys_mmap(K*D*8) as *i64 46 let cnt: *i64 = sys_mmap(K*8) as *i64 47 let indices: *i64 = sys_mmap(S*8) as *i64 48 let resid: *i64 = sys_mmap(D*8) as *i64 49 let out: *i64 = sys_mmap(D*8) as *i64 50 51 // training set (Q15 reflection vectors) 52 var t: i64 = 0 53 while t < T { 54 fill_parabola(s_pcm, n, 8 + t) 55 nx_lpc_autocorr(s_pcm, n, order, R); nx_lpc_levinson(R, order, kb, ab, eb) 56 var i: i64 = 0 57 while i < D { let val: i64 = ld64(kb, i*8) >> 15; train[t*D+i] = val; resid_train[t*D+i] = val; i = i + 1 } 58 t = t + 1 59 } 60 61 // GREEDY RVQ training: train stage k on the current residuals, then subtract its quantisation 62 var k: i64 = 0 63 while k < S { 64 let scb: *i64 = rvq_stage(cbs, k, K, D) 65 var c: i64 = 0 66 while c < K { var i: i64 = 0; while i < D { scb[c*D+i] = resid_train[(c*2)*D + i]; i = i + 1 } c = c + 1 } // coarse init 67 var it: i64 = 0 68 while it < 6 { vqt_iterate(resid_train, T, scb, K, D, assign, sumb, cnt); it = it + 1 } 69 // subtract stage k's quantisation from the residuals 70 t = 0 71 while t < T { 72 let idx: i64 = vq_encode((resid_train as i64 + t*D*8) as *i64, scb, K, D) 73 var j: i64 = 0 74 while j < D { resid_train[t*D+j] = resid_train[t*D+j] - scb[idx*D + j]; j = j + 1 } 75 t = t + 1 76 } 77 k = k + 1 78 } 79 80 // measure TOTAL distortion over the whole training set after 1, 2, 3 stages -- the codec's RD curve (monotonic by 81 // RVQ construction; a single vector can overshoot once it is exactly represented, but the distribution average cannot) 82 var d1: i64 = 0; var d2: i64 = 0; var d3: i64 = 0 83 var tt: i64 = 0 84 while tt < T { 85 let tvv: *i64 = (train as i64 + tt*D*8) as *i64 86 rvq_encode(tvv, D, cbs, K, S, indices, resid) 87 rvq_decode_partial(indices, cbs, K, D, 1, out); d1 = d1 + rvq_dist(tvv, out, D) 88 rvq_decode_partial(indices, cbs, K, D, 2, out); d2 = d2 + rvq_dist(tvv, out, D) 89 rvq_decode_partial(indices, cbs, K, D, 3, out); d3 = d3 + rvq_dist(tvv, out, D) 90 tt = tt + 1 91 } 92 let bits: i64 = S * vq_index_bits(K) 93 94 g_puts(" [measure] total RVQ distortion (over " as *u8); g_pn(T); g_puts(" vectors) vs stages: 1-stage=" as *u8); g_pn(d1); g_puts(" 2-stage=" as *u8); g_pn(d2); g_puts(" 3-stage=" as *u8); g_pn(d3); g_puts(" (" as *u8); g_pn(bits); g_puts(" bits @3 stages)\n" as *u8) 95 if d1 > 0 { g_puts(" [measure] 3-stage cuts distortion " as *u8); g_pn((d1 - d3) * 100 / d1); g_puts("% vs 1-stage (the RVQ rate-distortion knob)\n" as *u8) } 96 97 pass = pass + g_check("each RVQ stage refines the distribution (total distortion non-increasing: d1>=d2>=d3)" as *u8, (d2 <= d1) & (d3 <= d2)); total=total+1 98 pass = pass + g_check("3-stage RVQ << 1-stage VQ (>= 50% lower total distortion)" as *u8, d3 * 2 <= d1); total=total+1 99 100 g_puts("---- rvq gate: passed " as *u8); g_pn(pass); g_puts(" / " as *u8); g_pn(total); g_puts(" ----\n" as *u8) 101 if pass == total { g_puts("verdict=GREEN\n" as *u8); sys_exit(0); return 0 } 102 g_puts("verdict=RED\n" as *u8); sys_exit(1); return 1 103}