nx_safetensors_load_gate.nx source
↩ module page · 75 lines · 4183 B
1// nx_safetensors_load_gate.nx -- proof of the safetensors->f32 loader on a synthetic in-memory blob (header JSON +
2// F32 and BF16 tensors), verifying header/offset parse + exact value reads. expect_exit: 0
3import "nx_syscalls.nx"
4import "nx_f32_cvt.nx"
5import "nx_safetensors_load.nx"
6
7func gp(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } return sys_write(1, s, n) }
8func gn(v: i64) -> i64 {
9 let bb: *u8 = sys_mmap(28); var m: i64 = v
10 if m < 0 { sys_write(1, "-" as *u8, 1); m = 0 - m }
11 let t: *u8 = sys_mmap(28); var k: i64 = 0
12 if m == 0 { t[0] = 48 as u8; k = 1 }
13 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
14 var i: i64 = 0; while i < k { bb[i] = t[k - 1 - i]; i = i + 1 }
15 return sys_write(1, bb, k)
16}
17func put_u64le(buf: *u8, off: i64, v: i64) -> i64 { var i: i64=0; while i<8 { buf[off+i]=((v>>(i*8))&0xff) as u8; i=i+1 } return 0 }
18func put_u32le(buf: *u8, off: i64, v: i64) -> i64 { var i: i64=0; while i<4 { buf[off+i]=((v>>(i*8))&0xff) as u8; i=i+1 } return 0 }
19func put_str(buf: *u8, off: i64, s: *u8) -> i64 { var i: i64=0; while s[i]!=(0 as u8){ buf[off+i]=s[i]; i=i+1 } return off+i }
20
21func main(argc: i64, argv: *i64) -> i64 {
22 var pass: i64 = 0
23 let blob: *u8 = sys_mmap(512)
24 // header JSON: w F32[2]@[0,8], b F32[1]@[8,12], h BF16[1]@[12,14]
25 let jstr: *u8 = "{\"w\":{\"dtype\":\"F32\",\"shape\":[2],\"data_offsets\":[0,8]},\"b\":{\"dtype\":\"F32\",\"shape\":[1],\"data_offsets\":[8,12]},\"h\":{\"dtype\":\"BF16\",\"shape\":[1],\"data_offsets\":[12,14]}}" as *u8
26 let jend: i64 = put_str(blob, 8, jstr)
27 let jlen: i64 = jend - 8
28 put_u64le(blob, 0, jlen)
29 let ds: i64 = 8 + jlen
30 put_u32le(blob, ds + 0, nx_i32_to_f32(1)) // w[0]=1.0
31 put_u32le(blob, ds + 4, nx_i32_to_f32(2)) // w[1]=2.0
32 put_u32le(blob, ds + 8, nx_i32_to_f32(3)) // b[0]=3.0
33 blob[ds + 12] = 0x80 as u8; blob[ds + 13] = 0x40 as u8 // h[0] = bf16(4.0) = 0x4080
34
35 // S1 header/data-start parse
36 let hlen: i64 = stl_header_len(blob)
37 let dstart: i64 = stl_data_start(blob)
38 if hlen == jlen { if dstart == ds { pass = pass + 1; gp("S1 header_len + data_start parse OK\n" as *u8) } }
39 if hlen != jlen { gp("S1 FAIL hlen=" as *u8); gn(hlen); gp(" want=" as *u8); gn(jlen); gp("\n" as *u8) }
40
41 let dtype: *u8 = sys_mmap(16)
42 let offs: *i64 = sys_mmap(16) as *i64
43 let out: *i64 = sys_mmap(8 * 4) as *i64
44
45 // S2 lookup w -> F32, offs [0,8]
46 let f2: i64 = stl_tensor(blob, hlen, "\"w\":" as *u8, dtype, offs)
47 if f2 == 1 { if stl_streq(dtype, "F32" as *u8) == 1 { if offs[0]==0 { if offs[1]==8 {
48 pass = pass + 1; gp("S2 tensor w -> F32 offs[0,8] OK\n" as *u8)
49 } } } }
50 if f2 != 1 { gp("S2 FAIL found=" as *u8); gn(f2); gp(" dtype-offs0=" as *u8); gn(offs[0]); gp("\n" as *u8) }
51
52 // S3 read w -> [1.0, 2.0]
53 let n3: i64 = stl_read_f32(blob, dstart, offs, dtype, out)
54 if n3 == 2 { if out[0]==nx_i32_to_f32(1) { if out[1]==nx_i32_to_f32(2) {
55 pass = pass + 1; gp("S3 read w = [1.0,2.0] exact OK\n" as *u8)
56 } } }
57 if n3 != 2 { gp("S3 FAIL n=" as *u8); gn(n3); gp(" o0=" as *u8); gn(out[0]); gp("\n" as *u8) }
58
59 // S4 lookup+read b -> [3.0]
60 stl_tensor(blob, hlen, "\"b\":" as *u8, dtype, offs)
61 let n4: i64 = stl_read_f32(blob, dstart, offs, dtype, out)
62 if n4 == 1 { if out[0]==nx_i32_to_f32(3) { pass = pass + 1; gp("S4 read b = [3.0] exact OK\n" as *u8) } }
63 if n4 != 1 { gp("S4 FAIL n=" as *u8); gn(n4); gp(" o0=" as *u8); gn(out[0]); gp("\n" as *u8) }
64
65 // S5 BF16 h -> 4.0 (0x40800000)
66 stl_tensor(blob, hlen, "\"h\":" as *u8, dtype, offs)
67 let n5: i64 = stl_read_f32(blob, dstart, offs, dtype, out)
68 if n5 == 1 { if out[0]==nx_i32_to_f32(4) { pass = pass + 1; gp("S5 BF16 read h = 4.0 OK\n" as *u8) } }
69 if n5 != 1 { gp("S5 FAIL n=" as *u8); gn(n5); gp(" o0=" as *u8); gn(out[0]); gp(" want=" as *u8); gn(nx_i32_to_f32(4)); gp("\n" as *u8) }
70
71 gp("SAFETENSORS-LOAD-GATE pass=" as *u8); gn(pass); gp("/5\n" as *u8)
72 if pass == 5 { gp("SAFETENSORS-LOAD-GATE GREEN 5/5 (header + tensor lookup + F32/BF16 value reads)\n" as *u8); sys_exit(0) }
73 sys_exit(1)
74 return 0
75}