nx_dequant_iter_bench_test.nx source
↩ module page · 201 lines · 7019 B
1// nx_dequant_iter_bench_test.nx -- iter vs whole-block bench.
2//
3// Measures elapsed ns for:
4// path A: nx_gguf_dequant_q4_k (always 256 values)
5// path B: nx_q4k_iter_init + K * iter_value (K of 256)
6// for K in {1, 8, 32, 256}.
7//
8// Writes results to /tmp/nx_dequant_iter_bench.tsv for the smoke
9// wrapper to cat. Smoke exits 0 if all paths complete without
10// stale-INF saturation (sanity), regardless of measured numbers.
11
12import "nx_syscalls.nx"
13import "nx_tier.nx"
14import "nx_le.nx"
15import "nx_strconv.nx"
16import "nx_gguf.nx"
17import "nx_gguf_load.nx"
18import "nx_dequant_iter.nx"
19import "nx_dequant_iter_bench.nx"
20
21// Emit "label\tns_per_block\n" line.
22func _emit(fd: i64, label: *u8, label_len: nx_int,
23 k_per_block: i64, ns_per_block: i64) -> i64 {
24 let line: *u8 = sys_mmap(128)
25 var lo: i64 = 0
26 var i: nx_int = 0
27 while i < label_len { line[lo] = label[i]; lo = lo + 1; i = i + 1 }
28 line[lo] = 0x09; lo = lo + 1
29
30 let dec: *u8 = sys_mmap(32)
31 var n: i64 = nx_strconv_format_i64(k_per_block, dec)
32 var j: i64 = 0
33 while j < n { line[lo] = dec[j]; lo = lo + 1; j = j + 1 }
34 line[lo] = 0x09; lo = lo + 1
35
36 n = nx_strconv_format_i64(ns_per_block, dec)
37 j = 0
38 while j < n { line[lo] = dec[j]; lo = lo + 1; j = j + 1 }
39 line[lo] = 0x0A; lo = lo + 1
40
41 return sys_write(fd, line, lo)
42}
43
44// Build N pseudo-random Q4_K blocks at offset 0 in buf.
45func _build_blocks(buf: *u8, n: i64) -> nx_int {
46 let buf_bytes: i64 = n * NX_GL_Q4_K_BPB
47 var b: i64 = 0
48 while b < n {
49 let base: i64 = b * NX_GL_Q4_K_BPB
50 let pat: i64 = b - (b / 3) * 3
51 var d_raw: i64 = 0x3C00
52 if pat == 0 { d_raw = 0x3800 }
53 if pat == 2 { d_raw = 0x4000 }
54 nx_le_write_u16(buf, base + 0, d_raw)
55 nx_le_write_u16(buf, base + 2, 0x0000)
56 var i: i64 = 0
57 while i < 12 {
58 buf[base + 4 + i] = (b * 7 + i * 3) & 0xFF
59 i = i + 1
60 }
61 var jj: i64 = 0
62 while jj < 128 {
63 buf[base + 16 + jj] = (b * 11 + jj * 5 + 1) & 0xFF
64 jj = jj + 1
65 }
66 b = b + 1
67 }
68 return 0
69}
70
71// Run path A: whole-block dequant. Returns total ns.
72func _time_path_a(buf: *u8, n_blocks: i64, runs: i64,
73 out: *i64, n_values: i64) -> i64 {
74 let t0: *i64 = sys_mmap(16) as *i64
75 sys_clock_gettime_mono(t0)
76 let t0_s: i64 = t0[0]
77 let t0_n: i64 = t0[1]
78
79 var r: i64 = 0
80 while r < runs {
81 nx_gguf_dequant_q4_k(buf, 0, n_values, out)
82 r = r + 1
83 }
84
85 let t1: *i64 = sys_mmap(16) as *i64
86 sys_clock_gettime_mono(t1)
87 return (t1[0] - t0_s) * 1000000000 + (t1[1] - t0_n)
88}
89
90// Run path B: iter init + K iter_value reads per block. K in [1, 256].
91// Returns total ns.
92func _time_path_b(buf: *u8, n_blocks: i64, runs: i64, k: i64) -> i64 {
93 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
94 var dummy: i64 = 0 // prevents the compiler from eliding the work
95
96 let t0: *i64 = sys_mmap(16) as *i64
97 sys_clock_gettime_mono(t0)
98 let t0_s: i64 = t0[0]
99 let t0_n: i64 = t0[1]
100
101 var r: i64 = 0
102 while r < runs {
103 var b: i64 = 0
104 while b < n_blocks {
105 nx_q4k_iter_init(buf, b * NX_GL_Q4_K_BPB, it)
106 var v: i64 = 0
107 while v < k {
108 dummy = dummy + nx_q4k_iter_value(it, v)
109 v = v + 1
110 }
111 b = b + 1
112 }
113 r = r + 1
114 }
115
116 let t1: *i64 = sys_mmap(16) as *i64
117 sys_clock_gettime_mono(t1)
118 let elapsed: i64 = (t1[0] - t0_s) * 1000000000 + (t1[1] - t0_n)
119
120 // Touch dummy so the compiler can't elide the loop.
121 if dummy == 0x7FFFFFFFFFFFFFFE { return 0 - 1 }
122 return elapsed
123}
124
125func main() -> i64 {
126 let buf_bytes: i64 = NX_DIB_N_BLOCKS * NX_GL_Q4_K_BPB
127 let buf: *u8 = sys_mmap(buf_bytes + 64)
128 _build_blocks(buf, NX_DIB_N_BLOCKS)
129
130 let n_values: i64 = NX_DIB_N_BLOCKS * 256
131 let out: *i64 = sys_mmap(n_values * 8) as *i64
132
133 // ----- Warmup both paths -----
134 nx_gguf_dequant_q4_k(buf, 0, n_values, out)
135 let it_warm: *NxQ4KBlockIter = nx_q4k_iter_alloc()
136 nx_q4k_iter_init(buf, 0, it_warm)
137 let _ww: i64 = nx_q4k_iter_value(it_warm, 0)
138
139 // ----- Time path A -----
140 let ns_a: i64 = _time_path_a(buf, NX_DIB_N_BLOCKS, NX_DIB_K_RUNS, out, n_values)
141 if ns_a <= 0 { return 30 }
142 let ns_a_per_block: i64 = ns_a / (NX_DIB_N_BLOCKS * NX_DIB_K_RUNS)
143
144 // ----- Time path B for K in {1, 8, 32, 256} -----
145 let ns_b1: i64 = _time_path_b(buf, NX_DIB_N_BLOCKS, NX_DIB_K_RUNS, 1)
146 let ns_b8: i64 = _time_path_b(buf, NX_DIB_N_BLOCKS, NX_DIB_K_RUNS, 8)
147 let ns_b32: i64 = _time_path_b(buf, NX_DIB_N_BLOCKS, NX_DIB_K_RUNS, 32)
148 let ns_b256: i64 = _time_path_b(buf, NX_DIB_N_BLOCKS, NX_DIB_K_RUNS, 256)
149 if ns_b1 < 0 { return 40 }
150 if ns_b256 < 0 { return 41 }
151 let ns_b1_per: i64 = ns_b1 / (NX_DIB_N_BLOCKS * NX_DIB_K_RUNS)
152 let ns_b8_per: i64 = ns_b8 / (NX_DIB_N_BLOCKS * NX_DIB_K_RUNS)
153 let ns_b32_per: i64 = ns_b32 / (NX_DIB_N_BLOCKS * NX_DIB_K_RUNS)
154 let ns_b256_per: i64 = ns_b256 / (NX_DIB_N_BLOCKS * NX_DIB_K_RUNS)
155
156 // ----- Write TSV -----
157 let path: *u8 = sys_mmap(64)
158 path[0]=0x2F; path[1]=0x74; path[2]=0x6D; path[3]=0x70
159 path[4]=0x2F
160 path[5]=0x6E; path[6]=0x78; path[7]=0x5F
161 path[8]=0x64; path[9]=0x65; path[10]=0x71; path[11]=0x75
162 path[12]=0x61; path[13]=0x6E; path[14]=0x74; path[15]=0x5F
163 path[16]=0x69; path[17]=0x74; path[18]=0x65; path[19]=0x72
164 path[20]=0x5F; path[21]=0x62; path[22]=0x65; path[23]=0x6E
165 path[24]=0x63; path[25]=0x68; path[26]=0x2E; path[27]=0x74
166 path[28]=0x73; path[29]=0x76 // ".tsv"
167 path[30]=0
168
169 let fd: i64 = sys_openat_wr(path, 0x1A4)
170 if fd < 0 { return 80 }
171
172 // Header line: "path\tk_per_block\tns_per_block"
173 let hdr: *u8 = sys_mmap(64)
174 hdr[0]=0x70; hdr[1]=0x61; hdr[2]=0x74; hdr[3]=0x68 // "path"
175 hdr[4]=0x09
176 hdr[5]=0x6B; hdr[6]=0x5F; hdr[7]=0x70; hdr[8]=0x65 // "k_pe"
177 hdr[9]=0x72; hdr[10]=0x5F; hdr[11]=0x62; hdr[12]=0x6C
178 hdr[13]=0x6F; hdr[14]=0x63; hdr[15]=0x6B // "r_block"
179 hdr[16]=0x09
180 hdr[17]=0x6E; hdr[18]=0x73; hdr[19]=0x5F // "ns_"
181 hdr[20]=0x70; hdr[21]=0x65; hdr[22]=0x72; hdr[23]=0x5F
182 hdr[24]=0x62; hdr[25]=0x6C; hdr[26]=0x6F; hdr[27]=0x63
183 hdr[28]=0x6B // "per_block"
184 hdr[29]=0x0A
185 sys_write(fd, hdr, 30)
186
187 let lab_a: *u8 = sys_mmap(16)
188 lab_a[0]=0x77; lab_a[1]=0x68; lab_a[2]=0x6F; lab_a[3]=0x6C
189 lab_a[4]=0x65 // "whole"
190 _emit(fd, lab_a, 5, 256, ns_a_per_block)
191
192 let lab_b: *u8 = sys_mmap(16)
193 lab_b[0]=0x69; lab_b[1]=0x74; lab_b[2]=0x65; lab_b[3]=0x72 // "iter"
194 _emit(fd, lab_b, 4, 1, ns_b1_per)
195 _emit(fd, lab_b, 4, 8, ns_b8_per)
196 _emit(fd, lab_b, 4, 32, ns_b32_per)
197 _emit(fd, lab_b, 4, 256, ns_b256_per)
198
199 sys_close(fd)
200 return 0
201}