lib/accy/src/profiling/versus/runner.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const accy = @import("accy");
4 const bench = @import("bench");
5 const sys = @import("sys");
6
7 const jsonl = @import("jsonl.zig");
8 const workload_mod = @import("workload.zig");
9
10 const coz = bench.coz;
11
12 const Timer = struct {
13 start_ns: i128,
14
15 fn start() Timer {
16 return .{ .start_ns = sys.time.nanoTimestamp() };
17 }
18
19 fn reset(self: *Timer) void {
20 self.start_ns = sys.time.nanoTimestamp();
21 }
22
23 fn read(self: *const Timer) u64 {
24 const end = sys.time.nanoTimestamp();
25 if (end <= self.start_ns) return 0;
26 const elapsed = end - self.start_ns;
27 if (elapsed > std.math.maxInt(u64)) return std.math.maxInt(u64);
28 return @intCast(elapsed);
29 }
30 };
31
32 const Allocator = std.mem.Allocator;
33 const Workload = workload_mod.Workload;
34
35 pub const system_name = "accy";
36
37 pub const Options = struct {
38 samples: u32 = 20,
39 dispatch_samples: u32 = 50,
40 pipeline_batch: u32 = 16,
41 pipeline_rounds: u32 = 5,
42 warmup: u32 = 3,
43 compile_repeats: u32 = 4,
44 workload: ?[]const u8 = null,
45 math_tier: gpu.BackendMathTier = .exact,
46 verify: bool = true,
47 dump_ptx_dir: ?[]const u8 = null,
48 };
49
50 pub const oracle_cells = 256;
51
52 fn compileAndLoadModule(
53 allocator: Allocator,
54 handle: gpu.BackendHandle,
55 module: *accy.choir.SemanticModule,
56 options: accy.executable.FragmentCompilerOptions,
57 ) !*accy.executable.LoadedFragment {
58 const compiled = try accy.executable.compileFragmentFromSemanticModule(allocator, handle, module, options);
59 return try accy.executable.loadFragment(allocator, handle, compiled, options);
60 }
61
62 pub fn run(arena: Allocator, backing: Allocator, out: *std.Io.Writer, options: Options) !u8 {
63 const system = systemName(options.math_tier);
64 var state = gpu.cuda.State.initDevice(backing, 0) catch |err| switch (err) {
65 error.RuntimeUnavailable => {
66 try jsonl.writeMeta(out, system, "unavailable", "CUDA driver or device unavailable");
67 try out.flush();
68 return 3;
69 },
70 else => return err,
71 };
72 defer state.deinit();
73
74 try jsonl.writeMeta(out, system, "cuda:0", "compile=compileFragmentFromSemanticModule load=loadFragment");
75
76 for (workload_mod.battery) |workload| {
77 if (options.workload) |filter| {
78 if (!std.mem.eql(u8, filter, workload.name)) continue;
79 }
80 runOne(arena, backing, out, &state, workload, options) catch |err| {
81 var note_buffer: [96]u8 = undefined;
82 const note = std.fmt.bufPrint(note_buffer[0..], "error: {t}", .{err}) catch "error";
83 try jsonl.writeRow(out, .{
84 .system = system,
85 .workload = workload.name,
86 .metric = .latency_ns,
87 .status = .failed,
88 .oracle = .skipped,
89 .samples = 0,
90 .median_ns = 0,
91 .p10_ns = 0,
92 .p90_ns = 0,
93 .flops = workload.flops(),
94 .moved_bytes = workload.movedBytes(),
95 .note = note,
96 });
97 };
98 try out.flush();
99 }
100 return 0;
101 }
102
103 fn systemName(math_tier: gpu.BackendMathTier) []const u8 {
104 return switch (math_tier) {
105 .exact => system_name,
106 .tf32_tensor => "accy_tf32",
107 };
108 }
109
110 pub fn buildModule(allocator: Allocator, workload: Workload, scene: workload_mod.Scene) !*accy.choir.SemanticModule {
111 var builder = try accy.choir.SemanticBuilder.init(allocator, accy.choir.SemanticBuilder.ContextLimits.standard);
112 defer builder.deinit();
113
114 switch (workload.kind) {
115 .matmul => {
116 const lhs_ty = try builder.tensor(.f32, &.{ workload.m, workload.k });
117 const rhs_ty = try builder.tensor(.f32, &.{ workload.k, workload.n });
118 const out_ty = try builder.tensor(.f32, &.{ workload.m, workload.n });
119 var function = try builder.beginFunction("matmul", &.{ lhs_ty, rhs_ty }, &.{out_ty});
120 const product = try function.dotGeneral(function.parameter(0), function.parameter(1), out_ty, &.{1}, &.{0}, &.{}, &.{});
121 try function.return_(&.{product});
122 try function.finish();
123 },
124 .dense => {
125 const lhs_ty = try builder.tensor(.f32, &.{ workload.m, workload.k });
126 const rhs_ty = try builder.tensor(.f32, &.{ workload.k, workload.n });
127 const bias_ty = try builder.tensor(.f32, &.{workload.n});
128 const out_ty = try builder.tensor(.f32, &.{ workload.m, workload.n });
129 var function = try builder.beginFunction("dense", &.{ lhs_ty, rhs_ty, bias_ty }, &.{out_ty});
130 const product = try function.dotGeneral(function.parameter(0), function.parameter(1), out_ty, &.{1}, &.{0}, &.{}, &.{});
131 const bias = try function.broadcastInDim(function.parameter(2), out_ty, &.{ workload.m, workload.n }, &.{1});
132 const sum = try function.add(product, bias);
133 const activated = try function.tanh(sum);
134 try function.return_(&.{activated});
135 try function.finish();
136 },
137 .ewchain => {
138 const vec_ty = try builder.tensor(.f32, &.{@as(i64, @intCast(workload.elements))});
139 var function = try builder.beginFunction("ewchain", &.{vec_ty}, &.{vec_ty});
140 const x = function.parameter(0);
141 var value = x;
142 var link: u32 = 0;
143 while (link < workload_mod.ewchain_links) : (link += 1) {
144 const squared = try function.mul(value, value);
145 const shifted = try function.add(squared, x);
146 value = try function.tanh(shifted);
147 }
148 try function.return_(&.{value});
149 try function.finish();
150 },
151 .reduce => {
152 const vec_ty = try builder.tensor(.f32, &.{@as(i64, @intCast(workload.elements))});
153 const scalar_ty = try builder.tensor(.f32, &.{});
154 var function = try builder.beginFunction("reduce", &.{vec_ty}, &.{scalar_ty});
155 const zero = try function.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
156 const total = try function.reduce(function.parameter(0), zero, scalar_ty, "sum", &.{0});
157 try function.return_(&.{total});
158 try function.finish();
159 },
160 .scan => {
161 const count: i64 = @intCast(workload.elements);
162 const vec_ty = try builder.tensor(.f32, &.{count});
163 var function = try builder.beginFunction("scan", &.{vec_ty}, &.{vec_ty});
164 const out = try function.cumsum(function.parameter(0), vec_ty, 0);
165 try function.return_(&.{out});
166 try function.finish();
167 },
168 .softmax => {
169 const rows: i64 = @intCast(workload.m);
170 const cols: i64 = @intCast(workload.n);
171 const x_ty = try builder.tensor(.f32, &.{ rows, cols });
172 const row_ty = try builder.tensor(.f32, &.{rows});
173 const scalar_ty = try builder.tensor(.f32, &.{});
174 var function = try builder.beginFunction("softmax", &.{x_ty}, &.{x_ty});
175 const x = function.parameter(0);
176 const lowest = try function.constant(scalar_ty, std.mem.asBytes(&@as(f32, -std.math.floatMax(f32))));
177 const row_max = try function.reduce(x, lowest, row_ty, "max", &.{1});
178 const shifted = try function.sub(x, try function.broadcastInDim(row_max, x_ty, &.{ rows, cols }, &.{0}));
179 const exps = try function.exp(shifted);
180 const zero = try function.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
181 const row_sum = try function.reduce(exps, zero, row_ty, "sum", &.{1});
182 const out = try function.div(exps, try function.broadcastInDim(row_sum, x_ty, &.{ rows, cols }, &.{0}));
183 try function.return_(&.{out});
184 try function.finish();
185 },
186 .layernorm => {
187 const rows: i64 = @intCast(workload.m);
188 const cols: i64 = @intCast(workload.n);
189 const x_ty = try builder.tensor(.f32, &.{ rows, cols });
190 const row_ty = try builder.tensor(.f32, &.{rows});
191 const col_ty = try builder.tensor(.f32, &.{cols});
192 const scalar_ty = try builder.tensor(.f32, &.{});
193 var function = try builder.beginFunction("layernorm", &.{ x_ty, col_ty, col_ty }, &.{x_ty});
194 const x = function.parameter(0);
195 const gamma = function.parameter(1);
196 const beta = function.parameter(2);
197
198 const zero = try function.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
199 const row_sum = try function.reduce(x, zero, row_ty, "sum", &.{1});
200 var inv_cols_value: f32 = 1.0 / @as(f32, @floatFromInt(cols));
201 const inv_cols_values = try allocator.alloc(f32, @intCast(rows));
202 defer allocator.free(inv_cols_values);
203 @memset(inv_cols_values, inv_cols_value);
204 _ = &inv_cols_value;
205 const inv_cols = try function.constant(row_ty, std.mem.sliceAsBytes(inv_cols_values));
206 const mean = try function.mul(row_sum, inv_cols);
207 const mean_b = try function.broadcastInDim(mean, x_ty, &.{ rows, cols }, &.{0});
208 const centered = try function.sub(x, mean_b);
209 const centered_squared = try function.mul(centered, centered);
210 const square_sum = try function.reduce(centered_squared, zero, row_ty, "sum", &.{1});
211 const variance = try function.mul(square_sum, inv_cols);
212 const eps_values = try allocator.alloc(f32, @intCast(rows));
213 defer allocator.free(eps_values);
214 @memset(eps_values, layernorm_epsilon);
215 const eps = try function.constant(row_ty, std.mem.sliceAsBytes(eps_values));
216 const variance_eps = try function.add(variance, eps);
217 const stddev = try function.sqrt(variance_eps);
218 const one_values = try allocator.alloc(f32, @intCast(rows));
219 defer allocator.free(one_values);
220 @memset(one_values, 1.0);
221 const one = try function.constant(row_ty, std.mem.sliceAsBytes(one_values));
222 const inv_std = try function.div(one, stddev);
223 const inv_std_b = try function.broadcastInDim(inv_std, x_ty, &.{ rows, cols }, &.{0});
224 const normalized = try function.mul(centered, inv_std_b);
225 const gamma_b = try function.broadcastInDim(gamma, x_ty, &.{ rows, cols }, &.{1});
226 const scaled = try function.mul(normalized, gamma_b);
227 const beta_b = try function.broadcastInDim(beta, x_ty, &.{ rows, cols }, &.{1});
228 const out = try function.add(scaled, beta_b);
229 try function.return_(&.{out});
230 try function.finish();
231 },
232 .rmsnorm => {
233 const rows: i64 = @intCast(workload.m);
234 const cols: i64 = @intCast(workload.n);
235 const x_ty = try builder.tensor(.f32, &.{ rows, cols });
236 const row_ty = try builder.tensor(.f32, &.{rows});
237 const col_ty = try builder.tensor(.f32, &.{cols});
238 const scalar_ty = try builder.tensor(.f32, &.{});
239 var function = try builder.beginFunction("rmsnorm", &.{ x_ty, col_ty }, &.{x_ty});
240 const x = function.parameter(0);
241 const gamma = function.parameter(1);
242
243 const zero = try function.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
244 const squared = try function.mul(x, x);
245 const square_sum = try function.reduce(squared, zero, row_ty, "sum", &.{1});
246 const inv_cols_values = try allocator.alloc(f32, @intCast(rows));
247 defer allocator.free(inv_cols_values);
248 @memset(inv_cols_values, 1.0 / @as(f32, @floatFromInt(cols)));
249 const inv_cols = try function.constant(row_ty, std.mem.sliceAsBytes(inv_cols_values));
250 const mean_sq = try function.mul(square_sum, inv_cols);
251 const eps_values = try allocator.alloc(f32, @intCast(rows));
252 defer allocator.free(eps_values);
253 @memset(eps_values, layernorm_epsilon);
254 const eps = try function.constant(row_ty, std.mem.sliceAsBytes(eps_values));
255 const mean_eps = try function.add(mean_sq, eps);
256 const rms = try function.sqrt(mean_eps);
257 const one_values = try allocator.alloc(f32, @intCast(rows));
258 defer allocator.free(one_values);
259 @memset(one_values, 1.0);
260 const one = try function.constant(row_ty, std.mem.sliceAsBytes(one_values));
261 const inv_rms = try function.div(one, rms);
262 const inv_rms_b = try function.broadcastInDim(inv_rms, x_ty, &.{ rows, cols }, &.{0});
263 const normalized = try function.mul(x, inv_rms_b);
264 const gamma_b = try function.broadcastInDim(gamma, x_ty, &.{ rows, cols }, &.{1});
265 const out = try function.mul(normalized, gamma_b);
266 try function.return_(&.{out});
267 try function.finish();
268 },
269 .attention => {
270 const seq: i64 = @intCast(workload.m);
271 const dim: i64 = @intCast(workload.k);
272 const qkv_ty = try builder.tensor(.f32, &.{ seq, dim });
273 const kt_ty = try builder.tensor(.f32, &.{ dim, seq });
274 const scores_ty = try builder.tensor(.f32, &.{ seq, seq });
275 const row_ty = try builder.tensor(.f32, &.{seq});
276 const scalar_ty = try builder.tensor(.f32, &.{});
277 var function = try builder.beginFunction("attention", &.{ qkv_ty, kt_ty, qkv_ty }, &.{qkv_ty});
278 const q = function.parameter(0);
279 const kt = function.parameter(1);
280 const v = function.parameter(2);
281
282 const raw_scores = try function.dotGeneral(q, kt, scores_ty, &.{1}, &.{0}, &.{}, &.{});
283 const scale_value: f32 = 1.0 / @sqrt(@as(f32, @floatFromInt(dim)));
284 const scale_scalar = try function.constant(scalar_ty, std.mem.asBytes(&scale_value));
285 const scale = try function.broadcast(scale_scalar, scores_ty, &.{ seq, seq });
286 const scores = try function.mul(raw_scores, scale);
287
288 const lowest = try function.constant(scalar_ty, std.mem.asBytes(&@as(f32, -std.math.floatMax(f32))));
289 const row_max = try function.reduce(scores, lowest, row_ty, "max", &.{1});
290 const shifted = try function.sub(scores, try function.broadcastInDim(row_max, scores_ty, &.{ seq, seq }, &.{0}));
291 const exps = try function.exp(shifted);
292 const zero = try function.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
293 const row_sum = try function.reduce(exps, zero, row_ty, "sum", &.{1});
294 const probs = try function.div(exps, try function.broadcastInDim(row_sum, scores_ty, &.{ seq, seq }, &.{0}));
295
296 const out = try function.dotGeneral(probs, v, qkv_ty, &.{1}, &.{0}, &.{}, &.{});
297 try function.return_(&.{out});
298 try function.finish();
299 },
300 .nbody => {
301 const n: i64 = @intCast(workload.m);
302 const vec_ty = try builder.tensor(.f32, &.{n});
303 const pair_ty = try builder.tensor(.f32, &.{ n, n });
304 const out_ty = try builder.tensor(.f32, &.{3 * n});
305 const scalar_ty = try builder.tensor(.f32, &.{});
306 var function = try builder.beginFunction("nbody", &.{ vec_ty, vec_ty, vec_ty, vec_ty }, &.{out_ty});
307 const px = function.parameter(0);
308 const py = function.parameter(1);
309 const pz = function.parameter(2);
310 const mass = function.parameter(3);
311 const dx = try function.sub(
312 try function.broadcastInDim(px, pair_ty, &.{ n, n }, &.{1}),
313 try function.broadcastInDim(px, pair_ty, &.{ n, n }, &.{0}),
314 );
315 const dy = try function.sub(
316 try function.broadcastInDim(py, pair_ty, &.{ n, n }, &.{1}),
317 try function.broadcastInDim(py, pair_ty, &.{ n, n }, &.{0}),
318 );
319 const dz = try function.sub(
320 try function.broadcastInDim(pz, pair_ty, &.{ n, n }, &.{1}),
321 try function.broadcastInDim(pz, pair_ty, &.{ n, n }, &.{0}),
322 );
323 const r2 = try function.add(
324 try function.add(try function.mul(dx, dx), try function.mul(dy, dy)),
325 try function.mul(dz, dz),
326 );
327 const eps = try function.constant(scalar_ty, std.mem.asBytes(&workload_mod.nbody_softening));
328 const r2e = try function.add(r2, try function.broadcast(eps, pair_ty, &.{ n, n }));
329 const denom = try function.mul(r2e, try function.sqrt(r2e));
330 const mass_j = try function.broadcastInDim(mass, pair_ty, &.{ n, n }, &.{1});
331 const weight = try function.div(mass_j, denom);
332 const zero = try function.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
333 const ax = try function.reduce(try function.mul(dx, weight), zero, vec_ty, "sum", &.{1});
334 const ay = try function.reduce(try function.mul(dy, weight), zero, vec_ty, "sum", &.{1});
335 const az = try function.reduce(try function.mul(dz, weight), zero, vec_ty, "sum", &.{1});
336 const out = try function.concatenate(&.{ ax, ay, az }, out_ty, 0);
337 try function.return_(&.{out});
338 try function.finish();
339 },
340 .stencil => {
341 const rows: i64 = @intCast(workload.m);
342 const cols: i64 = @intCast(workload.n);
343 const x_ty = try builder.tensor(.f32, &.{ rows, cols });
344 const pad_ty = try builder.tensor(.f32, &.{ rows + 2, cols + 2 });
345 const scalar_ty = try builder.tensor(.f32, &.{});
346 var function = try builder.beginFunction("stencil", &.{x_ty}, &.{x_ty});
347 const x = function.parameter(0);
348 const zero = try function.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.0)));
349 const padded = try function.pad(x, zero, pad_ty, &.{ 1, 1 }, &.{ 1, 1 }, &.{ 0, 0 });
350 const up = try function.slice(padded, x_ty, &.{ 0, 1 }, &.{ rows, cols + 1 }, &.{ 1, 1 });
351 const down = try function.slice(padded, x_ty, &.{ 2, 1 }, &.{ rows + 2, cols + 1 }, &.{ 1, 1 });
352 const left = try function.slice(padded, x_ty, &.{ 1, 0 }, &.{ rows + 1, cols }, &.{ 1, 1 });
353 const right = try function.slice(padded, x_ty, &.{ 1, 2 }, &.{ rows + 1, cols + 2 }, &.{ 1, 1 });
354 const quarter = try function.constant(scalar_ty, std.mem.asBytes(&@as(f32, 0.25)));
355 const sum = try function.add(try function.add(up, down), try function.add(left, right));
356 const out = try function.mul(sum, try function.broadcast(quarter, x_ty, &.{ rows, cols }));
357 try function.return_(&.{out});
358 try function.finish();
359 },
360 .warp2d => {
361 const rows: i64 = @intCast(workload.m);
362 const cols: i64 = @intCast(workload.n);
363 const img_ty = try builder.tensor(.f32, &.{ rows, cols });
364 const flat_ty = try builder.tensor(.f32, &.{rows * cols});
365 const idx_ty = try builder.tensor(.i32, &.{ rows, cols });
366 const scalar_ty = try builder.tensor(.f32, &.{});
367 var function = try builder.beginFunction("warp2d", &.{flat_ty}, &.{img_ty});
368
369 const Consts = struct {
370 function: *@TypeOf(function),
371 img_ty: accy.choir.ir.Type,
372 scalar_ty: accy.choir.ir.Type,
373 rows: i64,
374 cols: i64,
375
376 fn splat(self: *@This(), value: f32) !*accy.choir.ir.Value {
377 const scalar = try self.function.constant(self.scalar_ty, std.mem.asBytes(&value));
378 return self.function.broadcast(scalar, self.img_ty, &.{ self.rows, self.cols });
379 }
380 };
381 var consts = Consts{ .function = &function, .img_ty = img_ty, .scalar_ty = scalar_ty, .rows = rows, .cols = cols };
382
383 const coeffs = workload_mod.warp2dCoeffs(workload.m, workload.n);
384 const ix = try function.iota(img_ty, 1);
385 const iy = try function.iota(img_ty, 0);
386 const sx = try function.add(
387 try function.add(
388 try function.mul(ix, try consts.splat(coeffs.ca)),
389 try function.mul(iy, try consts.splat(coeffs.sa)),
390 ),
391 try consts.splat(coeffs.tx),
392 );
393 const sy = try function.add(
394 try function.sub(
395 try function.mul(iy, try consts.splat(coeffs.ca)),
396 try function.mul(ix, try consts.splat(coeffs.sa)),
397 ),
398 try consts.splat(coeffs.ty),
399 );
400 const zero = try consts.splat(0.0);
401 const cxx = try function.min(try function.max(sx, zero), try consts.splat(@floatFromInt(cols - 1)));
402 const cyy = try function.min(try function.max(sy, zero), try consts.splat(@floatFromInt(rows - 1)));
403 const x0 = try function.min(try function.floor(cxx), try consts.splat(@floatFromInt(cols - 2)));
404 const y0 = try function.min(try function.floor(cyy), try consts.splat(@floatFromInt(rows - 2)));
405 const fx = try function.sub(cxx, x0);
406 const fy = try function.sub(cyy, y0);
407 const one = try consts.splat(1.0);
408 const gx = try function.sub(one, fx);
409 const gy = try function.sub(one, fy);
410 const width = try consts.splat(@floatFromInt(cols));
411 const base = try function.add(try function.mul(y0, width), x0);
412 const idx00 = try function.convert(base, idx_ty, .i32);
413 const idx01 = try function.convert(try function.add(base, one), idx_ty, .i32);
414 const below = try function.add(base, width);
415 const idx10 = try function.convert(below, idx_ty, .i32);
416 const idx11 = try function.convert(try function.add(below, one), idx_ty, .i32);
417 const src = function.parameter(0);
418 const g00 = try function.gather(src, idx00, img_ty, 0);
419 const g01 = try function.gather(src, idx01, img_ty, 0);
420 const g10 = try function.gather(src, idx10, img_ty, 0);
421 const g11 = try function.gather(src, idx11, img_ty, 0);
422 const top = try function.add(try function.mul(g00, gx), try function.mul(g01, fx));
423 const bottom = try function.add(try function.mul(g10, gx), try function.mul(g11, fx));
424 const out = try function.add(try function.mul(top, gy), try function.mul(bottom, fy));
425 try function.return_(&.{out});
426 try function.finish();
427 },
428 .raymarch => {
429 const rows: i64 = @intCast(workload.m);
430 const cols: i64 = @intCast(workload.n);
431 const img_ty = try builder.tensor(.f32, &.{ rows, cols });
432 const scalar_ty = try builder.tensor(.f32, &.{});
433 var function = try builder.beginFunction("raymarch", &.{img_ty}, &.{img_ty});
434
435 const Consts = struct {
436 function: *@TypeOf(function),
437 img_ty: accy.choir.ir.Type,
438 scalar_ty: accy.choir.ir.Type,
439 rows: i64,
440 cols: i64,
441
442 fn splat(self: *@This(), value: f32) !*accy.choir.ir.Value {
443 const scalar = try self.function.constant(self.scalar_ty, std.mem.asBytes(&value));
444 return self.function.broadcast(scalar, self.img_ty, &.{ self.rows, self.cols });
445 }
446 };
447 var consts = Consts{ .function = &function, .img_ty = img_ty, .scalar_ty = scalar_ty, .rows = rows, .cols = cols };
448
449 const ix = try function.iota(img_ty, 1);
450 const iy = try function.iota(img_ty, 0);
451 const half = try consts.splat(0.5);
452 const one = try consts.splat(1.0);
453 const u = try function.sub(try function.mul(try function.add(ix, half), try consts.splat(2.0 / @as(f32, @floatFromInt(cols)))), one);
454 const v = try function.sub(one, try function.mul(try function.add(iy, half), try consts.splat(2.0 / @as(f32, @floatFromInt(rows)))));
455 const raw_z = try consts.splat(1.4);
456 const norm2 = try function.add(
457 try function.add(try function.mul(u, u), try function.mul(v, v)),
458 try function.mul(raw_z, raw_z),
459 );
460 const inv_len = try function.div(one, try function.sqrt(norm2));
461 const dx = try function.mul(u, inv_len);
462 const dy = try function.mul(v, inv_len);
463 const dz = try function.mul(raw_z, inv_len);
464 var it = try function.beginIterate(&.{function.parameter(0)}, workload_mod.raymarch_steps);
465 const body = it.inner();
466 const BodyConsts = struct {
467 body: *@TypeOf(function),
468 img_ty: accy.choir.ir.Type,
469 scalar_ty: accy.choir.ir.Type,
470 rows: i64,
471 cols: i64,
472
473 fn splat(self: *@This(), value: f32) !*accy.choir.ir.Value {
474 const scalar = try self.body.constant(self.scalar_ty, std.mem.asBytes(&value));
475 return self.body.broadcast(scalar, self.img_ty, &.{ self.rows, self.cols });
476 }
477 };
478 var body_consts = BodyConsts{ .body = body, .img_ty = img_ty, .scalar_ty = scalar_ty, .rows = rows, .cols = cols };
479 const pred_ty = try builder.tensor(.i1, &.{ rows, cols });
480
481 const oy = try body_consts.splat(1.2);
482 const zero = try body_consts.splat(0.0);
483 const far = try body_consts.splat(workload_mod.raymarch_scene.max_distance);
484 const eps = try body_consts.splat(workload_mod.raymarch_epsilon);
485
486 const t = it.carry(0);
487 const px = try body.mul(dx, t);
488 const py = try body.add(oy, try body.mul(dy, t));
489 const pz = try body.mul(dz, t);
490 const d = try emitSceneDistance(body, img_ty, scalar_ty, rows, cols, scene, px, py, pz);
491 const advanced = try body.min(try body.add(t, try body.max(d, zero)), far);
492 const margin = try body.min(try body.sub(d, eps), try body.sub(far, advanced));
493 const active = try body.compare(margin, zero, pred_ty, .gt);
494 try it.yield_(active, &.{advanced});
495
496 try function.return_(&.{it.result(0)});
497 try function.finish();
498 },
499 .mandelbrot => {
500 const rows: i64 = @intCast(workload.m);
501 const cols: i64 = @intCast(workload.n);
502 const img_ty = try builder.tensor(.f32, &.{ rows, cols });
503 const scalar_ty = try builder.tensor(.f32, &.{});
504 const pred_ty = try builder.tensor(.i1, &.{ rows, cols });
505 var function = try builder.beginFunction("mandelbrot", &.{img_ty}, &.{img_ty});
506
507 const Consts = struct {
508 function: *@TypeOf(function),
509 img_ty: accy.choir.ir.Type,
510 scalar_ty: accy.choir.ir.Type,
511 rows: i64,
512 cols: i64,
513
514 fn splat(self: *@This(), value: f32) !*accy.choir.ir.Value {
515 const scalar = try self.function.constant(self.scalar_ty, std.mem.asBytes(&value));
516 return self.function.broadcast(scalar, self.img_ty, &.{ self.rows, self.cols });
517 }
518 };
519 var consts = Consts{ .function = &function, .img_ty = img_ty, .scalar_ty = scalar_ty, .rows = rows, .cols = cols };
520
521 const ix = try function.iota(img_ty, 1);
522 const iy = try function.iota(img_ty, 0);
523 const x_step = try consts.splat(scene.mandelbrot_span / @as(f32, @floatFromInt(cols)));
524 const y_step = try consts.splat(scene.mandelbrot_span / @as(f32, @floatFromInt(rows)));
525 const cx = try function.add(try consts.splat(scene.mandelbrot_x0), try function.mul(ix, x_step));
526 const cy = try function.add(try consts.splat(scene.mandelbrot_y0), try function.mul(iy, y_step));
527 const zeros = try function.mul(cx, try consts.splat(0.0));
528
529 var it = try function.beginIterate(&.{ cx, cy, zeros }, workload_mod.mandelbrot_max_iters);
530 const body = it.inner();
531 const BodyConsts = struct {
532 body: *@TypeOf(function),
533 img_ty: accy.choir.ir.Type,
534 scalar_ty: accy.choir.ir.Type,
535 rows: i64,
536 cols: i64,
537
538 fn splat(self: *@This(), value: f32) !*accy.choir.ir.Value {
539 const scalar = try self.body.constant(self.scalar_ty, std.mem.asBytes(&value));
540 return self.body.broadcast(scalar, self.img_ty, &.{ self.rows, self.cols });
541 }
542 };
543 var body_consts = BodyConsts{ .body = body, .img_ty = img_ty, .scalar_ty = scalar_ty, .rows = rows, .cols = cols };
544 const zx = it.carry(0);
545 const zy = it.carry(1);
546 const count = it.carry(2);
547 const zx2 = try body.mul(zx, zx);
548 const zy2 = try body.mul(zy, zy);
549 const zxzy = try body.mul(zx, zy);
550 const new_zx = try body.add(try body.sub(zx2, zy2), cx);
551 const new_zy = try body.add(try body.add(zxzy, zxzy), cy);
552 const one_body = try body_consts.splat(1.0);
553 const new_count = try body.add(count, one_body);
554 const mag = try body.add(try body.mul(new_zx, new_zx), try body.mul(new_zy, new_zy));
555 const escape = try body_consts.splat(workload_mod.mandelbrot_escape);
556 const active = try body.compare(mag, escape, pred_ty, .lt);
557 try it.yield_(active, &.{ new_zx, new_zy, new_count });
558
559 try function.return_(&.{it.result(2)});
560 try function.finish();
561 },
562 }
563 return try builder.finish();
564 }
565
566 pub fn emitSplat(
567 function: *accy.choir.semantic.FunctionBuilder,
568 img_ty: accy.choir.ir.Type,
569 scalar_ty: accy.choir.ir.Type,
570 rows: i64,
571 cols: i64,
572 value: f32,
573 ) !*accy.choir.ir.Value {
574 const scalar = try function.constant(scalar_ty, std.mem.asBytes(&value));
575 return function.broadcast(scalar, img_ty, &.{ rows, cols });
576 }
577
578 pub fn emitSceneDistance(
579 function: *accy.choir.semantic.FunctionBuilder,
580 img_ty: accy.choir.ir.Type,
581 scalar_ty: accy.choir.ir.Type,
582 rows: i64,
583 cols: i64,
584 scene: workload_mod.Scene,
585 px: *accy.choir.ir.Value,
586 py: *accy.choir.ir.Value,
587 pz: *accy.choir.ir.Value,
588 ) !*accy.choir.ir.Value {
589 var d = py;
590 for (scene.spheres) |sphere| {
591 const sx = try function.sub(px, try emitSplat(function, img_ty, scalar_ty, rows, cols, sphere[0]));
592 const sy = try function.sub(py, try emitSplat(function, img_ty, scalar_ty, rows, cols, sphere[1]));
593 const sz = try function.sub(pz, try emitSplat(function, img_ty, scalar_ty, rows, cols, sphere[2]));
594 const dist2 = try function.add(
595 try function.add(try function.mul(sx, sx), try function.mul(sy, sy)),
596 try function.mul(sz, sz),
597 );
598 const ds = try function.sub(try function.sqrt(dist2), try emitSplat(function, img_ty, scalar_ty, rows, cols, sphere[3]));
599 d = try function.min(d, ds);
600 }
601 return d;
602 }
603
604 pub const HostBuffers = struct {
605 inputs: [4][]f32 = .{ &.{}, &.{}, &.{}, &.{} },
606 input_count: usize = 0,
607 output: []f32 = &.{},
608
609 pub fn deinit(self: *HostBuffers, allocator: Allocator) void {
610 for (self.inputs[0..self.input_count]) |input| allocator.free(input);
611 allocator.free(self.output);
612 self.* = .{};
613 }
614 };
615
616 pub fn allocHostBuffers(allocator: Allocator, workload: Workload) !HostBuffers {
617 var buffers = HostBuffers{};
618 errdefer buffers.deinit(allocator);
619
620 switch (workload.kind) {
621 .matmul => {
622 buffers.inputs[0] = try allocator.alloc(f32, @as(usize, workload.m) * workload.k);
623 buffers.inputs[1] = try allocator.alloc(f32, @as(usize, workload.k) * workload.n);
624 buffers.input_count = 2;
625 buffers.output = try allocator.alloc(f32, @as(usize, workload.m) * workload.n);
626 },
627 .dense => {
628 buffers.inputs[0] = try allocator.alloc(f32, @as(usize, workload.m) * workload.k);
629 buffers.inputs[1] = try allocator.alloc(f32, @as(usize, workload.k) * workload.n);
630 buffers.inputs[2] = try allocator.alloc(f32, workload.n);
631 buffers.input_count = 3;
632 buffers.output = try allocator.alloc(f32, @as(usize, workload.m) * workload.n);
633 },
634 .ewchain, .reduce, .scan => {
635 buffers.inputs[0] = try allocator.alloc(f32, workload.elements);
636 buffers.input_count = 1;
637 buffers.output = try allocator.alloc(f32, if (workload.kind == .reduce) 1 else workload.elements);
638 },
639 .softmax, .stencil, .raymarch, .mandelbrot, .warp2d => {
640 const cells = @as(usize, workload.m) * workload.n;
641 buffers.inputs[0] = try allocator.alloc(f32, cells);
642 buffers.input_count = 1;
643 buffers.output = try allocator.alloc(f32, cells);
644 },
645 .layernorm => {
646 const cells = @as(usize, workload.m) * workload.n;
647 buffers.inputs[0] = try allocator.alloc(f32, cells);
648 buffers.inputs[1] = try allocator.alloc(f32, workload.n);
649 buffers.inputs[2] = try allocator.alloc(f32, workload.n);
650 buffers.input_count = 3;
651 buffers.output = try allocator.alloc(f32, cells);
652 },
653 .rmsnorm => {
654 const cells = @as(usize, workload.m) * workload.n;
655 buffers.inputs[0] = try allocator.alloc(f32, cells);
656 buffers.inputs[1] = try allocator.alloc(f32, workload.n);
657 buffers.input_count = 2;
658 buffers.output = try allocator.alloc(f32, cells);
659 },
660 .attention => {
661 const cells = @as(usize, workload.m) * workload.k;
662 buffers.inputs[0] = try allocator.alloc(f32, cells);
663 buffers.inputs[1] = try allocator.alloc(f32, cells);
664 buffers.inputs[2] = try allocator.alloc(f32, cells);
665 buffers.input_count = 3;
666 buffers.output = try allocator.alloc(f32, cells);
667 },
668 .nbody => {
669 const bodies: usize = workload.m;
670 for (buffers.inputs[0..4]) |*input| {
671 input.* = try allocator.alloc(f32, bodies);
672 buffers.input_count += 1;
673 }
674 buffers.output = try allocator.alloc(f32, 3 * bodies);
675 },
676 }
677
678 var offset: u64 = 0;
679 for (buffers.inputs[0..buffers.input_count]) |input| {
680 workload_mod.fill(input, offset);
681 offset += input.len;
682 }
683 return buffers;
684 }
685
686 fn dumpKernelArtifacts(
687 arena: Allocator,
688 fragment: *accy.executable.LoadedFragment,
689 workload: Workload,
690 dir: []const u8,
691 ) !void {
692 var kernel_index: usize = 0;
693 while (kernel_index < fragment.kernelCount()) : (kernel_index += 1) {
694 var artifact = try fragment.copyKernelArtifact(arena, kernel_index);
695 defer artifact.deinit();
696 const text = switch (artifact.payload) {
697 .text => |text| text,
698 else => continue,
699 };
700 const file_path = try std.fmt.allocPrint(arena, "{s}/{s}_{d}_{s}.ptx", .{ dir, workload.name, kernel_index, artifact.entry_name });
701 sys.fs.writeFile(file_path, text) catch continue;
702 }
703 }
704
705 fn oracleCellIndex(sample: usize, total: u64) u64 {
706 if (total <= 1) return 0;
707 return switch (sample) {
708 0 => 0,
709 1 => total - 1,
710 2 => total / 2,
711 3 => total / 3,
712 4 => (total / 3) * 2,
713 else => oracleMix(@as(u64, @intCast(sample - 5)) +% total) % total,
714 };
715 }
716
717 fn oracleMix(value: u64) u64 {
718 var mixed = value +% 0x9E3779B97F4A7C15;
719 mixed ^= mixed >> 30;
720 mixed *%= 0xBF58476D1CE4E5B9;
721 mixed ^= mixed >> 27;
722 mixed *%= 0x94D049BB133111EB;
723 mixed ^= mixed >> 31;
724 return mixed;
725 }
726
727 fn matmulOracle(workload: Workload, math_tier: gpu.BackendMathTier, lhs: []const f32, rhs: []const f32, cell: u64) f64 {
728 const row = cell / workload.n;
729 const col = cell % workload.n;
730 var acc: f64 = 0;
731 var k: usize = 0;
732 while (k < workload.k) : (k += 1) {
733 const a = oracleInput(math_tier, lhs[row * workload.k + k]);
734 const b = oracleInput(math_tier, rhs[k * workload.n + col]);
735 acc += @as(f64, a) * @as(f64, b);
736 }
737 return acc;
738 }
739
740 fn oracleInput(math_tier: gpu.BackendMathTier, value: f32) f32 {
741 return switch (math_tier) {
742 .exact => value,
743 .tf32_tensor => accy.kernel.oracle.tf32RoundF32(value),
744 };
745 }
746
747 fn ewchainOracle(x: f64) f64 {
748 var value = x;
749 var link: u32 = 0;
750 while (link < workload_mod.ewchain_links) : (link += 1) {
751 value = std.math.tanh(value * value + x);
752 }
753 return value;
754 }
755
756 pub fn verifyOutput(workload: Workload, math_tier: gpu.BackendMathTier, scene: workload_mod.Scene, buffers: *const HostBuffers) !void {
757 switch (workload.kind) {
758 .matmul => {
759 const total = @as(u64, workload.m) * workload.n;
760 var sample: usize = 0;
761 while (sample < oracle_cells) : (sample += 1) {
762 const cell = oracleCellIndex(sample, total);
763 const expected = matmulOracle(workload, math_tier, buffers.inputs[0], buffers.inputs[1], cell);
764 try expectClose(expected, buffers.output[@intCast(cell)], 1e-2);
765 }
766 },
767 .dense => {
768 const total = @as(u64, workload.m) * workload.n;
769 var sample: usize = 0;
770 while (sample < oracle_cells) : (sample += 1) {
771 const cell = oracleCellIndex(sample, total);
772 const col = cell % workload.n;
773 const product = matmulOracle(workload, math_tier, buffers.inputs[0], buffers.inputs[1], cell);
774 const expected = std.math.tanh(product + @as(f64, buffers.inputs[2][@intCast(col)]));
775 try expectClose(expected, buffers.output[@intCast(cell)], 1e-2);
776 }
777 },
778 .ewchain => {
779 var sample: usize = 0;
780 while (sample < oracle_cells) : (sample += 1) {
781 const cell = oracleCellIndex(sample, workload.elements);
782 const expected = ewchainOracle(@as(f64, buffers.inputs[0][@intCast(cell)]));
783 try expectClose(expected, buffers.output[@intCast(cell)], 1e-3);
784 }
785 },
786 .reduce => {
787 var acc: f64 = 0;
788 for (buffers.inputs[0]) |value| acc += @as(f64, value);
789 try expectClose(acc, buffers.output[0], 5e-2);
790 },
791 .scan => {
792 const x = buffers.inputs[0];
793 var acc: f64 = 0;
794 var next_check: u64 = 1;
795 var sample: usize = 0;
796 var checked: usize = 0;
797 _ = &sample;
798 var index: u64 = 0;
799 while (index < workload.elements) : (index += 1) {
800 acc += @as(f64, x[@intCast(index)]);
801 if (index + 1 == next_check or index + 1 == workload.elements) {
802 try expectCloseAbsRel(acc, buffers.output[@intCast(index)], 1e-2, 1e-4);
803 next_check = next_check * 7 / 2 + 13;
804 checked += 1;
805 }
806 }
807 if (checked == 0) return error.OracleMismatch;
808 },
809 .softmax => {
810 const cols: u64 = workload.n;
811 var sample: usize = 0;
812 while (sample < oracle_cells) : (sample += 1) {
813 const cell = oracleCellIndex(sample, @as(u64, workload.m) * workload.n);
814 const row = cell / cols;
815 const x = buffers.inputs[0];
816 var row_max: f64 = -std.math.inf(f64);
817 var col: usize = 0;
818 while (col < cols) : (col += 1) {
819 row_max = @max(row_max, @as(f64, x[@intCast(row * cols + col)]));
820 }
821 var row_sum: f64 = 0;
822 col = 0;
823 while (col < cols) : (col += 1) {
824 row_sum += @exp(@as(f64, x[@intCast(row * cols + col)]) - row_max);
825 }
826 const expected = @exp(@as(f64, x[@intCast(cell)]) - row_max) / row_sum;
827 try expectCloseAbsRel(expected, buffers.output[@intCast(cell)], 1e-7, 1e-3);
828 }
829 },
830 .layernorm => {
831 const cols: u64 = workload.n;
832 var sample: usize = 0;
833 while (sample < oracle_cells) : (sample += 1) {
834 const cell = oracleCellIndex(sample, @as(u64, workload.m) * workload.n);
835 const row = cell / cols;
836 const x = buffers.inputs[0];
837 const gamma = buffers.inputs[1];
838 const beta = buffers.inputs[2];
839 var mean_acc: f64 = 0;
840 var col: usize = 0;
841 while (col < cols) : (col += 1) {
842 mean_acc += @as(f64, x[@intCast(row * cols + col)]);
843 }
844 const mean = mean_acc / @as(f64, @floatFromInt(cols));
845 var var_acc: f64 = 0;
846 col = 0;
847 while (col < cols) : (col += 1) {
848 const centered = @as(f64, x[@intCast(row * cols + col)]) - mean;
849 var_acc += centered * centered;
850 }
851 const inv_std = 1.0 / @sqrt(var_acc / @as(f64, @floatFromInt(cols)) + @as(f64, layernorm_epsilon));
852 const col_index: usize = @intCast(cell % cols);
853 const expected = (@as(f64, x[@intCast(cell)]) - mean) * inv_std * @as(f64, gamma[col_index]) + @as(f64, beta[col_index]);
854 try expectCloseAbsRel(expected, buffers.output[@intCast(cell)], 1e-3, 1e-3);
855 }
856 },
857 .rmsnorm => {
858 const cols: u64 = workload.n;
859 var sample: usize = 0;
860 while (sample < oracle_cells) : (sample += 1) {
861 const cell = oracleCellIndex(sample, @as(u64, workload.m) * workload.n);
862 const row = cell / cols;
863 const x = buffers.inputs[0];
864 const gamma = buffers.inputs[1];
865 var sq_acc: f64 = 0;
866 var col: usize = 0;
867 while (col < cols) : (col += 1) {
868 const value = @as(f64, x[@intCast(row * cols + col)]);
869 sq_acc += value * value;
870 }
871 const inv_rms = 1.0 / @sqrt(sq_acc / @as(f64, @floatFromInt(cols)) + @as(f64, layernorm_epsilon));
872 const col_index: usize = @intCast(cell % cols);
873 const expected = @as(f64, x[@intCast(cell)]) * inv_rms * @as(f64, gamma[col_index]);
874 try expectCloseAbsRel(expected, buffers.output[@intCast(cell)], 1e-3, 1e-3);
875 }
876 },
877 .attention => {
878 const seq: u64 = workload.m;
879 const dim: u64 = workload.k;
880 var sample: usize = 0;
881 while (sample < oracle_cells) : (sample += 1) {
882 const cell = oracleCellIndex(sample, seq * dim);
883 const row = cell / dim;
884 const out_col = cell % dim;
885 const q = buffers.inputs[0];
886 const kt = buffers.inputs[1];
887 const v = buffers.inputs[2];
888 const scale = 1.0 / @sqrt(@as(f64, @floatFromInt(dim)));
889
890 var row_max: f64 = -std.math.inf(f64);
891 var j: usize = 0;
892 while (j < seq) : (j += 1) {
893 var dot: f64 = 0;
894 var d: usize = 0;
895 while (d < dim) : (d += 1) {
896 dot += @as(f64, q[@intCast(row * dim + d)]) * @as(f64, kt[@intCast(d * seq + j)]);
897 }
898 row_max = @max(row_max, dot * scale);
899 }
900 var row_sum: f64 = 0;
901 var acc: f64 = 0;
902 j = 0;
903 while (j < seq) : (j += 1) {
904 var dot: f64 = 0;
905 var d: usize = 0;
906 while (d < dim) : (d += 1) {
907 dot += @as(f64, q[@intCast(row * dim + d)]) * @as(f64, kt[@intCast(d * seq + j)]);
908 }
909 const p = @exp(dot * scale - row_max);
910 row_sum += p;
911 acc += p * @as(f64, v[@intCast(j * dim + out_col)]);
912 }
913 const expected = acc / row_sum;
914 try expectCloseAbsRel(expected, buffers.output[@intCast(cell)], 1e-3, 5e-3);
915 }
916 },
917 .nbody => try verifyNbodyAccels(workload, buffers),
918 .stencil => {
919 const rows: u64 = workload.m;
920 const cols: u64 = workload.n;
921 const x = buffers.inputs[0];
922 var sample: usize = 0;
923 while (sample < oracle_cells) : (sample += 1) {
924 const cell = oracleCellIndex(sample, rows * cols);
925 const row = cell / cols;
926 const col = cell % cols;
927 const up: f64 = if (row > 0) x[@intCast((row - 1) * cols + col)] else 0;
928 const down: f64 = if (row + 1 < rows) x[@intCast((row + 1) * cols + col)] else 0;
929 const left: f64 = if (col > 0) x[@intCast(row * cols + col - 1)] else 0;
930 const right: f64 = if (col + 1 < cols) x[@intCast(row * cols + col + 1)] else 0;
931 const expected = 0.25 * (up + down + left + right);
932 try expectCloseAbsRel(expected, buffers.output[@intCast(cell)], 1e-5, 1e-5);
933 }
934 },
935 .raymarch => {
936 const rows: u64 = workload.m;
937 const cols: u64 = workload.n;
938 var sample: usize = 0;
939 while (sample < oracle_cells) : (sample += 1) {
940 const cell = oracleCellIndex(sample, rows * cols);
941 const row = cell / cols;
942 const col = cell % cols;
943 const expected = raymarchOracle(rows, cols, row, col, buffers.inputs[0][@intCast(cell)]);
944 try expectCloseAbsRel(expected, buffers.output[@intCast(cell)], 4e-3, 2e-3);
945 }
946 },
947 .warp2d => {
948 const rows: u64 = workload.m;
949 const cols: u64 = workload.n;
950 const coeffs = workload_mod.warp2dCoeffs(rows, cols);
951 const ca: f64 = coeffs.ca;
952 const sa: f64 = coeffs.sa;
953 const tx: f64 = coeffs.tx;
954 const ty: f64 = coeffs.ty;
955 const src = buffers.inputs[0];
956 var sample: usize = 0;
957 while (sample < oracle_cells) : (sample += 1) {
958 const cell = oracleCellIndex(sample, rows * cols);
959 const row: f64 = @floatFromInt(cell / cols);
960 const col: f64 = @floatFromInt(cell % cols);
961 const sx = col * ca + row * sa + tx;
962 const sy = row * ca - col * sa + ty;
963 const cxx = @min(@max(sx, 0.0), @as(f64, @floatFromInt(cols - 1)));
964 const cyy = @min(@max(sy, 0.0), @as(f64, @floatFromInt(rows - 1)));
965 const x0 = @min(@floor(cxx), @as(f64, @floatFromInt(cols - 2)));
966 const y0 = @min(@floor(cyy), @as(f64, @floatFromInt(rows - 2)));
967 const fx = cxx - x0;
968 const fy = cyy - y0;
969 const xi: u64 = @intFromFloat(x0);
970 const yi: u64 = @intFromFloat(y0);
971 const g00: f64 = src[@intCast(yi * cols + xi)];
972 const g01: f64 = src[@intCast(yi * cols + xi + 1)];
973 const g10: f64 = src[@intCast((yi + 1) * cols + xi)];
974 const g11: f64 = src[@intCast((yi + 1) * cols + xi + 1)];
975 const expected = (g00 * (1.0 - fx) + g01 * fx) * (1.0 - fy) + (g10 * (1.0 - fx) + g11 * fx) * fy;
976 try expectCloseAbsRel(expected, buffers.output[@intCast(cell)], 5e-3, 5e-3);
977 }
978 },
979 .mandelbrot => {
980 if (mandelbrotMismatches(workload, scene, buffers.output) != 0) return error.OracleMismatch;
981 },
982 }
983 }
984
985 pub const layernorm_epsilon: f32 = 1e-5;
986
987 pub const mandelbrot_frame_mismatch_budget: usize = 6;
988
989 pub fn verifyMandelbrotFrame(workload: Workload, scene: workload_mod.Scene, output: []const f32) !void {
990 const mismatches = mandelbrotMismatches(workload, scene, output);
991 if (mismatches > mandelbrot_frame_mismatch_budget) return error.OracleMismatch;
992 }
993
994 fn mandelbrotMismatches(workload: Workload, scene: workload_mod.Scene, output: []const f32) usize {
995 const rows: u64 = workload.m;
996 const cols: u64 = workload.n;
997 var mismatches: usize = 0;
998 var sample: usize = 0;
999 while (sample < oracle_cells) : (sample += 1) {
1000 const cell = oracleCellIndex(sample, rows * cols);
1001 const row = cell / cols;
1002 const col = cell % cols;
1003 const expected = mandelbrotOracle(rows, cols, row, col, scene);
1004 expectCloseAbsRel(expected, output[@intCast(cell)], 1.05, 0.0) catch {
1005 mismatches += 1;
1006 };
1007 }
1008 return mismatches;
1009 }
1010
1011 pub fn verifyNbodyAccels(workload: Workload, buffers: *const HostBuffers) !void {
1012 const bodies: u64 = workload.m;
1013 const px = buffers.inputs[0];
1014 const py = buffers.inputs[1];
1015 const pz = buffers.inputs[2];
1016 const mass = buffers.inputs[3];
1017 var sample: usize = 0;
1018 while (sample < oracle_cells) : (sample += 1) {
1019 const flat = oracleCellIndex(sample, 3 * bodies);
1020 const axis = flat / bodies;
1021 const body: usize = @intCast(flat % bodies);
1022 var acc: f64 = 0;
1023 var other: usize = 0;
1024 while (other < bodies) : (other += 1) {
1025 const dx = @as(f64, px[other]) - @as(f64, px[body]);
1026 const dy = @as(f64, py[other]) - @as(f64, py[body]);
1027 const dz = @as(f64, pz[other]) - @as(f64, pz[body]);
1028 const r2 = dx * dx + dy * dy + dz * dz + workload_mod.nbody_softening;
1029 const weight = @as(f64, mass[other]) / (r2 * @sqrt(r2));
1030 const delta = switch (axis) {
1031 0 => dx,
1032 1 => dy,
1033 else => dz,
1034 };
1035 acc += delta * weight;
1036 }
1037 try expectCloseAbsRel(acc, buffers.output[@intCast(flat)], 5e-2, 5e-3);
1038 }
1039 }
1040
1041 fn mandelbrotOracle(rows: u64, cols: u64, row: u64, col: u64, scene: workload_mod.Scene) f64 {
1042 const cx: f32 = scene.mandelbrot_x0 + @as(f32, @floatFromInt(col)) * (scene.mandelbrot_span / @as(f32, @floatFromInt(cols)));
1043 const cy: f32 = scene.mandelbrot_y0 + @as(f32, @floatFromInt(row)) * (scene.mandelbrot_span / @as(f32, @floatFromInt(rows)));
1044 var zx: f32 = cx;
1045 var zy: f32 = cy;
1046 var count: f32 = 0;
1047 var iters: usize = 0;
1048 while (iters < workload_mod.mandelbrot_max_iters) : (iters += 1) {
1049 const new_zx = zx * zx - zy * zy + cx;
1050 const new_zy = 2.0 * zx * zy + cy;
1051 count += 1;
1052 zx = new_zx;
1053 zy = new_zy;
1054 if (!(zx * zx + zy * zy < workload_mod.mandelbrot_escape)) break;
1055 }
1056 return count;
1057 }
1058
1059 fn raymarchOracle(rows: u64, cols: u64, row: u64, col: u64, t0: f32) f64 {
1060 const scene = workload_mod.raymarch_scene;
1061 const u = (@as(f64, @floatFromInt(col)) + 0.5) * (2.0 / @as(f64, @floatFromInt(cols))) - 1.0;
1062 const v = 1.0 - (@as(f64, @floatFromInt(row)) + 0.5) * (2.0 / @as(f64, @floatFromInt(rows)));
1063 const raw_z = 1.4;
1064 const inv_len = 1.0 / @sqrt(u * u + v * v + raw_z * raw_z);
1065 const dx = u * inv_len;
1066 const dy = v * inv_len;
1067 const dz = raw_z * inv_len;
1068 var t: f64 = @floatCast(t0);
1069 var step: u32 = 0;
1070 while (step < workload_mod.raymarch_steps) : (step += 1) {
1071 const px = dx * t;
1072 const py = 1.2 + dy * t;
1073 const pz = dz * t;
1074 var d = py;
1075 for (scene.spheres) |sphere| {
1076 const sx = px - sphere[0];
1077 const sy = py - sphere[1];
1078 const sz = pz - sphere[2];
1079 const ds = @sqrt(sx * sx + sy * sy + sz * sz) - sphere[3];
1080 d = @min(d, ds);
1081 }
1082 t = @min(t + @max(d, 0), scene.max_distance);
1083 const margin = @min(d - workload_mod.raymarch_epsilon, scene.max_distance - t);
1084 if (!(margin > 0)) break;
1085 }
1086 return t;
1087 }
1088
1089 fn expectCloseAbsRel(expected: f64, actual: f32, abs_tolerance: f64, rel_tolerance: f64) !void {
1090 const actual_wide: f64 = @floatCast(actual);
1091 if (@abs(expected - actual_wide) > abs_tolerance + rel_tolerance * @abs(expected)) {
1092 last_mismatch = .{ .expected = expected, .actual = actual_wide };
1093 return error.OracleMismatch;
1094 }
1095 }
1096
1097 const Mismatch = struct {
1098 expected: f64 = 0,
1099 actual: f64 = 0,
1100 };
1101
1102 var last_mismatch: Mismatch = .{};
1103
1104 fn expectClose(expected: f64, actual: f32, tolerance: f64) !void {
1105 const actual_wide: f64 = @floatCast(actual);
1106 const magnitude = @max(@abs(expected), 1.0);
1107 if (@abs(expected - actual_wide) > tolerance * magnitude) {
1108 last_mismatch = .{ .expected = expected, .actual = actual_wide };
1109 return error.OracleMismatch;
1110 }
1111 }
1112
1113 pub fn inputBytes(buffers: *const HostBuffers, storage: *[4][]const u8) []const []const u8 {
1114 for (buffers.inputs[0..buffers.input_count], 0..) |input, index| {
1115 storage[index] = std.mem.sliceAsBytes(input);
1116 }
1117 return storage[0..buffers.input_count];
1118 }
1119
1120 fn runOne(
1121 arena: Allocator,
1122 backing: Allocator,
1123 out: *std.Io.Writer,
1124 state: *gpu.cuda.State,
1125 workload: Workload,
1126 options: Options,
1127 ) !void {
1128 const system = systemName(options.math_tier);
1129 const flops = workload.flops();
1130 const moved_bytes = workload.movedBytes();
1131
1132 var timer = Timer.start();
1133
1134 timer.reset();
1135 const first_module = try buildModule(backing, workload, workload_mod.default_scene);
1136 var fragment = try compileAndLoadModule(backing, state.handle(), first_module, .{ .math_tier = options.math_tier });
1137 const compile_cold_ns = timer.read();
1138 defer fragment.deinit();
1139
1140 var compile_warm_stats: ?jsonl.Samples = null;
1141
1142 if (options.compile_repeats > 0) {
1143 const warm_samples = try arena.alloc(u64, options.compile_repeats);
1144 for (warm_samples) |*sample| {
1145 timer.reset();
1146 const module = try buildModule(backing, workload, workload_mod.default_scene);
1147 var warm_fragment = try compileAndLoadModule(backing, state.handle(), module, .{ .math_tier = options.math_tier });
1148 sample.* = timer.read();
1149 warm_fragment.deinit();
1150 }
1151 compile_warm_stats = jsonl.Samples.init(warm_samples);
1152 }
1153
1154 if (options.dump_ptx_dir) |dir| try dumpKernelArtifacts(arena, fragment, workload, dir);
1155
1156 var buffers = try allocHostBuffers(backing, workload);
1157 defer buffers.deinit(backing);
1158
1159 var input_storage: [4][]const u8 = undefined;
1160 const inputs = inputBytes(&buffers, &input_storage);
1161
1162 var mismatch_note: [128]u8 = undefined;
1163 if (options.verify) {
1164 var invocation = try accy.executable.prepareInvocation(fragment, backing, inputs);
1165 defer invocation.deinit();
1166 try invocation.launch(backing);
1167 try invocation.readOutput(0, std.mem.sliceAsBytes(buffers.output));
1168 verifyOutput(workload, options.math_tier, workload_mod.default_scene, &buffers) catch {
1169 try jsonl.writeRow(out, .{
1170 .system = system,
1171 .workload = workload.name,
1172 .metric = .kernel_ns,
1173 .status = .failed,
1174 .oracle = .failed,
1175 .samples = 0,
1176 .median_ns = 0,
1177 .p10_ns = 0,
1178 .p90_ns = 0,
1179 .flops = flops,
1180 .moved_bytes = moved_bytes,
1181 .note = std.fmt.bufPrint(mismatch_note[0..], "oracle mismatch: expected {e:.6} got {e:.6}, all timing skipped", .{ last_mismatch.expected, last_mismatch.actual }) catch "oracle mismatch: all timing skipped",
1182 });
1183 return;
1184 };
1185 }
1186
1187 const oracle: jsonl.Oracle = if (options.verify) .passed else .skipped;
1188 try jsonl.writeRow(out, .{
1189 .system = system,
1190 .workload = workload.name,
1191 .metric = .compile_ns,
1192 .oracle = oracle,
1193 .samples = 1,
1194 .median_ns = compile_cold_ns,
1195 .p10_ns = compile_cold_ns,
1196 .p90_ns = compile_cold_ns,
1197 .flops = flops,
1198 .moved_bytes = moved_bytes,
1199 });
1200 if (compile_warm_stats) |stats| {
1201 try jsonl.writeRow(out, .{
1202 .system = system,
1203 .workload = workload.name,
1204 .metric = .compile_warm_ns,
1205 .oracle = oracle,
1206 .samples = options.compile_repeats,
1207 .median_ns = stats.median(),
1208 .p10_ns = stats.p10(),
1209 .p90_ns = stats.p90(),
1210 .flops = flops,
1211 .moved_bytes = moved_bytes,
1212 });
1213 }
1214 if (options.compile_repeats > 0) {
1215 try jsonl.writeRow(out, .{
1216 .system = system,
1217 .workload = workload.name,
1218 .metric = .compile_incremental_ns,
1219 .status = .failed,
1220 .oracle = .skipped,
1221 .samples = 0,
1222 .median_ns = 0,
1223 .p10_ns = 0,
1224 .p90_ns = 0,
1225 .flops = flops,
1226 .moved_bytes = moved_bytes,
1227 .note = "MissingWorkContract: native artifact and load reuse is unavailable",
1228 });
1229 }
1230
1231 const kernel_count = fragment.kernelCount();
1232 var candidate_invocation = try accy.executable.prepareInvocation(fragment, backing, inputs);
1233 defer candidate_invocation.deinit();
1234 var kernel_index: usize = 0;
1235 while (kernel_index < kernel_count) : (kernel_index += 1) {
1236 const summary = try fragment.kernelSummary(kernel_index);
1237 if (summary.launch_candidate_count == 0) continue;
1238
1239 const records = try candidate_invocation.measureLaunchCandidateRecords(arena, backing, kernel_index, .{
1240 .warmup = options.warmup,
1241 .samples = options.samples,
1242 .synchronize = .device,
1243 });
1244 defer arena.free(records);
1245
1246 var best = records[0];
1247 for (records[1..]) |record| {
1248 if (record.median_ns < best.median_ns) best = record;
1249 }
1250 try fragment.recordLaunchCandidateRecords(records);
1251
1252 const note = try std.fmt.allocPrint(arena, "grid={d}x{d}x{d} block={d}x{d}x{d} candidates={d}", .{
1253 best.geometry.grid[0], best.geometry.grid[1], best.geometry.grid[2],
1254 best.geometry.threadgroup[0], best.geometry.threadgroup[1], best.geometry.threadgroup[2],
1255 records.len,
1256 });
1257 try jsonl.writeRow(out, .{
1258 .system = system,
1259 .workload = workload.name,
1260 .metric = .candidate_ns,
1261 .oracle = oracle,
1262 .samples = options.samples,
1263 .median_ns = best.median_ns,
1264 .p10_ns = best.median_ns,
1265 .p90_ns = best.median_ns,
1266 .flops = flops,
1267 .moved_bytes = moved_bytes,
1268 .kernel = best.kernel.entry_name,
1269 .note = note,
1270 });
1271 }
1272
1273 var kernel_warmup_index: u32 = 0;
1274 while (kernel_warmup_index < options.warmup) : (kernel_warmup_index += 1) {
1275 var invocation = try accy.executable.prepareInvocation(fragment, backing, inputs);
1276 defer invocation.deinit();
1277 try invocation.launch(backing);
1278 }
1279
1280 const kernel_samples = try arena.alloc(u64, options.samples);
1281 for (kernel_samples) |*sample| {
1282 var invocation = try accy.executable.prepareInvocation(fragment, backing, inputs);
1283 defer invocation.deinit();
1284 timer.reset();
1285 try invocation.launch(backing);
1286 sample.* = timer.read();
1287 coz.progressNamed("accy.versus.kernel_sample");
1288 }
1289 const kernel_stats = jsonl.Samples.init(kernel_samples);
1290 const kernel_note = try std.fmt.allocPrint(arena, "prepared launch+synchronize kernels={d}", .{kernel_count});
1291 try jsonl.writeRow(out, .{
1292 .system = system,
1293 .workload = workload.name,
1294 .metric = .kernel_ns,
1295 .oracle = oracle,
1296 .samples = options.samples,
1297 .median_ns = kernel_stats.median(),
1298 .p10_ns = kernel_stats.p10(),
1299 .p90_ns = kernel_stats.p90(),
1300 .flops = flops,
1301 .moved_bytes = moved_bytes,
1302 .note = kernel_note,
1303 });
1304
1305 var graph = try fragment.createLaunchGraphPlan(backing, .{});
1306 defer graph.deinit();
1307
1308 var warmup_index: u32 = 0;
1309 while (warmup_index < options.warmup) : (warmup_index += 1) {
1310 var invocation = try accy.executable.prepareInvocation(fragment, backing, inputs);
1311 defer invocation.deinit();
1312 try invocation.launchWithGraph(backing, graph.plan());
1313 }
1314
1315 const latency_samples = try arena.alloc(u64, options.samples);
1316 for (latency_samples) |*sample| {
1317 var invocation = try accy.executable.prepareInvocation(fragment, backing, inputs);
1318 defer invocation.deinit();
1319 timer.reset();
1320 try invocation.launchWithGraph(backing, graph.plan());
1321 sample.* = timer.read();
1322 coz.progressNamed("accy.versus.latency_sample");
1323 }
1324 const latency_stats = jsonl.Samples.init(latency_samples);
1325 try jsonl.writeRow(out, .{
1326 .system = system,
1327 .workload = workload.name,
1328 .metric = .latency_ns,
1329 .oracle = oracle,
1330 .samples = options.samples,
1331 .median_ns = latency_stats.median(),
1332 .p10_ns = latency_stats.p10(),
1333 .p90_ns = latency_stats.p90(),
1334 .flops = flops,
1335 .moved_bytes = moved_bytes,
1336 });
1337
1338 const dispatch_samples = try arena.alloc(u64, options.dispatch_samples);
1339 for (dispatch_samples) |*sample| {
1340 var invocation = try accy.executable.prepareInvocation(fragment, backing, inputs);
1341 defer invocation.deinit();
1342 timer.reset();
1343 try invocation.launchWithGraph(backing, graph.plan());
1344 sample.* = timer.read();
1345 }
1346 const dispatch_stats = jsonl.Samples.init(dispatch_samples);
1347 try jsonl.writeRow(out, .{
1348 .system = system,
1349 .workload = workload.name,
1350 .metric = .dispatch_ns,
1351 .oracle = oracle,
1352 .samples = options.dispatch_samples,
1353 .median_ns = dispatch_stats.median(),
1354 .p10_ns = dispatch_stats.p10(),
1355 .p90_ns = dispatch_stats.p90(),
1356 .flops = flops,
1357 .moved_bytes = moved_bytes,
1358 .note = "synchronous invocation, kernel count in kernel rows",
1359 });
1360
1361 const pipeline_samples = try arena.alloc(u64, options.pipeline_rounds);
1362 for (pipeline_samples) |*sample| {
1363 timer.reset();
1364 var launch_index: u32 = 0;
1365 while (launch_index < options.pipeline_batch) : (launch_index += 1) {
1366 var invocation = try accy.executable.prepareInvocation(fragment, backing, inputs);
1367 defer invocation.deinit();
1368 try invocation.launchWithGraph(backing, graph.plan());
1369 }
1370 sample.* = timer.read() / options.pipeline_batch;
1371 }
1372 const pipeline_stats = jsonl.Samples.init(pipeline_samples);
1373 try jsonl.writeRow(out, .{
1374 .system = system,
1375 .workload = workload.name,
1376 .metric = .pipeline_ns,
1377 .oracle = oracle,
1378 .samples = options.pipeline_rounds,
1379 .median_ns = pipeline_stats.median(),
1380 .p10_ns = pipeline_stats.p10(),
1381 .p90_ns = pipeline_stats.p90(),
1382 .flops = flops,
1383 .moved_bytes = moved_bytes,
1384 });
1385 }
1386
1387 test "module builders produce verified semantic modules" {
1388 for (workload_mod.battery) |workload| {
1389 if (workload.kind == .matmul and workload.m > 512) continue;
1390 if (workload.kind == .dense and workload.m > 512) continue;
1391 var scaled = workload;
1392 if (scaled.elements > 4096) scaled.elements = 4096;
1393 switch (scaled.kind) {
1394 .softmax, .stencil, .raymarch => {
1395 scaled.m = @min(scaled.m, 64);
1396 scaled.n = @min(scaled.n, 64);
1397 },
1398 .nbody => scaled.m = @min(scaled.m, 64),
1399 else => {},
1400 }
1401 const module = try buildModule(std.testing.allocator, scaled, workload_mod.default_scene);
1402 defer module.deinit();
1403 try module.verify();
1404 }
1405 }
1406
1407 test "raymarch oracle marches toward scene surfaces" {
1408 const center = raymarchOracle(1024, 1024, 512, 512, 0.0);
1409 try std.testing.expect(center > 0.5);
1410 try std.testing.expect(center <= workload_mod.raymarch_scene.max_distance);
1411 const sky = raymarchOracle(1024, 1024, 0, 0, 0.0);
1412 try std.testing.expectEqual(@as(f64, workload_mod.raymarch_scene.max_distance), sky);
1413 }
1414
1415 test "oracle helpers agree with direct evaluation" {
1416 const workload = Workload{ .name = "matmul_f32_2", .kind = .matmul, .m = 2, .n = 2, .k = 2 };
1417 const lhs = [_]f32{ 1, 2, 3, 4 };
1418 const rhs = [_]f32{ 5, 6, 7, 8 };
1419 try std.testing.expectEqual(@as(f64, 19), matmulOracle(workload, .exact, lhs[0..], rhs[0..], 0));
1420 try std.testing.expectEqual(@as(f64, 22), matmulOracle(workload, .exact, lhs[0..], rhs[0..], 1));
1421 }
1422
1423 test "oracle cell sampler includes boundaries" {
1424 const total: u64 = 1000;
1425 try std.testing.expectEqual(@as(u64, 0), oracleCellIndex(0, total));
1426 try std.testing.expectEqual(total - 1, oracleCellIndex(1, total));
1427 try std.testing.expectEqual(total / 2, oracleCellIndex(2, total));
1428 var sample: usize = 0;
1429 while (sample < oracle_cells) : (sample += 1) {
1430 try std.testing.expect(oracleCellIndex(sample, total) < total);
1431 }
1432 }
1433
1434 test "expectClose accepts close values and rejects far values" {
1435 try expectClose(1.0, 1.0005, 1e-2);
1436 try std.testing.expectError(error.OracleMismatch, expectClose(1.0, 1.5, 1e-2));
1437 }
1438
1439 test "mandelbrot frame gate accepts oracle-built frames and rejects shifted scenes" {
1440 const workload = workload_mod.byName("mandelbrot_f32_2048x2048") orelse return error.TestUnexpectedResult;
1441 var scene = workload_mod.default_scene;
1442 scene.mandelbrot_span = 0.02;
1443 scene.mandelbrot_x0 = -0.74364388 - 0.01;
1444 scene.mandelbrot_y0 = 0.13182590 - 0.01;
1445
1446 const rows: u64 = workload.m;
1447 const cols: u64 = workload.n;
1448 const output = try std.testing.allocator.alloc(f32, @intCast(rows * cols));
1449 defer std.testing.allocator.free(output);
1450 @memset(output, 0.0);
1451 var sample: usize = 0;
1452 while (sample < oracle_cells) : (sample += 1) {
1453 const cell = oracleCellIndex(sample, rows * cols);
1454 output[@intCast(cell)] = @floatCast(mandelbrotOracle(rows, cols, cell / cols, cell % cols, scene));
1455 }
1456 try verifyMandelbrotFrame(workload, scene, output);
1457
1458 var shifted = scene;
1459 shifted.mandelbrot_x0 += scene.mandelbrot_span * 0.5;
1460 try std.testing.expectError(error.OracleMismatch, verifyMandelbrotFrame(workload, shifted, output));
1461 }