From 5ddd034d2cba7cade919e104a2ae318fb0b89316 Mon Sep 17 00:00:00 2001 From: Tarry Singh Date: Wed, 20 Dec 2023 17:18:02 +0000 Subject: [PATCH] pjrt: Fix profiler by allowing i64 resource IDs and reserving memory when creating array lists. --- pjrt/convert/trace_container.zig | 155 +++++++++++++------------------ pjrt/convert/xplane_visitor.zig | 36 +++++-- pjrt/profiler.zig | 3 +- 3 files changed, 95 insertions(+), 99 deletions(-) diff --git a/pjrt/convert/trace_container.zig b/pjrt/convert/trace_container.zig index a967087..7156a1b 100644 --- a/pjrt/convert/trace_container.zig +++ b/pjrt/convert/trace_container.zig @@ -32,62 +32,43 @@ pub const TraceContainer = struct { devices: std.AutoArrayHashMapUnmanaged(u32, Device) = .{}, pub const Device = struct { - name: []const u8 = &[_]u8{}, - device_id: u32 = 0, - resources: std.AutoArrayHashMapUnmanaged(u32, trace_events_proto.Resource) = .{}, + name: []const u8, + device_id: u32, + resources: std.AutoArrayHashMapUnmanaged(i64, Resource) = .{}, + }; + + pub const Resource = struct { + name: []const u8, + sort_index: i64, }; pub const TraceEvent = struct { device_id: u32 = 0, - resource_id: u32 = 0, + resource_id: i64 = 0, name: []const u8 = &[_]u8{}, timestamp_ps: u128 = 0, duration_ps: u64 = 0, args: std.StringArrayHashMapUnmanaged([]const u8) = .{}, }; - pub fn init(allocator: std.mem.Allocator, pb_buffer: []const u8, max_events: ?usize) !TraceContainer { - var self: TraceContainer = .{ + pub fn init(allocator: std.mem.Allocator) TraceContainer { + return .{ .arena = std.heap.ArenaAllocator.init(allocator), }; - const arena = self.arena.allocator(); - - const xspace = try xplane_proto.XSpace.decode(pb_buffer, arena); - try self.fromXSpace(arena, &xspace, max_events); - return self; } pub fn deinit(self: *TraceContainer) void { self.arena.deinit(); } - fn picoToMicro(p: anytype) f64 { - return @as(f64, @floatFromInt(p)) / 1E6; + pub fn parseXSpaceBytes(self: *TraceContainer, pb_buffer: []const u8, max_events: ?usize) !void { + const arena = self.arena.allocator(); + + const xspace = try xplane_proto.XSpace.decode(pb_buffer, arena); + try self.fromXSpace(arena, xspace, max_events); } - fn xLineDisplayId(xline: *const xplane_proto.XLine) i64 { - return if (xline.display_id != 0) xline.display_id else xline.id; - } - - fn xLineDisplayName(xline: *const xplane_proto.XLine) []const u8 { - return switch (xline.display_name) { - .Empty => xline.name.getSlice(), - else => xline.display_name.getSlice(), - }; - } - - fn xstatValueToString(stat: *const xplane_proto.XStat, plane: *const xplane_visitor.XPlaneVisitor, writer: anytype) !void { - if (stat.value) |val| { - switch (val) { - inline .int64_value, .uint64_value, .double_value => |v| try writer.print("{d}", .{v}), - .str_value => |*v| try writer.writeAll(v.getSlice()), - .bytes_value => try writer.writeAll(""), - .ref_value => |v| try writer.writeAll(plane.getStatMetadataName(@intCast(v))), - } - } else return; - } - - fn findPlaneWithName(space: *const xplane_proto.XSpace, name: []const u8) ?*xplane_proto.XPlane { + fn findPlaneWithName(space: xplane_proto.XSpace, name: []const u8) ?*xplane_proto.XPlane { for (space.planes.items) |*v| { if (std.mem.eql(u8, v.name.getSlice(), name)) return v; } @@ -95,49 +76,37 @@ pub const TraceContainer = struct { } fn findPlanesWithPrefix( - allocator: std.mem.Allocator, - space: *const xplane_proto.XSpace, + out: *std.ArrayList(*const xplane_proto.XPlane), + space: xplane_proto.XSpace, prefix: []const u8, - ) ![]*const xplane_proto.XPlane { - var res = std.ArrayList(*const xplane_proto.XPlane).init(allocator); + ) !void { for (space.planes.items) |*p| { - if (std.mem.startsWith(u8, p.name.getSlice(), prefix)) try res.append(p); - } - return res.toOwnedSlice(); - } - - fn buildDeviceAndResources(allocator: std.mem.Allocator, device_id: u32, plane: *const xplane_visitor.XPlaneVisitor, device: *Device) !void { - device.name = plane.name(); - device.device_id = device_id; - const sort_by_ordinal = (device_id == host_threads_device_id); - var ordinal: u32 = 0; - for (plane.plane.lines.items) |*xline| { - const resource_id: u32 = @intCast(xLineDisplayId(xline)); - var resource: trace_events_proto.Resource = .{ - .resource_id = resource_id, - .name = .{ .Const = xLineDisplayName(xline) }, - }; - - if (sort_by_ordinal) { - ordinal += 1; - resource.sort_index = ordinal; + if (std.mem.startsWith(u8, p.name.getSlice(), prefix)) { + try out.append(p); } - try device.resources.put(allocator, resource_id, resource); } } fn xplaneToTraceEvents(self: *TraceContainer, allocator: std.mem.Allocator, device_id: u32, xplane: *const xplane_visitor.XPlaneVisitor) !void { // Convert devices and resources. - const device_entry = try self.devices.getOrPut(allocator, device_id); - if (!device_entry.found_existing) device_entry.value_ptr.* = .{}; + const device_entry = try self.devices.getOrPutValue(allocator, device_id, .{ .name = xplane.name(), .device_id = device_id }); + var device = device_entry.value_ptr.*; + defer device_entry.value_ptr.* = device; - try buildDeviceAndResources(allocator, device_id, xplane, device_entry.value_ptr); + try device.resources.ensureUnusedCapacity(allocator, xplane.plane.lines.items.len); + const sort_by_ordinal = (device_id == host_threads_device_id); // Convert events. - for (xplane.plane.lines.items) |*xline| { - const resource_id: u32 = @intCast(xLineDisplayId(xline)); + for (xplane.plane.lines.items, 0..) |*xline, ordinal| { + const resource_id = if (xline.display_id != 0) xline.display_id else xline.id; + const resource_name = if (xline.display_name.isEmpty()) xline.name.getSlice() else xline.display_name.getSlice(); + device.resources.putAssumeCapacity(resource_id, .{ + .name = resource_name, + .sort_index = if (sort_by_ordinal) @intCast(ordinal) else resource_id, + }); + + if (std.mem.eql(u8, resource_name, xla_async_op_line_name)) continue; - if (std.mem.eql(u8, xLineDisplayName(xline), xla_async_op_line_name)) continue; for (xline.events.items) |*xevent| { const event_type = xplane.getEventType(xevent.metadata_id); if (event_type.isInternalEvent()) continue; @@ -146,9 +115,11 @@ pub const TraceContainer = struct { event.resource_id = resource_id; if (xplane.event_metadata_by_id.get(xevent.metadata_id)) |metadata| { + try event.args.ensureUnusedCapacity(allocator, 1 + metadata.stats.items.len); + if (metadata.display_name != .Empty) { event.name = metadata.display_name.getSlice(); - try event.args.put(allocator, "long_name", metadata.name.getSlice()); + event.args.putAssumeCapacity("long_name", metadata.name.getSlice()); } else { event.name = metadata.name.getSlice(); } @@ -156,51 +127,54 @@ pub const TraceContainer = struct { event.timestamp_ps = (@as(u128, @intCast(xline.timestamp_ns)) * 1000) + @as(u128, @intCast(xevent.data.?.offset_ps)); event.duration_ps = @intCast(xevent.duration_ps); - for (metadata.stats.items) |*xstat| { + for (metadata.stats.items) |xstat| { if (xstat.value == null) continue; var stat_buffer = std.ArrayList(u8).init(allocator); - try xstatValueToString(xstat, xplane, stat_buffer.writer().any()); + try xplane.xstatToString(xstat, stat_buffer.writer().any()); const stat_str = try stat_buffer.toOwnedSlice(); const stat_type = xplane.getStatType(xstat.metadata_id); if (stat_type.isInternalStat()) continue; if (stat_type == .step_name) event.name = stat_str; - try event.args.put(allocator, xplane.getStatMetadataName(xstat.metadata_id), stat_str); + event.args.putAssumeCapacity(xplane.getStatMetadataName(xstat.metadata_id), stat_str); } } - for (xevent.stats.items) |*xstat| { + try event.args.ensureUnusedCapacity(allocator, xevent.stats.items.len); + for (xevent.stats.items) |xstat| { if (xstat.value == null) continue; var stat_buffer = std.ArrayList(u8).init(allocator); - try xstatValueToString(xstat, xplane, stat_buffer.writer().any()); + try xplane.xstatToString(xstat, stat_buffer.writer().any()); const stat_str = try stat_buffer.toOwnedSlice(); const stat_type = xplane.getStatType(xstat.metadata_id); if (stat_type.isInternalStat()) continue; if (stat_type == .step_name) event.name = stat_str; - try event.args.put(allocator, xplane.getStatMetadataName(xstat.metadata_id), stat_str); + event.args.putAssumeCapacity(xplane.getStatMetadataName(xstat.metadata_id), stat_str); } } } } - fn fromXSpace(self: *TraceContainer, allocator: std.mem.Allocator, xspace: *const xplane_proto.XSpace, max_events: ?usize) !void { + fn fromXSpace(self: *TraceContainer, allocator: std.mem.Allocator, xspace: xplane_proto.XSpace, max_events: ?usize) !void { if (findPlaneWithName(xspace, host_threads_plane_name)) |hp| { const xplane = try xplane_visitor.XPlaneVisitor.init(allocator, hp); try self.xplaneToTraceEvents(allocator, host_threads_device_id, &xplane); } - var device_planes = try findPlanesWithPrefix(allocator, xspace, gpu_plane_prefix); + var device_planes = std.ArrayList(*const xplane_proto.XPlane).init(allocator); + defer device_planes.deinit(); - // We don't expect GPU and TPU planes and custom devices to be present in the - // same XSpace. - if (device_planes.len == 0) { - device_planes = try findPlanesWithPrefix(allocator, xspace, tpu_plane_prefix); + try findPlanesWithPrefix(&device_planes, xspace, gpu_plane_prefix); + // We don't expect GPU and TPU planes and custom devices to be present in the same XSpace. + if (device_planes.items.len == 0) { + try findPlanesWithPrefix(&device_planes, xspace, tpu_plane_prefix); } - if (device_planes.len == 0) { - device_planes = try findPlanesWithPrefix(allocator, xspace, custom_plane_prefix); + if (device_planes.items.len == 0) { + try findPlanesWithPrefix(&device_planes, xspace, custom_plane_prefix); } - for (device_planes) |dp| { - const xplane = try xplane_visitor.XPlaneVisitor.init(allocator, dp); + for (device_planes.items) |dp| { + var xplane = try xplane_visitor.XPlaneVisitor.init(allocator, dp); + defer xplane.deinit(allocator); const device_id: u32 = first_device_id + @as(u32, @intCast(xplane.plane.id)); try self.xplaneToTraceEvents(allocator, device_id, &xplane); } @@ -257,26 +231,25 @@ pub const TraceContainer = struct { }); device.resources.sort(struct { - keys: []const u32, + keys: []const i64, pub fn lessThan(ctx: @This(), lhs: usize, rhs: usize) bool { return ctx.keys[lhs] < ctx.keys[rhs]; } }{ .keys = device.resources.keys() }); for (device.resources.keys(), device.resources.values()) |resource_id, resource| { - if (resource.name.getSlice().len != 0) { + if (resource.name.len != 0) { try writer.print( \\{{"ph":"M","pid":{d},"tid":{d},"name":"thread_name","args":{{"name":"{s}"}}}}, , .{ device_id, resource_id, - resource.name.getSlice(), + resource.name, }); } - const sort_index = if (resource.sort_index != 0) resource.sort_index else resource_id; try writer.print( \\{{"ph":"M","pid":{d},"tid":{d},"name":"thread_sort_index","args":{{"sort_index":{d}}}}}, - , .{ device_id, resource_id, sort_index }); + , .{ device_id, resource_id, resource.sort_index }); } } @@ -321,3 +294,7 @@ pub const TraceContainer = struct { try writer.writeAll("{}]}"); } }; + +fn picoToMicro(p: anytype) f64 { + return @as(f64, @floatFromInt(p)) / 1E6; +} diff --git a/pjrt/convert/xplane_visitor.zig b/pjrt/convert/xplane_visitor.zig index 091bec9..a2f8eba 100644 --- a/pjrt/convert/xplane_visitor.zig +++ b/pjrt/convert/xplane_visitor.zig @@ -13,38 +13,56 @@ pub const XPlaneVisitor = struct { ) !XPlaneVisitor { var res: XPlaneVisitor = .{ .plane = plane }; + try res.event_metadata_by_id.ensureUnusedCapacity(allocator, @intCast(plane.event_metadata.items.len)); // build event metadata map for (plane.event_metadata.items) |*event_metadata| { - try res.event_metadata_by_id.put(allocator, event_metadata.key, &event_metadata.value.?); + res.event_metadata_by_id.putAssumeCapacity(event_metadata.key, &event_metadata.value.?); } // build stat metadata map + try res.stat_metadata_by_id.ensureUnusedCapacity(allocator, @intCast(plane.stat_metadata.items.len)); for (plane.stat_metadata.items) |*stat_metadata| { - try res.stat_metadata_by_id.put(allocator, stat_metadata.key, &stat_metadata.value.?); + res.stat_metadata_by_id.putAssumeCapacity(stat_metadata.key, &stat_metadata.value.?); } return res; } - pub fn getEventType(self: *const XPlaneVisitor, event_metadata_id: i64) xplane_schema.HostEventType { + pub fn deinit(self: *XPlaneVisitor, allocator: std.mem.Allocator) void { + self.stat_metadata_by_id.deinit(allocator); + self.event_metadata_by_id.deinit(allocator); + } + + pub fn name(self: XPlaneVisitor) []const u8 { + return self.plane.name.getSlice(); + } + + pub fn getEventType(self: XPlaneVisitor, event_metadata_id: i64) xplane_schema.HostEventType { if (self.event_metadata_by_id.get(event_metadata_id)) |v| { return xplane_schema.HostEventType.fromString(v.name.getSlice()); } else return .unknown; } - pub fn name(self: *const XPlaneVisitor) []const u8 { - return self.plane.name.getSlice(); - } - - pub fn getStatMetadataName(self: *const XPlaneVisitor, stat_metadata_id: i64) []const u8 { + pub fn getStatMetadataName(self: XPlaneVisitor, stat_metadata_id: i64) []const u8 { if (self.stat_metadata_by_id.get(stat_metadata_id)) |v| { return v.name.getSlice(); } else return &[_]u8{}; } - pub fn getStatType(self: *const XPlaneVisitor, stat_metadata_id: i64) xplane_schema.StatType { + pub fn getStatType(self: XPlaneVisitor, stat_metadata_id: i64) xplane_schema.StatType { if (self.stat_metadata_by_id.get(stat_metadata_id)) |v| { return xplane_schema.StatType.fromString(v.name.getSlice()); } else return .unknown; } + + pub fn xstatToString(self: XPlaneVisitor, stat: xplane_proto.XStat, writer: anytype) !void { + if (stat.value == null) return; + + switch (stat.value.?) { + inline .int64_value, .uint64_value, .double_value => |v| try writer.print("{d}", .{v}), + .str_value => |*v| try writer.writeAll(v.getSlice()), + .bytes_value => try writer.writeAll(""), + .ref_value => |v| try writer.writeAll(self.getStatMetadataName(@intCast(v))), + } + } }; diff --git a/pjrt/profiler.zig b/pjrt/profiler.zig index 1e6e2f0..1b0d1a6 100644 --- a/pjrt/profiler.zig +++ b/pjrt/profiler.zig @@ -135,8 +135,9 @@ pub const Profiler = struct { return; } - var converter = try TraceContainer.init(allocator, profile_data.items(), 1_000_000); + var converter = TraceContainer.init(allocator); defer converter.deinit(); + try converter.parseXSpaceBytes(profile_data.items(), 1_000_000); var output_file = try dir.createFile(file_name, .{}); defer output_file.close();