lib/accy/src/profiling/versus/baseline/warplang.py
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 """Nvidia Warp 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 launch + synchronize, median over samples. Compile
5 rows time warp kernel module load (cold and cached).
6
7 Run inside a venv with warp-lang installed:
8 python warp.py [workload-name] > warp.jsonl
9 """
10
11 import json
12 import sys
13 import time
14
15 import numpy as np
16 import warp as wp
17
18 SCHEMA = "accy.versus/v2"
19 ORACLE_CELLS = 256
20
21 SAMPLES = 20
22 DISPATCH_SAMPLES = 50
23 PIPELINE_ROUNDS = 5
24 PIPELINE_BATCH = 16
25 WARMUP = 3
26
27 RAYMARCH_STEPS = 48
28 RAYMARCH_FAR = 24.0
29 NBODY_SOFTENING = 1e-3
30 SPHERES = [
31 (0.0, 1.0, 6.0, 1.0),
32 (-2.2, 0.6, 5.0, 0.6),
33 (1.9, 0.8, 7.5, 0.8),
34 ]
35
36 BATTERY = [
37 {"name": "softmax_f32_4096x4096", "kind": "softmax", "m": 4096, "n": 4096},
38 {"name": "layernorm_f32_4096x4096", "kind": "layernorm", "m": 4096, "n": 4096},
39 {"name": "nbody_f32_4096", "kind": "nbody", "m": 4096, "n": 1},
40 {"name": "stencil5_f32_4096x4096", "kind": "stencil", "m": 4096, "n": 4096},
41 {"name": "raymarch_f32_1024x1024", "kind": "raymarch", "m": 1024, "n": 1024},
42 {"name": "mandelbrot_f32_2048x2048", "kind": "mandelbrot", "m": 2048, "n": 2048},
43 {"name": "warp2d_f32_2048x2048", "kind": "warp2d", "m": 2048, "n": 2048},
44 {"name": "rmsnorm_f32_4096x4096", "kind": "rmsnorm", "m": 4096, "n": 4096},
45 {"name": "ewchain_f32_16m", "kind": "ewchain", "m": 1, "n": 1, "elements": 16 * 1024 * 1024},
46 ]
47
48 EWCHAIN_LINKS = 4
49
50
51 def flops(w):
52 if w["kind"] == "warp2d":
53 return 24 * w["m"] * w["n"]
54 if w["kind"] == "softmax":
55 return 5 * w["m"] * w["n"]
56 if w["kind"] == "layernorm":
57 return 8 * w["m"] * w["n"]
58 if w["kind"] == "mandelbrot":
59 return 10 * 64 * w["m"] * w["n"]
60 if w["kind"] == "rmsnorm":
61 return 5 * w["m"] * w["n"]
62 if w["kind"] == "nbody":
63 return 20 * w["m"] * w["m"]
64 if w["kind"] == "stencil":
65 return 4 * w["m"] * w["n"]
66 if w["kind"] == "raymarch":
67 return (12 * 3 + 14) * RAYMARCH_STEPS * w["m"] * w["n"]
68 return 3 * EWCHAIN_LINKS * w["elements"]
69
70
71 def moved_bytes(w):
72 s = 4
73 if w["kind"] == "warp2d":
74 return s * 2 * w["m"] * w["n"]
75 if w["kind"] == "softmax":
76 return s * 2 * w["m"] * w["n"]
77 if w["kind"] == "layernorm":
78 return s * (2 * w["m"] * w["n"] + 2 * w["n"])
79 if w["kind"] == "mandelbrot":
80 return s * w["m"] * w["n"]
81 if w["kind"] == "rmsnorm":
82 return s * (2 * w["m"] * w["n"] + w["n"])
83 if w["kind"] == "nbody":
84 return s * (4 * w["m"] + 3 * w["m"])
85 if w["kind"] == "stencil":
86 return s * 2 * w["m"] * w["n"]
87 if w["kind"] == "raymarch":
88 return s * w["m"] * w["n"]
89 return s * 2 * w["elements"]
90
91
92 def fill(count, offset):
93 index = np.arange(offset, offset + count, dtype=np.uint64)
94 state = index.astype(np.uint32)
95 state = state * np.uint32(1664525) + np.uint32(1013904223)
96 state = state * np.uint32(1664525) + np.uint32(1013904223)
97 return (state >> np.uint32(8)).astype(np.float32) / np.float32(1 << 24) - np.float32(0.5)
98
99
100 def emit(w, metric, samples, values_ns, kernel=None, note=None):
101 ordered = sorted(values_ns)
102 row = {
103 "schema": SCHEMA,
104 "system": "warp",
105 "workload": w["name"],
106 "metric": metric,
107 "status": "ok",
108 "oracle": "passed",
109 "samples": samples,
110 "median_ns": int(ordered[len(ordered) // 2]),
111 "p10_ns": int(ordered[len(ordered) // 10]),
112 "p90_ns": int(ordered[(len(ordered) * 9) // 10]),
113 "flops": flops(w),
114 "moved_bytes": moved_bytes(w),
115 }
116 if kernel:
117 row["kernel"] = kernel
118 if note:
119 row["note"] = note
120 print(json.dumps(row), flush=True)
121
122
123 def oracle_cells(total):
124 return [oracle_cell_index(s, total) for s in range(ORACLE_CELLS)]
125
126
127 def oracle_cell_index(sample, total):
128 if total <= 1:
129 return 0
130 if sample == 0:
131 return 0
132 if sample == 1:
133 return total - 1
134 if sample == 2:
135 return total // 2
136 if sample == 3:
137 return total // 3
138 if sample == 4:
139 return (total // 3) * 2
140 return oracle_mix(sample - 5 + total) % total
141
142
143 def oracle_mix(value):
144 mask = (1 << 64) - 1
145 mixed = (value + 0x9E3779B97F4A7C15) & mask
146 mixed = ((mixed ^ (mixed >> 30)) * 0xBF58476D1CE4E5B9) & mask
147 mixed = ((mixed ^ (mixed >> 27)) * 0x94D049BB133111EB) & mask
148 return mixed ^ (mixed >> 31)
149
150
151 def close_abs_rel(expected, actual, abs_tol, rel_tol):
152 return abs(expected - float(actual)) <= abs_tol + rel_tol * abs(expected)
153
154
155 def measure(launch, w, kernel_name, note=None):
156 for _ in range(WARMUP):
157 launch()
158 wp.synchronize()
159 latency = []
160 for _ in range(SAMPLES):
161 start = time.perf_counter_ns()
162 launch()
163 wp.synchronize()
164 latency.append(time.perf_counter_ns() - start)
165 emit(w, "kernel_ns", SAMPLES, latency, kernel=kernel_name, note=note)
166 emit(w, "latency_ns", SAMPLES, latency)
167
168 dispatch = []
169 for _ in range(DISPATCH_SAMPLES):
170 start = time.perf_counter_ns()
171 launch()
172 dispatch.append(time.perf_counter_ns() - start)
173 wp.synchronize()
174 emit(w, "dispatch_ns", DISPATCH_SAMPLES, dispatch)
175
176 pipeline = []
177 for _ in range(PIPELINE_ROUNDS):
178 wp.synchronize()
179 start = time.perf_counter_ns()
180 for _ in range(PIPELINE_BATCH):
181 launch()
182 wp.synchronize()
183 pipeline.append((time.perf_counter_ns() - start) // PIPELINE_BATCH)
184 emit(w, "pipeline_ns", PIPELINE_ROUNDS, pipeline)
185
186
187 @wp.kernel
188 def softmax_rows(x: wp.array2d(dtype=float), out: wp.array2d(dtype=float)):
189 row = wp.tid()
190 cols = x.shape[1]
191 row_max = float(-3.4e38)
192 for col in range(cols):
193 row_max = wp.max(row_max, x[row, col])
194 row_sum = float(0.0)
195 for col in range(cols):
196 row_sum += wp.exp(x[row, col] - row_max)
197 inv = 1.0 / row_sum
198 for col in range(cols):
199 out[row, col] = wp.exp(x[row, col] - row_max) * inv
200
201
202 @wp.kernel
203 def rmsnorm_rows(
204 x: wp.array2d(dtype=float),
205 gamma: wp.array(dtype=float),
206 out: wp.array2d(dtype=float),
207 ):
208 row = wp.tid()
209 cols = x.shape[1]
210 sq_acc = float(0.0)
211 for col in range(cols):
212 sq_acc += x[row, col] * x[row, col]
213 inv_rms = 1.0 / wp.sqrt(sq_acc / float(cols) + 1.0e-5)
214 for col in range(cols):
215 out[row, col] = x[row, col] * inv_rms * gamma[col]
216
217
218 @wp.kernel
219 def warp2d_bilinear(
220 src: wp.array(dtype=float),
221 out: wp.array2d(dtype=float),
222 ca: float,
223 sa: float,
224 tx: float,
225 ty: float,
226 ):
227 row, col = wp.tid()
228 rows = out.shape[0]
229 cols = out.shape[1]
230 sx = float(col) * ca + float(row) * sa + tx
231 sy = float(row) * ca - float(col) * sa + ty
232 cxx = wp.min(wp.max(sx, 0.0), float(cols - 1))
233 cyy = wp.min(wp.max(sy, 0.0), float(rows - 1))
234 x0 = wp.min(wp.floor(cxx), float(cols - 2))
235 y0 = wp.min(wp.floor(cyy), float(rows - 2))
236 fx = cxx - x0
237 fy = cyy - y0
238 xi = int(x0)
239 yi = int(y0)
240 g00 = src[yi * cols + xi]
241 g01 = src[yi * cols + xi + 1]
242 g10 = src[(yi + 1) * cols + xi]
243 g11 = src[(yi + 1) * cols + xi + 1]
244 top = g00 * (1.0 - fx) + g01 * fx
245 bottom = g10 * (1.0 - fx) + g11 * fx
246 out[row, col] = top * (1.0 - fy) + bottom * fy
247
248
249 @wp.kernel
250 def mandelbrot_escape(out: wp.array2d(dtype=float)):
251 row, col = wp.tid()
252 rows = out.shape[0]
253 cols = out.shape[1]
254 cx = -2.0 + float(col) * (2.5 / float(cols))
255 cy = -1.25 + float(row) * (2.5 / float(rows))
256 zx = cx
257 zy = cy
258 count = float(0.0)
259 for _ in range(64):
260 nzx = zx * zx - zy * zy + cx
261 nzy = 2.0 * zx * zy + cy
262 count += 1.0
263 zx = nzx
264 zy = nzy
265 if not (zx * zx + zy * zy < 4.0):
266 break
267 out[row, col] = count
268
269
270 @wp.kernel
271 def layernorm_rows(
272 x: wp.array2d(dtype=float),
273 gamma: wp.array(dtype=float),
274 beta: wp.array(dtype=float),
275 out: wp.array2d(dtype=float),
276 ):
277 row = wp.tid()
278 cols = x.shape[1]
279 mean_acc = float(0.0)
280 for col in range(cols):
281 mean_acc += x[row, col]
282 mean = mean_acc / float(cols)
283 var_acc = float(0.0)
284 for col in range(cols):
285 centered = x[row, col] - mean
286 var_acc += centered * centered
287 inv_std = 1.0 / wp.sqrt(var_acc / float(cols) + 1.0e-5)
288 for col in range(cols):
289 out[row, col] = (x[row, col] - mean) * inv_std * gamma[col] + beta[col]
290
291
292 @wp.kernel
293 def nbody_accel(
294 px: wp.array(dtype=float),
295 py: wp.array(dtype=float),
296 pz: wp.array(dtype=float),
297 mass: wp.array(dtype=float),
298 out: wp.array(dtype=float),
299 ):
300 body = wp.tid()
301 bodies = px.shape[0]
302 xi = px[body]
303 yi = py[body]
304 zi = pz[body]
305 ax = float(0.0)
306 ay = float(0.0)
307 az = float(0.0)
308 for other in range(bodies):
309 dx = px[other] - xi
310 dy = py[other] - yi
311 dz = pz[other] - zi
312 r2 = dx * dx + dy * dy + dz * dz + 1e-3
313 inv = 1.0 / wp.sqrt(r2)
314 weight = mass[other] * inv * inv * inv
315 ax += dx * weight
316 ay += dy * weight
317 az += dz * weight
318 out[body] = ax
319 out[bodies + body] = ay
320 out[2 * bodies + body] = az
321
322
323 @wp.kernel
324 def stencil5(x: wp.array2d(dtype=float), out: wp.array2d(dtype=float)):
325 row, col = wp.tid()
326 rows = x.shape[0]
327 cols = x.shape[1]
328 up = wp.where(row > 0, x[wp.max(row - 1, 0), col], 0.0)
329 down = wp.where(row + 1 < rows, x[wp.min(row + 1, rows - 1), col], 0.0)
330 left = wp.where(col > 0, x[row, wp.max(col - 1, 0)], 0.0)
331 right = wp.where(col + 1 < cols, x[row, wp.min(col + 1, cols - 1)], 0.0)
332 out[row, col] = 0.25 * (up + down + left + right)
333
334
335 @wp.kernel
336 def raymarch(out: wp.array2d(dtype=float)):
337 row, col = wp.tid()
338 rows = out.shape[0]
339 cols = out.shape[1]
340 u = (float(col) + 0.5) * (2.0 / float(cols)) - 1.0
341 v = 1.0 - (float(row) + 0.5) * (2.0 / float(rows))
342 raw_z = float(1.4)
343 inv_len = 1.0 / wp.sqrt(u * u + v * v + raw_z * raw_z)
344 dx = u * inv_len
345 dy = v * inv_len
346 dz = raw_z * inv_len
347 t = float(0.0)
348 for _step in range(48):
349 p = wp.vec3(dx * t, 1.2 + dy * t, dz * t)
350 d = p[1]
351 d = wp.min(d, wp.length(p - wp.vec3(0.0, 1.0, 6.0)) - 1.0)
352 d = wp.min(d, wp.length(p - wp.vec3(-2.2, 0.6, 5.0)) - 0.6)
353 d = wp.min(d, wp.length(p - wp.vec3(1.9, 0.8, 7.5)) - 0.8)
354 t = wp.min(t + wp.max(d, 0.0), 24.0)
355 if d < 1e-4 or t >= 24.0:
356 break
357 out[row, col] = t
358
359
360 @wp.kernel
361 def ewchain(x: wp.array(dtype=float), out: wp.array(dtype=float)):
362 i = wp.tid()
363 v = x[i]
364 r = v
365 for _link in range(4):
366 r = wp.tanh(r * r + v)
367 out[i] = r
368
369
370 def run_softmax(w):
371 rows, cols = w["m"], w["n"]
372 x = fill(rows * cols, 0).reshape(rows, cols)
373 dev_x = wp.array(x, dtype=float)
374 dev_out = wp.zeros((rows, cols), dtype=float)
375
376 def launch():
377 wp.launch(softmax_rows, dim=rows, inputs=[dev_x, dev_out])
378
379 launch()
380 wp.synchronize()
381 out = dev_out.numpy().reshape(-1)
382 x64 = x.astype(np.float64)
383 for cell in oracle_cells(rows * cols):
384 r, c = divmod(cell, cols)
385 shifted = x64[r] - x64[r].max()
386 expected = float(np.exp(shifted[c]) / np.exp(shifted).sum())
387 if not close_abs_rel(expected, out[cell], 1e-7, 1e-3):
388 print(f"verify failed for {w['name']}", file=sys.stderr)
389 return
390 measure(launch, w, "softmax_rows", note="thread-per-row three-pass")
391
392
393 def run_rmsnorm(w):
394 rows, cols = w["m"], w["n"]
395 cells = rows * cols
396 x = fill(cells, 0).reshape(rows, cols)
397 gamma = fill(cols, cells)
398 dev_x = wp.array(x, dtype=float)
399 dev_gamma = wp.array(gamma, dtype=float)
400 dev_out = wp.zeros((rows, cols), dtype=float)
401
402 def launch():
403 wp.launch(rmsnorm_rows, dim=rows, inputs=[dev_x, dev_gamma, dev_out])
404
405 launch()
406 wp.synchronize()
407 out = dev_out.numpy().reshape(-1)
408 x64 = x.astype(np.float64)
409 g64 = gamma.astype(np.float64)
410 for cell in oracle_cells(cells):
411 r, c = divmod(cell, cols)
412 mean_sq = (x64[r] ** 2).mean()
413 expected = float(x64[r, c] / np.sqrt(mean_sq + 1e-5) * g64[c])
414 if not close_abs_rel(expected, out[cell], 1e-3, 1e-3):
415 print(f"verify failed for {w['name']}", file=sys.stderr)
416 return
417 measure(launch, w, "rmsnorm_rows", note="thread-per-row two-pass")
418
419
420 def run_mandelbrot(w):
421 rows, cols = w["m"], w["n"]
422 dev_out = wp.zeros((rows, cols), dtype=float)
423
424 def launch():
425 wp.launch(mandelbrot_escape, dim=(rows, cols), inputs=[dev_out])
426
427 launch()
428 wp.synchronize()
429 out = dev_out.numpy().reshape(-1)
430 for cell in oracle_cells(rows * cols):
431 r, c = divmod(cell, cols)
432 cx = np.float32(-2.0 + c * (2.5 / cols))
433 cy = np.float32(-1.25 + r * (2.5 / rows))
434 zx, zy, count = cx, cy, 0.0
435 for _ in range(64):
436 nzx = zx * zx - zy * zy + cx
437 nzy = 2 * zx * zy + cy
438 count += 1.0
439 zx, zy = np.float32(nzx), np.float32(nzy)
440 if not (zx * zx + zy * zy < 4.0):
441 break
442 if not close_abs_rel(count, out[cell], 1.05, 0.0):
443 print(f"verify failed for {w['name']}", file=sys.stderr)
444 return
445 measure(launch, w, "mandelbrot_escape", note="thread-per-pixel early exit")
446
447
448 def run_warp2d(w):
449 import math
450 rows, cols = w["m"], w["n"]
451 cells = rows * cols
452 x = fill(cells, 0)
453 dev_x = wp.array(x, dtype=float)
454 dev_out = wp.zeros((rows, cols), dtype=float)
455
456 inv_scale = 1.0 / 1.15
457 ca_d = math.cos(0.35) * inv_scale
458 sa_d = math.sin(0.35) * inv_scale
459 cx_d = 0.5 * (cols - 1)
460 cy_d = 0.5 * (rows - 1)
461 ca = float(np.float32(ca_d))
462 sa = float(np.float32(sa_d))
463 tx = float(np.float32(cx_d - ca_d * cx_d - sa_d * cy_d))
464 ty = float(np.float32(cy_d + sa_d * cx_d - ca_d * cy_d))
465
466 def launch():
467 wp.launch(warp2d_bilinear, dim=(rows, cols), inputs=[dev_x, dev_out, ca, sa, tx, ty])
468
469 launch()
470 wp.synchronize()
471 out = dev_out.numpy().reshape(-1)
472 x64 = x.astype(np.float64)
473 for cell in oracle_cells(cells):
474 r, c = divmod(cell, cols)
475 sx = c * ca + r * sa + tx
476 sy = r * ca - c * sa + ty
477 cxx = min(max(sx, 0.0), float(cols - 1))
478 cyy = min(max(sy, 0.0), float(rows - 1))
479 x0 = min(math.floor(cxx), cols - 2)
480 y0 = min(math.floor(cyy), rows - 2)
481 fx = cxx - x0
482 fy = cyy - y0
483 g00 = x64[y0 * cols + x0]
484 g01 = x64[y0 * cols + x0 + 1]
485 g10 = x64[(y0 + 1) * cols + x0]
486 g11 = x64[(y0 + 1) * cols + x0 + 1]
487 expected = (g00 * (1.0 - fx) + g01 * fx) * (1.0 - fy) + (g10 * (1.0 - fx) + g11 * fx) * fy
488 if not close_abs_rel(expected, out[cell], 5e-3, 5e-3):
489 print(f"verify failed for {w['name']}", file=sys.stderr)
490 return
491 measure(launch, w, "warp2d_bilinear", note="thread-per-pixel bilinear")
492
493
494 def run_layernorm(w):
495 rows, cols = w["m"], w["n"]
496 cells = rows * cols
497 x = fill(cells, 0).reshape(rows, cols)
498 gamma = fill(cols, cells)
499 beta = fill(cols, cells + cols)
500 dev_x = wp.array(x, dtype=float)
501 dev_gamma = wp.array(gamma, dtype=float)
502 dev_beta = wp.array(beta, dtype=float)
503 dev_out = wp.zeros((rows, cols), dtype=float)
504
505 def launch():
506 wp.launch(layernorm_rows, dim=rows, inputs=[dev_x, dev_gamma, dev_beta, dev_out])
507
508 launch()
509 wp.synchronize()
510 out = dev_out.numpy().reshape(-1)
511 x64 = x.astype(np.float64)
512 g64 = gamma.astype(np.float64)
513 b64 = beta.astype(np.float64)
514 for cell in oracle_cells(cells):
515 r, c = divmod(cell, cols)
516 mean = x64[r].mean()
517 var = ((x64[r] - mean) ** 2).mean()
518 expected = float((x64[r, c] - mean) / np.sqrt(var + 1e-5) * g64[c] + b64[c])
519 if not close_abs_rel(expected, out[cell], 1e-3, 1e-3):
520 print(f"verify failed for {w['name']}", file=sys.stderr)
521 return
522 measure(launch, w, "layernorm_rows", note="thread-per-row three-pass")
523
524
525 def run_nbody(w):
526 bodies = w["m"]
527 px = fill(bodies, 0)
528 py = fill(bodies, bodies)
529 pz = fill(bodies, 2 * bodies)
530 mass = fill(bodies, 3 * bodies)
531 dev = [wp.array(v, dtype=float) for v in (px, py, pz, mass)]
532 dev_out = wp.zeros(3 * bodies, dtype=float)
533
534 def launch():
535 wp.launch(nbody_accel, dim=bodies, inputs=[*dev, dev_out])
536
537 launch()
538 wp.synchronize()
539 out = dev_out.numpy()
540 px64, py64, pz64, mass64 = (v.astype(np.float64) for v in (px, py, pz, mass))
541 for flat in oracle_cells(3 * bodies):
542 axis, body = divmod(flat, bodies)
543 dx = px64 - px64[body]
544 dy = py64 - py64[body]
545 dz = pz64 - pz64[body]
546 r2 = dx * dx + dy * dy + dz * dz + NBODY_SOFTENING
547 weight = mass64 / (r2 * np.sqrt(r2))
548 delta = dx if axis == 0 else dy if axis == 1 else dz
549 expected = float(np.sum(delta * weight))
550 if not close_abs_rel(expected, out[flat], 5e-2, 5e-3):
551 print(f"verify failed for {w['name']}", file=sys.stderr)
552 return
553 measure(launch, w, "nbody_accel", note="thread-per-body serial pairs")
554
555
556 def run_stencil(w):
557 rows, cols = w["m"], w["n"]
558 x = fill(rows * cols, 0).reshape(rows, cols)
559 dev_x = wp.array(x, dtype=float)
560 dev_out = wp.zeros((rows, cols), dtype=float)
561
562 def launch():
563 wp.launch(stencil5, dim=(rows, cols), inputs=[dev_x, dev_out])
564
565 launch()
566 wp.synchronize()
567 out = dev_out.numpy().reshape(-1)
568 x64 = x.astype(np.float64)
569 for cell in oracle_cells(rows * cols):
570 r, c = divmod(cell, cols)
571 up = x64[r - 1, c] if r > 0 else 0.0
572 down = x64[r + 1, c] if r + 1 < rows else 0.0
573 left = x64[r, c - 1] if c > 0 else 0.0
574 right = x64[r, c + 1] if c + 1 < cols else 0.0
575 if not close_abs_rel(0.25 * (up + down + left + right), out[cell], 1e-5, 1e-5):
576 print(f"verify failed for {w['name']}", file=sys.stderr)
577 return
578 measure(launch, w, "stencil5", note="thread-per-cell global loads")
579
580
581 def run_raymarch(w):
582 rows, cols = w["m"], w["n"]
583 dev_out = wp.zeros((rows, cols), dtype=float)
584
585 def launch():
586 wp.launch(raymarch, dim=(rows, cols), inputs=[dev_out])
587
588 launch()
589 wp.synchronize()
590 out = dev_out.numpy().reshape(-1)
591 for cell in oracle_cells(rows * cols):
592 r, c = divmod(cell, cols)
593 u = (c + 0.5) * (2.0 / cols) - 1.0
594 v = 1.0 - (r + 0.5) * (2.0 / rows)
595 raw_z = 1.4
596 inv_len = 1.0 / np.sqrt(u * u + v * v + raw_z * raw_z)
597 dx, dy, dz = u * inv_len, v * inv_len, raw_z * inv_len
598 t = 0.0
599 for _ in range(RAYMARCH_STEPS):
600 p = (dx * t, 1.2 + dy * t, dz * t)
601 d = p[1]
602 for cx, cy, cz, radius in SPHERES:
603 d = min(d, np.sqrt((p[0] - cx) ** 2 + (p[1] - cy) ** 2 + (p[2] - cz) ** 2) - radius)
604 t = min(t + max(d, 0.0), RAYMARCH_FAR)
605 if not close_abs_rel(t, out[cell], 2e-3, 1e-3):
606 print(f"verify failed for {w['name']}", file=sys.stderr)
607 return
608 measure(launch, w, "raymarch", note="per-ray early exit")
609
610
611 def run_ewchain(w):
612 x = fill(w["elements"], 0)
613 dev_x = wp.array(x, dtype=float)
614 dev_out = wp.zeros(w["elements"], dtype=float)
615
616 def launch():
617 wp.launch(ewchain, dim=w["elements"], inputs=[dev_x, dev_out])
618
619 launch()
620 wp.synchronize()
621 out = dev_out.numpy()
622 x64 = x.astype(np.float64)
623 for cell in oracle_cells(w["elements"]):
624 expected = x64[cell]
625 value = expected
626 for _ in range(EWCHAIN_LINKS):
627 value = np.tanh(value * value + expected)
628 if not close_abs_rel(float(value), out[cell], 1e-3, 1e-3):
629 print(f"verify failed for {w['name']}", file=sys.stderr)
630 return
631 measure(launch, w, "ewchain")
632
633
634 def main():
635 name_filter = sys.argv[1] if len(sys.argv) > 1 else None
636 wp.init()
637 print(json.dumps({
638 "schema": SCHEMA,
639 "system": "warp",
640 "meta": True,
641 "device": str(wp.get_device()),
642 "detail": f"warp={wp.__version__}, numpy={np.__version__}, timing=wall(launch+synchronize)",
643 }), flush=True)
644
645 runners = {
646 "warp2d": run_warp2d,
647 "softmax": run_softmax,
648 "layernorm": run_layernorm,
649 "mandelbrot": run_mandelbrot,
650 "rmsnorm": run_rmsnorm,
651 "nbody": run_nbody,
652 "stencil": run_stencil,
653 "raymarch": run_raymarch,
654 "ewchain": run_ewchain,
655 }
656 for w in BATTERY:
657 if name_filter and w["name"] != name_filter:
658 continue
659 start = time.perf_counter_ns()
660 runners[w["kind"]](w)
661 _ = start
662
663
664 if __name__ == "__main__":
665 main()