lib/accy/src/preparation/target.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir = @import("choir");
4 const accy_root = @import("../root.zig");
5
6 const ir = choir.ir;
7
8 pub const backend_kind_attr_name = "accy.backend.kind";
9 pub const artifact_format_attr_name = "accy.backend.artifact_format";
10 pub const math_tier_attr_name = "accy.backend.math_tier";
11 pub const dtype_bits_attr_name = "accy.backend.dtype_bits";
12 pub const feature_bits_attr_name = "accy.backend.feature_bits";
13 pub const generated_scan_schedule_attr_name = "accy.backend.scan_schedule";
14
15 pub const GeneratedScanSchedule = struct {
16 threads: u32,
17 items: u32,
18
19 pub fn eql(self: GeneratedScanSchedule, other: GeneratedScanSchedule) bool {
20 return self.threads == other.threads and self.items == other.items;
21 }
22 };
23
24 pub const GeneratedScanScheduleDecision = struct {
25 total: ?u64 = null,
26 schedule: GeneratedScanSchedule,
27 };
28
29 pub fn setGeneratedScanSchedules(
30 allocator: std.mem.Allocator,
31 ctx: *ir.Context,
32 module: *ir.Operation,
33 decisions: []const GeneratedScanScheduleDecision,
34 ) !void {
35 const encoded = try encodeGeneratedScanSchedules(allocator, decisions);
36 defer allocator.free(encoded);
37 try module.setAttr(generated_scan_schedule_attr_name, try ctx.getStringAttr(encoded));
38 }
39
40 /// The recipe calls this function to record caller decisions for one generated kernel, such as its
41 /// thread count, optionally tied to a problem size, as a request's scan schedule choices in the
42 /// same bytes the module stores. The recipe encodes a stage's options into its stage record. The
43 /// function writes the choices as comma-separated entries: `threads`x`items` or
44 /// `total`=`threads`x`items` when a choice is tied to a total length. The bytes match the module
45 /// attribute the same choices are stored in, and the caller owns them. The call returns
46 /// `error.InvalidArtifact` for an empty list.
47 pub fn encodeGeneratedScanSchedules(
48 allocator: std.mem.Allocator,
49 decisions: []const GeneratedScanScheduleDecision,
50 ) ![]u8 {
51 if (decisions.len == 0) return error.InvalidArtifact;
52 var text = std.ArrayListUnmanaged(u8).empty;
53 errdefer text.deinit(allocator);
54 for (decisions, 0..) |decision, index| {
55 if (index != 0) try text.append(allocator, ',');
56 var entry_buffer: [64]u8 = undefined;
57 const entry = if (decision.total) |total|
58 try std.fmt.bufPrint(&entry_buffer, "{d}={d}x{d}", .{ total, decision.schedule.threads, decision.schedule.items })
59 else
60 try std.fmt.bufPrint(&entry_buffer, "{d}x{d}", .{ decision.schedule.threads, decision.schedule.items });
61 try text.appendSlice(allocator, entry);
62 }
63 return text.toOwnedSlice(allocator);
64 }
65
66 pub fn readGeneratedScanSchedules(module: *const ir.Operation) ?[]const u8 {
67 const attr = module.getAttrAs(ir.Attribute.StringAttr, generated_scan_schedule_attr_name) orelse return null;
68 return attr.getValue();
69 }
70
71 pub fn resolveGeneratedScanSchedule(
72 encoded: []const u8,
73 total: u64,
74 ) error{InvalidArtifact}!?GeneratedScanSchedule {
75 var fallback: ?GeneratedScanSchedule = null;
76 var entries = std.mem.splitScalar(u8, encoded, ',');
77 while (entries.next()) |entry| {
78 if (std.mem.indexOfScalar(u8, entry, '=')) |split_index| {
79 const entry_total = std.fmt.parseInt(u64, entry[0..split_index], 10) catch return error.InvalidArtifact;
80 const schedule = try parseGeneratedScanSchedule(entry[split_index + 1 ..]);
81 if (entry_total == total) return schedule;
82 } else {
83 fallback = try parseGeneratedScanSchedule(entry);
84 }
85 }
86 return fallback;
87 }
88
89 fn parseGeneratedScanSchedule(text: []const u8) error{InvalidArtifact}!GeneratedScanSchedule {
90 const split_index = std.mem.indexOfScalar(u8, text, 'x') orelse return error.InvalidArtifact;
91 const threads = std.fmt.parseInt(u32, text[0..split_index], 10) catch return error.InvalidArtifact;
92 const items = std.fmt.parseInt(u32, text[split_index + 1 ..], 10) catch return error.InvalidArtifact;
93 if (threads == 0 or items == 0) return error.InvalidArtifact;
94 return .{ .threads = threads, .items = items };
95 }
96
97 pub const generated_row_pipeline_schedule_attr_name = "accy.backend.row_pipeline_schedule";
98
99 pub const GeneratedRowPipelineSchedule = struct {
100 threads: u32,
101
102 pub fn eql(self: GeneratedRowPipelineSchedule, other: GeneratedRowPipelineSchedule) bool {
103 return self.threads == other.threads;
104 }
105 };
106
107 pub const GeneratedRowPipelineShape = struct {
108 rows: u64,
109 cols: u64,
110 };
111
112 pub const GeneratedRowPipelineScheduleDecision = struct {
113 shape: ?GeneratedRowPipelineShape = null,
114 schedule: GeneratedRowPipelineSchedule,
115 };
116
117 pub fn setGeneratedRowPipelineSchedules(
118 allocator: std.mem.Allocator,
119 ctx: *ir.Context,
120 module: *ir.Operation,
121 decisions: []const GeneratedRowPipelineScheduleDecision,
122 ) !void {
123 const encoded = try encodeGeneratedRowPipelineSchedules(allocator, decisions);
124 defer allocator.free(encoded);
125 try module.setAttr(generated_row_pipeline_schedule_attr_name, try ctx.getStringAttr(encoded));
126 }
127
128 /// The recipe calls this function to record a request's row schedule choices in the same bytes the
129 /// module stores. The function writes the choices as comma-separated entries: `threads` or
130 /// `rows`x`cols`=`threads` when a choice is tied to a row shape. The bytes match the module
131 /// attribute the same choices are stored in, and the caller owns them. The call returns
132 /// `error.InvalidArtifact` for an empty list.
133 pub fn encodeGeneratedRowPipelineSchedules(
134 allocator: std.mem.Allocator,
135 decisions: []const GeneratedRowPipelineScheduleDecision,
136 ) ![]u8 {
137 if (decisions.len == 0) return error.InvalidArtifact;
138 var text = std.ArrayListUnmanaged(u8).empty;
139 errdefer text.deinit(allocator);
140 for (decisions, 0..) |decision, index| {
141 if (index != 0) try text.append(allocator, ',');
142 var entry_buffer: [96]u8 = undefined;
143 const entry = if (decision.shape) |shape|
144 try std.fmt.bufPrint(&entry_buffer, "{d}x{d}={d}", .{ shape.rows, shape.cols, decision.schedule.threads })
145 else
146 try std.fmt.bufPrint(&entry_buffer, "{d}", .{decision.schedule.threads});
147 try text.appendSlice(allocator, entry);
148 }
149 return text.toOwnedSlice(allocator);
150 }
151
152 pub fn readGeneratedRowPipelineSchedules(module: *const ir.Operation) ?[]const u8 {
153 const attr = module.getAttrAs(ir.Attribute.StringAttr, generated_row_pipeline_schedule_attr_name) orelse return null;
154 return attr.getValue();
155 }
156
157 pub fn resolveGeneratedRowPipelineSchedule(
158 encoded: []const u8,
159 rows: u64,
160 cols: u64,
161 ) error{InvalidArtifact}!?GeneratedRowPipelineSchedule {
162 var fallback: ?GeneratedRowPipelineSchedule = null;
163 var entries = std.mem.splitScalar(u8, encoded, ',');
164 while (entries.next()) |entry| {
165 if (std.mem.indexOfScalar(u8, entry, '=')) |split_index| {
166 const shape_text = entry[0..split_index];
167 const shape_split = std.mem.indexOfScalar(u8, shape_text, 'x') orelse return error.InvalidArtifact;
168 const entry_rows = std.fmt.parseInt(u64, shape_text[0..shape_split], 10) catch return error.InvalidArtifact;
169 const entry_cols = std.fmt.parseInt(u64, shape_text[shape_split + 1 ..], 10) catch return error.InvalidArtifact;
170 const schedule = try parseGeneratedRowPipelineSchedule(entry[split_index + 1 ..]);
171 if (entry_rows == rows and entry_cols == cols) return schedule;
172 } else {
173 fallback = try parseGeneratedRowPipelineSchedule(entry);
174 }
175 }
176 return fallback;
177 }
178
179 fn parseGeneratedRowPipelineSchedule(text: []const u8) error{InvalidArtifact}!GeneratedRowPipelineSchedule {
180 const threads = std.fmt.parseInt(u32, text, 10) catch return error.InvalidArtifact;
181 if (threads == 0) return error.InvalidArtifact;
182 return .{ .threads = threads };
183 }
184
185 pub const BackendTargetProfile = accy_root.choir.record.target.BackendTargetProfile;
186
187 pub fn setBackendTargetProfile(
188 ctx: *ir.Context,
189 module: *ir.Operation,
190 profile: BackendTargetProfile,
191 ) !void {
192 try module.setAttr(backend_kind_attr_name, try ctx.getStringAttr(@tagName(profile.backend_kind)));
193 try module.setAttr(artifact_format_attr_name, try ctx.getStringAttr(@tagName(profile.artifact_format)));
194 try module.setAttr(math_tier_attr_name, try ctx.getStringAttr(@tagName(profile.math_tier)));
195 try module.setAttr(dtype_bits_attr_name, try ctx.getIntegerAttr(@bitCast(profile.dtype_bits), 64, false));
196 try module.setAttr(feature_bits_attr_name, try ctx.getIntegerAttr(@bitCast(profile.feature_bits), 64, false));
197 }
198
199 pub fn readBackendTargetProfile(module: *const ir.Operation) ?BackendTargetProfile {
200 const backend_kind = readEnumAttr(gpu.BackendKind, module, backend_kind_attr_name) orelse return null;
201 const artifact_format = readEnumAttr(gpu.ArtifactFormat, module, artifact_format_attr_name) orelse return null;
202 const math_tier = readEnumAttr(gpu.BackendMathTier, module, math_tier_attr_name) orelse return null;
203 const dtype_bits = readU64Attr(module, dtype_bits_attr_name) orelse return null;
204 const feature_bits = readU64Attr(module, feature_bits_attr_name) orelse return null;
205 return .{
206 .backend_kind = backend_kind,
207 .artifact_format = artifact_format,
208 .math_tier = math_tier,
209 .dtype_bits = dtype_bits,
210 .feature_bits = feature_bits,
211 };
212 }
213
214 fn readEnumAttr(comptime E: type, module: *const ir.Operation, attr_name: []const u8) ?E {
215 const string_attr = module.getAttrAs(ir.Attribute.StringAttr, attr_name) orelse return null;
216 return std.meta.stringToEnum(E, string_attr.getValue());
217 }
218
219 fn readU64Attr(module: *const ir.Operation, attr_name: []const u8) ?u64 {
220 const int_attr = module.getAttrAs(ir.Attribute.IntegerAttr, attr_name) orelse return null;
221 return int_attr.getUnsignedValue();
222 }
223
224 const testing = std.testing;
225
226 test "backend target profile round trips through prepared module attrs" {
227 var arena = std.heap.ArenaAllocator.init(testing.allocator);
228 defer arena.deinit();
229 const allocator = arena.allocator();
230
231 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
232 defer ctx.deinit(allocator);
233
234 const module = try choir.dialects.builtin.BuiltinDialect.ModuleOp.create(&ctx, ir.Location.getUnknown());
235 const profile = try BackendTargetProfile.init(.{
236 .identity = .{
237 .backend = .vulkan,
238 .family = .vulkan,
239 },
240 .dtypes = gpu.DTypeSet.init(&.{ .i32, .f32 }),
241 .features = .{ .atomic_i32 = true },
242 .artifact_formats = gpu.ArtifactFormatSet.init(&.{.vulkan_spirv}),
243 }, .vulkan, .vulkan_spirv);
244
245 try setBackendTargetProfile(&ctx, module.op, profile);
246 const read = readBackendTargetProfile(module.op) orelse return error.MissingTargetProfile;
247
248 try testing.expectEqual(profile.backend_kind, read.backend_kind);
249 try testing.expectEqual(profile.artifact_format, read.artifact_format);
250 try testing.expectEqual(profile.math_tier, read.math_tier);
251 try testing.expectEqual(profile.dtype_bits, read.dtype_bits);
252 try testing.expectEqual(profile.feature_bits, read.feature_bits);
253 try testing.expect(read.supportsDType(.f32));
254 try testing.expect(!read.supportsDType(.f64));
255 }
256
257 test "generated scan schedule decisions round trip through module attrs" {
258 var arena = std.heap.ArenaAllocator.init(testing.allocator);
259 defer arena.deinit();
260 const allocator = arena.allocator();
261
262 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
263 defer ctx.deinit(allocator);
264
265 const module = try choir.dialects.builtin.BuiltinDialect.ModuleOp.create(&ctx, ir.Location.getUnknown());
266 try setGeneratedScanSchedules(allocator, &ctx, module.op, &.{
267 .{ .total = 16777216, .schedule = .{ .threads = 256, .items = 16 } },
268 .{ .schedule = .{ .threads = 512, .items = 16 } },
269 });
270
271 const encoded = readGeneratedScanSchedules(module.op) orelse return error.MissingScanSchedules;
272 try testing.expectEqualStrings("16777216=256x16,512x16", encoded);
273
274 const exact = (try resolveGeneratedScanSchedule(encoded, 16777216)) orelse return error.MissingScanSchedule;
275 try testing.expect(exact.eql(.{ .threads = 256, .items = 16 }));
276
277 const wildcard = (try resolveGeneratedScanSchedule(encoded, 8192)) orelse return error.MissingScanSchedule;
278 try testing.expect(wildcard.eql(.{ .threads = 512, .items = 16 }));
279 }
280
281 test "generated scan schedule resolution rejects malformed encodings" {
282 try testing.expectError(error.InvalidArtifact, resolveGeneratedScanSchedule("", 4096));
283 try testing.expectError(error.InvalidArtifact, resolveGeneratedScanSchedule("512", 4096));
284 try testing.expectError(error.InvalidArtifact, resolveGeneratedScanSchedule("4096=0x16", 4096));
285 try testing.expectError(error.InvalidArtifact, resolveGeneratedScanSchedule("abc=512x16", 4096));
286 try testing.expectEqual(@as(?GeneratedScanSchedule, null), try resolveGeneratedScanSchedule("8192=512x16", 4096));
287 }
288
289 test "generated row pipeline schedule decisions round trip through module attrs" {
290 var arena = std.heap.ArenaAllocator.init(testing.allocator);
291 defer arena.deinit();
292 const allocator = arena.allocator();
293
294 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
295 defer ctx.deinit(allocator);
296
297 const module = try choir.dialects.builtin.BuiltinDialect.ModuleOp.create(&ctx, ir.Location.getUnknown());
298 try setGeneratedRowPipelineSchedules(allocator, &ctx, module.op, &.{
299 .{ .shape = .{ .rows = 4096, .cols = 4096 }, .schedule = .{ .threads = 512 } },
300 .{ .schedule = .{ .threads = 128 } },
301 });
302
303 const encoded = readGeneratedRowPipelineSchedules(module.op) orelse return error.MissingRowPipelineSchedules;
304 try testing.expectEqualStrings("4096x4096=512,128", encoded);
305
306 const exact = (try resolveGeneratedRowPipelineSchedule(encoded, 4096, 4096)) orelse return error.MissingRowPipelineSchedule;
307 try testing.expect(exact.eql(.{ .threads = 512 }));
308
309 const wildcard = (try resolveGeneratedRowPipelineSchedule(encoded, 8, 2048)) orelse return error.MissingRowPipelineSchedule;
310 try testing.expect(wildcard.eql(.{ .threads = 128 }));
311 }
312
313 test "generated row pipeline schedule resolution rejects malformed encodings" {
314 try testing.expectError(error.InvalidArtifact, resolveGeneratedRowPipelineSchedule("", 4, 1024));
315 try testing.expectError(error.InvalidArtifact, resolveGeneratedRowPipelineSchedule("4096=512", 4, 1024));
316 try testing.expectError(error.InvalidArtifact, resolveGeneratedRowPipelineSchedule("4x1024=0", 4, 1024));
317 try testing.expectEqual(@as(?GeneratedRowPipelineSchedule, null), try resolveGeneratedRowPipelineSchedule("8x2048=512", 4, 1024));
318 }
319
320 test "backend target profile gates tf32 tensor math on cuda tensor cores" {
321 const caps = gpu.BackendCapabilities{
322 .identity = .{
323 .backend = .cuda,
324 .family = .nvidia_cuda,
325 },
326 .dtypes = gpu.DTypeSet.init(&.{.f32}),
327 .features = .{ .tensor_cores = true },
328 .artifact_formats = gpu.ArtifactFormatSet.init(&.{.cuda_ptx}),
329 };
330 const profile = try BackendTargetProfile.initWithMathTier(caps, .cuda, .cuda_ptx, .tf32_tensor);
331 try testing.expectEqual(gpu.BackendMathTier.tf32_tensor, profile.math_tier);
332 try testing.expect(profile.isSupportedBy(caps));
333
334 var no_tensor = caps;
335 no_tensor.features.tensor_cores = false;
336 try testing.expectError(error.CapabilityMismatch, BackendTargetProfile.initWithMathTier(no_tensor, .cuda, .cuda_ptx, .tf32_tensor));
337 }