From b3b2b915c42a79c5f752b2ada51a63dc6fd9b722 Mon Sep 17 00:00:00 2001 From: Ben Sinclair Date: Wed, 20 Mar 2024 20:42:57 +1100 Subject: [PATCH] Add memory allocation tracking Adds a custom Allocator wrapper which keeps track of the allocated memory, then record the maximum allocation readings in the Runner alongside timing readings and carry them into the Result. --- examples/json.zig | 9 +++- util/runner.zig | 69 ++++++++++++++++++-------- util/runner/types.zig | 5 +- util/statistics.zig | 1 + util/tracking_allocator.zig | 81 +++++++++++++++++++++++++++++++ zbench.zig | 97 ++++++++++++++++++++++++++++++++----- 6 files changed, 227 insertions(+), 35 deletions(-) create mode 100644 util/tracking_allocator.zig diff --git a/examples/json.zig b/examples/json.zig index 577776d..71e345b 100644 --- a/examples/json.zig +++ b/examples/json.zig @@ -17,7 +17,14 @@ test "bench test json" { var bench = zbench.Benchmark.init(test_allocator, .{}); defer bench.deinit(); - try bench.add("My Benchmark 1", myBenchmark, .{ .iterations = 10 }); + try bench.add("My Benchmark 1", myBenchmark, .{ + .iterations = 10, + .track_allocations = false, + }); + try bench.add("My Benchmark 2", myBenchmark, .{ + .iterations = 10, + .track_allocations = true, + }); try stdout.writeAll("["); var iter = try bench.iterator(); diff --git a/util/runner.zig b/util/runner.zig index c3f46bf..f09703d 100644 --- a/util/runner.zig +++ b/util/runner.zig @@ -47,10 +47,14 @@ const State = union(enum) { /// Array of timings collected. timings_ns: []u64, + + /// Array of memory usages collected. + allocations: ?[]usize, }; }; allocator: std.mem.Allocator, +track_allocations: bool, state: State, pub fn init( @@ -58,19 +62,26 @@ pub fn init( iterations: u16, max_iterations: u16, time_budget_ns: u64, + track_allocations: bool, ) Error!Runner { return if (iterations == 0) .{ .allocator = allocator, + .track_allocations = track_allocations, .state = .{ .preparing = .{ .max_iterations = max_iterations, .time_budget_ns = time_budget_ns, } }, } else .{ .allocator = allocator, + .track_allocations = track_allocations, .state = .{ .running = .{ .iterations_count = iterations, .iterations_remaining = iterations, .timings_ns = try allocator.alloc(u64, iterations), + .allocations = if (track_allocations) + try allocator.alloc(usize, iterations) + else + null, } }, }; } @@ -102,6 +113,10 @@ pub fn next(self: *Runner, reading: Reading) Error!?Step { .iterations_count = N, .iterations_remaining = N, .timings_ns = try self.allocator.alloc(u64, N), + .allocations = if (self.track_allocations) + try self.allocator.alloc(usize, N) + else + null, } }; } return .more; @@ -110,6 +125,12 @@ pub fn next(self: *Runner, reading: Reading) Error!?Step { if (0 < st.iterations_remaining) { const i = st.timings_ns.len - st.iterations_remaining; st.timings_ns[i] = reading.timing_ns; + if (st.allocations) |allocs| { + if (reading.max_allocated) |m| allocs[i] = m else { + self.allocator.free(allocs); + st.allocations = null; + } + } st.iterations_remaining -= 1; } return if (st.iterations_remaining == 0) null else .more; @@ -123,9 +144,11 @@ pub fn finish(self: *Runner) Error!Readings { return switch (self.state) { .preparing => .{ .timings_ns = &.{}, + .max_allocations = null, }, .running => |st| .{ .timings_ns = st.timings_ns, + .max_allocations = st.allocations, }, }; } @@ -134,7 +157,10 @@ pub fn finish(self: *Runner) Error!Readings { pub fn abort(self: *Runner) void { return switch (self.state) { .preparing => {}, - .running => |st| self.allocator.free(st.timings_ns), + .running => |st| { + self.allocator.free(st.timings_ns); + if (st.allocations) |allocs| self.allocator.free(allocs); + }, }; } @@ -157,31 +183,32 @@ pub fn status(self: Runner) Status { } test "Runner" { - var r = try Runner.init(std.testing.allocator, 0, 16384, 2e9); + var r = try Runner.init(std.testing.allocator, 0, 16384, 2e9, false); { errdefer r.abort(); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(Step.more, try r.next(Reading.init(200_000_000))); - try expectEq(@as(?Step, null), try r.next(Reading.init(200_000_000))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(Step.more, try r.next(Reading.init(200_000_000, null))); + try expectEq(@as(?Step, null), try r.next(Reading.init(200_000_000, null))); } const result = try r.finish(); defer std.testing.allocator.free(result.timings_ns); + defer if (result.max_allocations) |m| std.testing.allocator.free(m); try expectEqSlices(u64, &.{ 200_000_000, 200_000_000, 200_000_000, 200_000_000, 200_000_000, 200_000_000, 200_000_000, 200_000_000, diff --git a/util/runner/types.zig b/util/runner/types.zig index 61eb7a4..72f4fc9 100644 --- a/util/runner/types.zig +++ b/util/runner/types.zig @@ -6,14 +6,17 @@ pub const Step = enum { more }; pub const Reading = struct { timing_ns: u64, + max_allocated: ?usize, - pub fn init(timing_ns: u64) Reading { + pub fn init(timing_ns: u64, max_allocated: ?usize) Reading { return .{ .timing_ns = timing_ns, + .max_allocated = max_allocated, }; } }; pub const Readings = struct { timings_ns: []u64, + max_allocations: ?[]usize, }; diff --git a/util/statistics.zig b/util/statistics.zig index b084a62..7be963b 100644 --- a/util/statistics.zig +++ b/util/statistics.zig @@ -1,5 +1,6 @@ const std = @import("std"); +/// Collect common statistical calculations together. pub fn Statistics(comptime T: type) type { return struct { total: T, diff --git a/util/tracking_allocator.zig b/util/tracking_allocator.zig new file mode 100644 index 0000000..6af83b6 --- /dev/null +++ b/util/tracking_allocator.zig @@ -0,0 +1,81 @@ +const std = @import("std"); +const Allocator = std.mem.Allocator; + +const TrackingAllocator = @This(); + +parent_allocator: Allocator, +current_allocated: usize = 0, +max_allocated: usize = 0, + +pub fn init(parent_allocator: Allocator) TrackingAllocator { + return .{ + .parent_allocator = parent_allocator, + }; +} + +pub fn allocator(self: *TrackingAllocator) Allocator { + return .{ + .ptr = self, + .vtable = &.{ + .alloc = alloc, + .resize = resize, + .free = free, + }, + }; +} + +pub fn maxAllocated(self: TrackingAllocator) usize { + return self.max_allocated; +} + +fn alloc( + ctx: *anyopaque, + len: usize, + log2_ptr_align: u8, + ra: usize, +) ?[*]u8 { + const self: *TrackingAllocator = @ptrCast(@alignCast(ctx)); + const result = self.parent_allocator.rawAlloc(len, log2_ptr_align, ra); + if (result) |_| { + self.current_allocated += len; + if (self.max_allocated < self.current_allocated) + self.max_allocated = self.current_allocated; + } + return result; +} + +fn resize( + ctx: *anyopaque, + buf: []u8, + log2_buf_align: u8, + new_len: usize, + ra: usize, +) bool { + const self: *TrackingAllocator = @ptrCast(@alignCast(ctx)); + const result = self.parent_allocator.rawResize(buf, log2_buf_align, new_len, ra); + if (result) { + if (buf.len < new_len) { + self.current_allocated += new_len - buf.len; + if (self.max_allocated < self.current_allocated) + self.max_allocated = self.current_allocated; + } else self.current_allocated -= buf.len - new_len; + } + return result; +} + +fn free( + ctx: *anyopaque, + buf: []u8, + log2_buf_align: u8, + ra: usize, +) void { + const self: *TrackingAllocator = @ptrCast(@alignCast(ctx)); + self.parent_allocator.rawFree(buf, log2_buf_align, ra); + self.current_allocated -= buf.len; +} + +/// This allocator is used in front of another allocator and tracks the maximum +/// memory usage on every call to the allocator. +pub fn trackingAllocator(parent_allocator: Allocator) TrackingAllocator { + return TrackingAllocator.init(parent_allocator); +} diff --git a/zbench.zig b/zbench.zig index ded27db..2469afa 100644 --- a/zbench.zig +++ b/zbench.zig @@ -13,6 +13,7 @@ const Optional = @import("./util/optional.zig").Optional; const optional = @import("./util/optional.zig").optional; const platform = @import("./util/platform.zig"); const Runner = @import("./util/runner.zig"); +const TrackingAllocator = @import("./util/tracking_allocator.zig"); /// Hooks containing optional hooks for lifecycle events in benchmarking. /// Each field in this struct is a nullable function pointer. @@ -48,6 +49,10 @@ pub const Config = struct { /// This allows for custom setup and teardown operations, as well as fine-grained control /// over the environment in which benchmarks are run. hooks: Hooks = .{}, + + /// Track memory allocations made using the Allocator provided to + /// benchmarks. + track_allocations: bool = false, }; /// A benchmark definition. @@ -68,13 +73,23 @@ const Definition = struct { if (self.config.hooks.before_each) |hook| hook(); defer if (self.config.hooks.after_each) |hook| hook(); + var tracking = TrackingAllocator.init(allocator); + const alloc = if (self.config.track_allocations) + tracking.allocator() + else + allocator; + var t = try std.time.Timer.start(); switch (self.defn) { - .simple => |func| func(allocator), - .parameterised => |x| x.func(@ptrCast(x.context), allocator), + .simple => |func| func(alloc), + .parameterised => |x| x.func(@ptrCast(x.context), alloc), } return Runner.Reading{ .timing_ns = t.read(), + .max_allocated = if (self.config.track_allocations) + tracking.maxAllocated() + else + null, }; } }; @@ -184,6 +199,7 @@ pub const Benchmark = struct { config.iterations, config.max_iterations, config.time_budget_ns, + config.track_allocations, ); break :blk &self.runner.?; }; @@ -281,6 +297,7 @@ pub const Result = struct { allocator: std.mem.Allocator, name: []const u8, timings_ns: []const u64, + max_allocations: ?[]const usize, pub fn init( allocator: std.mem.Allocator, @@ -288,15 +305,19 @@ pub const Result = struct { readings: Runner.Readings, ) !Result { std.sort.heap(u64, readings.timings_ns, {}, std.sort.asc(u64)); + if (readings.max_allocations) |allocs| + std.sort.heap(usize, allocs, {}, std.sort.asc(usize)); return Result{ .allocator = allocator, .name = name, .timings_ns = readings.timings_ns, + .max_allocations = readings.max_allocations, }; } pub fn deinit(self: Result) void { self.allocator.free(self.timings_ns); + if (self.max_allocations) |allocs| self.allocator.free(allocs); } /// Formats and prints the benchmark result in a human readable format. @@ -339,6 +360,41 @@ pub const Result = struct { // End of line try setColor(colors, writer, Color.reset); try writer.writeAll("\n"); + + if (self.max_allocations) |allocs| { + const m = Statistics(usize).init(allocs); + // Benchmark name + const name = try std.fmt.bufPrint(&buf, "{s} [MEMORY]", .{ + self.name, + }); + try writer.print("{s:<46} ", .{name}); + // Mean + standard deviation + try setColor(colors, writer, Color.green); + try writer.print("{s:<23}", .{ + try std.fmt.bufPrint(&buf, "{:.3} ± {:.3}", .{ + std.fmt.fmtIntSizeBin(m.mean), + std.fmt.fmtIntSizeBin(m.stddev), + }), + }); + // Minimum and maximum + try setColor(colors, writer, Color.blue); + try writer.print("{s:<29}", .{ + try std.fmt.bufPrint(&buf, "({:.3} ... {:.3})", .{ + std.fmt.fmtIntSizeBin(m.min), + std.fmt.fmtIntSizeBin(m.max), + }), + }); + // Percentiles + try setColor(colors, writer, Color.cyan); + try writer.print("{:<10} {:<10} {:<10}", .{ + std.fmt.fmtIntSizeBin(m.percentiles.p75), + std.fmt.fmtIntSizeBin(m.percentiles.p99), + std.fmt.fmtIntSizeBin(m.percentiles.p995), + }); + // End of line + try setColor(colors, writer, Color.reset); + try writer.writeAll("\n"); + } } fn setColor(colors: bool, writer: anytype, color: Color) !void { @@ -347,15 +403,32 @@ pub const Result = struct { pub fn writeJSON(self: Result, writer: anytype) !void { const timings_ns_stats = Statistics(u64).init(self.timings_ns); - try writer.print( - \\{{ "name": "{s}", - \\ "timing_statistics": {}, "timings": {} }} - , - .{ - std.fmt.fmtSliceEscapeLower(self.name), - statistics.fmtJSON(u64, "nanoseconds", timings_ns_stats), - format.fmtJSONArray(u64, self.timings_ns), - }, - ); + if (self.max_allocations) |allocs| { + const max_allocated_stats = Statistics(usize).init(allocs); + try writer.print( + \\{{ "name": "{s}", + \\ "timing_statistics": {}, "timings": {}, + \\ "max_allocation_statistics": {}, "max_allocations": {} }} + , + .{ + std.fmt.fmtSliceEscapeLower(self.name), + statistics.fmtJSON(u64, "nanoseconds", timings_ns_stats), + format.fmtJSONArray(u64, self.timings_ns), + statistics.fmtJSON(usize, "bytes", max_allocated_stats), + format.fmtJSONArray(usize, allocs), + }, + ); + } else { + try writer.print( + \\{{ "name": "{s}", + \\ "timing_statistics": {}, "timings": {} }} + , + .{ + std.fmt.fmtSliceEscapeLower(self.name), + statistics.fmtJSON(u64, "nanoseconds", timings_ns_stats), + format.fmtJSONArray(u64, self.timings_ns), + }, + ); + } } };