lib/accy/src/profiling/versus/baseline/torchlane.py
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 """PyTorch 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 + torch.cuda.synchronize, median over
5 samples. Eager production kernels only (ATen, cuDNN, SDPA); tf32 is
6 disabled to stay in the exact fp32 math tier. Workloads with no
7 production torch op (nbody, raymarch, mandelbrot) are not covered.
8
9 Run inside a venv with torch installed:
10 python torchlane.py [workload-name] > torch.jsonl
11 """
12
13 import json
14 import math
15 import sys
16 import time
17
18 import numpy as np
19 import torch
20 import torch.nn.functional as F
21
22 EWCHAIN_LINKS = 4
23 SCHEMA = "accy.versus/v2"
24 ORACLE_CELLS = 256
25
26 SAMPLES = 20
27 DISPATCH_SAMPLES = 50
28 PIPELINE_ROUNDS = 5
29 PIPELINE_BATCH = 16
30 WARMUP = 3
31
32 BATTERY = [
33 {"name": "matmul_f32_256", "kind": "matmul", "m": 256, "n": 256, "k": 256},
34 {"name": "matmul_f32_512", "kind": "matmul", "m": 512, "n": 512, "k": 512},
35 {"name": "matmul_f32_1024", "kind": "matmul", "m": 1024, "n": 1024, "k": 1024},
36 {"name": "matmul_f32_2048", "kind": "matmul", "m": 2048, "n": 2048, "k": 2048},
37 {"name": "matmul_f32_4096", "kind": "matmul", "m": 4096, "n": 4096, "k": 4096},
38 {"name": "dense_f32_1024", "kind": "dense", "m": 1024, "n": 1024, "k": 1024},
39 {"name": "ewchain_f32_16m", "kind": "ewchain", "elements": 16 * 1024 * 1024},
40 {"name": "reduce_f32_16m", "kind": "reduce", "elements": 16 * 1024 * 1024},
41 {"name": "scan_f32_16m", "kind": "scan", "elements": 16 * 1024 * 1024},
42 {"name": "softmax_f32_4096x4096", "kind": "softmax", "m": 4096, "n": 4096},
43 {"name": "layernorm_f32_4096x4096", "kind": "layernorm", "m": 4096, "n": 4096},
44 {"name": "rmsnorm_f32_4096x4096", "kind": "rmsnorm", "m": 4096, "n": 4096},
45 {"name": "attention_f32_4096x64", "kind": "attention", "m": 4096, "k": 64},
46 {"name": "stencil5_f32_4096x4096", "kind": "stencil", "m": 4096, "n": 4096},
47 {"name": "warp2d_f32_2048x2048", "kind": "warp2d", "m": 2048, "n": 2048},
48 ]
49
50
51 def flops(w):
52 if w["kind"] == "matmul":
53 return 2 * w["m"] * w["n"] * w["k"]
54 if w["kind"] == "dense":
55 return 2 * w["m"] * w["n"] * w["k"] + 2 * w["m"] * w["n"]
56 if w["kind"] == "ewchain":
57 return 3 * EWCHAIN_LINKS * w["elements"]
58 if w["kind"] == "softmax":
59 return 5 * w["m"] * w["n"]
60 if w["kind"] == "scan":
61 return w["elements"]
62 if w["kind"] == "layernorm":
63 return 8 * w["m"] * w["n"]
64 if w["kind"] == "rmsnorm":
65 return 5 * w["m"] * w["n"]
66 if w["kind"] == "attention":
67 return 4 * w["m"] * w["m"] * w["k"] + 5 * w["m"] * w["m"]
68 if w["kind"] == "stencil":
69 return 4 * w["m"] * w["n"]
70 if w["kind"] == "warp2d":
71 return 24 * w["m"] * w["n"]
72 return w["elements"]
73
74
75 def moved_bytes(w):
76 s = 4
77 if w["kind"] == "matmul":
78 return s * (w["m"] * w["k"] + w["k"] * w["n"] + w["m"] * w["n"])
79 if w["kind"] == "dense":
80 return s * (w["m"] * w["k"] + w["k"] * w["n"] + w["n"] + w["m"] * w["n"])
81 if w["kind"] == "ewchain":
82 return s * 2 * w["elements"]
83 if w["kind"] == "softmax":
84 return s * 2 * w["m"] * w["n"]
85 if w["kind"] == "scan":
86 return s * 2 * w["elements"]
87 if w["kind"] == "layernorm":
88 return s * (2 * w["m"] * w["n"] + 2 * w["n"])
89 if w["kind"] == "rmsnorm":
90 return s * (2 * w["m"] * w["n"] + w["n"])
91 if w["kind"] == "attention":
92 return s * 4 * w["m"] * w["k"]
93 if w["kind"] == "stencil":
94 return s * 2 * w["m"] * w["n"]
95 if w["kind"] == "warp2d":
96 return s * 2 * w["m"] * w["n"]
97 return s * (w["elements"] + 1)
98
99
100 def fill(count, offset):
101 index = np.arange(offset, offset + count, dtype=np.uint64)
102 state = index.astype(np.uint32)
103 state = state * np.uint32(1664525) + np.uint32(1013904223)
104 state = state * np.uint32(1664525) + np.uint32(1013904223)
105 return (state >> np.uint32(8)).astype(np.float32) / np.float32(1 << 24) - np.float32(0.5)
106
107
108 def make_row(system, w, metric, samples, values_ns, kernel=None, note=None):
109 ordered = sorted(values_ns)
110 row = {
111 "schema": SCHEMA,
112 "system": system,
113 "workload": w["name"],
114 "metric": metric,
115 "status": "ok",
116 "oracle": "passed",
117 "samples": samples,
118 "median_ns": int(ordered[len(ordered) // 2]),
119 "p10_ns": int(ordered[len(ordered) // 10]),
120 "p90_ns": int(ordered[(len(ordered) * 9) // 10]),
121 "flops": flops(w),
122 "moved_bytes": moved_bytes(w),
123 }
124 if kernel:
125 row["kernel"] = kernel
126 if note:
127 row["note"] = note
128 return row
129
130
131 def emit(w, metric, samples, values_ns, kernel=None, note=None):
132 print(json.dumps(make_row("torch", w, metric, samples, values_ns,
133 kernel=kernel, note=note)), flush=True)
134
135
136 def oracle_cells(total):
137 return [oracle_cell_index(s, total) for s in range(ORACLE_CELLS)]
138
139
140 def oracle_cell_index(sample, total):
141 if total <= 1:
142 return 0
143 if sample == 0:
144 return 0
145 if sample == 1:
146 return total - 1
147 if sample == 2:
148 return total // 2
149 if sample == 3:
150 return total // 3
151 if sample == 4:
152 return (total // 3) * 2
153 return oracle_mix(sample - 5 + total) % total
154
155
156 def oracle_mix(value):
157 mask = (1 << 64) - 1
158 mixed = (value + 0x9E3779B97F4A7C15) & mask
159 mixed = ((mixed ^ (mixed >> 30)) * 0xBF58476D1CE4E5B9) & mask
160 mixed = ((mixed ^ (mixed >> 27)) * 0x94D049BB133111EB) & mask
161 return mixed ^ (mixed >> 31)
162
163
164 def close(expected, actual, tolerance):
165 magnitude = max(abs(expected), 1.0)
166 return abs(expected - float(actual)) <= tolerance * magnitude
167
168
169 def close_abs_rel(expected, actual, abs_tol, rel_tol):
170 return abs(expected - float(actual)) <= abs_tol + rel_tol * abs(expected)
171
172
173 def measure(fn, w, kernel_name, note=None):
174 for _ in range(WARMUP):
175 fn()
176 torch.cuda.synchronize()
177 latency = []
178 for _ in range(SAMPLES):
179 start = time.perf_counter_ns()
180 fn()
181 torch.cuda.synchronize()
182 latency.append(time.perf_counter_ns() - start)
183 emit(w, "kernel_ns", SAMPLES, latency, kernel=kernel_name, note=note)
184 emit(w, "latency_ns", SAMPLES, latency)
185
186 dispatch = []
187 for _ in range(DISPATCH_SAMPLES):
188 start = time.perf_counter_ns()
189 fn()
190 dispatch.append(time.perf_counter_ns() - start)
191 torch.cuda.synchronize()
192 emit(w, "dispatch_ns", DISPATCH_SAMPLES, dispatch)
193
194 pipeline = []
195 for _ in range(PIPELINE_ROUNDS):
196 fn()
197 torch.cuda.synchronize()
198 start = time.perf_counter_ns()
199 for _ in range(PIPELINE_BATCH):
200 fn()
201 torch.cuda.synchronize()
202 pipeline.append((time.perf_counter_ns() - start) // PIPELINE_BATCH)
203 emit(w, "pipeline_ns", PIPELINE_ROUNDS, pipeline)
204
205
206 def fail(w):
207 print(f"verify failed for {w['name']}", file=sys.stderr)
208
209
210 def run_matmul(w, device):
211 a = fill(w["m"] * w["k"], 0).reshape(w["m"], w["k"])
212 b = fill(w["k"] * w["n"], w["m"] * w["k"]).reshape(w["k"], w["n"])
213 dev_a = torch.from_numpy(a).to(device)
214 dev_b = torch.from_numpy(b).to(device)
215
216 out = torch.mm(dev_a, dev_b).cpu().numpy().reshape(-1)
217 a64 = a.astype(np.float64)
218 b64 = b.astype(np.float64)
219 for cell in oracle_cells(w["m"] * w["n"]):
220 row, col = divmod(cell, w["n"])
221 if not close(float(a64[row] @ b64[:, col]), out[cell], 1e-2):
222 return fail(w)
223 measure(lambda: torch.mm(dev_a, dev_b), w, "aten mm", note="cublas fp32")
224
225
226 def run_dense(w, device):
227 a = fill(w["m"] * w["k"], 0).reshape(w["m"], w["k"])
228 b = fill(w["k"] * w["n"], w["m"] * w["k"]).reshape(w["k"], w["n"])
229 bias = fill(w["n"], w["m"] * w["k"] + w["k"] * w["n"])
230 dev_a = torch.from_numpy(a).to(device)
231 dev_b = torch.from_numpy(b).to(device)
232 dev_bias = torch.from_numpy(bias).to(device)
233
234 out = torch.tanh(torch.addmm(dev_bias, dev_a, dev_b)).cpu().numpy().reshape(-1)
235 a64 = a.astype(np.float64)
236 b64 = b.astype(np.float64)
237 for cell in oracle_cells(w["m"] * w["n"]):
238 row, col = divmod(cell, w["n"])
239 expected = np.tanh(float(a64[row] @ b64[:, col]) + float(bias[col]))
240 if not close(expected, out[cell], 1e-2):
241 return fail(w)
242 measure(lambda: torch.tanh(torch.addmm(dev_bias, dev_a, dev_b)), w,
243 "aten addmm+tanh", note="two launches per iteration")
244
245
246 def run_ewchain(w, device):
247 x = fill(w["elements"], 0)
248 dev_x = torch.from_numpy(x).to(device)
249
250 def f():
251 r = dev_x
252 for _ in range(EWCHAIN_LINKS):
253 r = torch.tanh(r * r + dev_x)
254 return r
255
256 out = f().cpu().numpy()
257 x64 = x.astype(np.float64)
258 for cell in oracle_cells(w["elements"]):
259 expected = x64[cell]
260 value = expected
261 for _ in range(EWCHAIN_LINKS):
262 value = np.tanh(value * value + expected)
263 if not close(float(value), out[cell], 1e-3):
264 return fail(w)
265 measure(f, w, "eager op chain", note="unfused, 12 kernels")
266
267
268 def run_reduce(w, device):
269 x = fill(w["elements"], 0)
270 dev_x = torch.from_numpy(x).to(device)
271
272 result = float(torch.sum(dev_x))
273 expected = float(np.sum(x.astype(np.float64)))
274 if not close(expected, result, 5e-2):
275 return fail(w)
276 measure(lambda: torch.sum(dev_x), w, "aten sum")
277
278
279 def run_scan(w, device):
280 n = w["elements"]
281 x = fill(n, 0)
282 dev_x = torch.from_numpy(x).to(device)
283
284 out = torch.cumsum(dev_x, 0).cpu().numpy()
285 prefix = np.cumsum(x.astype(np.float64))
286 next_check = 1
287 while next_check <= n:
288 i = next_check - 1
289 if not close_abs_rel(float(prefix[i]), out[i], 1e-2, 1e-4):
290 return fail(w)
291 next_check = next_check * 7 // 2 + 13
292 if not close_abs_rel(float(prefix[n - 1]), out[n - 1], 1e-2, 1e-4):
293 return fail(w)
294 measure(lambda: torch.cumsum(dev_x, 0), w, "aten cumsum")
295
296
297 def run_softmax(w, device):
298 rows, cols = w["m"], w["n"]
299 x = fill(rows * cols, 0).reshape(rows, cols)
300 dev_x = torch.from_numpy(x).to(device)
301
302 out = F.softmax(dev_x, dim=1).cpu().numpy().reshape(-1)
303 x64 = x.astype(np.float64)
304 for cell in oracle_cells(rows * cols):
305 row, col = divmod(cell, cols)
306 shifted = x64[row] - x64[row].max()
307 expected = float(np.exp(shifted[col]) / np.exp(shifted).sum())
308 if not close_abs_rel(expected, out[cell], 1e-7, 1e-3):
309 return fail(w)
310 measure(lambda: F.softmax(dev_x, dim=1), w, "aten softmax")
311
312
313 def run_layernorm(w, device):
314 rows, cols = w["m"], w["n"]
315 cells = rows * cols
316 x = fill(cells, 0).reshape(rows, cols)
317 gamma = fill(cols, cells)
318 beta = fill(cols, cells + cols)
319 epsilon = 1e-5
320 dev_x = torch.from_numpy(x).to(device)
321 dev_gamma = torch.from_numpy(gamma).to(device)
322 dev_beta = torch.from_numpy(beta).to(device)
323
324 out = F.layer_norm(dev_x, (cols,), dev_gamma, dev_beta, epsilon).cpu().numpy().reshape(-1)
325 x64 = x.astype(np.float64)
326 g64 = gamma.astype(np.float64)
327 b64 = beta.astype(np.float64)
328 for cell in oracle_cells(cells):
329 row, col = divmod(cell, cols)
330 mean = x64[row].mean()
331 var = ((x64[row] - mean) ** 2).mean()
332 expected = float((x64[row, col] - mean) / np.sqrt(var + epsilon) * g64[col] + b64[col])
333 if not close_abs_rel(expected, out[cell], 1e-3, 1e-3):
334 return fail(w)
335 measure(lambda: F.layer_norm(dev_x, (cols,), dev_gamma, dev_beta, epsilon),
336 w, "aten layer_norm")
337
338
339 def run_rmsnorm(w, device):
340 if not hasattr(F, "rms_norm"):
341 print("torch has no rms_norm; skipping", file=sys.stderr)
342 return
343 rows, cols = w["m"], w["n"]
344 cells = rows * cols
345 x = fill(cells, 0).reshape(rows, cols)
346 gamma = fill(cols, cells)
347 epsilon = 1e-5
348 dev_x = torch.from_numpy(x).to(device)
349 dev_gamma = torch.from_numpy(gamma).to(device)
350
351 out = F.rms_norm(dev_x, (cols,), dev_gamma, epsilon).cpu().numpy().reshape(-1)
352 x64 = x.astype(np.float64)
353 g64 = gamma.astype(np.float64)
354 for cell in oracle_cells(cells):
355 row, col = divmod(cell, cols)
356 mean_sq = (x64[row] ** 2).mean()
357 expected = float(x64[row, col] / np.sqrt(mean_sq + epsilon) * g64[col])
358 if not close_abs_rel(expected, out[cell], 1e-3, 1e-3):
359 return fail(w)
360 measure(lambda: F.rms_norm(dev_x, (cols,), dev_gamma, epsilon), w, "aten rms_norm")
361
362
363 def run_attention(w, device):
364 seq, dim = w["m"], w["k"]
365 q = fill(seq * dim, 0).reshape(seq, dim)
366 kt = fill(seq * dim, seq * dim).reshape(dim, seq)
367 v = fill(seq * dim, 2 * seq * dim).reshape(seq, dim)
368 scale = 1.0 / np.sqrt(dim)
369 dev_q = torch.from_numpy(q).to(device).view(1, 1, seq, dim)
370 dev_k = torch.from_numpy(kt.T.copy()).to(device).view(1, 1, seq, dim)
371 dev_v = torch.from_numpy(v).to(device).view(1, 1, seq, dim)
372
373 out = F.scaled_dot_product_attention(dev_q, dev_k, dev_v).cpu().numpy().reshape(-1)
374 q64 = q.astype(np.float64)
375 kt64 = kt.astype(np.float64)
376 v64 = v.astype(np.float64)
377 for cell in oracle_cells(seq * dim):
378 row, out_col = divmod(cell, dim)
379 scores = (q64[row] @ kt64) * scale
380 shifted = np.exp(scores - scores.max())
381 probs = shifted / shifted.sum()
382 expected = float(probs @ v64[:, out_col])
383 if not close_abs_rel(expected, out[cell], 1e-3, 5e-3):
384 return fail(w)
385 measure(lambda: F.scaled_dot_product_attention(dev_q, dev_k, dev_v), w,
386 "sdpa", note="auto backend, fp32")
387
388
389 def run_stencil(w, device):
390 rows, cols = w["m"], w["n"]
391 x = fill(rows * cols, 0).reshape(rows, cols)
392 dev_x = torch.from_numpy(x).to(device).view(1, 1, rows, cols)
393 weight = torch.tensor(
394 [[0.0, 0.25, 0.0], [0.25, 0.0, 0.25], [0.0, 0.25, 0.0]],
395 dtype=torch.float32, device=device,
396 ).view(1, 1, 3, 3)
397
398 out = F.conv2d(dev_x, weight, padding=1).cpu().numpy().reshape(-1)
399 x64 = x.astype(np.float64)
400 for cell in oracle_cells(rows * cols):
401 row, col = divmod(cell, cols)
402 up = x64[row - 1, col] if row > 0 else 0.0
403 down = x64[row + 1, col] if row + 1 < rows else 0.0
404 left = x64[row, col - 1] if col > 0 else 0.0
405 right = x64[row, col + 1] if col + 1 < cols else 0.0
406 if not close_abs_rel(0.25 * (up + down + left + right), out[cell], 1e-5, 1e-5):
407 return fail(w)
408 measure(lambda: F.conv2d(dev_x, weight, padding=1), w,
409 "aten conv2d", note="cudnn, 3x3 cross weights")
410
411
412 def run_warp2d(w, device):
413 rows, cols = w["m"], w["n"]
414 x = fill(rows * cols, 0)
415 dev_x = torch.from_numpy(x.reshape(rows, cols)).to(device).view(1, 1, rows, cols)
416
417 inv_scale = 1.0 / 1.15
418 ca_d = math.cos(0.35) * inv_scale
419 sa_d = math.sin(0.35) * inv_scale
420 cx_d = 0.5 * (cols - 1)
421 cy_d = 0.5 * (rows - 1)
422 ca = np.float32(ca_d)
423 sa = np.float32(sa_d)
424 tx = np.float32(cx_d - ca_d * cx_d - sa_d * cy_d)
425 ty = np.float32(cy_d + sa_d * cx_d - ca_d * cy_d)
426
427 iy, ix = np.meshgrid(
428 np.arange(rows, dtype=np.float32),
429 np.arange(cols, dtype=np.float32),
430 indexing="ij",
431 )
432 sx = ix * ca + iy * sa + tx
433 sy = iy * ca - ix * sa + ty
434 grid = np.stack(
435 [sx / np.float32(cols - 1) * 2.0 - 1.0, sy / np.float32(rows - 1) * 2.0 - 1.0],
436 axis=-1,
437 ).astype(np.float32)
438 dev_grid = torch.from_numpy(grid).to(device).view(1, rows, cols, 2)
439
440 def f():
441 return F.grid_sample(dev_x, dev_grid, mode="bilinear",
442 padding_mode="border", align_corners=True)
443
444 out = f().cpu().numpy().reshape(-1)
445 x64 = x.astype(np.float64)
446 for cell in oracle_cells(rows * cols):
447 row, col = divmod(cell, cols)
448 sxx = col * float(ca) + row * float(sa) + float(tx)
449 syy = row * float(ca) - col * float(sa) + float(ty)
450 cxx = min(max(sxx, 0.0), float(cols - 1))
451 cyy = min(max(syy, 0.0), float(rows - 1))
452 x0 = min(math.floor(cxx), cols - 2)
453 y0 = min(math.floor(cyy), rows - 2)
454 fx = cxx - x0
455 fy = cyy - y0
456 g00 = x64[y0 * cols + x0]
457 g01 = x64[y0 * cols + x0 + 1]
458 g10 = x64[(y0 + 1) * cols + x0]
459 g11 = x64[(y0 + 1) * cols + x0 + 1]
460 expected = (g00 * (1.0 - fx) + g01 * fx) * (1.0 - fy) + (g10 * (1.0 - fx) + g11 * fx) * fy
461 if not close_abs_rel(expected, out[cell], 5e-3, 5e-3):
462 return fail(w)
463 measure(f, w, "aten grid_sample", note="precomputed affine grid input, border clamp")
464
465
466 def main():
467 name_filter = sys.argv[1] if len(sys.argv) > 1 else None
468 assert torch.cuda.is_available()
469 device = torch.device("cuda:0")
470 torch.backends.cuda.matmul.allow_tf32 = False
471 torch.backends.cudnn.allow_tf32 = False
472 print(json.dumps({
473 "schema": SCHEMA,
474 "system": "torch",
475 "meta": True,
476 "device": torch.cuda.get_device_name(device),
477 "detail": f"torch={torch.__version__}, eager fp32, tf32 off, "
478 "timing=wall(dispatch+synchronize)",
479 }), flush=True)
480
481 runners = {
482 "matmul": run_matmul,
483 "dense": run_dense,
484 "ewchain": run_ewchain,
485 "reduce": run_reduce,
486 "scan": run_scan,
487 "softmax": run_softmax,
488 "layernorm": run_layernorm,
489 "rmsnorm": run_rmsnorm,
490 "attention": run_attention,
491 "stencil": run_stencil,
492 "warp2d": run_warp2d,
493 }
494 for w in BATTERY:
495 if name_filter and w["name"] != name_filter:
496 continue
497 runners[w["kind"]](w, device)
498
499
500 if __name__ == "__main__":
501 main()