nx_vitpose_patchembed.nx source
↩ module page · 101 lines · 5836 B
1// nx_vitpose_patchembed.nx -- FIRST assembly rung of the faithful f32 ViTPose forward: load the REAL patch-embed
2// weights (projection.weight [768,3,16,16] + bias [768]) and position_embeddings [1,193,768] from the downloaded
3// safetensors (targeted seek+read), run the patch-embed Conv2d(3->768,k16,s16) on a 256x192x3 input, transpose the
4// [768,16,12] feature map into 192 tokens x 768, and add the position embeddings (index 1.. -> skip the cls slot).
5// Structural gate: token block is 192x768, finite (no NaN/Inf), non-trivial. (No PyTorch reference in-env -> the
6// assembly is built rung-by-rung with structural/property gates; final validation = keypoints on a real image.)
7// expect_exit: 0
8import "nx_syscalls.nx"
9import "nx_f32.nx"
10import "nx_f32_cvt.nx"
11import "nx_f32_conv2d.nx"
12import "nx_safetensors_load.nx"
13const K_MAGIC_4096: i64 = 4096
14const K_MAGIC_589824: i64 = 589824
15const K_MAGIC_148224: i64 = 148224
16const K_MAGIC_100000: i64 = 100000
17
18func w(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
19func wn(v: i64) -> i64 { var m: i64=v; if m<0{w("-" as *u8);m=0-m} let t:*u8=sys_mmap(24); 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; let o:*u8=sys_mmap(24); while i<k{o[i]=t[k-1-i];i=i+1} sys_write(1,o,k); return 0 }
20func rdfile(path: *u8, buf: *u8, cap: i64) -> i64 { let fd: i64=sys_openat_rd(path); if fd<0{return 0-1} var n: i64=sys_read(fd,buf,cap-1); sys_close(fd); if n<0{return 0-1} var go: i64=1; while go==1{go=0; if n>0{ let c: i64=buf[n-1] as i64; if c==10{n=n-1;go=1} if c==13{n=n-1;go=1} if c==32{n=n-1;go=1} }} buf[n]=0 as u8; return n }
21func u64le(buf: *u8, off: i64) -> i64 { var v: i64=0; var i: i64=0; while i<8 { v=v|((buf[off+i]&0xff)<<(i*8)); i=i+1 } return v }
22
23// load a named tensor via targeted seek+read into out (f32 bits). Returns element count or -1.
24func load_named(fd: i64, hbuf: *u8, hlen: i64, data_start: i64, name: *u8, out: *i64) -> i64 {
25 let dtype: *u8 = sys_mmap(16); let offs: *i64 = sys_mmap(16) as *i64
26 if stl_tensor(hbuf, hlen, name, dtype, offs) == 0 { return 0 - 1 }
27 let nbytes: i64 = offs[1] - offs[0]
28 let dbuf: *u8 = sys_mmap(nbytes + 16)
29 sys_lseek(fd, data_start + offs[0], 0)
30 var got: i64 = 0
31 while got < nbytes { let n: i64=sys_read(fd, ((dbuf as i64)+got) as *u8, nbytes-got); if n<=0{got=nbytes}else{got=got+n} }
32 let lo: *i64 = sys_mmap(16) as *i64; lo[0]=0; lo[1]=nbytes
33 return stl_read_f32(dbuf, 0, lo, dtype, out)
34}
35
36func main() -> i64 {
37 let path: *u8 = sys_mmap(K_MAGIC_4096)
38 if rdfile("data/mp_fetch_dest.txt" as *u8, path, K_MAGIC_4096) < 0 { w("no dest\n" as *u8); sys_exit(1); return 1 }
39 let fd: i64 = sys_openat_rd(path); if fd<0{ w("open fail\n" as *u8); sys_exit(1); return 1 }
40 let pre: *u8 = sys_mmap(16); sys_read(fd, pre, 8)
41 let hlen: i64 = u64le(pre, 0)
42 let hbuf: *u8 = sys_mmap(hlen + 16)
43 var i: i64 = 0; while i<8 { hbuf[i]=pre[i]; i=i+1 }
44 var got: i64 = 0; while got<hlen { let n: i64=sys_read(fd, ((hbuf as i64)+8+got) as *u8, hlen-got); if n<=0{got=hlen}else{got=got+n} }
45 let data_start: i64 = 8 + hlen
46
47 // load weights
48 let pw: *i64 = sys_mmap(8 * 768 * 3 * 16 * 16) as *i64 // projection.weight [768,3,16,16]
49 let pb: *i64 = sys_mmap(8 * 768) as *i64 // projection.bias [768]
50 let pe: *i64 = sys_mmap(8 * 193 * 768) as *i64 // position_embeddings [1,193,768]
51 let nw: i64 = load_named(fd, hbuf, hlen, data_start, "\"backbone.embeddings.patch_embeddings.projection.weight\":" as *u8, pw)
52 let nb: i64 = load_named(fd, hbuf, hlen, data_start, "\"backbone.embeddings.patch_embeddings.projection.bias\":" as *u8, pb)
53 let ne: i64 = load_named(fd, hbuf, hlen, data_start, "\"backbone.embeddings.position_embeddings\":" as *u8, pe)
54 sys_close(fd)
55 w("loaded pw=" as *u8); wn(nw); w(" pb=" as *u8); wn(nb); w(" pe=" as *u8); wn(ne); w("\n" as *u8)
56
57 // synthetic input image [1,3,256,192] (a small varied pattern in a sane normalized range)
58 let H: i64 = 256; let W: i64 = 192
59 let img: *i64 = sys_mmap(8 * 3 * H * W) as *i64
60 i = 0; while i < 3*H*W { img[i] = nx_i32_to_f32((i % 7) - 3); i = i + 1 }
61
62 // patch embed: Conv2d(3->768, k16, s16, pad0) -> [768, 16, 12]
63 let feat: *i64 = sys_mmap(8 * 768 * 16 * 12) as *i64
64 let rc: i64 = nx_f32_conv2d_forward(img, 1, 3, H, W, pw, 768, 16, 16, 16, 0, pb, feat)
65 w("patch-embed conv rc=" as *u8); wn(rc); w(" feat[768,16,12]\n" as *u8)
66
67 // transpose [768,16,12] -> tokens [192,768]; token p (=h*12+w) gets channel c from feat[c*192 + p]
68 let tok: *i64 = sys_mmap(8 * 192 * 768) as *i64
69 var p: i64 = 0
70 while p < 192 {
71 var c: i64 = 0
72 while c < 768 { tok[p*768 + c] = feat[c*192 + p]; c = c + 1 }
73 p = p + 1
74 }
75 // add position embeddings (skip cls slot: use pe index 1+p)
76 p = 0
77 while p < 192 {
78 var c: i64 = 0
79 while c < 768 { tok[p*768 + c] = nx_f32_add(tok[p*768 + c], pe[(1+p)*768 + c]); c = c + 1 }
80 p = p + 1
81 }
82
83 // structural checks: finite + non-trivial
84 var bad: i64 = 0; var nz: i64 = 0
85 i = 0
86 while i < 192*768 {
87 if nx_f32_is_nan(tok[i]) == 1 { bad = bad + 1 }
88 if nx_f32_is_inf(tok[i]) == 1 { bad = bad + 1 }
89 if nx_f32_is_zero(tok[i]) == 0 { nz = nz + 1 }
90 i = i + 1
91 }
92 w("tokens: NaN/Inf=" as *u8); wn(bad); w(" nonzero=" as *u8); wn(nz); w("/147456\n" as *u8)
93
94 if nw == K_MAGIC_589824 { if nb == 768 { if ne == K_MAGIC_148224 { if rc == 0 { if bad == 0 { if nz > K_MAGIC_100000 {
95 w("PATCHEMBED OK: real weights loaded, conv ran, 192x768 tokens finite + non-trivial\n" as *u8)
96 sys_exit(0); return 0
97 } } } } } }
98 w("PATCHEMBED FAIL\n" as *u8)
99 sys_exit(1)
100 return 1
101}