lib/accy/src/profiling/wos/suite.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const accy = @import("accy");
5 const bench = @import("bench");
6
7 const options_mod = @import("options.zig");
8 const stats_mod = @import("stats.zig");
9 const workload = @import("workload.zig");
10
11 const Allocator = std.mem.Allocator;
12 const coz = bench.coz;
13 const pretty_json = bench.pretty.json;
14
15 const analytic_bias_limit: f64 = 0.02;
16
17 pub fn run(arena: Allocator, backing_allocator: Allocator, out: *std.Io.Writer, options: options_mod.Options) !void {
18 var run_start_stream = pretty_json.Writer.init(out, .minified);
19 const run_start = try run_start_stream.object();
20 try run_start.field("kind", "run_start");
21 try run_start.field("benchmark", "wos-while-falsifier");
22 try run_start.field("walks", options.walks);
23 try run_start.field("warmup", options.warmup);
24 try run_start.field("samples", options.samples);
25 try run_start.endLine();
26
27 var cuda_state = gpu.cuda.State.initDevice(backing_allocator, 0) catch |err| switch (err) {
28 error.RuntimeUnavailable => {
29 var skip_stream = pretty_json.Writer.init(out, .minified);
30 const skip = try skip_stream.object();
31 try skip.field("kind", "skip");
32 try skip.field("reason", "cuda runtime unavailable");
33 try skip.endLine();
34 try writeRunEnd(out);
35 return;
36 },
37 else => return err,
38 };
39 defer cuda_state.deinit();
40 const handle = cuda_state.handle();
41
42 if (options.customSweep()) |sweep| {
43 try runSweep(arena, backing_allocator, out, options, handle, sweep);
44 } else {
45 for (options.sweeps()) |sweep| {
46 try runSweep(arena, backing_allocator, out, options, handle, sweep);
47 }
48 }
49 try writeRunEnd(out);
50 }
51
52 fn writeRunEnd(out: *std.Io.Writer) !void {
53 var stream = pretty_json.Writer.init(out, .minified);
54 const object = try stream.object();
55 try object.field("kind", "run_end");
56 try object.field("benchmark", "wos-while-falsifier");
57 try object.endLine();
58 }
59
60 const VariantRun = struct {
61 artifact: gpu.KernelArtifact,
62 loaded: gpu.LoadedArtifact,
63 dst: gpu.BufferBinding,
64 steps: gpu.BufferBinding,
65 ptx_bytes: usize,
66 compile_ns: u64,
67 load_ns: u64,
68 host_dst: []f32,
69 host_steps: []i32,
70 timing: bench.SampleStats,
71
72 fn deinit(self: *VariantRun, handle: gpu.BackendHandle) void {
73 handle.destroyObject(self.loaded.id);
74 handle.destroyObject(self.dst.handle.id);
75 handle.destroyObject(self.steps.handle.id);
76 self.artifact.deinit();
77 }
78 };
79
80 fn runSweep(
81 arena: Allocator,
82 backing_allocator: Allocator,
83 out: *std.Io.Writer,
84 options: options_mod.Options,
85 handle: gpu.BackendHandle,
86 sweep: options_mod.Sweep,
87 ) !void {
88 const sweep_phase = bench.phaseAt("accy.wos.sweep", @src());
89 defer sweep_phase.end();
90
91 const walks = options.walks;
92 const xs = try arena.alloc(f32, walks);
93 const ys = try arena.alloc(f32, walks);
94 workload.fillWalkOrigins(xs, ys);
95
96 const xs_binding = try writeDeviceBuffer(handle, .f32, std.mem.sliceAsBytes(xs), walks, .read_only);
97 defer handle.destroyObject(xs_binding.handle.id);
98 const ys_binding = try writeDeviceBuffer(handle, .f32, std.mem.sliceAsBytes(ys), walks, .read_only);
99 defer handle.destroyObject(ys_binding.handle.id);
100
101 const geometry = workload.launchGeometry(walks);
102 const seeds = workload.seedArguments();
103 const variants = [_]workload.Variant{ .fold, .while_loop };
104 var runs: [variants.len]VariantRun = undefined;
105 var prepared: usize = 0;
106 defer for (runs[0..prepared]) |*variant_run| variant_run.deinit(handle);
107
108 for (variants, 0..) |variant, variant_index| {
109 const compile_start = bench.nowNs();
110 var graph = try workload.buildWalkGraph(backing_allocator, variant, sweep.epsilon, sweep.cap);
111 defer graph.deinit();
112 var artifact = try accy.kernel.createKernelArtifact(backing_allocator, handle, &graph, .{
113 .authored_kernel_diagnostic_id = "profiling/wos-while-falsifier",
114 });
115 var artifact_owned = true;
116 errdefer if (artifact_owned) artifact.deinit();
117 const compile_ns: u64 = @intCast(bench.nowNs() - compile_start);
118
119 const load_start = bench.nowNs();
120 const loaded = try handle.loadArtifact(&artifact);
121 errdefer if (artifact_owned) handle.destroyObject(loaded.id);
122 const load_ns: u64 = @intCast(bench.nowNs() - load_start);
123
124 const dst_binding = try allocateDeviceBuffer(handle, .f32, walks * @sizeOf(f32), walks);
125 errdefer if (artifact_owned) handle.destroyObject(dst_binding.handle.id);
126 const steps_binding = try allocateDeviceBuffer(handle, .i32, walks * @sizeOf(i32), walks);
127 errdefer if (artifact_owned) handle.destroyObject(steps_binding.handle.id);
128
129 runs[variant_index] = .{
130 .artifact = artifact,
131 .loaded = loaded,
132 .dst = dst_binding,
133 .steps = steps_binding,
134 .ptx_bytes = switch (artifact.payload) {
135 .text => |text| text.len,
136 else => 0,
137 },
138 .compile_ns = compile_ns,
139 .load_ns = load_ns,
140 .host_dst = try arena.alloc(f32, walks),
141 .host_steps = try arena.alloc(i32, walks),
142 .timing = undefined,
143 };
144 artifact_owned = false;
145 prepared += 1;
146
147 var artifact_stream = pretty_json.Writer.init(out, .minified);
148 const artifact_object = try artifact_stream.object();
149 try artifact_object.field("kind", "artifact");
150 try artifact_object.field("sweep", sweep.name);
151 try artifact_object.field("variant", variant.label());
152 try artifact_object.field("ptx_bytes", runs[variant_index].ptx_bytes);
153 try artifact_object.field("compile_ns", compile_ns);
154 try artifact_object.field("load_ns", load_ns);
155 try artifact_object.endLine();
156 }
157
158 for (&runs) |*variant_run| {
159 try launchVariant(handle, variant_run, xs_binding, ys_binding, seeds, geometry);
160 try handle.synchronize(.{ .scope = .device });
161 try handle.readBuffer(.{ .handle = variant_run.dst.handle, .bytes = std.mem.sliceAsBytes(variant_run.host_dst) });
162 try handle.readBuffer(.{ .handle = variant_run.steps.handle, .bytes = std.mem.sliceAsBytes(variant_run.host_steps) });
163 }
164
165 const dst_equal = std.mem.eql(
166 u8,
167 std.mem.sliceAsBytes(runs[0].host_dst),
168 std.mem.sliceAsBytes(runs[1].host_dst),
169 );
170 const steps_equal = std.mem.eql(i32, runs[0].host_steps, runs[1].host_steps);
171 const bias = analyticBias(xs, ys, runs[1].host_dst);
172 var equality_stream = pretty_json.Writer.init(out, .minified);
173 const equality_gate = try equality_stream.object();
174 try equality_gate.field("kind", "equality_gate");
175 try equality_gate.field("sweep", sweep.name);
176 try equality_gate.field("dst_equal", dst_equal);
177 try equality_gate.field("steps_equal", steps_equal);
178 try equality_gate.print("analytic_bias", "{d}", .{bias});
179 try equality_gate.print("bias_limit", "{d}", .{analytic_bias_limit});
180 try equality_gate.endLine();
181 if (!dst_equal or !steps_equal or @abs(bias) > analytic_bias_limit) {
182 var failed_stream = pretty_json.Writer.init(out, .minified);
183 const failed = try failed_stream.object();
184 try failed.field("kind", "gate_failed");
185 try failed.field("sweep", sweep.name);
186 try failed.endLine();
187 return;
188 }
189
190 const walk_stats = try stats_mod.computeWalkLengthStats(arena, runs[1].host_steps, sweep.cap);
191 var lengths_stream = pretty_json.Writer.init(out, .minified);
192 const lengths = try lengths_stream.object();
193 try lengths.field("kind", "walk_lengths");
194 try lengths.field("sweep", sweep.name);
195 try lengths.print("epsilon", "{d}", .{sweep.epsilon});
196 try lengths.field("cap", sweep.cap);
197 try lengths.field("walks", walk_stats.walks);
198 try lengths.print("mean", "{d}", .{walk_stats.mean});
199 try lengths.field("p50", walk_stats.p50);
200 try lengths.field("p90", walk_stats.p90);
201 try lengths.field("p99", walk_stats.p99);
202 try lengths.field("max", walk_stats.max);
203 try lengths.field("cap_hits", walk_stats.cap_hits);
204 try lengths.print("warp_max_mean", "{d}", .{walk_stats.warp_max_mean});
205 try lengths.print("fold_cost", "{d}", .{walk_stats.foldCost(sweep.cap)});
206 try lengths.print("while_cost", "{d}", .{walk_stats.whileCost()});
207 try lengths.print("persistent_cost", "{d}", .{walk_stats.persistentCost()});
208 try lengths.print("predicted_while_speedup", "{d}", .{walk_stats.predictedWhileSpeedup(sweep.cap)});
209 try lengths.print("persistent_headroom", "{d}", .{walk_stats.persistentHeadroom()});
210 try lengths.endLine();
211
212 const samples = try arena.alloc(u64, options.samples);
213 var statistics = try bench.StatisticsStorage.init(arena, .{ .samples = samples.len });
214 defer statistics.deinit(arena);
215 statistics.activate();
216 for (&runs, 0..) |*variant_run, variant_index| {
217 const variant = variants[variant_index];
218 const timing_phase = bench.phaseAt("accy.wos.timing", @src());
219 defer timing_phase.end();
220
221 for (0..options.warmup) |_| {
222 try launchVariant(handle, variant_run, xs_binding, ys_binding, seeds, geometry);
223 try handle.synchronize(.{ .scope = .device });
224 }
225 for (samples) |*sample| {
226 const start = bench.nowNs();
227 try launchVariant(handle, variant_run, xs_binding, ys_binding, seeds, geometry);
228 try handle.synchronize(.{ .scope = .device });
229 sample.* = @intCast(bench.nowNs() - start);
230 coz.progressNamed("accy.wos.sample");
231 }
232 variant_run.timing = try bench.computeSampleStatsWithBootstrap(
233 &statistics,
234 samples,
235 .{},
236 );
237
238 const intervals = variant_run.timing.confidence_intervals.?;
239 var timing_stream = pretty_json.Writer.init(out, .minified);
240 const timing = try timing_stream.object();
241 try timing.field("kind", "timing");
242 try timing.field("sweep", sweep.name);
243 try timing.field("variant", variant.label());
244 try timing.field("samples", samples.len);
245 try timing.field("min_ns", variant_run.timing.min_ns);
246 try timing.field("median_ns", variant_run.timing.median_ns);
247 try timing.field("mean_ns", variant_run.timing.mean_ns);
248 try timing.field("p95_ns", variant_run.timing.p95_ns);
249 try timing.field("max_ns", variant_run.timing.max_ns);
250 try timing.field("ci_median_low_ns", intervals.median_ns.low_ns);
251 try timing.field("ci_median_high_ns", intervals.median_ns.high_ns);
252 try timing.endLine();
253 }
254
255 const fold_ci = runs[0].timing.confidence_intervals.?.median_ns;
256 const while_ci = runs[1].timing.confidence_intervals.?.median_ns;
257 const verdict = if (while_ci.high_ns < fold_ci.low_ns)
258 "improved"
259 else if (while_ci.low_ns > fold_ci.high_ns)
260 "regressed"
261 else
262 "within-noise";
263 const measured_speedup = @as(f64, @floatFromInt(runs[0].timing.median_ns)) /
264 @as(f64, @floatFromInt(runs[1].timing.median_ns));
265 var verdict_stream = pretty_json.Writer.init(out, .minified);
266 const verdict_object = try verdict_stream.object();
267 try verdict_object.field("kind", "verdict");
268 try verdict_object.field("sweep", sweep.name);
269 try verdict_object.field("verdict", verdict);
270 try verdict_object.field("fold_median_ns", runs[0].timing.median_ns);
271 try verdict_object.field("while_median_ns", runs[1].timing.median_ns);
272 try verdict_object.print("measured_speedup", "{d}", .{measured_speedup});
273 try verdict_object.print("predicted_speedup", "{d}", .{walk_stats.predictedWhileSpeedup(sweep.cap)});
274 try verdict_object.endLine();
275 coz.progressNamed("accy.wos.sweep.complete");
276 }
277
278 fn launchVariant(
279 handle: gpu.BackendHandle,
280 variant_run: *VariantRun,
281 xs_binding: gpu.BufferBinding,
282 ys_binding: gpu.BufferBinding,
283 seeds: [2]choir_abi.ScalarArgument,
284 geometry: choir_abi.LaunchGeometry,
285 ) !void {
286 const buffers = [_]gpu.BufferBinding{ xs_binding, ys_binding, variant_run.dst, variant_run.steps };
287 try handle.launch(.{
288 .artifact = &variant_run.artifact,
289 .loaded_artifact = variant_run.loaded,
290 .buffers = buffers[0..],
291 .scalar_arguments = seeds[0..],
292 .geometry = geometry,
293 });
294 }
295
296 fn allocateDeviceBuffer(
297 handle: gpu.BackendHandle,
298 dtype: choir_abi.DType,
299 byte_size: usize,
300 element_count: usize,
301 ) !gpu.BufferBinding {
302 const device_buffer = try handle.allocateBuffer(.{
303 .byte_size = byte_size,
304 .alignment = 256,
305 .dtype = dtype,
306 .element_count = element_count,
307 });
308 return .{
309 .handle = device_buffer,
310 .access = .read_write,
311 .ownership = device_buffer.ownership,
312 .byte_size = device_buffer.byte_size,
313 };
314 }
315
316 fn writeDeviceBuffer(
317 handle: gpu.BackendHandle,
318 dtype: choir_abi.DType,
319 bytes: []const u8,
320 element_count: usize,
321 access: gpu.BufferAccess,
322 ) !gpu.BufferBinding {
323 var binding = try allocateDeviceBuffer(handle, dtype, bytes.len, element_count);
324 binding.access = access;
325 try handle.writeBuffer(.{ .handle = binding.handle, .bytes = bytes });
326 return binding;
327 }
328
329 fn analyticBias(xs: []const f32, ys: []const f32, estimates: []const f32) f64 {
330 var total: f64 = 0.0;
331 for (xs, ys, estimates) |x, y, estimate| {
332 const analytic = @as(f64, x) * x * x - 3.0 * @as(f64, x) * y * y;
333 total += @as(f64, estimate) - analytic;
334 }
335 return total / @as(f64, @floatFromInt(estimates.len));
336 }
337
338 test "analytic bias vanishes when estimates match the analytic field" {
339 const xs = [_]f32{ 0.1, -0.4, 0.7 };
340 const ys = [_]f32{ 0.3, 0.2, -0.5 };
341 var estimates: [3]f32 = undefined;
342 for (&estimates, xs, ys) |*estimate, x, y| {
343 estimate.* = x * x * x - 3.0 * x * y * y;
344 }
345
346 const bias = analyticBias(xs[0..], ys[0..], estimates[0..]);
347 try std.testing.expect(@abs(bias) < 0.0000001);
348 }