lib/accy/src/profiling/scan/root.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 const graph_mod = @import("graph.zig");
7 pub const floor_mod = @import("floor.zig");
8 pub const gate = @import("gate.zig");
9
10 const Allocator = std.mem.Allocator;
11 const tensor = accy.tensor;
12 const buildScanProgram = graph_mod.buildScanProgram;
13 const buildUnrolledProgram = graph_mod.buildUnrolledProgram;
14 const countProgramOperationTree = graph_mod.countProgramOperationTree;
15 const countProgramValueTree = graph_mod.countProgramValueTree;
16 const pretty_json = bench.pretty.json;
17
18 const benchmark_name = "accy.scan.iterate";
19
20 const usage_text =
21 \\usage: accy-scan-bench [--steps N] [--width N] [--warmup N] [--samples N]
22 \\
23 \\Measure tensor scan lowering through accy.iterate against an equivalent hand-unrolled tensor graph.
24 \\
25 \\Options:
26 \\ --steps N Scan or unrolled recurrence length. Default: 32.
27 \\ --width N Vector lane count. Default: 1024.
28 \\ --warmup N Untimed launches per variant. Default: 2.
29 \\ --samples N Timed launches per variant. Default: 10.
30 \\ -h, --help Print this message.
31 ;
32
33 const Options = struct {
34 steps: u32 = 32,
35 width: u32 = 1024,
36 warmup: u32 = 2,
37 samples: u32 = 10,
38 help: bool = false,
39 };
40
41 const Variant = enum {
42 scan,
43 unrolled,
44
45 fn label(self: Variant) []const u8 {
46 return switch (self) {
47 .scan => "scan",
48 .unrolled => "unrolled",
49 };
50 }
51
52 /// Returns the scope for this variant, the name a gate line reports:
53 /// "accy.scan.iterate.scan" or "accy.scan.iterate.unrolled", so the
54 /// benchmark passes this name to the gate to label each variant's verdict.
55 fn scope(self: Variant) []const u8 {
56 return switch (self) {
57 .scan => benchmark_name ++ ".scan",
58 .unrolled => benchmark_name ++ ".unrolled",
59 };
60 }
61 };
62
63 const PreparedVariant = struct {
64 variant: Variant,
65 graph: tensor.Graph,
66 prepared: tensor.BackendPreparedJob,
67 fragment: *tensor.LoadedFragment,
68 inputs: [2][]const u8,
69 source_top_ops: usize,
70 source_tree_ops: usize,
71 source_top_values: usize,
72 source_tree_values: usize,
73 artifact_bytes: usize,
74 generated_kernels: usize,
75
76 fn deinit(self: *PreparedVariant) void {
77 self.fragment.deinit();
78 self.prepared.deinit();
79 self.graph.deinit();
80 self.* = undefined;
81 }
82 };
83
84 const SourceShape = struct {
85 top_ops: usize,
86 tree_ops: usize,
87 top_values: usize,
88 tree_values: usize,
89 };
90
91 pub fn main(init: sys.process.Init) !void {
92 var gpa: std.heap.DebugAllocator(.{}) = .{};
93 const backing_allocator = gpa.allocator();
94
95 const exit_code: u8 = blk: {
96 var arena_state = std.heap.ArenaAllocator.init(backing_allocator);
97 defer arena_state.deinit();
98 const arena = arena_state.allocator();
99
100 const args = init.minimal.args.toSlice(arena) catch |err| {
101 try writeError(err);
102 break :blk 1;
103 };
104
105 break :blk run(arena, backing_allocator, args[1..]) catch |err| switch (err) {
106 error.InvalidArguments => 2,
107 else => {
108 try writeError(err);
109 break :blk 1;
110 },
111 };
112 };
113
114 if (gpa.deinit() == .leak) sys.process.exit(1);
115 sys.process.exit(exit_code);
116 }
117
118 pub fn run(arena: Allocator, backing_allocator: Allocator, args: []const []const u8) !u8 {
119 var stdout_buffer: [8192]u8 = undefined;
120 var stdout = sys.stdio.stdout().writer(sys.stdio.debugIo(), &stdout_buffer);
121 defer stdout.interface.flush() catch {};
122 const out = &stdout.interface;
123
124 const options = try parseArgs(args);
125 if (options.help) {
126 try bench.pretty.write(out, .{ .text = usage_text }, .{});
127 return 0;
128 }
129
130 try runBenchmark(arena, backing_allocator, out, options);
131 return 0;
132 }
133
134 pub fn parseArgs(args: []const []const u8) !Options {
135 var options = Options{};
136 var index: usize = 0;
137 while (index < args.len) : (index += 1) {
138 const arg = args[index];
139 if (std.mem.eql(u8, arg, "-h") or std.mem.eql(u8, arg, "--help")) {
140 options.help = true;
141 } else if (std.mem.eql(u8, arg, "--steps")) {
142 index += 1;
143 if (index >= args.len) return error.InvalidArguments;
144 options.steps = try bench.parse.parsePositiveU32(args[index]);
145 } else if (std.mem.eql(u8, arg, "--width")) {
146 index += 1;
147 if (index >= args.len) return error.InvalidArguments;
148 options.width = try bench.parse.parsePositiveU32(args[index]);
149 } else if (std.mem.eql(u8, arg, "--warmup")) {
150 index += 1;
151 if (index >= args.len) return error.InvalidArguments;
152 options.warmup = try bench.parse.parseU32(args[index]);
153 } else if (std.mem.eql(u8, arg, "--samples")) {
154 index += 1;
155 if (index >= args.len) return error.InvalidArguments;
156 options.samples = try bench.parse.parsePositiveU32(args[index]);
157 } else {
158 return error.InvalidArguments;
159 }
160 }
161 return options;
162 }
163
164 fn runBenchmark(arena: Allocator, backing_allocator: Allocator, out: *std.Io.Writer, options: Options) !void {
165 var run_start_stream = pretty_json.Writer.init(out, .minified);
166 const run_start = try run_start_stream.object();
167 try run_start.field("kind", "run_start");
168 try run_start.field("benchmark", benchmark_name);
169 try run_start.field("steps", options.steps);
170 try run_start.field("width", options.width);
171 try run_start.field("warmup", options.warmup);
172 try run_start.field("samples", options.samples);
173 try run_start.endLine();
174
175 const width: usize = options.width;
176 const xs = try arena.alloc(f32, width);
177 const cs = try arena.alloc(f32, width);
178 fillInputs(xs, cs);
179
180 var state = gpu.cpu.State.init(backing_allocator);
181 defer state.deinit();
182
183 const variants = [_]Variant{ .scan, .unrolled };
184 var prepared_variants: [variants.len]PreparedVariant = undefined;
185 var prepared_count: usize = 0;
186 defer for (prepared_variants[0..prepared_count]) |*prepared| prepared.deinit();
187
188 for (variants, 0..) |variant, variant_index| {
189 var graph = switch (variant) {
190 .scan => try buildScanProgram(backing_allocator, options.width, options.steps),
191 .unrolled => try buildUnrolledProgram(backing_allocator, options.width, options.steps),
192 };
193 const shape = sourceShape(graph);
194 var failure: accy.preparation.BackendPreparationFailure = .{};
195 defer failure.deinit(backing_allocator);
196
197 prepared_variants[variant_index] = prepareVariant(
198 backing_allocator,
199 state.handle(),
200 variant,
201 &graph,
202 shape,
203 xs,
204 cs,
205 &failure,
206 ) catch |err| switch (err) {
207 error.OutOfMemory => {
208 graph.deinit();
209 return err;
210 },
211 else => {
212 try writeVariantFailure(out, variant, shape, err, &failure);
213 graph.deinit();
214 continue;
215 },
216 };
217 prepared_count += 1;
218 try writePrepared(out, &prepared_variants[variant_index]);
219 }
220
221 for (prepared_variants[0..prepared_count]) |*prepared| {
222 try runVariant(arena, backing_allocator, out, options, xs, cs, prepared);
223 }
224
225 var run_end_stream = pretty_json.Writer.init(out, .minified);
226 const run_end = try run_end_stream.object();
227 try run_end.field("kind", "run_end");
228 try run_end.field("benchmark", benchmark_name);
229 try run_end.endLine();
230 }
231
232 fn prepareVariant(
233 backing_allocator: Allocator,
234 handle: gpu.BackendHandle,
235 variant: Variant,
236 graph: *tensor.Graph,
237 shape: SourceShape,
238 xs: []const f32,
239 cs: []const f32,
240 failure: *accy.preparation.BackendPreparationFailure,
241 ) !PreparedVariant {
242 var prepared = try tensor.prepareFragment(backing_allocator, handle, graph, .{
243 .artifact_format = .cpu_object,
244 .preparation_failure = failure,
245 });
246 errdefer prepared.deinit();
247
248 const options = tensor.FragmentCompilerOptions{
249 .artifact_format = .cpu_object,
250 };
251 const compiled = try tensor.compileFragmentFromPreparedJob(backing_allocator, handle, &prepared, options);
252 var fragment = try accy.executable.loadFragment(backing_allocator, handle, compiled, options);
253 errdefer fragment.deinit();
254
255 const inputs = [_][]const u8{ std.mem.sliceAsBytes(xs), std.mem.sliceAsBytes(cs) };
256
257 return .{
258 .variant = variant,
259 .graph = graph.*,
260 .prepared = prepared,
261 .fragment = fragment,
262 .inputs = inputs,
263 .source_top_ops = shape.top_ops,
264 .source_tree_ops = shape.tree_ops,
265 .source_top_values = shape.top_values,
266 .source_tree_values = shape.tree_values,
267 .artifact_bytes = try fragmentArtifactBytes(backing_allocator, fragment),
268 .generated_kernels = try prepared.generatedKernelCount(),
269 };
270 }
271
272 fn runVariant(
273 arena: Allocator,
274 scratch: Allocator,
275 out: *std.Io.Writer,
276 options: Options,
277 xs: []const f32,
278 cs: []const f32,
279 prepared: *PreparedVariant,
280 ) !void {
281 const phase = bench.phaseAt("accy.scan.iterate", @src());
282 defer phase.end();
283
284 var validation_invocation = try accy.executable.prepareInvocation(prepared.fragment, scratch, &prepared.inputs);
285 defer validation_invocation.deinit();
286 try validation_invocation.launch(scratch);
287 try validateOutputs(arena, options, xs, cs, validation_invocation);
288 var validation_stream = pretty_json.Writer.init(out, .minified);
289 const validation = try validation_stream.object();
290 try validation.field("kind", "validation");
291 try validation.field("benchmark", benchmark_name);
292 try validation.field("variant", prepared.variant.label());
293 try validation.field("passed", true);
294 try validation.endLine();
295
296 for (0..options.warmup) |_| {
297 var invocation = try accy.executable.prepareInvocation(prepared.fragment, scratch, &prepared.inputs);
298 defer invocation.deinit();
299 try invocation.launch(scratch);
300 bench.coz.progressNamed("accy.scan.warmup");
301 }
302
303 const samples = try arena.alloc(u64, options.samples);
304 var cycles = gate.Cycles{};
305 for (samples, 0..) |*sample, sample_index| {
306 var region = gate.open();
307 defer gate.close(®ion);
308 const start = bench.nowNs();
309 var invocation = try accy.executable.prepareInvocation(prepared.fragment, scratch, &prepared.inputs);
310 defer invocation.deinit();
311 try invocation.launch(scratch);
312 sample.* = @intCast(bench.nowNs() - start);
313 cycles.record(®ion, sample.*);
314 bench.coz.progressNamed("accy.scan.sample");
315 var sample_stream = pretty_json.Writer.init(out, .minified);
316 const sample_object = try sample_stream.object();
317 try sample_object.field("kind", "sample");
318 try sample_object.field("benchmark", benchmark_name);
319 try sample_object.field("variant", prepared.variant.label());
320 try sample_object.field("index", sample_index);
321 try sample_object.field("runtime_ns", sample.*);
322 try sample_object.endLine();
323 }
324
325 try writeTiming(arena, out, prepared.variant, samples);
326 const shape = floor_mod.Shape{ .lanes = options.width, .steps = options.steps };
327 const scope = prepared.variant.scope();
328 if (try gate.write(arena, out, scope, shape, samples, cycles.clock())) {
329 return error.FloorAboveMeasurement;
330 }
331 }
332
333 /// Prints one JSON line with the spread of one variant's samples: count,
334 /// minimum, median, mean, the 75th, 95th and 99th percentiles, maximum and
335 /// total, all in nanoseconds, for the benchmark after its timed launches. The
336 /// statistics run with no bootstrap resampling, in storage taken from `arena`.
337 fn writeTiming(
338 arena: Allocator,
339 out: *std.Io.Writer,
340 variant: Variant,
341 samples: []const u64,
342 ) !void {
343 var statistics = try bench.StatisticsStorage.init(arena, .{
344 .samples = samples.len,
345 .bootstrap_iterations = 0,
346 });
347 defer statistics.deinit(arena);
348 statistics.activate();
349 const timing = try bench.computeSampleStatsWithBootstrap(
350 &statistics,
351 samples,
352 .{ .iterations = 0 },
353 );
354 var timing_stream = pretty_json.Writer.init(out, .minified);
355 const timing_object = try timing_stream.object();
356 try timing_object.field("kind", "timing");
357 try timing_object.field("benchmark", benchmark_name);
358 try timing_object.field("variant", variant.label());
359 try timing_object.field("samples", samples.len);
360 try timing_object.field("min_ns", timing.min_ns);
361 try timing_object.field("median_ns", timing.median_ns);
362 try timing_object.field("mean_ns", timing.mean_ns);
363 try timing_object.field("p75_ns", timing.p75_ns);
364 try timing_object.field("p95_ns", timing.p95_ns);
365 try timing_object.field("p99_ns", timing.p99_ns);
366 try timing_object.field("max_ns", timing.max_ns);
367 try timing_object.field("total_ns", timing.total_ns);
368 try timing_object.endLine();
369 }
370
371 fn writePrepared(out: *std.Io.Writer, prepared: *const PreparedVariant) !void {
372 var stream = pretty_json.Writer.init(out, .minified);
373 const object = try stream.object();
374 try object.field("kind", "prepared");
375 try object.field("benchmark", benchmark_name);
376 try object.field("variant", prepared.variant.label());
377 try object.field("source_top_ops", prepared.source_top_ops);
378 try object.field("source_tree_ops", prepared.source_tree_ops);
379 try object.field("source_top_values", prepared.source_top_values);
380 try object.field("source_tree_values", prepared.source_tree_values);
381 try object.field("initial_choir_ops", prepared.prepared.run.initial_choir_ops);
382 try object.field("final_choir_ops", prepared.prepared.run.final_choir_ops);
383 try object.field("generated_kernels", prepared.generated_kernels);
384 try object.field("executable_kernels", prepared.fragment.kernelCount());
385 try object.field("artifact_bytes", prepared.artifact_bytes);
386 try object.field("preparation_ns", prepared.prepared.run.total_ns);
387 try object.endLine();
388 }
389
390 fn writeVariantFailure(
391 out: *std.Io.Writer,
392 variant: Variant,
393 shape: SourceShape,
394 err: anyerror,
395 failure: *const accy.preparation.BackendPreparationFailure,
396 ) !void {
397 const failure_kind: ?[]const u8 = if (failure.failure_kind) |kind| @tagName(kind) else null;
398 var stream = pretty_json.Writer.init(out, .minified);
399 const object = try stream.object();
400 try object.field("kind", "variant_failure");
401 try object.field("benchmark", benchmark_name);
402 try object.field("variant", variant.label());
403 try object.field("source_top_ops", shape.top_ops);
404 try object.field("source_tree_ops", shape.tree_ops);
405 try object.field("source_top_values", shape.top_values);
406 try object.field("source_tree_values", shape.tree_values);
407 try object.field("error", @errorName(err));
408 try object.field("pipeline_name", failure.pipeline_name);
409 try object.field("failure_kind", failure_kind);
410 try object.field("pass_name", failure.pass_name);
411 try object.field("target_op_name", failure.target_op_name);
412 try object.field("target_symbol_name", failure.target_symbol_name);
413 try object.field("verifier_error_name", failure.verifier_error_name);
414 try object.field("worker_count", failure.worker_count);
415 try object.endLine();
416 }
417
418 fn sourceShape(graph: tensor.Graph) SourceShape {
419 return .{
420 .top_ops = graph.operationCount(),
421 .tree_ops = countProgramOperationTree(graph),
422 .top_values = graph.valueCount(),
423 .tree_values = countProgramValueTree(graph),
424 };
425 }
426
427 fn fillInputs(xs: []f32, cs: []f32) void {
428 for (xs, cs, 0..) |*x, *c, index| {
429 const x_mod: u32 = @intCast(index % 17);
430 const c_mod: i32 = @as(i32, @intCast(index % 7)) - 3;
431 x.* = 0.05 + @as(f32, @floatFromInt(x_mod)) * 0.001;
432 c.* = @as(f32, @floatFromInt(c_mod)) * 0.001;
433 }
434 }
435
436 fn validateOutputs(
437 arena: Allocator,
438 options: Options,
439 xs: []const f32,
440 cs: []const f32,
441 invocation: *accy.executable.Invocation,
442 ) !void {
443 const width: usize = options.width;
444 const actual_x = try arena.alloc(f32, width);
445 const actual_acc = try arena.alloc(f32, width);
446 try invocation.readOutput(0, std.mem.sliceAsBytes(actual_x));
447 try invocation.readOutput(1, std.mem.sliceAsBytes(actual_acc));
448
449 const expected_x = try arena.dupe(f32, xs);
450 const expected_acc = try arena.alloc(f32, width);
451 @memset(expected_acc, 0);
452 for (0..options.steps) |_| {
453 for (expected_x, expected_acc, cs) |*x, *acc, c| {
454 x.* = x.* * x.* * 0.5 + c;
455 acc.* += x.*;
456 }
457 }
458
459 for (actual_x, actual_acc, expected_x, expected_acc) |got_x, got_acc, want_x, want_acc| {
460 if (@abs(got_x - want_x) > 1e-5) return error.ValidationFailed;
461 if (@abs(got_acc - want_acc) > 1e-4) return error.ValidationFailed;
462 }
463 }
464
465 fn fragmentArtifactBytes(allocator: Allocator, fragment: *tensor.LoadedFragment) !usize {
466 var total: usize = 0;
467 var kernel_index: usize = 0;
468 while (kernel_index < fragment.kernelCount()) : (kernel_index += 1) {
469 var artifact = try fragment.copyKernelArtifact(allocator, kernel_index);
470 defer artifact.deinit();
471 total += try artifactPayloadBytes(artifact.payload);
472 }
473 return total;
474 }
475
476 fn artifactPayloadBytes(payload: gpu.ArtifactPayload) !usize {
477 return switch (payload) {
478 .text => |text| text.len,
479 .bytes => |bytes| bytes.len,
480 .words_u32 => |words| words.len * @sizeOf(u32),
481 .none, .external => error.InvalidArtifact,
482 };
483 }
484
485 fn writeError(err: anyerror) !void {
486 try bench.writeStderrFmt("error: {t}\n", .{err});
487 }