tiny.accy.artifact.summary

1pub const summary = @import("summary.zig");

1const std = @import("std");2const gpu = @import("gpu");3const choir_abi = @import("choir_abi");4const artifact = @import("root.zig");56pub const KernelSource = enum {7    tensor,8    kernel_call,9    choir_kernel,10};1112pub const KernelSummary = struct {13    source: KernelSource,14    kernel_id: usize,15    work_item_id: usize,16    output_layout_fingerprint: u64,17    input_layout_fingerprint: u64,18    element_count: u64,19    op_count: usize,20    entry_name: []const u8,21    artifact_format: gpu.ArtifactFormat,22    artifact_payload_bytes: usize,23    compile_argument_count: u32,24    compile_required_dtype_bits: u64,25    compile_required_features: choir_abi.Features,26    compile_required_subgroup: choir_abi.SubgroupRequirements,27    compile_payload: artifact.PlannedKernelCompilePayload,28    compile_payload_bytes: usize,29    compile_launch: artifact.PlannedKernelCompileLaunch,30    runtime_scalar_argument_count: u32,31    launch_geometry: choir_abi.LaunchGeometry,32    launch_candidate_count: usize,33    launch_resource_class: []const u8,34    fixed_threadgroup: bool,35    subgroup_aligned: bool,36    subgroup_size: ?u32,37    static_bytes_complete: bool,38};3940pub const KernelSummaries = struct {41    allocator: std.mem.Allocator,42    items: []KernelSummary,4344    pub fn deinit(self: *KernelSummaries) void {45        for (self.items) |item| {46            self.allocator.free(item.entry_name);47        }48        self.allocator.free(self.items);49        self.* = undefined;50    }5152    pub fn len(self: *const KernelSummaries) usize {53        return self.items.len;54    }5556    pub fn summary(self: *const KernelSummaries, kernel_index: usize) gpu.BackendError!KernelSummary {57        if (kernel_index >= self.items.len) return error.InvalidArtifact;58        return self.items[kernel_index];59    }6061    pub fn summaryForWork(self: *const KernelSummaries, work_item_id: usize) gpu.BackendError!KernelSummary {62        for (self.items) |item| {63            if (item.work_item_id == work_item_id) return item;64        }65        return error.InvalidArtifact;66    }67};6869pub fn copyKernelSummary(allocator: std.mem.Allocator, summary: KernelSummary) gpu.BackendError!KernelSummary {70    var copied = summary;71    copied.entry_name = allocator.dupe(u8, summary.entry_name) catch return error.OutOfMemory;72    return copied;73}7475pub fn kernelSource(source: artifact.PlannedKernelSource) KernelSource {76    return switch (source) {77        .tensor => .tensor,78        .kernel_call => .kernel_call,79        .choir_kernel => .choir_kernel,80    };81}8283pub fn artifactPayloadByteCount(kernel_artifact: gpu.KernelArtifact) gpu.BackendError!usize {84    return switch (kernel_artifact.payload) {85        .bytes => |bytes| bytes.len,86        .words_u32 => |words| words.len * @sizeOf(u32),87        .text => |text| text.len,88        .none, .external => error.InvalidArtifact,89    };90}9192pub fn launchResourceClassName(class: artifact.LaunchResourceClass) []const u8 {93    return switch (class) {94        .unknown => "unknown",95        .memory_bound => "memory_bound",96        .balanced => "balanced",97        .compute_weighted => "compute_weighted",98    };99}100101pub fn summarizePlannedKernel(planned: artifact.PlannedKernel) gpu.BackendError!KernelSummary {102    return .{103        .source = kernelSource(planned.compile.source),104        .kernel_id = planned.kernel_id,105        .work_item_id = planned.work_item_id,106        .output_layout_fingerprint = planned.output_layout_fingerprint,107        .input_layout_fingerprint = planned.input_layout_fingerprint,108        .element_count = planned.element_count,109        .op_count = planned.op_count,110        .entry_name = planned.artifact.entry_name,111        .artifact_format = planned.artifact.format,112        .artifact_payload_bytes = try artifactPayloadByteCount(planned.artifact),113        .compile_argument_count = planned.compile.argument_count,114        .compile_required_dtype_bits = planned.compile.required_dtypes.bits,115        .compile_required_features = planned.compile.required_features,116        .compile_required_subgroup = planned.compile.required_subgroup,117        .compile_payload = planned.compile.payload,118        .compile_payload_bytes = planned.compile.payload_byte_count,119        .compile_launch = planned.compile.launch,120        .runtime_scalar_argument_count = planned.runtime_scalar_argument_count,121        .launch_geometry = planned.launch_resources.geometry,122        .launch_candidate_count = planned.launch_resources.candidate_count,123        .launch_resource_class = launchResourceClassName(planned.launch_resources.resource_class),124        .fixed_threadgroup = planned.launch_resources.fixed_threadgroup,125        .subgroup_aligned = planned.launch_resources.subgroup_aligned,126        .subgroup_size = planned.launch_resources.subgroup_size,127        .static_bytes_complete = planned.launch_resources.static_bytes_complete,128    };129}130131pub fn kernelSummariesEqual(lhs: KernelSummary, rhs: KernelSummary) bool {132    return lhs.source == rhs.source and133        lhs.kernel_id == rhs.kernel_id and134        lhs.work_item_id == rhs.work_item_id and135        lhs.output_layout_fingerprint == rhs.output_layout_fingerprint and136        lhs.input_layout_fingerprint == rhs.input_layout_fingerprint and137        lhs.element_count == rhs.element_count and138        lhs.op_count == rhs.op_count and139        std.mem.eql(u8, lhs.entry_name, rhs.entry_name) and140        lhs.artifact_format == rhs.artifact_format and141        lhs.artifact_payload_bytes == rhs.artifact_payload_bytes and142        lhs.compile_argument_count == rhs.compile_argument_count and143        lhs.compile_required_dtype_bits == rhs.compile_required_dtype_bits and144        std.meta.eql(lhs.compile_required_features, rhs.compile_required_features) and145        std.meta.eql(lhs.compile_required_subgroup, rhs.compile_required_subgroup) and146        lhs.compile_payload == rhs.compile_payload and147        lhs.compile_payload_bytes == rhs.compile_payload_bytes and148        lhs.compile_launch == rhs.compile_launch and149        lhs.runtime_scalar_argument_count == rhs.runtime_scalar_argument_count and150        launchGeometriesEqual(lhs.launch_geometry, rhs.launch_geometry) and151        lhs.launch_candidate_count == rhs.launch_candidate_count and152        std.mem.eql(u8, lhs.launch_resource_class, rhs.launch_resource_class) and153        lhs.fixed_threadgroup == rhs.fixed_threadgroup and154        lhs.subgroup_aligned == rhs.subgroup_aligned and155        lhs.subgroup_size == rhs.subgroup_size and156        lhs.static_bytes_complete == rhs.static_bytes_complete;157}158159pub fn launchGeometriesEqual(lhs: choir_abi.LaunchGeometry, rhs: choir_abi.LaunchGeometry) bool {160    return lhs.grid[0] == rhs.grid[0] and161        lhs.grid[1] == rhs.grid[1] and162        lhs.grid[2] == rhs.grid[2] and163        lhs.threadgroup[0] == rhs.threadgroup[0] and164        lhs.threadgroup[1] == rhs.threadgroup[1] and165        lhs.threadgroup[2] == rhs.threadgroup[2] and166        lhs.dynamic_shared_memory_bytes == rhs.dynamic_shared_memory_bytes;167}168169fn testKernelSummary() KernelSummary {170    return .{171        .source = .tensor,172        .kernel_id = 7,173        .work_item_id = 3,174        .output_layout_fingerprint = 11,175        .input_layout_fingerprint = 13,176        .element_count = 1024,177        .op_count = 4,178        .entry_name = "accy_test_kernel",179        .artifact_format = .cuda_ptx,180        .artifact_payload_bytes = 128,181        .compile_argument_count = 3,182        .compile_required_dtype_bits = gpu.DTypeSet.init(&.{ .f32, .i32 }).bits,183        .compile_required_features = .{},184        .compile_required_subgroup = .{},185        .compile_payload = .text,186        .compile_payload_bytes = 96,187        .compile_launch = .generic,188        .runtime_scalar_argument_count = 0,189        .launch_geometry = .{190            .grid = .{ 8, 2, 1 },191            .threadgroup = .{ 128, 1, 1 },192            .dynamic_shared_memory_bytes = 256,193        },194        .launch_candidate_count = 2,195        .launch_resource_class = "balanced",196        .fixed_threadgroup = false,197        .subgroup_aligned = true,198        .subgroup_size = 32,199        .static_bytes_complete = true,200    };201}202203test "artifact kernel summary equality includes compile contract" {204    const base = testKernelSummary();205    try std.testing.expect(kernelSummariesEqual(base, testKernelSummary()));206207    var changed = testKernelSummary();208    changed.compile_argument_count += 1;209    try std.testing.expect(!kernelSummariesEqual(base, changed));210211    changed = testKernelSummary();212    changed.compile_required_dtype_bits ^= gpu.DTypeSet.init(&.{.f16}).bits;213    try std.testing.expect(!kernelSummariesEqual(base, changed));214215    changed = testKernelSummary();216    changed.compile_required_features = .{ .tensor_cores = true };217    try std.testing.expect(!kernelSummariesEqual(base, changed));218219    changed = testKernelSummary();220    changed.compile_required_subgroup = .{ .supported = true, .shuffle = true };221    try std.testing.expect(!kernelSummariesEqual(base, changed));222223    changed = testKernelSummary();224    changed.compile_payload = .bytes;225    try std.testing.expect(!kernelSummariesEqual(base, changed));226227    changed = testKernelSummary();228    changed.compile_payload_bytes += 1;229    try std.testing.expect(!kernelSummariesEqual(base, changed));230231    changed = testKernelSummary();232    changed.compile_launch = .dot_general;233    try std.testing.expect(!kernelSummariesEqual(base, changed));234235    changed = testKernelSummary();236    changed.runtime_scalar_argument_count += 1;237    try std.testing.expect(!kernelSummariesEqual(base, changed));238239    changed = testKernelSummary();240    changed.launch_geometry.dynamic_shared_memory_bytes += 1;241    try std.testing.expect(!kernelSummariesEqual(base, changed));242}