lib/accy/src/profiling/versus/baseline/xla.py
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 """XLA (JAX) side of the accy.versus/v2 battle harness.
2
3 Emits the shared JSONL schema on stdout. Timing mirrors the accy runner:
4 host wall clock around dispatch + block_until_ready, median over samples.
5
6 Run inside a venv with jax[cuda] installed:
7 python xla.py [workload-name] > xla.jsonl
8 """
9
10 import json
11 import sys
12 import time
13
14 import jax
15 import jaxlib
16 import jax.numpy as jnp
17 import numpy as np
18
19 EWCHAIN_LINKS = 4
20 SCHEMA = "accy.versus/v2"
21 ORACLE_CELLS = 256
22
23 SAMPLES = 20
24 DISPATCH_SAMPLES = 50
25 PIPELINE_ROUNDS = 5
26 PIPELINE_BATCH = 16
27 WARMUP = 3
28
29 BATTERY = [
30 {"name": "matmul_f32_256", "kind": "matmul", "m": 256, "n": 256, "k": 256},
31 {"name": "matmul_f32_512", "kind": "matmul", "m": 512, "n": 512, "k": 512},
32 {"name": "matmul_f32_1024", "kind": "matmul", "m": 1024, "n": 1024, "k": 1024},
33 {"name": "matmul_f32_2048", "kind": "matmul", "m": 2048, "n": 2048, "k": 2048},
34 {"name": "matmul_f32_4096", "kind": "matmul", "m": 4096, "n": 4096, "k": 4096},
35 {"name": "dense_f32_1024", "kind": "dense", "m": 1024, "n": 1024, "k": 1024},
36 {"name": "ewchain_f32_16m", "kind": "ewchain", "elements": 16 * 1024 * 1024},
37 {"name": "reduce_f32_16m", "kind": "reduce", "elements": 16 * 1024 * 1024},
38 {"name": "scan_f32_16m", "kind": "scan", "elements": 16 * 1024 * 1024},
39 {"name": "softmax_f32_4096x4096", "kind": "softmax", "m": 4096, "n": 4096},
40 {"name": "layernorm_f32_4096x4096", "kind": "layernorm", "m": 4096, "n": 4096},
41 {"name": "nbody_f32_4096", "kind": "nbody", "m": 4096},
42 {"name": "stencil5_f32_4096x4096", "kind": "stencil", "m": 4096, "n": 4096},
43 {"name": "raymarch_f32_1024x1024", "kind": "raymarch", "m": 1024, "n": 1024},
44 {"name": "mandelbrot_f32_2048x2048", "kind": "mandelbrot", "m": 2048, "n": 2048},
45 {"name": "warp2d_f32_2048x2048", "kind": "warp2d", "m": 2048, "n": 2048},
46 {"name": "rmsnorm_f32_4096x4096", "kind": "rmsnorm", "m": 4096, "n": 4096},
47 {"name": "attention_f32_4096x64", "kind": "attention", "m": 4096, "k": 64},
48 ]
49
50 RAYMARCH_STEPS = 48
51 RAYMARCH_FAR = 24.0
52 NBODY_SOFTENING = 1e-3
53 SPHERES = [
54 (0.0, 1.0, 6.0, 1.0),
55 (-2.2, 0.6, 5.0, 0.6),
56 (1.9, 0.8, 7.5, 0.8),
57 ]
58
59
60 def flops(w):
61 if w["kind"] == "matmul":
62 return 2 * w["m"] * w["n"] * w["k"]
63 if w["kind"] == "dense":
64 return 2 * w["m"] * w["n"] * w["k"] + 2 * w["m"] * w["n"]
65 if w["kind"] == "ewchain":
66 return 3 * EWCHAIN_LINKS * w["elements"]
67 if w["kind"] == "softmax":
68 return 5 * w["m"] * w["n"]
69 if w["kind"] == "scan":
70 return w["elements"]
71 if w["kind"] == "layernorm":
72 return 8 * w["m"] * w["n"]
73 if w["kind"] == "rmsnorm":
74 return 5 * w["m"] * w["n"]
75 if w["kind"] == "attention":
76 return 4 * w["m"] * w["m"] * w["k"] + 5 * w["m"] * w["m"]
77 if w["kind"] == "nbody":
78 return 20 * w["m"] * w["m"]
79 if w["kind"] == "stencil":
80 return 4 * w["m"] * w["n"]
81 if w["kind"] == "raymarch":
82 return (12 * 3 + 14) * RAYMARCH_STEPS * w["m"] * w["n"]
83 if w["kind"] == "mandelbrot":
84 return 10 * 64 * w["m"] * w["n"]
85 if w["kind"] == "warp2d":
86 return 24 * w["m"] * w["n"]
87 return w["elements"]
88
89
90 def moved_bytes(w):
91 s = 4
92 if w["kind"] == "matmul":
93 return s * (w["m"] * w["k"] + w["k"] * w["n"] + w["m"] * w["n"])
94 if w["kind"] == "dense":
95 return s * (w["m"] * w["k"] + w["k"] * w["n"] + w["n"] + w["m"] * w["n"])
96 if w["kind"] == "ewchain":
97 return s * 2 * w["elements"]
98 if w["kind"] == "softmax":
99 return s * 2 * w["m"] * w["n"]
100 if w["kind"] == "scan":
101 return s * 2 * w["elements"]
102 if w["kind"] == "layernorm":
103 return s * (2 * w["m"] * w["n"] + 2 * w["n"])
104 if w["kind"] == "rmsnorm":
105 return s * (2 * w["m"] * w["n"] + w["n"])
106 if w["kind"] == "attention":
107 return s * 4 * w["m"] * w["k"]
108 if w["kind"] == "nbody":
109 return s * (4 * w["m"] + 3 * w["m"])
110 if w["kind"] == "stencil":
111 return s * 2 * w["m"] * w["n"]
112 if w["kind"] == "raymarch":
113 return s * w["m"] * w["n"]
114 if w["kind"] == "mandelbrot":
115 return s * w["m"] * w["n"]
116 if w["kind"] == "warp2d":
117 return s * 2 * w["m"] * w["n"]
118 return s * (w["elements"] + 1)
119
120
121 def fill(count, offset):
122 index = np.arange(offset, offset + count, dtype=np.uint64)
123 state = index.astype(np.uint32)
124 state = state * np.uint32(1664525) + np.uint32(1013904223)
125 state = state * np.uint32(1664525) + np.uint32(1013904223)
126 return (state >> np.uint32(8)).astype(np.float32) / np.float32(1 << 24) - np.float32(0.5)
127
128
129 def make_row(system, w, metric, samples, values_ns, kernel=None, note=None):
130 ordered = sorted(values_ns)
131 row = {
132 "schema": SCHEMA,
133 "system": system,
134 "workload": w["name"],
135 "metric": metric,
136 "status": "ok",
137 "oracle": "passed",
138 "samples": samples,
139 "median_ns": int(ordered[len(ordered) // 2]),
140 "p10_ns": int(ordered[len(ordered) // 10]),
141 "p90_ns": int(ordered[(len(ordered) * 9) // 10]),
142 "flops": flops(w),
143 "moved_bytes": moved_bytes(w),
144 }
145 if kernel:
146 row["kernel"] = kernel
147 if note:
148 row["note"] = note
149 return row
150
151
152 def emit(system, w, metric, samples, values_ns, kernel=None, note=None):
153 emit_row(make_row(system, w, metric, samples, values_ns, kernel=kernel, note=note))
154
155
156 def emit_row(row):
157 print(json.dumps(row), flush=True)
158
159
160 def emit_rows(rows):
161 for row in rows:
162 emit_row(row)
163
164
165 def oracle_cells(total):
166 return [oracle_cell_index(s, total) for s in range(ORACLE_CELLS)]
167
168
169 def oracle_cell_index(sample, total):
170 if total <= 1:
171 return 0
172 if sample == 0:
173 return 0
174 if sample == 1:
175 return total - 1
176 if sample == 2:
177 return total // 2
178 if sample == 3:
179 return total // 3
180 if sample == 4:
181 return (total // 3) * 2
182 return oracle_mix(sample - 5 + total) % total
183
184
185 def oracle_mix(value):
186 mask = (1 << 64) - 1
187 mixed = (value + 0x9E3779B97F4A7C15) & mask
188 mixed = ((mixed ^ (mixed >> 30)) * 0xBF58476D1CE4E5B9) & mask
189 mixed = ((mixed ^ (mixed >> 27)) * 0x94D049BB133111EB) & mask
190 return mixed ^ (mixed >> 31)
191
192
193 def measure(fn, args, w, note=None):
194 # latency: dispatch + block, per sample
195 for _ in range(WARMUP):
196 fn(*args).block_until_ready()
197 latency = []
198 for _ in range(SAMPLES):
199 start = time.perf_counter_ns()
200 fn(*args).block_until_ready()
201 latency.append(time.perf_counter_ns() - start)
202 emit("xla", w, "kernel_ns", SAMPLES, latency, note=note)
203 emit("xla", w, "latency_ns", SAMPLES, latency)
204
205 dispatch = []
206 for _ in range(DISPATCH_SAMPLES):
207 start = time.perf_counter_ns()
208 result = fn(*args)
209 dispatch.append(time.perf_counter_ns() - start)
210 result.block_until_ready()
211 emit("xla", w, "dispatch_ns", DISPATCH_SAMPLES, dispatch)
212
213 pipeline = []
214 for _ in range(PIPELINE_ROUNDS):
215 jax.block_until_ready(fn(*args))
216 start = time.perf_counter_ns()
217 result = None
218 for _ in range(PIPELINE_BATCH):
219 result = fn(*args)
220 result.block_until_ready()
221 pipeline.append((time.perf_counter_ns() - start) // PIPELINE_BATCH)
222 emit("xla", w, "pipeline_ns", PIPELINE_ROUNDS, pipeline)
223
224
225 def compile_timing(fn, args, w):
226 rows = []
227 start = time.perf_counter_ns()
228 lowered = jax.jit(fn).lower(*args)
229 compiled = lowered.compile()
230 cold = time.perf_counter_ns() - start
231 rows.append(make_row("xla", w, "compile_ns", 1, [cold], note="jit lower+compile"))
232
233 warm = []
234 for _ in range(4):
235 jax.clear_caches()
236 start = time.perf_counter_ns()
237 jax.jit(fn).lower(*args).compile()
238 warm.append(time.perf_counter_ns() - start)
239 rows.append(make_row("xla", w, "compile_warm_ns", 4, warm, note="jit lower+compile, cleared caches"))
240
241 incremental = []
242 for _ in range(4):
243 fresh = _fresh_copy(fn)
244 start = time.perf_counter_ns()
245 jax.jit(fresh).lower(*args).compile()
246 incremental.append(time.perf_counter_ns() - start)
247 rows.append(make_row("xla", w, "compile_incremental_ns", 4, incremental,
248 note="fresh identical function, in-process caches warm"))
249 return compiled, rows
250
251
252 def _fresh_copy(fn):
253 import types
254 return types.FunctionType(fn.__code__, fn.__globals__, fn.__name__ + "_fresh",
255 fn.__defaults__, fn.__closure__)
256
257
258 def close(expected, actual, tolerance):
259 magnitude = max(abs(expected), 1.0)
260 return abs(expected - float(actual)) <= tolerance * magnitude
261
262
263 def verify_matmul(w, a, b, out):
264 total = w["m"] * w["n"]
265 a64 = a.astype(np.float64)
266 b64 = b.astype(np.float64)
267 flat = np.asarray(out).reshape(-1)
268 for cell in oracle_cells(total):
269 row, col = divmod(cell, w["n"])
270 expected = float(a64[row] @ b64[:, col])
271 if not close(expected, flat[cell], 1e-2):
272 return False
273 return True
274
275
276 def run_matmul(w):
277 a = fill(w["m"] * w["k"], 0).reshape(w["m"], w["k"])
278 b = fill(w["k"] * w["n"], w["m"] * w["k"]).reshape(w["k"], w["n"])
279
280 def f(x, y):
281 return jnp.dot(x, y, precision=jax.lax.Precision.HIGHEST)
282
283 dev_a = jax.device_put(a)
284 dev_b = jax.device_put(b)
285 compiled, compile_rows = compile_timing(f, (dev_a, dev_b), w)
286 out = compiled(dev_a, dev_b)
287 out.block_until_ready()
288 if not verify_matmul(w, a, b, out):
289 print(f"verify failed for {w['name']}", file=sys.stderr)
290 return
291 emit_rows(compile_rows)
292 measure(compiled, (dev_a, dev_b), w, note="precision=HIGHEST")
293
294
295 def run_dense(w):
296 a = fill(w["m"] * w["k"], 0).reshape(w["m"], w["k"])
297 b = fill(w["k"] * w["n"], w["m"] * w["k"]).reshape(w["k"], w["n"])
298 bias = fill(w["n"], w["m"] * w["k"] + w["k"] * w["n"])
299
300 def f(x, y, z):
301 return jnp.tanh(jnp.dot(x, y, precision=jax.lax.Precision.HIGHEST) + z)
302
303 dev = (jax.device_put(a), jax.device_put(b), jax.device_put(bias))
304 compiled, compile_rows = compile_timing(f, dev, w)
305 out = np.asarray(compiled(*dev)).reshape(-1)
306
307 a64 = a.astype(np.float64)
308 b64 = b.astype(np.float64)
309 for cell in oracle_cells(w["m"] * w["n"]):
310 row, col = divmod(cell, w["n"])
311 expected = np.tanh(float(a64[row] @ b64[:, col]) + float(bias[col]))
312 if not close(expected, out[cell], 1e-2):
313 print(f"verify failed for {w['name']}", file=sys.stderr)
314 return
315 emit_rows(compile_rows)
316 measure(compiled, dev, w)
317
318
319 def run_rmsnorm(w):
320 rows, cols = w["m"], w["n"]
321 cells = rows * cols
322 x = fill(cells, 0).reshape(rows, cols)
323 gamma = fill(cols, cells)
324 epsilon = 1e-5
325
326 def f(v, g):
327 mean_sq = (v * v).mean(axis=1, keepdims=True)
328 return v / jnp.sqrt(mean_sq + epsilon) * g[None, :]
329
330 dev_x = jax.device_put(x)
331 dev_gamma = jax.device_put(gamma)
332 compiled, compile_rows = compile_timing(f, (dev_x, dev_gamma), w)
333 out = np.asarray(compiled(dev_x, dev_gamma)).reshape(-1)
334
335 x64 = x.astype(np.float64)
336 g64 = gamma.astype(np.float64)
337 for cell in oracle_cells(cells):
338 row, col = divmod(cell, cols)
339 mean_sq = (x64[row] ** 2).mean()
340 expected = float(x64[row, col] / np.sqrt(mean_sq + epsilon) * g64[col])
341 if not close_abs_rel(expected, out[cell], 1e-3, 1e-3):
342 print(f"verify failed for {w['name']}", file=sys.stderr)
343 return
344 emit_rows(compile_rows)
345 measure(compiled, (dev_x, dev_gamma), w, note="jax rmsnorm")
346
347
348 def run_attention(w):
349 seq, dim = w["m"], w["k"]
350 q = fill(seq * dim, 0).reshape(seq, dim)
351 kt = fill(seq * dim, seq * dim).reshape(dim, seq)
352 v = fill(seq * dim, 2 * seq * dim).reshape(seq, dim)
353 scale = 1.0 / np.sqrt(dim)
354
355 def f(qq, kk, vv):
356 scores = (qq @ kk) * scale
357 probs = jax.nn.softmax(scores, axis=1)
358 return probs @ vv
359
360 dev_q = jax.device_put(q)
361 dev_kt = jax.device_put(kt)
362 dev_v = jax.device_put(v)
363 compiled, compile_rows = compile_timing(f, (dev_q, dev_kt, dev_v), w)
364 out = np.asarray(compiled(dev_q, dev_kt, dev_v)).reshape(-1)
365
366 q64 = q.astype(np.float64)
367 kt64 = kt.astype(np.float64)
368 v64 = v.astype(np.float64)
369 for cell in oracle_cells(seq * dim):
370 row, out_col = divmod(cell, dim)
371 scores = (q64[row] @ kt64) * scale
372 shifted = np.exp(scores - scores.max())
373 probs = shifted / shifted.sum()
374 expected = float(probs @ v64[:, out_col])
375 if not close_abs_rel(expected, out[cell], 1e-3, 5e-3):
376 print(f"verify failed for {w['name']}", file=sys.stderr)
377 return
378 emit_rows(compile_rows)
379 measure(compiled, (dev_q, dev_kt, dev_v), w, note="jax matmul softmax matmul")
380
381
382 def run_mandelbrot(w):
383 rows, cols = w["m"], w["n"]
384 max_iters = 64
385
386 def f():
387 col = jnp.arange(cols, dtype=jnp.float32)[None, :]
388 row = jnp.arange(rows, dtype=jnp.float32)[:, None]
389 cx = -2.0 + col * (2.5 / cols)
390 cy = -1.25 + row * (2.5 / rows)
391 cx = jnp.broadcast_to(cx, (rows, cols))
392 cy = jnp.broadcast_to(cy, (rows, cols))
393
394 def body(_, state):
395 zx, zy, count, active = state
396 nzx = zx * zx - zy * zy + cx
397 nzy = 2.0 * zx * zy + cy
398 count = jnp.where(active, count + 1.0, count)
399 zx = jnp.where(active, nzx, zx)
400 zy = jnp.where(active, nzy, zy)
401 active = jnp.logical_and(active, zx * zx + zy * zy < 4.0)
402 return zx, zy, count, active
403
404 state = (cx, cy, jnp.zeros((rows, cols), jnp.float32), jnp.ones((rows, cols), bool))
405 zx, zy, count, active = jax.lax.fori_loop(0, max_iters, body, state)
406 return count
407
408 compiled, compile_rows = compile_timing(f, (), w)
409 out = np.asarray(compiled()).reshape(-1)
410
411 for cell in oracle_cells(rows * cols):
412 r, c = divmod(cell, cols)
413 cx = np.float32(-2.0 + c * (2.5 / cols))
414 cy = np.float32(-1.25 + r * (2.5 / rows))
415 zx, zy, count = cx, cy, 0.0
416 for _ in range(max_iters):
417 nzx = zx * zx - zy * zy + cx
418 nzy = 2 * zx * zy + cy
419 count += 1.0
420 zx, zy = np.float32(nzx), np.float32(nzy)
421 if not (zx * zx + zy * zy < 4.0):
422 break
423 if not close_abs_rel(count, out[cell], 1.05, 0.0):
424 print(f"verify failed for {w['name']}", file=sys.stderr)
425 return
426 emit_rows(compile_rows)
427 measure(compiled, (), w, note="fori_loop masked, no early exit")
428
429
430 def run_ewchain(w):
431 x = fill(w["elements"], 0)
432
433 def f(v):
434 r = v
435 for _ in range(EWCHAIN_LINKS):
436 r = jnp.tanh(r * r + v)
437 return r
438
439 dev_x = jax.device_put(x)
440 compiled, compile_rows = compile_timing(f, (dev_x,), w)
441 out = np.asarray(compiled(dev_x))
442
443 x64 = x.astype(np.float64)
444 for cell in oracle_cells(w["elements"]):
445 expected = x64[cell]
446 value = expected
447 for _ in range(EWCHAIN_LINKS):
448 value = np.tanh(value * value + expected)
449 if not close(float(value), out[cell], 1e-3):
450 print(f"verify failed for {w['name']}", file=sys.stderr)
451 return
452 emit_rows(compile_rows)
453 measure(compiled, (dev_x,), w)
454
455
456 def run_reduce(w):
457 x = fill(w["elements"], 0)
458
459 def f(v):
460 return jnp.sum(v)
461
462 dev_x = jax.device_put(x)
463 compiled, compile_rows = compile_timing(f, (dev_x,), w)
464 result = float(compiled(dev_x))
465 expected = float(np.sum(x.astype(np.float64)))
466 if not close(expected, result, 5e-2):
467 print(f"verify failed for {w['name']}: {result} vs {expected}", file=sys.stderr)
468 return
469 emit_rows(compile_rows)
470 measure(compiled, (dev_x,), w)
471
472
473
474
475 def close_abs_rel(expected, actual, abs_tol, rel_tol):
476 return abs(expected - float(actual)) <= abs_tol + rel_tol * abs(expected)
477
478
479 def run_scan(w):
480 n = w["elements"]
481 x = fill(n, 0)
482
483 def f(v):
484 return jnp.cumsum(v)
485
486 dev_x = jax.device_put(x)
487 compiled, compile_rows = compile_timing(f, (dev_x,), w)
488 out = np.asarray(compiled(dev_x))
489
490 acc = 0.0
491 next_check = 1
492 x64 = x.astype(np.float64)
493 i = 0
494 prefix = np.cumsum(x64)
495 while next_check <= n:
496 i = next_check - 1
497 if not close_abs_rel(float(prefix[i]), out[i], 1e-2, 1e-4):
498 print(f"verify failed for {w['name']}", file=sys.stderr)
499 return
500 next_check = next_check * 7 // 2 + 13
501 if not close_abs_rel(float(prefix[n - 1]), out[n - 1], 1e-2, 1e-4):
502 print(f"verify failed for {w['name']}", file=sys.stderr)
503 return
504 emit_rows(compile_rows)
505 measure(compiled, (dev_x,), w, note="jnp.cumsum")
506
507
508 def run_softmax(w):
509 rows, cols = w["m"], w["n"]
510 x = fill(rows * cols, 0).reshape(rows, cols)
511
512 def f(v):
513 return jax.nn.softmax(v, axis=1)
514
515 dev_x = jax.device_put(x)
516 compiled, compile_rows = compile_timing(f, (dev_x,), w)
517 out = np.asarray(compiled(dev_x)).reshape(-1)
518
519 x64 = x.astype(np.float64)
520 for cell in oracle_cells(rows * cols):
521 row, col = divmod(cell, cols)
522 shifted = x64[row] - x64[row].max()
523 expected = float(np.exp(shifted[col]) / np.exp(shifted).sum())
524 if not close_abs_rel(expected, out[cell], 1e-7, 1e-3):
525 print(f"verify failed for {w['name']}", file=sys.stderr)
526 return
527 emit_rows(compile_rows)
528 measure(compiled, (dev_x,), w, note="jax.nn.softmax")
529
530
531 def run_layernorm(w):
532 rows, cols = w["m"], w["n"]
533 cells = rows * cols
534 x = fill(cells, 0).reshape(rows, cols)
535 gamma = fill(cols, cells)
536 beta = fill(cols, cells + cols)
537 epsilon = 1e-5
538
539 def f(v, g, b):
540 mean = v.mean(axis=1, keepdims=True)
541 var = ((v - mean) ** 2).mean(axis=1, keepdims=True)
542 return (v - mean) / jnp.sqrt(var + epsilon) * g[None, :] + b[None, :]
543
544 dev_x = jax.device_put(x)
545 dev_gamma = jax.device_put(gamma)
546 dev_beta = jax.device_put(beta)
547 compiled, compile_rows = compile_timing(f, (dev_x, dev_gamma, dev_beta), w)
548 out = np.asarray(compiled(dev_x, dev_gamma, dev_beta)).reshape(-1)
549
550 x64 = x.astype(np.float64)
551 g64 = gamma.astype(np.float64)
552 b64 = beta.astype(np.float64)
553 for cell in oracle_cells(cells):
554 row, col = divmod(cell, cols)
555 mean = x64[row].mean()
556 var = ((x64[row] - mean) ** 2).mean()
557 expected = float((x64[row, col] - mean) / np.sqrt(var + epsilon) * g64[col] + b64[col])
558 if not close_abs_rel(expected, out[cell], 1e-3, 1e-3):
559 print(f"verify failed for {w['name']}", file=sys.stderr)
560 return
561 emit_rows(compile_rows)
562 measure(compiled, (dev_x, dev_gamma, dev_beta), w, note="jax layernorm")
563
564
565 def run_nbody(w):
566 bodies = w["m"]
567 px = fill(bodies, 0)
568 py = fill(bodies, bodies)
569 pz = fill(bodies, 2 * bodies)
570 mass = fill(bodies, 3 * bodies)
571
572 def f(x, y, z, m):
573 dx = x[None, :] - x[:, None]
574 dy = y[None, :] - y[:, None]
575 dz = z[None, :] - z[:, None]
576 r2 = dx * dx + dy * dy + dz * dz + NBODY_SOFTENING
577 weight = m[None, :] / (r2 * jnp.sqrt(r2))
578 ax = jnp.sum(dx * weight, axis=1)
579 ay = jnp.sum(dy * weight, axis=1)
580 az = jnp.sum(dz * weight, axis=1)
581 return jnp.concatenate([ax, ay, az])
582
583 dev = tuple(jax.device_put(v) for v in (px, py, pz, mass))
584 compiled, compile_rows = compile_timing(f, dev, w)
585 out = np.asarray(compiled(*dev))
586
587 px64, py64, pz64, mass64 = (v.astype(np.float64) for v in (px, py, pz, mass))
588 for flat in oracle_cells(3 * bodies):
589 axis, body = divmod(flat, bodies)
590 dx = px64 - px64[body]
591 dy = py64 - py64[body]
592 dz = pz64 - pz64[body]
593 r2 = dx * dx + dy * dy + dz * dz + NBODY_SOFTENING
594 weight = mass64 / (r2 * np.sqrt(r2))
595 delta = dx if axis == 0 else dy if axis == 1 else dz
596 expected = float(np.sum(delta * weight))
597 if not close_abs_rel(expected, out[flat], 5e-2, 5e-3):
598 print(f"verify failed for {w['name']}", file=sys.stderr)
599 return
600 emit_rows(compile_rows)
601 measure(compiled, dev, w, note="same tensor formulation as accy")
602
603
604 def run_stencil(w):
605 rows, cols = w["m"], w["n"]
606 x = fill(rows * cols, 0).reshape(rows, cols)
607
608 def f(v):
609 padded = jnp.pad(v, 1)
610 up = padded[0:rows, 1 : cols + 1]
611 down = padded[2 : rows + 2, 1 : cols + 1]
612 left = padded[1 : rows + 1, 0:cols]
613 right = padded[1 : rows + 1, 2 : cols + 2]
614 return 0.25 * (up + down + left + right)
615
616 dev_x = jax.device_put(x)
617 compiled, compile_rows = compile_timing(f, (dev_x,), w)
618 out = np.asarray(compiled(dev_x)).reshape(-1)
619
620 x64 = x.astype(np.float64)
621 for cell in oracle_cells(rows * cols):
622 row, col = divmod(cell, cols)
623 up = x64[row - 1, col] if row > 0 else 0.0
624 down = x64[row + 1, col] if row + 1 < rows else 0.0
625 left = x64[row, col - 1] if col > 0 else 0.0
626 right = x64[row, col + 1] if col + 1 < cols else 0.0
627 if not close_abs_rel(0.25 * (up + down + left + right), out[cell], 1e-5, 1e-5):
628 print(f"verify failed for {w['name']}", file=sys.stderr)
629 return
630 emit_rows(compile_rows)
631 measure(compiled, (dev_x,), w, note="pad + shifted slices")
632
633
634 def run_warp2d(w):
635 import math
636 rows, cols = w["m"], w["n"]
637 x = fill(rows * cols, 0)
638
639 inv_scale = 1.0 / 1.15
640 ca_d = math.cos(0.35) * inv_scale
641 sa_d = math.sin(0.35) * inv_scale
642 cx_d = 0.5 * (cols - 1)
643 cy_d = 0.5 * (rows - 1)
644 ca = np.float32(ca_d)
645 sa = np.float32(sa_d)
646 tx = np.float32(cx_d - ca_d * cx_d - sa_d * cy_d)
647 ty = np.float32(cy_d + sa_d * cx_d - ca_d * cy_d)
648
649 def f(src):
650 i = jnp.arange(rows * cols, dtype=jnp.float32)
651 iy = jnp.floor(i * np.float32(1.0 / cols))
652 ix = i - iy * np.float32(cols)
653 sx = ix * ca + iy * sa + tx
654 sy = iy * ca - ix * sa + ty
655 cxx = jnp.clip(sx, 0.0, np.float32(cols - 1))
656 cyy = jnp.clip(sy, 0.0, np.float32(rows - 1))
657 x0 = jnp.minimum(jnp.floor(cxx), np.float32(cols - 2))
658 y0 = jnp.minimum(jnp.floor(cyy), np.float32(rows - 2))
659 fx = cxx - x0
660 fy = cyy - y0
661 base = (y0 * np.float32(cols) + x0).astype(jnp.int32)
662 g00 = src[base]
663 g01 = src[base + 1]
664 g10 = src[base + cols]
665 g11 = src[base + cols + 1]
666 top = g00 * (1.0 - fx) + g01 * fx
667 bottom = g10 * (1.0 - fx) + g11 * fx
668 return top * (1.0 - fy) + bottom * fy
669
670 dev_x = jax.device_put(x)
671 compiled, compile_rows = compile_timing(f, (dev_x,), w)
672 out = np.asarray(compiled(dev_x)).reshape(-1)
673
674 x64 = x.astype(np.float64)
675 ok = True
676 for cell in oracle_cells(rows * cols):
677 row, col = divmod(cell, cols)
678 sx = col * float(ca) + row * float(sa) + float(tx)
679 sy = row * float(ca) - col * float(sa) + float(ty)
680 cxx = min(max(sx, 0.0), float(cols - 1))
681 cyy = min(max(sy, 0.0), float(rows - 1))
682 x0 = min(math.floor(cxx), cols - 2)
683 y0 = min(math.floor(cyy), rows - 2)
684 fx = cxx - x0
685 fy = cyy - y0
686 g00 = x64[y0 * cols + x0]
687 g01 = x64[y0 * cols + x0 + 1]
688 g10 = x64[(y0 + 1) * cols + x0]
689 g11 = x64[(y0 + 1) * cols + x0 + 1]
690 expected = (g00 * (1.0 - fx) + g01 * fx) * (1.0 - fy) + (g10 * (1.0 - fx) + g11 * fx) * fy
691 if not close_abs_rel(expected, out[cell], 5e-3, 5e-3):
692 print(f"verify failed for {w['name']}", file=sys.stderr)
693 ok = False
694 break
695 if ok:
696 emit_rows(compile_rows)
697 measure(compiled, (dev_x,), w, note="fancy-index gathers")
698
699
700 def run_raymarch(w):
701 rows, cols = w["m"], w["n"]
702 seed = np.zeros((rows, cols), dtype=np.float32)
703
704 def f(t0):
705 col = (jnp.arange(cols, dtype=jnp.float32)[None, :] + 0.5) * (2.0 / cols) - 1.0
706 row = 1.0 - (jnp.arange(rows, dtype=jnp.float32)[:, None] + 0.5) * (2.0 / rows)
707 u = jnp.broadcast_to(col, (rows, cols))
708 v = jnp.broadcast_to(row, (rows, cols))
709 raw_z = 1.4
710 inv_len = 1.0 / jnp.sqrt(u * u + v * v + raw_z * raw_z)
711 dx = u * inv_len
712 dy = v * inv_len
713 dz = raw_z * inv_len
714 t = t0 * 0.0
715
716 def step(_, t):
717 px = dx * t
718 py = 1.2 + dy * t
719 pz = dz * t
720 d = py
721 for cx, cy, cz, radius in SPHERES:
722 sx = px - cx
723 sy = py - cy
724 sz = pz - cz
725 d = jnp.minimum(d, jnp.sqrt(sx * sx + sy * sy + sz * sz) - radius)
726 return jnp.minimum(t + jnp.maximum(d, 0.0), RAYMARCH_FAR)
727
728 return jax.lax.fori_loop(0, RAYMARCH_STEPS, step, t)
729
730 dev_seed = jax.device_put(seed)
731 compiled, compile_rows = compile_timing(f, (dev_seed,), w)
732 out = np.asarray(compiled(dev_seed)).reshape(-1)
733
734 for cell in oracle_cells(rows * cols):
735 row, col = divmod(cell, cols)
736 u = (col + 0.5) * (2.0 / cols) - 1.0
737 v = 1.0 - (row + 0.5) * (2.0 / rows)
738 raw_z = 1.4
739 inv_len = 1.0 / np.sqrt(u * u + v * v + raw_z * raw_z)
740 dx, dy, dz = u * inv_len, v * inv_len, raw_z * inv_len
741 t = 0.0
742 for _ in range(RAYMARCH_STEPS):
743 px, py, pz = dx * t, 1.2 + dy * t, dz * t
744 d = py
745 for cx, cy, cz, radius in SPHERES:
746 d = min(d, np.sqrt((px - cx) ** 2 + (py - cy) ** 2 + (pz - cz) ** 2) - radius)
747 t = min(t + max(d, 0.0), RAYMARCH_FAR)
748 if not close_abs_rel(t, out[cell], 2e-3, 1e-3):
749 print(f"verify failed for {w['name']}", file=sys.stderr)
750 return
751 emit_rows(compile_rows)
752 measure(compiled, (dev_seed,), w, note="fori_loop fixed 48 steps")
753
754
755 def main():
756 name_filter = sys.argv[1] if len(sys.argv) > 1 else None
757 device = jax.devices()[0]
758 print(json.dumps({
759 "schema": SCHEMA,
760 "system": "xla",
761 "meta": True,
762 "device": str(device.device_kind),
763 "detail": f"jax={jax.__version__}, jaxlib={jaxlib.__version__}, numpy={np.__version__}, timing=wall(dispatch+block)",
764 }), flush=True)
765
766 runners = {
767 "warp2d": run_warp2d,
768 "matmul": run_matmul,
769 "dense": run_dense,
770 "ewchain": run_ewchain,
771 "reduce": run_reduce,
772 "softmax": run_softmax,
773 "scan": run_scan,
774 "layernorm": run_layernorm,
775 "mandelbrot": run_mandelbrot,
776 "rmsnorm": run_rmsnorm,
777 "attention": run_attention,
778 "nbody": run_nbody,
779 "stencil": run_stencil,
780 "raymarch": run_raymarch,
781 }
782 for w in BATTERY:
783 if name_filter and w["name"] != name_filter:
784 continue
785 runners[w["kind"]](w)
786
787
788 if __name__ == "__main__":
789 main()