nx_vitpose_loadtest.nx source
↩ module page · 96 lines · 5411 B
1// nx_vitpose_loadtest.nx -- prove nx_safetensors_load reads REAL tensor values from the downloaded 343MB ViTPose
2// file, via TARGETED reads (header + seek to one tensor's bytes) so it never loads the whole body. Loads
3// head.conv.bias [17] and patch_embeddings.projection.bias [768], prints the count + first values as a sanity that
4// the assembly's weight-load path works on the real model. expect_exit: 0
5import "nx_syscalls.nx"
6import "nx_f32.nx"
7import "nx_f32_cvt.nx"
8import "nx_safetensors_load.nx"
9const K_MAGIC_10000: i64 = 10000
10const K_MAGIC_4096: i64 = 4096
11const K_MAGIC_1024: i64 = 1024
12
13func w(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
14func 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 }
15// print an f32 (i64-held bits) to ~4 decimals via a x10000 fixed-point path (sign + integer + 4 frac digits).
16func wf(x: i64) -> i64 {
17 let scaled: i64 = nx_f32_mul(x, nx_i32_to_f32(K_MAGIC_10000))
18 // round to nearest int via +/-0.5 then truncate through the exp; simplest: use nx_f32 compare loop is overkill.
19 // convert by repeated compare is heavy; instead print raw bits fallback is unreadable -> approximate: use sign+abs
20 let neg: i64 = nx_f32_lt(x, 0)
21 if neg == 1 { w("-" as *u8) }
22 var a: i64 = x
23 if neg == 1 { a = nx_f32_neg(x) }
24 let sc: i64 = nx_f32_mul(a, nx_i32_to_f32(K_MAGIC_10000))
25 // integer value of sc: subtract 1.0 repeatedly? too slow. Use exp-field decode:
26 let ip: i64 = f32_to_i64_trunc(sc)
27 wn(ip / K_MAGIC_10000); w("." as *u8)
28 var frac: i64 = ip - (ip/K_MAGIC_10000)*K_MAGIC_10000
29 if frac < 1000 { w("0" as *u8) }
30 if frac < 100 { w("0" as *u8) }
31 if frac < 10 { w("0" as *u8) }
32 wn(frac)
33 return 0
34}
35// truncate a NON-negative f32 to i64 (enough for our small magnitudes) via exp/mantissa decode.
36func f32_to_i64_trunc(x: i64) -> i64 {
37 if nx_f32_is_zero(x) == 1 { return 0 }
38 let exp: i64 = ((x >> 23) & 0xFF) - 127
39 if exp < 0 { return 0 }
40 let mant: i64 = (x & 0x7FFFFF) | 0x800000 // 24-bit mantissa with implicit 1
41 if exp >= 23 { return mant << (exp - 23) }
42 return mant >> (23 - exp)
43}
44
45func readfile(path: *u8, buf: *u8, cap: i64) -> i64 {
46 let fd: i64 = sys_openat_rd(path); if fd < 0 { return 0-1 }
47 var n: i64 = sys_read(fd, buf, cap-1); sys_close(fd); if n < 0 { return 0-1 }
48 var go: i64 = 1
49 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} } }
50 buf[n] = 0 as u8; return n
51}
52func 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 }
53
54// load ONE tensor by name from the file via targeted seek+read. Returns element count, fills out[].
55func load_one(fd: i64, hbuf: *u8, hlen: i64, data_start: i64, name: *u8, out: *i64) -> i64 {
56 let dtype: *u8 = sys_mmap(16)
57 let offs: *i64 = sys_mmap(16) as *i64
58 if stl_tensor(hbuf, hlen, name, dtype, offs) == 0 { return 0 - 1 }
59 let nbytes: i64 = offs[1] - offs[0]
60 let dbuf: *u8 = sys_mmap(nbytes + 16)
61 sys_lseek(fd, data_start + offs[0], 0)
62 var got: i64 = 0
63 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} }
64 let loffs: *i64 = sys_mmap(16) as *i64; loffs[0]=0; loffs[1]=nbytes
65 return stl_read_f32(dbuf, 0, loffs, dtype, out)
66}
67
68func main() -> i64 {
69 let path: *u8 = sys_mmap(K_MAGIC_4096)
70 if readfile("data/mp_fetch_dest.txt" as *u8, path, K_MAGIC_4096) < 0 { w("no dest\n" as *u8); sys_exit(1); return 1 }
71 let fd: i64 = sys_openat_rd(path); if fd < 0 { w("cannot open model\n" as *u8); sys_exit(1); return 1 }
72 let pre: *u8 = sys_mmap(16); sys_read(fd, pre, 8)
73 let hlen: i64 = u64le(pre, 0)
74 let hbuf: *u8 = sys_mmap(hlen + 16)
75 // hbuf must mirror the file layout for stl_tensor (searches [8..8+hlen)); put the 8-byte prefix then header.
76 var i: i64 = 0; while i < 8 { hbuf[i] = pre[i]; i = i + 1 }
77 var got: i64 = 0
78 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} }
79 let data_start: i64 = 8 + hlen
80 w("header_len=" as *u8); wn(hlen); w(" data_start=" as *u8); wn(data_start); w("\n" as *u8)
81
82 let out: *i64 = sys_mmap(8 * K_MAGIC_1024) as *i64
83 let n1: i64 = load_one(fd, hbuf, hlen, data_start, "\"head.conv.bias\":" as *u8, out)
84 w("head.conv.bias count=" as *u8); wn(n1); w(" vals[0..4]= " as *u8)
85 var j: i64 = 0; while j < 5 { if j < n1 { wf(out[j]); w(" " as *u8) } j = j + 1 } w("\n" as *u8)
86
87 let n2: i64 = load_one(fd, hbuf, hlen, data_start, "\"backbone.embeddings.patch_embeddings.projection.bias\":" as *u8, out)
88 w("patch_embed.projection.bias count=" as *u8); wn(n2); w(" vals[0..4]= " as *u8)
89 j = 0; while j < 5 { if j < n2 { wf(out[j]); w(" " as *u8) } j = j + 1 } w("\n" as *u8)
90 sys_close(fd)
91
92 if n1 == 17 { if n2 == 768 { w("LOADTEST OK: real ViTPose tensors read (head.conv.bias=17, patch bias=768)\n" as *u8); sys_exit(0); return 0 } }
93 w("LOADTEST FAIL: unexpected counts\n" as *u8)
94 sys_exit(1)
95 return 1
96}