code wiki / _hdl_build / nx_vcodec_heataq.nx
`nx_vcodec_heataq.nx` source

↩ module page · 235 lines · 13557 B
1// nx_vcodec_heataq.nx -- Adapts skip thresholds per macroblock based on past behavior to optimize bitrate allocation.
2import "nx_gate_base.nx"
3// nx_vcodec_heataq.nx -- GEN-2 P2 PROBE: past-propagating MB-heat adaptive skip-threshold (the RTC-safe cousin
4// of x264 mbtree, which needs future-lookahead=latency). Per luma MB, track PAST behavior across frames:
5//   heat  = consecutive-SKIP streak (stability: this content persists -> protect it from drift)
6//   churn = consecutive-CODED streak (transience: these bits die next frame -> spend less)
7// and modulate the SKIP THRESHOLD per MB (encoder-side only, decoder-transparent -- the wire carries a frame
8// qp and legal syntax either way). Variants measured MSU-style (BD-rate vs the SAME oracle points as the ship):
9//   argv[1]=0  FIDELITY: modulation OFF -> mirror must be BIT-IDENTICAL to vv_enc_rct8 (wire + recon) -- the
10//              trust gate for the mirror (pattern: rdskip/i16/cdef probes)
11//   argv[1]=1  CHURN arm:  churn>=3  -> thresh*2 (skip transient content earlier)
12//   argv[1]=2  STABLE arm: heat>=8   -> thresh/2 (protect persistent content from drift)
13//   argv[1]=3  BOTH arms
14// Emits OURS/OURSINTRA lines on /tmp/seq.y4m (48f CIF, emode 681 = the shipped RTC config) so the existing
15// msu_ours_splice.sh + msu_bdrate.sh tooling scores it unchanged. license: ORIGINAL
16import "nx_syscalls.nx"
17import "nx_video_codec_wasm.nx"
18const CHURN_MAGIC_2048: i64 = 2048
19const CHURN_MAGIC_4194304: i64 = 4194304
20const CHURN_MAGIC_5120: i64 = 5120
21const CHURN_MAGIC_8192: i64 = 8192
22
23const SW: i64 = 352
24const SH: i64 = 288
25const NF: i64 = 48
26const SKIPMUL: i64 = 30
27const CHURN_N: i64 = 3     // coded-streak arming the churn discount
28const STABLE_N: i64 = 8    // skip-streak arming the stability protection
29const SKIPBITS: i64 = 4    // CAVLC bits <= this across a block call ~= the MB skipped (mode+coeffs ride the rc section)
30
31func grow(name: *u8, ok: i64) -> i64 { if ok==1 { gw("  PASS " as *u8) } else { gw("  FAIL " as *u8) } gw(name); gw("
32" as *u8); return ok }
33func gn(v: i64) -> i64 {
  let b: *u8=sys_mmap(28); var m: i64=v; if m<0{sys_write(1,"-" as *u8,1);m=0-m}
  let t: *u8=sys_mmap(28); var k: i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1}
  var i: i64=0; while i<k{b[i]=t[k-1-i];i=i+1} sys_write(1,b,k); return 0 }
37func cpb(d: *u8, s: *u8, n: i64) -> i64 { var i: i64=0; while i<n { d[i]=s[i]; i=i+1 } return 0 }
38
39func frames(yuv: *u8, hi: i64, total: i64, fp: *i64, fdata: i64) -> i64 {
  var nf: i64=0; var p: i64=hi
  while nf < NF { if p>=total { break }
      var q: i64=p; while q<total { if (yuv[q]&0xff)==10 { break } q=q+1 } q=q+1
      if q+fdata>total { break } fp[nf]=(yuv as i64)+q; nf=nf+1; p=q+fdata }
  return nf }
45
46// per-MB threshold: the P2 policy under test (mode 0 = OFF = base, the fidelity arm).
47// mode packs sweep params (no new args -- 16-arg-cap respect): arm = mode&15, stable-N override =
48// (mode>>4)&255 (0 -> STABLE_N), divisor override = (mode>>12)&15 (0 -> 2).
49func hq_thr(base: i64, heat: i64, churn: i64, mode: i64) -> i64 {
  let arm: i64 = mode & 15
  if arm == 0 { return base }
  var sn: i64 = (mode >> 4) & 255
  if sn == 0 { sn = STABLE_N }
  var dv: i64 = (mode >> 12) & 15
  if dv == 0 { dv = 2 }
  if arm == 1 { if churn >= CHURN_N { return base * 2 } return base }
  if arm == 2 { if heat >= sn { return base / dv } return base }
  if churn >= CHURN_N { return base * 2 }
  if heat >= sn { return base / dv }
  return base
61}
62
63// MIRROR of vc_enc_frame_packed_rct8 (nx_vcodec.nx:1660) with ONE change: per-MB sad_thresh via hq_thr +
64// per-MB CAVLC-bit-delta skip flags out. Everything else byte-identical -- proven by the mode-0 fidelity arm.
65func hq_enc_frame(cur: *u8, prev: *u8, recon: *u8, W: i64, H: i64, qp: i64, keyframe: i64, sad_thresh: i64, buf: *u8, blk: *i64, mv: *i64, tf: i64, rctx: *i64, heat: *i64, churn: *i64, skipf: *i64, mode: i64) -> i64 {
  rctx[0] = rctx[0] | 8
  if (rctx[0] & 16) != 0 { rctx[0] = (rctx[0] | 4) - 4 }
  let est: *i64 = rctx[1] as *i64; let probs: *i64 = rctx[2] as *i64
  rc_enc_init(est)
  var ci: i64 = 0; while ci < RC_NCTX8 { probs[ci] = CHURN_MAGIC_2048; ci = ci + 1 }
  var bp: i64 = nx_bw_put(buf, 16, keyframe, 1)
  bp = nx_bw_put(buf, bp, 0, 1)
  let BW: i64 = W / 16; let BH: i64 = H / 16
  var by: i64 = 0
  while by < BH { var bx: i64 = 0
      while bx < BW {
          let lv: i64 = vc_aq_level(vc_mb_mad(cur, W, bx*16, by*16))
          bp = nx_bw_put(buf, bp, lv, 2)
          var t8: i64 = 0
          if (rctx[0] & 16) == 0 { if keyframe == 1 {
              if lv <= 1 { t8 = vc_mb_t8_ok(cur, W, bx*16, by*16) }
          } }
          rctx[5] = bp
          bp = nx_bw_put(buf, bp, t8, 1)
          let mi: i64 = by * BW + bx
          var th: i64 = sad_thresh
          if keyframe == 0 { th = hq_thr(sad_thresh, heat[mi], churn[mi], mode) }
          let bp0: i64 = bp
          bp = vc_enc_block_packed_e(cur, prev, recon, W, H, bx, by, vc_aq_qp(qp, lv), keyframe, th, buf, bp, blk, mv, tf + (t8 << 1), rctx)
          var sk: i64 = 0
          if keyframe == 0 { if bp - bp0 <= SKIPBITS { sk = 1 } }
          skipf[mi] = sk
          bx = bx + 1 } by = by + 1 }
  let rbytes: i64 = rc_enc_flush(est, rctx[3] as *u8)
  let cavlc_end: i64 = (bp + 7) / 8
  buf[0] = (cavlc_end & 0xff) as u8; buf[1] = ((cavlc_end >> 8) & 0xff) as u8
  let rcbuf: *u8 = rctx[3] as *u8; var i: i64 = 0
  while i < rbytes { buf[cavlc_end + i] = rcbuf[i]; i = i + 1 }
  return (cavlc_end + rbytes) * 8
100}
101
102// MIRROR of vv_enc_rct8 (nx_video_codec_wasm.nx:422): luma via hq_enc_frame (the hook), chroma via the REAL
103// frame fn, same headers + deblock + keyframe nf. After luma, fold skip flags into the heat/churn planes.
104func hq_enc_rct8(cur: *u8, prev: *u8, recon: *u8, W: i64, H: i64, qp: i64, keyframe: i64, sad: i64, out: *u8, out_cap: i64, blk: *i64, mv: *i64, rctx: *i64, heat: *i64, churn: *i64, skipf: *i64, mode: i64) -> i64 {
  let N: i64 = W * H
  let C: i64 = (W / 2) * (H / 2)
  rctx[0] = (rctx[0] | 16) - 16
  let yb: i64 = (hq_enc_frame(cur, prev, recon, W, H, qp, keyframe, sad, ((out as i64) + 12) as *u8, blk, mv, 1, rctx, heat, churn, skipf, mode) + 7) / 8
  if 12 + yb > out_cap { return 0 - 1 }
  // fold this frame's skip flags into the streak counters (used by the NEXT frame = past-propagation)
  let NMB: i64 = (W / 16) * (H / 16)
  var mi: i64 = 0
  while mi < NMB {
      if keyframe == 1 { heat[mi] = 0; churn[mi] = 0 } else {
          if skipf[mi] == 1 { heat[mi] = heat[mi] + 1; churn[mi] = 0 } else { churn[mi] = churn[mi] + 1; heat[mi] = 0 }
      }
      mi = mi + 1 }
  rctx[0] = rctx[0] | 16
  let ub: i64 = (vc_enc_frame_packed_rct8(((cur as i64) + N) as *u8, ((prev as i64) + N) as *u8, ((recon as i64) + N) as *u8, W / 2, H / 2, qp, keyframe, sad, ((out as i64) + 12 + yb) as *u8, blk, mv, 1, rctx) + 7) / 8
  if 12 + yb + ub > out_cap { return 0 - 1 }
  let vb: i64 = (vc_enc_frame_packed_rct8(((cur as i64) + N + C) as *u8, ((prev as i64) + N + C) as *u8, ((recon as i64) + N + C) as *u8, W / 2, H / 2, qp, keyframe, sad, ((out as i64) + 12 + yb + ub) as *u8, blk, mv, 1, rctx) + 7) / 8
  if 12 + yb + ub + vb > out_cap { return 0 - 1 }
  vv_wr_u32(out, 0, yb)
  vv_wr_u32(out, 4, ub)
  vv_wr_u32(out, 8, vb)
  vv_deblock(recon, W, H, qp, (rctx[0] >> 9) & 1)
  var tb2: i64 = 0
  if keyframe == 1 { if 12 + yb + ub + vb + (H/16 + 7)/8 <= out_cap {
      tb2 = vv_nf_enc(cur, recon, W, H, qp, ((out as i64) + 12) as *u8, 17, ((out as i64) + 12 + yb + ub + vb) as *u8, rctx)
  } }
  return 12 + yb + ub + vb + tb2
132}
133
134// FIDELITY (mode 0): mirror-with-modulation-OFF must equal vv_enc_rct8 EXACTLY -- wire bytes AND recon.
135func fidelity(fp: *i64, sz: i64) -> i64 {
  let prevA: *u8=sys_mmap(sz+64); let reconA: *u8=sys_mmap(sz+64); let wireA: *u8=sys_mmap(CHURN_MAGIC_4194304)
  let prevB: *u8=sys_mmap(sz+64); let reconB: *u8=sys_mmap(sz+64); let wireB: *u8=sys_mmap(CHURN_MAGIC_4194304)
  let blk: *i64=sys_mmap(512) as *i64; let mv: *i64=sys_mmap(128) as *i64
  let est: *i64=sys_mmap(64) as *i64; let probs: *i64=sys_mmap(32*8) as *i64; let rcbuf: *u8=sys_mmap(CHURN_MAGIC_4194304)
  let t8c: *i64=sys_mmap(CHURN_MAGIC_5120) as *i64; let rctx: *i64=sys_mmap(64) as *i64
  let heat: *i64=sys_mmap(CHURN_MAGIC_8192) as *i64; let churn: *i64=sys_mmap(CHURN_MAGIC_8192) as *i64; let skipf: *i64=sys_mmap(CHURN_MAGIC_8192) as *i64
  vc_t8_init(t8c)
  rctx[1]=est as i64; rctx[2]=probs as i64; rctx[3]=rcbuf as i64; rctx[4]=t8c as i64; rctx[5]=0; rctx[6]=0; rctx[7]=0
  var z: i64=0; while z<sz { prevA[z]=0 as u8; prevB[z]=0 as u8; z=z+1 }
  let qp: i64 = 20
  var f: i64 = 0
  var bad: i64 = 0
  while f < 12 {
      let cur: *u8 = fp[f] as *u8
      var key: i64 = 0; if f == 0 { key = 1 }
      rctx[0] = 681
      let na: i64 = vv_enc_rct8(cur, prevA, reconA, SW, SH, qp, key, qp*SKIPMUL, wireA, CHURN_MAGIC_4194304, blk, mv, rctx)
      rctx[0] = 681
      let nb: i64 = hq_enc_rct8(cur, prevB, reconB, SW, SH, qp, key, qp*SKIPMUL, wireB, CHURN_MAGIC_4194304, blk, mv, rctx, heat, churn, skipf, 0)
      if na != nb { bad = bad + 1 }
      var i: i64 = 0
      if na == nb { while i < na { if wireA[i] != wireB[i] { bad = bad + 1; break } i = i + 1 } }
      i = 0
      while i < sz { if reconA[i] != reconB[i] { bad = bad + 1; break } i = i + 1 }
      cpb(prevA, reconA, sz); cpb(prevB, reconB, sz)
      f = f + 1
  }
  if bad == 0 { gw("HEATAQ FIDELITY: GREEN (mirror == vv_enc_rct8 bit-exact, wire+recon, 12 frames)\n" as *u8); return 0 }
  gw("HEATAQ FIDELITY: RED bad=" as *u8); gn(bad); gw("\n" as *u8)
  return 1
166}
167
168func runpq(fp: *i64, kq: i64, pq: i64, emode: i64, sz: i64, allkey: i64, mode: i64) -> i64 {
  let prevE: *u8=sys_mmap(sz+64); let reconE: *u8=sys_mmap(sz+64)
  let wire: *u8=sys_mmap(CHURN_MAGIC_4194304); let blk: *i64=sys_mmap(512) as *i64; let mv: *i64=sys_mmap(128) as *i64
  let est: *i64=sys_mmap(64) as *i64; let probs: *i64=sys_mmap(32*8) as *i64; let rcbuf: *u8=sys_mmap(CHURN_MAGIC_4194304)
  let t8c: *i64=sys_mmap(CHURN_MAGIC_5120) as *i64; let rctx: *i64=sys_mmap(64) as *i64
  let heat: *i64=sys_mmap(CHURN_MAGIC_8192) as *i64; let churn: *i64=sys_mmap(CHURN_MAGIC_8192) as *i64; let skipf: *i64=sys_mmap(CHURN_MAGIC_8192) as *i64
  vc_t8_init(t8c)
  var z: i64=0; while z<sz { prevE[z]=0 as u8; z=z+1 }
  rctx[1]=est as i64; rctx[2]=probs as i64; rctx[3]=rcbuf as i64; rctx[4]=t8c as i64; rctx[5]=0; rctx[6]=0; rctx[7]=0
  var keyB: i64=0; var pB: i64=0; var sseY: i64=0; var keySSE: i64=0; var npix: i64=0
  let NL: i64=SW*SH
  var f: i64=0
  while f < NF {
      let cur: *u8=fp[f] as *u8
      var key: i64=0; if f==0 { key=1 } if allkey==1 { key=1 }
      var fq: i64 = kq; if key==0 { fq = pq }
      rctx[0]=emode
      let nb: i64 = hq_enc_rct8(cur, prevE, reconE, SW, SH, fq, key, fq*SKIPMUL, wire, CHURN_MAGIC_4194304, blk, mv, rctx, heat, churn, skipf, mode)
      if nb<=0 { gw("enc fail\n" as *u8); return 1 }
      if f==0 { keyB=nb } else { pB=pB+nb }
      var pxi: i64=0; var fsse: i64=0
      while pxi<NL { let d: i64=(reconE[pxi]&0xff)-(cur[pxi]&0xff); fsse=fsse+d*d; pxi=pxi+1 }
      sseY=sseY+fsse; if f==0 { keySSE=fsse }
      npix=npix+NL
      cpb(prevE, reconE, sz)
      f=f+1
  }
  if allkey==1 { gw("OURSINTRA kq=" as *u8) } else { gw("OURS kq=" as *u8) }
  gn(kq); gw(" pq=" as *u8); gn(pq); gw(" total=" as *u8); gn(keyB+pB); gw(" key=" as *u8); gn(keyB)
  gw(" ptot=" as *u8); gn(pB); gw(" sseY=" as *u8); gn(sseY); gw(" keySSE=" as *u8); gn(keySSE)
  gw(" npix=" as *u8); gn(npix); gw("\n" as *u8)
  return 0 }
200
201func main(argc: i64, argv: *i64) -> i64 {
  var mode: i64 = 0
  if argc >= 2 { let a: *u8 = argv[1] as *u8; mode = (a[0] & 0xff) - 48 }
  if mode < 0 { mode = 0 }
  if mode > 3 { mode = 0 }
  // optional sweep params: argv[2] = stable-N (1-2 digits), argv[3] = divisor (1 digit) -> packed into mode
  if argc >= 3 { let s: *u8 = argv[2] as *u8
      var sn: i64 = (s[0] & 0xff) - 48
      if s[1] != (0 as u8) { sn = sn * 10 + ((s[1] & 0xff) - 48) }
      if sn > 0 { if sn < 256 { mode = mode | (sn << 4) } } }
  if argc >= 4 { let d: *u8 = argv[3] as *u8
      let dv: i64 = (d[0] & 0xff) - 48
      if dv > 0 { if dv < 16 { mode = mode | (dv << 12) } } }
  let box: *i64 = sys_mmap(16) as *i64
  let yuv: *u8 = sys_read_file("/tmp/seq.y4m" as *u8, box)
  if (yuv as i64)==0 { gw("cannot read /tmp/seq.y4m -> RED\n" as *u8); return 1 }
  let total: i64=box[0]
  var hi: i64=0; while hi<total { if (yuv[hi]&0xff)==10 { break } hi=hi+1 } hi=hi+1
  let C2: i64=(SW/2)*(SH/2); let fdata: i64=SW*SH+2*C2; let sz: i64=fdata
  let fp: *i64=sys_mmap(128*8) as *i64
  let nf: i64=frames(yuv, hi, total, fp, fdata)
  if nf < NF { gw("too few frames -> RED\n" as *u8); return 1 }
  if mode == 0 { return fidelity(fp, sz) }
  runpq(fp, 8, 8, 681, sz, 0, mode)
  runpq(fp, 16, 16, 681, sz, 0, mode)
  runpq(fp, 24, 24, 681, sz, 0, mode)
  runpq(fp, 32, 32, 681, sz, 0, mode)
  runpq(fp, 44, 44, 681, sz, 0, mode)
  runpq(fp, 8, 8, 681, sz, 1, mode)
  runpq(fp, 16, 16, 681, sz, 1, mode)
  runpq(fp, 24, 24, 681, sz, 1, mode)
  runpq(fp, 32, 32, 681, sz, 1, mode)
  runpq(fp, 44, 44, 681, sz, 1, mode)
  gw("HEATAQ: DONE mode=" as *u8); gn(mode); gw("\n" as *u8)
  return 0 }
nx_vcodec_heataq.nx source

`nx_vcodec_heataq.nx` source
