lib/accy/src/artifact/job.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const artifact = @import("root.zig");
4 const model = @import("model/root.zig");
5 const target = @import("../target/root.zig");
6
7 const fingerprint_mod = artifact.fingerprint;
8 const plan = artifact.plan;
9
10 const ArtifactPlanOptions = model.ArtifactPlanOptions;
11 const BackendArtifactPlan = plan.BackendArtifactPlan;
12
13 /// A heap-allocated holder of one artifact plan, the compiled kernels and the
14 /// plan that describes them. A caller holds one of these between compiling a
15 /// prepared program to device code and copying that code into a runnable
16 /// program. `init` takes ownership of the plan and allocates the holder with
17 /// the given allocator, and `deinit` frees the plan and then the holder. The
18 /// holder lives only as long as the compile that made it and is never stored as
19 /// a stage record. The holder answers kernel counts and per-kernel summaries by
20 /// index or by work item, and it copies the plan or the summaries out for a
21 /// caller to keep.
22 pub const ArtifactJob = struct {
23 allocator: std.mem.Allocator,
24 artifact_plan: BackendArtifactPlan,
25
26 pub fn init(
27 allocator: std.mem.Allocator,
28 artifact_plan: BackendArtifactPlan,
29 ) !*ArtifactJob {
30 const artifact_job = try allocator.create(ArtifactJob);
31 errdefer allocator.destroy(artifact_job);
32 artifact_job.* = .{
33 .allocator = allocator,
34 .artifact_plan = artifact_plan,
35 };
36 return artifact_job;
37 }
38
39 pub fn deinit(self: *ArtifactJob) void {
40 self.artifact_plan.deinit();
41 const allocator = self.allocator;
42 allocator.destroy(self);
43 }
44
45 /// A 64-bit hash of the plan's target profile, kind, format, slots, input
46 /// and output ids, kernels and totals. A caller shows this number in
47 /// reports or stamps it beside compiled output as a label, and no identity
48 /// check or reuse decision compares it.
49 pub fn fingerprint(self: *const ArtifactJob) u64 {
50 return fingerprint_mod.artifactPlan(&self.artifact_plan);
51 }
52
53 pub fn kernelCount(self: *const ArtifactJob) usize {
54 return self.artifact_plan.kernelCount();
55 }
56
57 pub fn kernelSummary(self: *const ArtifactJob, kernel_index: usize) gpu.BackendError!artifact.KernelSummary {
58 if (kernel_index >= self.artifact_plan.kernels.items.len) return error.InvalidArtifact;
59 return try artifact.summarizePlannedKernel(self.artifact_plan.kernels.items[kernel_index]);
60 }
61
62 pub fn kernelSummaryForWork(self: *const ArtifactJob, work_item_id: usize) gpu.BackendError!artifact.KernelSummary {
63 for (self.artifact_plan.kernels.items) |kernel| {
64 if (kernel.work_item_id == work_item_id) return try artifact.summarizePlannedKernel(kernel);
65 }
66 return error.InvalidArtifact;
67 }
68
69 pub fn artifactPlan(self: *const ArtifactJob) *const BackendArtifactPlan {
70 return &self.artifact_plan;
71 }
72
73 pub fn copyKernelSummaries(self: *const ArtifactJob, result_allocator: std.mem.Allocator) gpu.BackendError!artifact.KernelSummaries {
74 const items = result_allocator.alloc(artifact.KernelSummary, self.artifact_plan.kernels.items.len) catch return error.OutOfMemory;
75 var copied: usize = 0;
76 errdefer {
77 for (items[0..copied]) |summary| {
78 result_allocator.free(summary.entry_name);
79 }
80 result_allocator.free(items);
81 }
82
83 for (items, self.artifact_plan.kernels.items) |*item, kernel| {
84 item.* = try artifact.copyKernelSummary(result_allocator, try artifact.summarizePlannedKernel(kernel));
85 copied += 1;
86 }
87
88 return .{
89 .allocator = result_allocator,
90 .items = items,
91 };
92 }
93
94 pub fn copyArtifactPlan(self: *const ArtifactJob, allocator: std.mem.Allocator) gpu.BackendError!BackendArtifactPlan {
95 return try self.artifact_plan.copy(allocator);
96 }
97 };
98
99 pub fn createArtifactJobFromTargetJob(
100 allocator: std.mem.Allocator,
101 handle: gpu.BackendHandle,
102 target_module: *target.TargetJob,
103 options: ArtifactPlanOptions,
104 ) !*ArtifactJob {
105 var artifact_plan = try plan.createBackendArtifactPlanFromTargetJob(
106 allocator,
107 handle,
108 target_module,
109 options,
110 );
111 var plan_owned = true;
112 errdefer if (plan_owned) artifact_plan.deinit();
113
114 const artifact_job = try ArtifactJob.init(
115 allocator,
116 artifact_plan,
117 );
118 plan_owned = false;
119 return artifact_job;
120 }
121
122 test "artifact job fingerprint includes artifact plan identity" {
123 const allocator = std.testing.allocator;
124
125 var first_plan = BackendArtifactPlan.init(allocator, .{
126 .backend_kind = .cuda,
127 .artifact_format = .cuda_ptx,
128 .dtype_bits = gpu.DTypeSet.init(&.{.f32}).bits,
129 });
130 var first_plan_owned = true;
131 errdefer if (first_plan_owned) first_plan.deinit();
132
133 const first = try ArtifactJob.init(allocator, first_plan);
134 first_plan_owned = false;
135 defer first.deinit();
136
137 var second_plan = BackendArtifactPlan.init(allocator, .{
138 .backend_kind = .cuda,
139 .artifact_format = .cuda_ptx,
140 .dtype_bits = gpu.DTypeSet.init(&.{ .f16, .f32 }).bits,
141 });
142 var second_plan_owned = true;
143 errdefer if (second_plan_owned) second_plan.deinit();
144
145 const second = try ArtifactJob.init(allocator, second_plan);
146 second_plan_owned = false;
147 defer second.deinit();
148
149 try std.testing.expect(first.fingerprint() != second.fingerprint());
150 }
151
152 test "artifact job exposes kernel count" {
153 const allocator = std.testing.allocator;
154
155 var artifact_plan = BackendArtifactPlan.init(allocator, .{
156 .backend_kind = .cuda,
157 .artifact_format = .cuda_ptx,
158 .dtype_bits = gpu.DTypeSet.init(&.{.f32}).bits,
159 });
160 var plan_owned = true;
161 errdefer if (plan_owned) artifact_plan.deinit();
162
163 const module = try ArtifactJob.init(allocator, artifact_plan);
164 plan_owned = false;
165 defer module.deinit();
166
167 try std.testing.expectEqual(@as(usize, 0), module.kernelCount());
168 }
169
170 test "artifact job exposes kernel summaries by index and work item" {
171 const allocator = std.testing.allocator;
172
173 var artifact_plan = BackendArtifactPlan.init(allocator, .{
174 .backend_kind = .cuda,
175 .artifact_format = .cuda_ptx,
176 .dtype_bits = gpu.DTypeSet.init(&.{.f32}).bits,
177 });
178 var plan_owned = true;
179 errdefer if (plan_owned) artifact_plan.deinit();
180
181 var kernel_artifact = try gpu.KernelArtifact.init(allocator, .{
182 .backend = .cuda,
183 .format = .cuda_ptx,
184 .entry_name = "artifact_summary_kernel",
185 .argument_count = 0,
186 });
187 var artifact_owned = true;
188 errdefer if (artifact_owned) kernel_artifact.deinit();
189 try kernel_artifact.setOwnedText("payload");
190
191 var compile = try artifact.PlannedKernelCompileContract.init(
192 allocator,
193 .choir_kernel,
194 .authored,
195 .cuda_ptx,
196 "artifact_summary_kernel",
197 0,
198 gpu.DTypeSet.init(&.{.f32}),
199 .{},
200 .{},
201 null,
202 .{ .text = "payload" },
203 );
204 var compile_owned = true;
205 errdefer if (compile_owned) compile.deinit(allocator);
206
207 try artifact_plan.addStandaloneKernel(
208 kernel_artifact,
209 .{
210 .format = .cuda_ptx,
211 .element_count = 1,
212 .geometry = .{
213 .grid = .{ 1, 1, 1 },
214 .threadgroup = .{ 1, 1, 1 },
215 },
216 },
217 compile,
218 .{},
219 );
220 artifact_owned = false;
221 compile_owned = false;
222
223 const module = try ArtifactJob.init(allocator, artifact_plan);
224 plan_owned = false;
225 defer module.deinit();
226
227 try std.testing.expectEqual(@as(usize, 1), module.kernelCount());
228 const indexed = try module.kernelSummary(0);
229 const by_work = try module.kernelSummaryForWork(indexed.work_item_id);
230 try std.testing.expectEqualStrings(indexed.entry_name, by_work.entry_name);
231 try std.testing.expectEqual(indexed.work_item_id, by_work.work_item_id);
232 try std.testing.expectEqual(artifact.KernelSource.choir_kernel, indexed.source);
233 try std.testing.expectError(error.InvalidArtifact, module.kernelSummary(1));
234 try std.testing.expectError(error.InvalidArtifact, module.kernelSummaryForWork(99));
235
236 var summaries = try module.copyKernelSummaries(allocator);
237 defer summaries.deinit();
238 try std.testing.expectEqual(@as(usize, 1), summaries.len());
239 const copied = try summaries.summaryForWork(indexed.work_item_id);
240 try std.testing.expectEqualStrings(indexed.entry_name, copied.entry_name);
241 try std.testing.expect(copied.entry_name.ptr != indexed.entry_name.ptr);
242 try std.testing.expectError(error.InvalidArtifact, summaries.summary(1));
243 try std.testing.expectError(error.InvalidArtifact, summaries.summaryForWork(99));
244 }