lib/accy/src/executable/candidate.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const gpu = @import("gpu");
  3 const choir_abi = @import("choir_abi");
  4 const accy_root = @import("../root.zig");
  5 const artifact_product = @import("../artifact/root.zig");
  6 const exec_product = @import("plan.zig");
  7 const tuning_mod = @import("tuning.zig");
  8 
  9 pub const LaunchCandidateMeasurement = tuning_mod.LaunchCandidateMeasurement;
 10 pub const LaunchOptions = exec_product.LaunchOptions;
 11 
 12 pub const LaunchCandidateSynchronization = enum {
 13     none,
 14     device,
 15     stream,
 16     event,
 17 };
 18 
 19 pub const LaunchCandidateBenchmarkOptions = struct {
 20     warmup: u32 = 1,
 21     samples: u32 = 5,
 22     base_options: LaunchOptions = .{},
 23     synchronize: LaunchCandidateSynchronization = .none,
 24 };
 25 
 26 pub const LaunchCandidateRecord = struct {
 27     kernel: artifact_product.KernelSummary,
 28     candidate_index: usize,
 29     geometry: choir_abi.LaunchGeometry,
 30     candidate_score: u32,
 31     estimated_static_bytes_per_threadgroup: u64,
 32     estimated_element_ops_per_threadgroup: u64,
 33     median_ns: u64,
 34     sample_count: u32,
 35 };
 36 
 37 pub fn launchCandidateRecord(
 38     planned: artifact_product.PlannedKernel,
 39     candidate: artifact_product.LaunchResourceCandidate,
 40     measurement: LaunchCandidateMeasurement,
 41 ) gpu.BackendError!LaunchCandidateRecord {
 42     return .{
 43         .kernel = try artifact_product.summarizePlannedKernel(planned),
 44         .candidate_index = measurement.candidate_index,
 45         .geometry = candidate.geometry,
 46         .candidate_score = candidate.score,
 47         .estimated_static_bytes_per_threadgroup = candidate.estimated_static_bytes_per_threadgroup,
 48         .estimated_element_ops_per_threadgroup = candidate.estimated_element_ops_per_threadgroup,
 49         .median_ns = measurement.median_ns,
 50         .sample_count = measurement.sample_count,
 51     };
 52 }
 53 
 54 pub fn plannedKernelForLaunchCandidateRecord(
 55     artifact_plan: *const artifact_product.BackendArtifactPlan,
 56     record: LaunchCandidateRecord,
 57 ) gpu.BackendError!artifact_product.PlannedKernel {
 58     for (artifact_plan.kernels.items) |planned| {
 59         if (planned.kernel_id != record.kernel.kernel_id) continue;
 60         if (!try launchCandidateRecordMatchesPlannedKernel(planned, record)) return error.LaunchArgumentMismatch;
 61         return planned;
 62     }
 63     return error.LaunchArgumentMismatch;
 64 }
 65 
 66 pub fn launchCandidateRecordMatchesPlannedKernel(
 67     planned: artifact_product.PlannedKernel,
 68     record: LaunchCandidateRecord,
 69 ) gpu.BackendError!bool {
 70     const planned_summary = try artifact_product.summarizePlannedKernel(planned);
 71     if (!artifact_product.kernelSummariesEqual(planned_summary, record.kernel)) return false;
 72     if (record.candidate_index >= planned.launch_resources.candidate_count) return false;
 73     const candidate = planned.launch_resources.candidates[record.candidate_index];
 74     return artifact_product.launchGeometriesEqual(candidate.geometry, record.geometry) and
 75         candidate.score == record.candidate_score and
 76         candidate.estimated_static_bytes_per_threadgroup == record.estimated_static_bytes_per_threadgroup and
 77         candidate.estimated_element_ops_per_threadgroup == record.estimated_element_ops_per_threadgroup and
 78         record.sample_count != 0;
 79 }
 80 
 81 pub fn copyKernelArtifactToAllocator(
 82     allocator: std.mem.Allocator,
 83     source: gpu.KernelArtifact,
 84 ) gpu.BackendError!gpu.KernelArtifact {
 85     var artifact = gpu.KernelArtifact.init(allocator, .{
 86         .backend = source.backend,
 87         .format = source.format,
 88         .entry_name = source.entry_name,
 89         .argument_count = source.argument_count,
 90         .scalar_argument_count = source.scalar_argument_count,
 91         .diagnostic_id = source.diagnostic_id,
 92         .interface = source.interface,
 93     }) catch return error.OutOfMemory;
 94     errdefer artifact.deinit();
 95 
 96     switch (source.payload) {
 97         .text => |text| try artifact.setOwnedText(text),
 98         .bytes => |bytes| try artifact.setOwnedBytes(bytes),
 99         .words_u32 => |words| try artifact.setOwnedWords(words),
100         .none, .external => return error.InvalidArtifact,
101     }
102 
103     return artifact;
104 }