lib/accy/src/kernel/call.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const choir = @import("choir");
5 const gpu_codegen = choir.backends.gpu;
6 const accy_root = @import("../root.zig");
7 const artifact_product = @import("../artifact/root.zig");
8 const preparation = @import("../preparation/root.zig");
9 const target = @import("../target/root.zig");
10 const kernel_model = @import("model/root.zig");
11 const builder = kernel_model.core.builder;
12 const plan_mod = kernel_model.plan;
13 const program_mod = @import("program/root.zig");
14
15 pub const product_name = "accy.kernel_call";
16
17 pub const Options = struct {
18 target: []const u8,
19 version: u32 = 1,
20 format: ?gpu.ArtifactFormat = null,
21 kernel_plan: plan_mod.Options = .{},
22 element_count_argument: artifact_product.ElementCountArgument = .none,
23 shape_family_fingerprint: ?u64 = null,
24 shape_profile: ?artifact_product.KernelCallShapeProfile = null,
25 launch: ?artifact_product.KernelCallLaunch = null,
26 runtime_scalar_argument_count: u32 = 0,
27 static_arguments: []const choir_abi.ScalarArgument = &.{},
28 };
29
30 pub const OwnedArtifact = struct {
31 allocator: std.mem.Allocator,
32 artifacts: [1]artifact_product.KernelCallArtifact,
33
34 pub fn entry(self: *const OwnedArtifact) artifact_product.KernelCallArtifact {
35 return self.artifacts[0];
36 }
37
38 pub fn registry(self: *const OwnedArtifact) artifact_product.KernelCallRegistry {
39 return .{ .entries = self.artifacts[0..] };
40 }
41
42 pub fn fingerprint(self: *const OwnedArtifact) u64 {
43 const registry_value = self.registry();
44 return artifact_product.kernelCallRegistryFingerprint(®istry_value);
45 }
46
47 pub fn productStamp(self: *const OwnedArtifact) choir.product.incremental.ProductStamp {
48 return choir.product.incremental.productStamp(product_name, self.fingerprint());
49 }
50
51 pub fn deinit(self: *OwnedArtifact) void {
52 const allocator = self.allocator;
53 const artifact = &self.artifacts[0];
54 allocator.free(@constCast(artifact.target));
55 allocator.free(@constCast(artifact.entry_name));
56 if (artifact.shape_profile) |profile| artifact_product.deinitKernelCallShapeProfile(allocator, profile);
57 if (artifact.static_arguments.len != 0) allocator.free(@constCast(artifact.static_arguments));
58 deinitCompilePayload(allocator, artifact.payload);
59 self.* = undefined;
60 }
61 };
62
63 pub fn createBackendArtifactFromEntry(
64 allocator: std.mem.Allocator,
65 handle: gpu.BackendHandle,
66 entry: artifact_product.KernelCallArtifact,
67 diagnostic_id: ?[]const u8,
68 ) !gpu.KernelArtifact {
69 const caps = try handle.queryCapabilities();
70 const scalar_argument_count = try entry.scalarArgumentCount();
71 var artifact = try gpu.KernelArtifact.init(allocator, .{
72 .backend = handle.backendKind() orelse caps.identity.backend,
73 .format = entry.format,
74 .entry_name = entry.entry_name,
75 .argument_count = entry.argument_count,
76 .scalar_argument_count = scalar_argument_count,
77 .diagnostic_id = diagnostic_id,
78 .interface = .{
79 .features = entry.required_features,
80 .subgroup = entry.required_subgroup,
81 .push_constants = entry.push_constants,
82 },
83 });
84 errdefer artifact.deinit();
85 switch (entry.payload) {
86 .text => |text| artifact.setBorrowedText(text),
87 .bytes => |bytes| artifact.setBorrowedBytes(bytes),
88 .words_u32 => |words| artifact.setBorrowedWords(words),
89 .none => return error.InvalidArtifact,
90 }
91 return artifact;
92 }
93
94 pub fn launchGeometryForEntry(
95 entry: artifact_product.KernelCallArtifact,
96 runtime_arguments: []const choir_abi.ScalarArgument,
97 ) gpu.BackendError!choir_abi.LaunchGeometry {
98 return switch (entry.launch) {
99 .fixed => |geometry| geometry,
100 .derived => |derived| try derived.geometry(runtime_arguments),
101 };
102 }
103
104 pub fn createArtifact(
105 allocator: std.mem.Allocator,
106 handle: gpu.BackendHandle,
107 program: *program_mod.Program,
108 options: Options,
109 ) !OwnedArtifact {
110 if (options.target.len == 0 or options.version == 0) return error.InvalidArtifact;
111
112 var authored_plan = try program.createCheckedPlan(allocator, options.kernel_plan);
113 defer authored_plan.deinit();
114
115 const caps = try handle.queryCapabilities();
116 const backend_kind = handle.backendKind() orelse caps.identity.backend;
117 const format = options.format orelse artifact_product.defaultArtifactFormat(backend_kind) orelse {
118 return error.UnsupportedOperation;
119 };
120 if (!caps.supportsArtifactFormat(format)) return error.UnsupportedArtifactFormat;
121 const geometry = launchGeometry(&authored_plan);
122 try caps.validateLaunchGeometry(geometry);
123
124 const compilation = try target.compileKernelForArtifactFormat(
125 allocator,
126 format,
127 authored_plan.entry_name,
128 program.kernelModule(),
129 compileOptions(format, geometry),
130 );
131 errdefer deinitCompilePayload(allocator, compilation.payload);
132
133 const owned_target = try allocator.dupe(u8, options.target);
134 errdefer allocator.free(owned_target);
135 const owned_entry_name = try allocator.dupe(u8, authored_plan.entry_name);
136 errdefer allocator.free(owned_entry_name);
137 const static_arguments = try compileStaticArguments(allocator, format, geometry, options.static_arguments);
138 errdefer if (static_arguments.len != 0) allocator.free(static_arguments);
139
140 const shape_family_fingerprint = options.shape_family_fingerprint orelse if (options.shape_profile) |profile| profile.fingerprint else null;
141 if (options.shape_profile) |profile| {
142 if (shape_family_fingerprint.? != profile.fingerprint) return error.InvalidArtifact;
143 try profile.validate(options.runtime_scalar_argument_count);
144 }
145 const shape_profile = if (options.shape_profile) |profile|
146 try artifact_product.duplicateKernelCallShapeProfile(allocator, profile)
147 else
148 null;
149 errdefer if (shape_profile) |profile| artifact_product.deinitKernelCallShapeProfile(allocator, profile);
150
151 return .{
152 .allocator = allocator,
153 .artifacts = .{.{
154 .target = owned_target,
155 .version = options.version,
156 .format = format,
157 .entry_name = owned_entry_name,
158 .argument_count = try compileArgumentCount(format, authored_plan.argument_count),
159 .shape_family_fingerprint = shape_family_fingerprint,
160 .shape_profile = shape_profile,
161 .required_dtypes = requiredDTypesForParams(authored_plan.params),
162 .required_features = gpu_codegen.featureRequirementsForModule(program.kernelModule()),
163 .required_subgroup = gpu_codegen.subgroupRequirementsForModule(program.kernelModule()),
164 .push_constants = compilation.push_constants,
165 .payload = compilation.payload,
166 .launch = options.launch orelse .{ .fixed = geometry },
167 .element_count_argument = options.element_count_argument,
168 .runtime_scalar_argument_count = options.runtime_scalar_argument_count,
169 .static_arguments = static_arguments,
170 }},
171 };
172 }
173
174 fn compileOptions(format: gpu.ArtifactFormat, geometry: choir_abi.LaunchGeometry) target.CompileOptions {
175 const element_count = geometry.threadCount() catch return .{};
176 return target.compileOptionsForArtifactFormat(format, element_count);
177 }
178
179 fn compileArgumentCount(format: gpu.ArtifactFormat, argument_count: u32) gpu.BackendError!u32 {
180 if (!gpu.artifactFormatUsesHostLoopLaunch(format)) return argument_count;
181 return choir_abi.kernelArgumentCount(argument_count);
182 }
183
184 fn compileStaticArguments(
185 allocator: std.mem.Allocator,
186 format: gpu.ArtifactFormat,
187 geometry: choir_abi.LaunchGeometry,
188 static_arguments: []const choir_abi.ScalarArgument,
189 ) gpu.BackendError![]choir_abi.ScalarArgument {
190 if (!gpu.artifactFormatUsesHostLoopLaunch(format)) {
191 return allocator.dupe(choir_abi.ScalarArgument, static_arguments) catch return error.OutOfMemory;
192 }
193 if (static_arguments.len != 0) return error.InvalidArtifact;
194 return choir_abi.launchShapeArguments(
195 allocator,
196 geometry.threadCount() catch return error.InvalidArtifact,
197 geometry,
198 );
199 }
200
201 fn requiredDTypesForParams(params: []const builder.Param) gpu.DTypeSet {
202 var dtypes: gpu.DTypeSet = .{};
203 for (params) |param| {
204 switch (param) {
205 .scalar => |dtype| dtypes.insert(dtype),
206 .buffer => |buffer| dtypes.insert(buffer.dtype),
207 }
208 }
209 return dtypes;
210 }
211
212 fn launchGeometry(authored_plan: *const plan_mod.Plan) choir_abi.LaunchGeometry {
213 return .{
214 .grid = authored_plan.launch.grid,
215 .threadgroup = authored_plan.launch.block,
216 };
217 }
218
219 fn deinitCompilePayload(allocator: std.mem.Allocator, payload: gpu.CompilePayload) void {
220 switch (payload) {
221 .bytes => |bytes| allocator.free(@constCast(bytes)),
222 .words_u32 => |words| allocator.free(@constCast(words)),
223 .text => |text| allocator.free(@constCast(text)),
224 .none => {},
225 }
226 }
227
228 test "kernel call artifact records atomic feature requirements" {
229 const allocator = std.testing.allocator;
230
231 var state = gpu.recording.BackendState{
232 .allocator = allocator,
233 .kind = .cuda,
234 .format = .cuda_ptx,
235 };
236
237 var builder_state = try program_mod.Builder.init(allocator, program_mod.Builder.Limits.testing, "kernel_call_atomic_add_i32", &.{
238 builder.dynamicBuffer(.i32),
239 builder.dynamicBuffer(.i32),
240 });
241 errdefer builder_state.deinit();
242
243 const axis = try builder_state.axis("i", 32);
244 try builder_state.bind(axis, .thread_x);
245 const acc = builder_state.argument(0);
246 const dst = builder_state.argument(1);
247 const index = try builder_state.globalId(.x);
248 const one = try builder_state.constantInt(.i32, 1);
249 const old = try builder_state.atomicRmw(.add, one, acc, index);
250 try builder_state.store(old, dst, index);
251 try builder_state.return_();
252
253 var program = try builder_state.finish();
254 defer program.deinit();
255
256 var artifact = try createArtifact(allocator, state.handle(), &program, .{
257 .target = "atomic-call",
258 .version = 1,
259 });
260 defer artifact.deinit();
261
262 try std.testing.expect(artifact.entry().required_features.atomic_i32);
263 }
264
265 test "kernel call entry materializes backend artifact and derived launch geometry" {
266 const allocator = std.testing.allocator;
267 var state = gpu.recording.BackendState{
268 .allocator = allocator,
269 .kind = .vulkan,
270 .format = .vulkan_spirv,
271 };
272 const words = [_]u32{0x07230203};
273 const static_arguments = [_]choir_abi.ScalarArgument{.{ .u32 = 11 }};
274 const entry = artifact_product.KernelCallArtifact{
275 .target = "image_resize_bilinear_family_16x16_rgba8",
276 .version = 1,
277 .format = .vulkan_spirv,
278 .entry_name = "accy_image_resize_bilinear_rgba8",
279 .argument_count = 10,
280 .payload = .{ .words_u32 = words[0..] },
281 .runtime_scalar_argument_count = 7,
282 .static_arguments = static_arguments[0..],
283 .launch = .{ .derived = .{
284 .grid = .{
285 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = 16 } },
286 .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = 16 } },
287 .{ .fixed = 1 },
288 },
289 .threadgroup = .{ 16, 16, 1 },
290 } },
291 };
292
293 var artifact = try createBackendArtifactFromEntry(allocator, state.handle(), entry, "accy/kernel-call/materialize-test");
294 defer artifact.deinit();
295
296 try std.testing.expectEqual(gpu.BackendKind.vulkan, artifact.backend);
297 try std.testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, artifact.format);
298 try std.testing.expectEqualStrings("accy_image_resize_bilinear_rgba8", artifact.entry_name);
299 try std.testing.expectEqual(@as(u32, 10), artifact.argument_count);
300 try std.testing.expectEqual(@as(u32, 8), artifact.scalar_argument_count);
301 try std.testing.expectEqualStrings("accy/kernel-call/materialize-test", artifact.diagnostic_id.?);
302 try std.testing.expectEqualSlices(u32, words[0..], artifact.payload.words_u32);
303
304 const runtime_arguments = [_]choir_abi.ScalarArgument{
305 .{ .u32 = 33 },
306 .{ .u32 = 17 },
307 };
308 const geometry = try launchGeometryForEntry(entry, runtime_arguments[0..]);
309 try std.testing.expectEqual(@as(u32, 3), geometry.grid[0]);
310 try std.testing.expectEqual(@as(u32, 2), geometry.grid[1]);
311 try std.testing.expectEqual(@as(u32, 1), geometry.grid[2]);
312 try std.testing.expectEqual(@as(u32, 16), geometry.threadgroup[0]);
313 try std.testing.expectEqual(@as(u32, 16), geometry.threadgroup[1]);
314 }