lib/accy/src/preparation/einsum.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const alloc_arena = @import("alloc_arena");
5 const choir = @import("choir");
6 const accy_root = @import("../root.zig");
7 const accy_choir = @import("../choir/root.zig");
8 const kernel_library = @import("../kernel/library/root.zig");
9 const kernel_selection = @import("../kernel/logical/selection/root.zig");
10 const call_preparation = @import("call.zig");
11 const library_preparation = @import("library.zig");
12 const dialect_mod = accy_choir.dialect;
13
14 const ir = choir.ir;
15 const rewrite = ir.rewrite;
16 const passes = choir.passes;
17 const work = passes.pass.work;
18 const Strategy = accy_choir.einsum.Strategy;
19
20 pub const KernelLibraryLowering = library_preparation.KernelLibraryLowering;
21
22 pub const Options = struct {
23 strategy: Strategy = .auto,
24 exact_state_limit: usize = accy_choir.einsum.default_exact_state_limit,
25 beam_width: usize = 64,
26 auto_beam_width: usize = accy_choir.einsum.default_auto_beam_width,
27 kernel_library: KernelLibraryLowering = .disabled,
28 matrix_product_schedule: ?kernel_library.MatrixProductSchedule = null,
29 matrix_product_tuning: ?kernel_library.linalg.MatrixProductScheduleReader = null,
30 family_tuning: ?*const kernel_library.tuning.FamilyTuningReader = null,
31
32 pub fn eql(self: Options, other: Options) bool {
33 return self.strategy == other.strategy and
34 self.exact_state_limit == other.exact_state_limit and
35 self.beam_width == other.beam_width and
36 self.auto_beam_width == other.auto_beam_width and
37 self.kernel_library == other.kernel_library and
38 std.meta.eql(self.matrix_product_schedule, other.matrix_product_schedule) and
39 matrixProductTuningEql(self.matrix_product_tuning, other.matrix_product_tuning) and
40 self.family_tuning == other.family_tuning;
41 }
42 };
43
44 fn matrixProductTuningEql(
45 lhs: ?kernel_library.linalg.MatrixProductScheduleReader,
46 rhs: ?kernel_library.linalg.MatrixProductScheduleReader,
47 ) bool {
48 if (lhs == null or rhs == null) return lhs == null and rhs == null;
49 return lhs.?.eql(rhs.?);
50 }
51
52 pub const einsum_lowering_pass_name = "accy-choir-einsum-lower";
53 pub const einsum_lowering_pass_description =
54 "Lower semantic Accy einsum operations into planned tensor contractions";
55
56 const strategy_option_choices = [_]passes.PassOptionChoice{
57 .{ .name = "auto" },
58 .{ .name = "left-to-right" },
59 .{ .name = "greedy" },
60 .{ .name = "beam" },
61 .{ .name = "anytime" },
62 .{ .name = "optimal" },
63 };
64
65 const kernel_library_option_choices = [_]passes.PassOptionChoice{
66 .{ .name = "disabled" },
67 .{ .name = "enabled" },
68 };
69
70 pub const einsum_lowering_pass_options = [_]passes.PassOptionSpec{
71 .{
72 .name = "strategy",
73 .description = "Einsum contraction planning strategy",
74 .kind = .choice,
75 .choices = &strategy_option_choices,
76 .default_value = "auto",
77 },
78 .{
79 .name = "exact-state-limit",
80 .description = "Maximum exact planner state count",
81 .kind = .unsigned,
82 .default_value = "262144",
83 },
84 .{
85 .name = "beam-width",
86 .description = "Beam planner width",
87 .kind = .unsigned,
88 .default_value = "64",
89 },
90 .{
91 .name = "auto-beam-width",
92 .description = "Beam width used by automatic planning",
93 .kind = .unsigned,
94 .default_value = "256",
95 },
96 .{
97 .name = "kernel-library",
98 .description = "Use kernel library calls for recognized contractions",
99 .kind = .choice,
100 .choices = &kernel_library_option_choices,
101 .default_value = "disabled",
102 },
103 };
104
105 pub fn einsumLoweringPass() passes.Pass {
106 return .{
107 .name = einsum_lowering_pass_name,
108 .description = einsum_lowering_pass_description,
109 .run_fn = runEinsumLoweringPass,
110 .work_contract = einsum_work_contract,
111 };
112 }
113
114 pub fn einsumLoweringPassWithOptions(options: *const Options) passes.Pass {
115 return .{
116 .name = einsum_lowering_pass_name,
117 .description = einsum_lowering_pass_description,
118 .state = @constCast(options),
119 .run_with_state_fn = runEinsumLoweringPassWithState,
120 .work_contract = einsum_work_contract,
121 };
122 }
123
124 pub fn einsumLoweringPassFromOptions(allocator: std.mem.Allocator, set: passes.PassOptionSet) anyerror!passes.Pass {
125 const options = try allocator.create(Options);
126 errdefer allocator.destroy(options);
127 options.* = .{
128 .strategy = try strategyFromText(set.choiceValue("strategy", "auto")),
129 .exact_state_limit = try set.unsignedValue(usize, "exact-state-limit", accy_choir.einsum.default_exact_state_limit),
130 .beam_width = try set.unsignedValue(usize, "beam-width", 64),
131 .auto_beam_width = try set.unsignedValue(usize, "auto-beam-width", accy_choir.einsum.default_auto_beam_width),
132 .kernel_library = try kernelLibraryLoweringFromText(set.choiceValue("kernel-library", "disabled")),
133 };
134 var pass = einsumLoweringPassWithOptions(options);
135 pass.state_deinit_fn = destroyOptions;
136 return pass;
137 }
138
139 const einsum_work_contract: work.Contract = .{
140 .identity = .{ .name = einsum_lowering_pass_name, .version = 1 },
141 .estimate = einsumWork,
142 };
143
144 const EinsumWork = struct {
145 root: *ir.Operation,
146 options: Options,
147 candidates: u64 = 0,
148 operands: u64 = 0,
149 type_bytes: u64 = 0,
150 scratch: u64 = 0,
151 visits: u64 = 0,
152 nodes: u64 = 0,
153 types: u64 = 0,
154 library_queries: u64 = 0,
155
156 fn visit(self: *EinsumWork, op: *ir.Operation) !ir.WalkResult {
157 if (op == self.root or op.getNumResults() != 1) return .advance;
158 if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.EinsumOp.operation_name)) {
159 const text = (dialect_mod.AccyDialect.EinsumOp{ .op = op }).getEquation() orelse
160 return .advance;
161 try self.equation(op, text.len);
162 } else if (self.options.kernel_library == .enabled and
163 (self.options.matrix_product_schedule != null or
164 self.options.matrix_product_tuning != null or
165 self.options.family_tuning != null) and
166 op.getNumOperands() == 2 and
167 std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.DotGeneralOp.operation_name))
168 {
169 try self.candidate(op, true);
170 self.library_queries = try work.add(self.library_queries, 1);
171 self.nodes = try work.add(self.nodes, 1);
172 }
173 return .advance;
174 }
175
176 fn candidate(self: *EinsumWork, op: *ir.Operation, result: bool) !void {
177 self.candidates = try work.add(self.candidates, 1);
178 self.operands = try work.add(self.operands, op.getNumOperands());
179 for (op.getOperandValues()) |operand| try self.typeBytes(operand.type);
180 if (result) try self.typeBytes(op.getResult(0).?.type);
181 }
182
183 fn typeBytes(self: *EinsumWork, typ: ir.Type) !void {
184 if (typ.getDialectParamKey()) |key| {
185 self.type_bytes = try work.add(self.type_bytes, key.len);
186 }
187 }
188
189 fn equation(self: *EinsumWork, op: *ir.Operation, text_bytes: usize) !void {
190 const library = self.options.kernel_library == .enabled;
191 try self.candidate(op, library);
192 const count = op.getNumOperands();
193 const parsing = try accy_choir.einsum.parseBounds(text_bytes, count);
194 self.scratch = try work.add(self.scratch, parsing.allocation_capacity);
195 self.visits = try work.add(self.visits, parsing.structural_visits);
196 if (count == 0 or count > 63) return;
197 const planning: accy_choir.einsum.Options = .{
198 .strategy = self.options.strategy,
199 .exact_state_limit = self.options.exact_state_limit,
200 .beam_width = self.options.beam_width,
201 .auto_beam_width = self.options.auto_beam_width,
202 };
203 const lowering = try accy_choir.einsum.loweringBounds(count);
204 self.scratch = try work.add(self.scratch, try work.add(
205 try planning.storageBound(count),
206 lowering.scratch_bytes,
207 ));
208 self.visits = try work.add(self.visits, try work.add(
209 try planning.workBound(count),
210 lowering.structural_visits,
211 ));
212 self.nodes = try work.add(self.nodes, lowering.operation_requests);
213 self.types = try work.add(self.types, lowering.tensor_type_requests);
214 if (library) {
215 self.library_queries = try work.add(self.library_queries, 1);
216 self.nodes = try work.add(self.nodes, 1);
217 self.scratch = try work.add(self.scratch, try work.multiply(
218 count,
219 @sizeOf(kernel_selection.EinsumOperand) + 62 * @sizeOf(i64) + 128,
220 ));
221 }
222 }
223 };
224
225 /// Sizes the storage needed for the owned description the kernel library returns for one request,
226 /// including its shapes, axes and schedule bindings (catalog descriptor). The pass declares its
227 /// costs before it runs so compilation can refuse a pass that would exceed the caller's limits, and
228 /// this work bound charges descriptor storage once per library request and once per tuning request.
229 /// The value covers the four matrix families the package authors: matrix product, matrix-vector
230 /// product, batched matrix product and outer product, including the names of their schedule tiles
231 /// and lanes and the copies each family makes. The largest descriptor is the batched matrix
232 /// product, which holds four shapes, six schedule bindings and ten axes. The storage adds room for
233 /// 64 expressions and 128 terms on top of that to cover the tensor dimensions, facts and runtime
234 /// sizes copied into a descriptor. The work bound multiplies the value by the number of library and
235 /// tuning requests.
236 fn einsumDescriptorStorage() u64 {
237 const entry = kernel_library.entry;
238 const shape = accy_choir.shape;
239 const arrays = 4 * std.ArrayList(shape.Symbol).growCapacity(4) * @sizeOf(shape.Symbol) +
240 4 * std.ArrayList(shape.Tensor).growCapacity(4) * @sizeOf(shape.Tensor) +
241 4 * std.ArrayList(shape.Fact).growCapacity(4) * @sizeOf(shape.Fact);
242 const records = 4 * @sizeOf(entry.Shape) + 10 * @sizeOf(entry.Axis) +
243 6 * @sizeOf(entry.ScheduleBinding) + @sizeOf(entry.Reduction) +
244 64 * @sizeOf(shape.Expression) + 128 * @sizeOf(shape.Term);
245 return 8 * (arrays + records + 2048 + 128 * 128) +
246 @sizeOf(alloc_arena.Arena) + @sizeOf(shape.Family) + 128;
247 }
248
249 const EinsumTuningWork = struct {
250 input_bytes: u64 = 0,
251 visits: u64 = 0,
252 family_queries: u64 = 0,
253 };
254
255 fn einsumTuningWork(options: Options, queries: u64) !EinsumTuningWork {
256 if (queries == 0 or options.matrix_product_schedule != null) return .{};
257 var result: EinsumTuningWork = .{};
258 if (options.matrix_product_tuning) |reader| {
259 var bytes = try work.add(@sizeOf(@TypeOf(reader)), reader.device.name.len);
260 if (reader.device.driver_version) |driver| bytes = try work.add(bytes, driver.len);
261 bytes = try work.add(bytes, try work.multiply(
262 reader.records.len,
263 @sizeOf(kernel_library.tuning.MatrixProductFamilyScheduleTuningRecord),
264 ));
265 result.input_bytes = try work.add(result.input_bytes, bytes);
266 }
267 if (options.family_tuning) |reader| {
268 result.family_queries = queries;
269 var bytes = try work.add(@sizeOf(@TypeOf(reader.*)), try work.multiply(
270 reader.table.records.len,
271 @sizeOf(kernel_library.tuning.FamilyTuningRecord),
272 ));
273 for (reader.table.records) |record| bytes = try work.add(bytes, record.target.len);
274 result.input_bytes = try work.add(result.input_bytes, bytes);
275 }
276 result.visits = try work.multiply(128, try work.multiply(
277 queries,
278 try work.add(result.input_bytes, 1),
279 ));
280 return result;
281 }
282
283 fn einsumWorkspace(facts: EinsumWork, tuning: EinsumTuningWork) !u64 {
284 const dimensions = try work.multiply(facts.type_bytes, 2 * @sizeOf(i64));
285 const shapes = try work.multiply(facts.operands, @sizeOf([]const u64) + 256);
286 const frames = try work.multiply(facts.candidates, 1024);
287 const decoding = try work.add(dimensions, try work.add(shapes, frames));
288 const arenas = try work.multiply(8, try work.add(facts.scratch, decoding));
289 const replacement = facts.options.matrix_product_schedule != null or
290 facts.options.matrix_product_tuning != null or facts.options.family_tuning != null;
291 const selections = try work.multiply(
292 facts.library_queries,
293 @as(u64, if (replacement) 8 else 4),
294 );
295 const descriptors = try work.multiply(
296 try work.add(selections, tuning.family_queries),
297 einsumDescriptorStorage(),
298 );
299 const spellings = try work.multiply(
300 tuning.family_queries,
301 kernel_library.geometry.max_thread_candidates * (128 + 8),
302 );
303 const calls = try work.multiply(
304 facts.library_queries,
305 2 * @sizeOf(accy_choir.semantics.KernelOperandEffect) + @sizeOf(?usize) + 256,
306 );
307 const queues = try work.add(
308 try work.arrayListGrowth(*ir.Operation, facts.nodes),
309 try work.arrayListGrowth(*ir.Operation, facts.candidates),
310 );
311 return work.add(
312 try work.add(arenas, descriptors),
313 try work.add(try work.add(spellings, calls), queues),
314 );
315 }
316
317 fn einsumWork(input: work.Input) !work.Bounds {
318 const options: *const Options = if (input.state) |state| @ptrCast(@alignCast(state)) else &.{};
319 const counts = try work.Census.inspect(input.operation);
320 var facts: EinsumWork = .{ .root = input.operation, .options = options.* };
321 _ = try input.operation.walk(.{ .order = .pre_order }, &facts, EinsumWork.visit);
322 const tuning = try einsumTuningWork(options.*, facts.library_queries);
323 const workspace = try einsumWorkspace(facts, tuning);
324 const nodes = try work.multiply(facts.nodes, @sizeOf(ir.Operation) + @sizeOf(ir.Value) +
325 2 * @sizeOf(ir.OpOperand) + 8 * @sizeOf(ir.NamedAttribute) + 4 * 62 * @sizeOf(i64) + 512);
326 const type_bytes = try work.multiply(facts.types, @sizeOf(ir.Type.DialectTypeStorage) +
327 62 * 21 + 128);
328 const emitted = try work.add(facts.nodes, 1);
329 const units = try work.add(try work.add(counts.atoms, counts.input_bytes), emitted);
330 const uses = try work.add(try work.add(counts.values, counts.operands), emitted);
331 const traversal = try work.multiply(128, try work.multiply(units, uses));
332 const catalog = try work.multiply(facts.library_queries, einsumDescriptorStorage());
333 const factories = try work.multiply(128, try work.add(type_bytes, catalog));
334 const processing = try work.add(facts.visits, try work.add(tuning.visits, factories));
335 return .{
336 .work = .{
337 .input_bytes = try work.add(counts.input_bytes, tuning.input_bytes),
338 .output_bytes = try work.add(nodes, type_bytes),
339 .structural_visits = try work.add(traversal, processing),
340 .allocation_capacity = workspace,
341 },
342 .workspace = workspace,
343 };
344 }
345
346 fn destroyOptions(raw: ?*anyopaque, allocator: std.mem.Allocator) void {
347 const options: *Options = @ptrCast(@alignCast(raw orelse return));
348 allocator.destroy(options);
349 }
350
351 fn strategyFromText(value: []const u8) !Strategy {
352 if (std.mem.eql(u8, value, "auto")) return .auto;
353 if (std.mem.eql(u8, value, "left-to-right")) return .left_to_right;
354 if (std.mem.eql(u8, value, "greedy")) return .greedy;
355 if (std.mem.eql(u8, value, "beam")) return .beam;
356 if (std.mem.eql(u8, value, "anytime")) return .anytime;
357 if (std.mem.eql(u8, value, "optimal")) return .optimal;
358 return error.InvalidPassOptionValue;
359 }
360
361 fn kernelLibraryLoweringFromText(value: []const u8) !KernelLibraryLowering {
362 if (std.mem.eql(u8, value, "disabled")) return .disabled;
363 if (std.mem.eql(u8, value, "enabled")) return .enabled;
364 return error.InvalidPassOptionValue;
365 }
366
367 fn runEinsumLoweringPass(pass_ctx: *passes.PassContext) passes.PassResult {
368 return runEinsumLoweringWithOptions(pass_ctx, .{});
369 }
370
371 fn runEinsumLoweringPassWithState(raw: ?*anyopaque, pass_ctx: *passes.PassContext) passes.PassResult {
372 const options: *const Options = @ptrCast(@alignCast(raw orelse return .failure));
373 return runEinsumLoweringWithOptions(pass_ctx, options.*);
374 }
375
376 fn runEinsumLoweringWithOptions(pass_ctx: *passes.PassContext, options: Options) passes.PassResult {
377 var rewriter = rewrite.PatternRewriter.init(pass_ctx.allocator, pass_ctx.ir_ctx);
378 defer rewriter.deinit();
379
380 var lowered_count: usize = 0;
381 lowerOnOp(pass_ctx.op, &rewriter, options, &lowered_count) catch return .failure;
382 if (lowered_count == 0) {
383 pass_ctx.preserveAllAnalyses();
384 } else {
385 rewriter.finalize(pass_ctx.op);
386 pass_ctx.markModified();
387 }
388 return .success;
389 }
390
391 fn lowerOnOp(
392 op: *ir.Operation,
393 rewriter: *rewrite.PatternRewriter,
394 options: Options,
395 lowered_count: *usize,
396 ) !void {
397 for (op.regions.items) |*region| {
398 var block_iter = region.getBlocks();
399 while (block_iter.next()) |block| {
400 var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
401 while (current) |current_op| {
402 const next = current_op.next_op;
403 if (current_op.regions.items.len > 0) {
404 try lowerOnOp(current_op, rewriter, options, lowered_count);
405 }
406 if (std.mem.eql(u8, current_op.name.name, dialect_mod.AccyDialect.EinsumOp.operation_name)) {
407 var guard = rewriter.insertionGuard();
408 defer guard.deinit();
409 rewriter.setInsertionPointBefore(current_op);
410 if (!(try lowerEinsumOp(current_op, rewriter, options))) return error.InvalidState;
411 lowered_count.* += 1;
412 } else if (options.kernel_library == .enabled and
413 (options.matrix_product_schedule != null or
414 options.matrix_product_tuning != null or
415 options.family_tuning != null) and
416 std.mem.eql(u8, current_op.name.name, dialect_mod.AccyDialect.DotGeneralOp.operation_name))
417 {
418 var guard = rewriter.insertionGuard();
419 defer guard.deinit();
420 rewriter.setInsertionPointBefore(current_op);
421 if (try lowerKnownKernelLibraryDotGeneral(current_op, rewriter, options)) {
422 lowered_count.* += 1;
423 }
424 }
425 current = next;
426 }
427 }
428 }
429 }
430
431 fn lowerEinsumOp(op: *ir.Operation, rewriter: *rewrite.PatternRewriter, options: Options) !bool {
432 if (!std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.EinsumOp.operation_name)) return false;
433 if (op.getNumResults() != 1) return false;
434 const equation_text = (dialect_mod.AccyDialect.EinsumOp{ .op = op }).getEquation() orelse return false;
435
436 var arena_state = alloc_arena.Arena.init(rewriter.allocator);
437 defer arena_state.deinit();
438 const arena = arena_state.allocator();
439
440 const operands = op.getOperandValues();
441 const shapes = try arena.alloc([]const u64, operands.len);
442 var dtype: ?choir_abi.DType = null;
443 for (operands, 0..) |operand, i| {
444 const tensor_type = try dialect_mod.decodeTensorType(arena, operand.type);
445 if (dtype) |known| {
446 if (known != tensor_type.dtype) return false;
447 } else {
448 dtype = tensor_type.dtype;
449 }
450 const dims = try arena.alloc(u64, tensor_type.dims.len);
451 for (tensor_type.dims, 0..) |dim, dim_index| {
452 if (dim < 0) return false;
453 dims[dim_index] = @intCast(dim);
454 }
455 shapes[i] = dims;
456 }
457
458 var equation = try accy_choir.einsum.parse(arena, equation_text, shapes);
459 defer equation.deinit();
460 if (options.kernel_library == .enabled) {
461 if (try lowerKnownKernelLibraryEinsum(op, rewriter, &equation, operands, dtype orelse return false, options)) |lowered| {
462 try rewriter.replaceOpWithValue(op, lowered);
463 return true;
464 }
465 }
466
467 var plan = try accy_choir.einsum.createPlan(arena, &equation, .{
468 .strategy = options.strategy,
469 .exact_state_limit = options.exact_state_limit,
470 .beam_width = options.beam_width,
471 .auto_beam_width = options.auto_beam_width,
472 });
473 defer plan.deinit();
474
475 const lowered = try accy_choir.einsum.lowerPlanWithRewriter(
476 arena,
477 rewriter,
478 &equation,
479 &plan,
480 operands,
481 dtype orelse return false,
482 );
483 try rewriter.replaceOpWithValue(op, lowered);
484 return true;
485 }
486
487 fn lowerKnownKernelLibraryEinsum(
488 op: *ir.Operation,
489 rewriter: *rewrite.PatternRewriter,
490 equation: *const accy_choir.einsum.Equation,
491 operands: []const *ir.Value,
492 dtype: choir_abi.DType,
493 options: Options,
494 ) !?*ir.Value {
495 const result = op.getResult(0) orelse return null;
496
497 var arena_state = alloc_arena.Arena.init(rewriter.allocator);
498 defer arena_state.deinit();
499 const arena = arena_state.allocator();
500 const output_type = try dialect_mod.decodeTensorType(arena, result.type);
501 if (output_type.dtype != dtype) return null;
502
503 const logical_inputs = try arena.alloc(kernel_selection.EinsumOperand, equation.inputs.len);
504 for (equation.inputs, logical_inputs) |input, *logical_input| {
505 logical_input.* = .{
506 .indices = input.indices,
507 .dims = try shapeToI64(arena, input.dims),
508 };
509 }
510
511 const request = kernel_selection.EinsumSelectionRequest{
512 .dtype = dtype,
513 .inputs = logical_inputs,
514 .output_indices = equation.output,
515 .output_dims = output_type.dims,
516 };
517 var selection = (try selectKnownKernelLibraryContraction(rewriter, request, options)) orelse return null;
518 defer selection.selected.deinit();
519
520 const result_types = [_]ir.Type{result.type};
521 const call = try call_preparation.insertCatalogCall(rewriter, .{
522 .descriptor = selection.selected.descriptor.descriptor,
523 .operands = operands,
524 .result_types = &result_types,
525 });
526 return call.getFirstResult();
527 }
528
529 const ContractionSelection = struct {
530 selected: kernel_selection.OwnedSelectedEinsumKernel,
531 scheduled: bool,
532 };
533
534 fn selectKnownKernelLibraryContraction(
535 rewriter: *rewrite.PatternRewriter,
536 initial_request: kernel_selection.EinsumSelectionRequest,
537 options: Options,
538 ) !?ContractionSelection {
539 var request = initial_request;
540 var selected = (try kernel_selection.selectOwnedEinsumCatalog(rewriter.allocator, request)) orelse return null;
541 var selected_owned = true;
542 defer if (selected_owned) selected.deinit();
543 var scheduled = false;
544 if (options.matrix_product_schedule) |matrix_product_schedule| {
545 if (selected.kind == .matrix_product) {
546 request.schedule = .{ .matrix_product = matrix_product_schedule };
547 const replacement = (try kernel_selection.selectOwnedEinsumCatalog(rewriter.allocator, request)) orelse return null;
548 selected.deinit();
549 selected = replacement;
550 scheduled = true;
551 }
552 } else if (options.matrix_product_tuning) |reader| {
553 if (selected.kind == .matrix_product) {
554 if (kernel_library.linalg.matrixProductInstanceFromSpecialization(
555 selected.descriptor.descriptor.metadata.specialization,
556 )) |instance| {
557 if (try reader.resolve(instance)) |thread_blocks| {
558 request.schedule = .{ .matrix_product = .{ .thread_blocks = thread_blocks } };
559 const replacement = (try kernel_selection.selectOwnedEinsumCatalog(rewriter.allocator, request)) orelse return null;
560 selected.deinit();
561 selected = replacement;
562 scheduled = true;
563 }
564 }
565 }
566 }
567 if (!scheduled and options.family_tuning != null) {
568 const family_tuning = options.family_tuning.?;
569 if (selected.kind == .matrix_product) {
570 if (kernel_library.linalg.matrixProductInstanceFromSpecialization(
571 selected.descriptor.descriptor.metadata.specialization,
572 )) |instance| {
573 if (try kernel_library.linalg.resolveMatrixProductSchedule(
574 rewriter.allocator,
575 family_tuning.*,
576 instance,
577 )) |thread_blocks| {
578 request.schedule = .{ .matrix_product = .{ .thread_blocks = thread_blocks } };
579 const replacement = (try kernel_selection.selectOwnedEinsumCatalog(rewriter.allocator, request)) orelse return null;
580 selected.deinit();
581 selected = replacement;
582 scheduled = true;
583 }
584 }
585 }
586 }
587 selected_owned = false;
588 return .{ .selected = selected, .scheduled = scheduled };
589 }
590
591 const dot_general_lhs_indices = "mk";
592 const dot_general_rhs_indices = "kn";
593 const dot_general_output_indices = "mn";
594
595 fn lowerKnownKernelLibraryDotGeneral(
596 op: *ir.Operation,
597 rewriter: *rewrite.PatternRewriter,
598 options: Options,
599 ) !bool {
600 if (op.getNumResults() != 1) return false;
601 const result = op.getResult(0) orelse return false;
602 const operands = op.getOperandValues();
603 if (operands.len != 2) return false;
604
605 const dot = dialect_mod.AccyDialect.DotGeneralOp{ .op = op };
606 const lhs_batch = dot.getLhsBatchPayload() orelse return false;
607 const rhs_batch = dot.getRhsBatchPayload() orelse return false;
608 const lhs_contract = dot.getLhsContractPayload() orelse return false;
609 const rhs_contract = dot.getRhsContractPayload() orelse return false;
610 if (lhs_batch.len != 0 or rhs_batch.len != 0) return false;
611 if (!std.mem.eql(u8, lhs_contract, std.mem.sliceAsBytes(&[_]i64{1}))) return false;
612 if (!std.mem.eql(u8, rhs_contract, std.mem.sliceAsBytes(&[_]i64{0}))) return false;
613
614 var arena_state = alloc_arena.Arena.init(rewriter.allocator);
615 defer arena_state.deinit();
616 const arena = arena_state.allocator();
617
618 const lhs_type = try dialect_mod.decodeTensorType(arena, operands[0].type);
619 const rhs_type = try dialect_mod.decodeTensorType(arena, operands[1].type);
620 const output_type = try dialect_mod.decodeTensorType(arena, result.type);
621 if (lhs_type.dtype != rhs_type.dtype or lhs_type.dtype != output_type.dtype) return false;
622 if (lhs_type.dims.len != 2 or rhs_type.dims.len != 2 or output_type.dims.len != 2) return false;
623
624 const logical_inputs = [_]kernel_selection.EinsumOperand{
625 .{ .indices = dot_general_lhs_indices, .dims = lhs_type.dims },
626 .{ .indices = dot_general_rhs_indices, .dims = rhs_type.dims },
627 };
628 const request = kernel_selection.EinsumSelectionRequest{
629 .dtype = lhs_type.dtype,
630 .inputs = logical_inputs[0..],
631 .output_indices = dot_general_output_indices,
632 .output_dims = output_type.dims,
633 };
634 var selection = (try selectKnownKernelLibraryContraction(rewriter, request, options)) orelse return false;
635 defer selection.selected.deinit();
636 if (selection.selected.kind != .matrix_product) return false;
637 if (options.matrix_product_schedule == null and !selection.scheduled) return false;
638
639 const result_types = [_]ir.Type{result.type};
640 const call = try call_preparation.insertCatalogCall(rewriter, .{
641 .descriptor = selection.selected.descriptor.descriptor,
642 .operands = operands,
643 .result_types = &result_types,
644 });
645 try rewriter.replaceOpWithValue(op, call.getFirstResult());
646 return true;
647 }
648
649 fn shapeToI64(allocator: std.mem.Allocator, shape: []const u64) ![]const i64 {
650 const dims = try allocator.alloc(i64, shape.len);
651 for (shape, 0..) |dim, index| {
652 if (dim > @as(u64, @intCast(std.math.maxInt(i64)))) return error.InvalidDimension;
653 dims[index] = @intCast(dim);
654 }
655 return dims;
656 }
657
658 fn findOpNamed(op: *ir.Operation, name: []const u8) ?*ir.Operation {
659 if (std.mem.eql(u8, op.name.name, name)) return op;
660 for (op.regions.items) |*region| {
661 var block_iter = region.getBlocks();
662 while (block_iter.next()) |block| {
663 var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
664 while (current) |current_op| {
665 if (findOpNamed(current_op, name)) |found| return found;
666 current = current_op.next_op;
667 }
668 }
669 }
670 return null;
671 }
672
673 const testing = std.testing;
674 const semantic = accy_choir.semantic;
675
676 test "einsum lowering pass replaces semantic einsum with planned contraction" {
677 const allocator = testing.allocator;
678
679 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
680 defer builder.deinit();
681 const lhs_ty = try builder.tensor(.f32, &.{ 5, 8 });
682 const rhs_ty = try builder.tensor(.f32, &.{ 8, 16 });
683 const out_ty = try builder.tensor(.f32, &.{ 5, 16 });
684 var fb = try builder.beginFunction("einsum_lowering_pass", &.{ lhs_ty, rhs_ty }, &.{out_ty});
685 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "ik,kj->ij");
686 try fb.return_(&.{out});
687 try fb.finish();
688 const module = try builder.finish();
689 defer module.deinit();
690
691 var pm = passes.PassManager.init(allocator);
692 defer pm.deinit();
693 try pm.addPass(einsumLoweringPass());
694
695 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
696 try module.verify();
697 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
698 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
699 }
700
701 test "einsum lowering pass selects kernel library matrix product" {
702 const allocator = testing.allocator;
703
704 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
705 defer builder.deinit();
706 const lhs_ty = try builder.tensor(.f32, &.{ 4, 8 });
707 const rhs_ty = try builder.tensor(.f32, &.{ 8, 16 });
708 const out_ty = try builder.tensor(.f32, &.{ 4, 16 });
709 var fb = try builder.beginFunction("einsum_kernel_library_lowering_pass", &.{ lhs_ty, rhs_ty }, &.{out_ty});
710 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "ik,kj->ij");
711 try fb.return_(&.{out});
712 try fb.finish();
713 const module = try builder.finish();
714 defer module.deinit();
715
716 var options = Options{ .kernel_library = .enabled };
717 var pm = passes.PassManager.init(allocator);
718 defer pm.deinit();
719 try pm.addPass(einsumLoweringPassWithOptions(&options));
720
721 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
722 try module.verify();
723 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
724 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
725 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
726 }
727
728 test "einsum lowering pass selects kernel library matrix product family" {
729 const allocator = testing.allocator;
730
731 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
732 defer builder.deinit();
733 const lhs_ty = try builder.tensor(.f32, &.{ 5, 3 });
734 const rhs_ty = try builder.tensor(.f32, &.{ 3, 7 });
735 const out_ty = try builder.tensor(.f32, &.{ 5, 7 });
736 var fb = try builder.beginFunction("einsum_kernel_library_matmul_family_lowering_pass", &.{ lhs_ty, rhs_ty }, &.{out_ty});
737 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "mk,kn->mn");
738 try fb.return_(&.{out});
739 try fb.finish();
740 const module = try builder.finish();
741 defer module.deinit();
742
743 var options = Options{ .kernel_library = .enabled };
744 var pm = passes.PassManager.init(allocator);
745 defer pm.deinit();
746 try pm.addPass(einsumLoweringPassWithOptions(&options));
747
748 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
749 try module.verify();
750 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
751 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
752 const kernel_call = findOpNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
753 return error.TestExpectedKernelCall;
754 };
755 const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
756 const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
757 try testing.expectEqualStrings("accy.kernel.linalg.matmul_family_7x5_f32", target.payload);
758 }
759
760 test "einsum lowering pass selects scheduled kernel library matrix product family" {
761 const allocator = testing.allocator;
762
763 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
764 defer builder.deinit();
765 const lhs_ty = try builder.tensor(.f32, &.{ 5, 3 });
766 const rhs_ty = try builder.tensor(.f32, &.{ 3, 7 });
767 const out_ty = try builder.tensor(.f32, &.{ 5, 7 });
768 var fb = try builder.beginFunction("einsum_kernel_library_matmul_family_schedule_lowering_pass", &.{ lhs_ty, rhs_ty }, &.{out_ty});
769 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "mk,kn->mn");
770 try fb.return_(&.{out});
771 try fb.finish();
772 const module = try builder.finish();
773 defer module.deinit();
774
775 var options = Options{
776 .kernel_library = .enabled,
777 .matrix_product_schedule = .{ .thread_blocks = .{ .x = 4, .y = 2 } },
778 };
779 var pm = passes.PassManager.init(allocator);
780 defer pm.deinit();
781 try pm.addPass(einsumLoweringPassWithOptions(&options));
782
783 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
784 try module.verify();
785 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
786 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
787 const kernel_call = findOpNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
788 return error.TestExpectedKernelCall;
789 };
790 const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
791 const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
792 try testing.expectEqualStrings("accy.kernel.linalg.matmul_family_4x2_f32", target.payload);
793 }
794
795 test "einsum lowering pass selects kernel library matrix vector product" {
796 const allocator = testing.allocator;
797
798 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
799 defer builder.deinit();
800 const matrix_ty = try builder.tensor(.f32, &.{ 4, 8 });
801 const vector_ty = try builder.tensor(.f32, &.{8});
802 const out_ty = try builder.tensor(.f32, &.{4});
803 var fb = try builder.beginFunction("einsum_kernel_library_matvec_lowering_pass", &.{ matrix_ty, vector_ty }, &.{out_ty});
804 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "mk,k->m");
805 try fb.return_(&.{out});
806 try fb.finish();
807 const module = try builder.finish();
808 defer module.deinit();
809
810 var options = Options{ .kernel_library = .enabled };
811 var pm = passes.PassManager.init(allocator);
812 defer pm.deinit();
813 try pm.addPass(einsumLoweringPassWithOptions(&options));
814
815 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
816 try module.verify();
817 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
818 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
819 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ReduceOp.operation_name));
820 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
821 }
822
823 test "einsum lowering pass selects kernel library outer product" {
824 const allocator = testing.allocator;
825
826 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
827 defer builder.deinit();
828 const lhs_ty = try builder.tensor(.f32, &.{4});
829 const rhs_ty = try builder.tensor(.f32, &.{3});
830 const out_ty = try builder.tensor(.f32, &.{ 4, 3 });
831 var fb = try builder.beginFunction("einsum_kernel_library_outer_lowering_pass", &.{ lhs_ty, rhs_ty }, &.{out_ty});
832 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "m,n->mn");
833 try fb.return_(&.{out});
834 try fb.finish();
835 const module = try builder.finish();
836 defer module.deinit();
837
838 var options = Options{ .kernel_library = .enabled };
839 var pm = passes.PassManager.init(allocator);
840 defer pm.deinit();
841 try pm.addPass(einsumLoweringPassWithOptions(&options));
842
843 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
844 try module.verify();
845 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
846 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name));
847 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.MulOp.operation_name));
848 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
849 }
850
851 test "einsum lowering pass selects kernel library transpose" {
852 const allocator = testing.allocator;
853
854 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
855 defer builder.deinit();
856 const input_ty = try builder.tensor(.f32, &.{ 8, 16 });
857 const out_ty = try builder.tensor(.f32, &.{ 16, 8 });
858 var fb = try builder.beginFunction("einsum_kernel_library_transpose_lowering_pass", &.{input_ty}, &.{out_ty});
859 const out = try fb.einsum(&.{fb.parameter(0)}, out_ty, "ij->ji");
860 try fb.return_(&.{out});
861 try fb.finish();
862 const module = try builder.finish();
863 defer module.deinit();
864
865 var options = Options{ .kernel_library = .enabled };
866 var pm = passes.PassManager.init(allocator);
867 defer pm.deinit();
868 try pm.addPass(einsumLoweringPassWithOptions(&options));
869
870 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
871 try module.verify();
872 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
873 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
874 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
875 }
876
877 test "einsum lowering pass selects kernel library scalar sum reduction" {
878 const allocator = testing.allocator;
879
880 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
881 defer builder.deinit();
882 const input_ty = try builder.tensor(.f32, &.{8});
883 const out_ty = try builder.tensor(.f32, &.{});
884 var fb = try builder.beginFunction("einsum_kernel_library_sum_lowering_pass", &.{input_ty}, &.{out_ty});
885 const out = try fb.einsum(&.{fb.parameter(0)}, out_ty, "i->");
886 try fb.return_(&.{out});
887 try fb.finish();
888 const module = try builder.finish();
889 defer module.deinit();
890
891 var options = Options{ .kernel_library = .enabled };
892 var pm = passes.PassManager.init(allocator);
893 defer pm.deinit();
894 try pm.addPass(einsumLoweringPassWithOptions(&options));
895
896 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
897 try module.verify();
898 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
899 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ReduceOp.operation_name));
900 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
901 }
902
903 test "einsum lowering pass selects kernel library scalar dot product" {
904 const allocator = testing.allocator;
905
906 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
907 defer builder.deinit();
908 const input_ty = try builder.tensor(.f32, &.{8});
909 const out_ty = try builder.tensor(.f32, &.{});
910 var fb = try builder.beginFunction("einsum_kernel_library_dot_lowering_pass", &.{ input_ty, input_ty }, &.{out_ty});
911 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "i,i->");
912 try fb.return_(&.{out});
913 try fb.finish();
914 const module = try builder.finish();
915 defer module.deinit();
916
917 var options = Options{ .kernel_library = .enabled };
918 var pm = passes.PassManager.init(allocator);
919 defer pm.deinit();
920 try pm.addPass(einsumLoweringPassWithOptions(&options));
921
922 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
923 try module.verify();
924 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
925 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.MulOp.operation_name));
926 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ReduceOp.operation_name));
927 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
928 }
929
930 test "einsum lowering pass keeps attention-shaped pure einsum generic with kernel library" {
931 const allocator = testing.allocator;
932
933 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
934 defer builder.deinit();
935 const query_ty = try builder.tensor(.f32, &.{ 2, 2, 2 });
936 const key_ty = try builder.tensor(.f32, &.{ 2, 3, 2 });
937 const value_ty = try builder.tensor(.f32, &.{ 2, 3, 2 });
938 const out_ty = try builder.tensor(.f32, &.{ 2, 2, 2 });
939 var fb = try builder.beginFunction("einsum_kernel_library_attention_shaped_generic_lowering_pass", &.{ query_ty, key_ty, value_ty }, &.{out_ty});
940 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1), fb.parameter(2) }, out_ty, "bqh,bkh,bkv->bqv");
941 try fb.return_(&.{out});
942 try fb.finish();
943 const module = try builder.finish();
944 defer module.deinit();
945
946 var options = Options{ .kernel_library = .enabled };
947 var pm = passes.PassManager.init(allocator);
948 defer pm.deinit();
949 try pm.addPass(einsumLoweringPassWithOptions(&options));
950
951 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
952 try module.verify();
953 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
954 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
955 try testing.expect(ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.MulOp.operation_name) > 0);
956 try testing.expect(ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ReduceOp.operation_name) > 0);
957 }
958
959 test "einsum lowering pass keeps registered matrix product generic by default" {
960 const allocator = testing.allocator;
961
962 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
963 defer builder.deinit();
964 const lhs_ty = try builder.tensor(.f32, &.{ 4, 8 });
965 const rhs_ty = try builder.tensor(.f32, &.{ 8, 16 });
966 const out_ty = try builder.tensor(.f32, &.{ 4, 16 });
967 var fb = try builder.beginFunction("einsum_registered_matmul_generic_lowering_pass", &.{ lhs_ty, rhs_ty }, &.{out_ty});
968 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "ik,kj->ij");
969 try fb.return_(&.{out});
970 try fb.finish();
971 const module = try builder.finish();
972 defer module.deinit();
973
974 var pm = passes.PassManager.init(allocator);
975 defer pm.deinit();
976 try pm.addPass(einsumLoweringPass());
977
978 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
979 try module.verify();
980 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
981 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
982 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
983 }
984
985 test "einsum lowering pass keeps registered transpose generic by default" {
986 const allocator = testing.allocator;
987
988 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
989 defer builder.deinit();
990 const input_ty = try builder.tensor(.f32, &.{ 8, 16 });
991 const out_ty = try builder.tensor(.f32, &.{ 16, 8 });
992 var fb = try builder.beginFunction("einsum_registered_transpose_generic_lowering_pass", &.{input_ty}, &.{out_ty});
993 const out = try fb.einsum(&.{fb.parameter(0)}, out_ty, "ij->ji");
994 try fb.return_(&.{out});
995 try fb.finish();
996 const module = try builder.finish();
997 defer module.deinit();
998
999 var pm = passes.PassManager.init(allocator);
1000 defer pm.deinit();
1001 try pm.addPass(einsumLoweringPass());
1002
1003 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1004 try module.verify();
1005 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
1006 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.TransposeOp.operation_name));
1007 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1008 }
1009
1010 test "einsum lowering pass preserves analyses when no einsum exists" {
1011 const allocator = testing.allocator;
1012
1013 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1014 defer builder.deinit();
1015 const ty = try builder.tensor(.f32, &.{4});
1016 var fb = try builder.beginFunction("einsum_lowering_noop", &.{ ty, ty }, &.{ty});
1017 const out = try fb.add(fb.parameter(0), fb.parameter(1));
1018 try fb.return_(&.{out});
1019 try fb.finish();
1020 const module = try builder.finish();
1021 defer module.deinit();
1022
1023 var pm = passes.PassManager.init(allocator);
1024 defer pm.deinit();
1025 try pm.addPass(einsumLoweringPass());
1026
1027 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1028 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
1029 }
1030
1031 test "einsum lowering pass accepts beam width options" {
1032 const allocator = testing.allocator;
1033
1034 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1035 defer builder.deinit();
1036 const a_ty = try builder.tensor(.f32, &.{ 32, 4 });
1037 const b_ty = try builder.tensor(.f32, &.{ 4, 64 });
1038 const c_ty = try builder.tensor(.f32, &.{ 4, 3 });
1039 const d_ty = try builder.tensor(.f32, &.{ 4, 16 });
1040 const out_ty = try builder.tensor(.f32, &.{ 32, 64, 3, 16 });
1041 var fb = try builder.beginFunction("einsum_lowering_beam_width", &.{ a_ty, b_ty, c_ty, d_ty }, &.{out_ty});
1042 const out = try fb.einsum(
1043 &.{ fb.parameter(0), fb.parameter(1), fb.parameter(2), fb.parameter(3) },
1044 out_ty,
1045 "ab,bc,bd,be->acde",
1046 );
1047 try fb.return_(&.{out});
1048 try fb.finish();
1049 const module = try builder.finish();
1050 defer module.deinit();
1051
1052 var options = Options{ .strategy = .beam, .beam_width = 2 };
1053 var pm = passes.PassManager.init(allocator);
1054 defer pm.deinit();
1055 try pm.addPass(einsumLoweringPassWithOptions(&options));
1056
1057 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1058 try module.verify();
1059 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.EinsumOp.operation_name));
1060 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
1061 try testing.expectEqual(@as(usize, 6), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name));
1062 try testing.expectEqual(@as(usize, 3), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.MulOp.operation_name));
1063 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ReduceOp.operation_name));
1064 }
1065
1066 fn familyTuningTestCapabilities() gpu.BackendCapabilities {
1067 return .{ .identity = .{
1068 .backend = .cuda,
1069 .family = .nvidia_cuda,
1070 .name = "pass-test-device",
1071 .vendor_id = 0x10de,
1072 .device_id = 0x2684,
1073 } };
1074 }
1075
1076 test "einsum lowering pass consults family tuning when schedule is null" {
1077 const allocator = testing.allocator;
1078
1079 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1080 defer builder.deinit();
1081 const lhs_ty = try builder.tensor(.f32, &.{ 5, 3 });
1082 const rhs_ty = try builder.tensor(.f32, &.{ 3, 7 });
1083 const out_ty = try builder.tensor(.f32, &.{ 5, 7 });
1084 var fb = try builder.beginFunction("einsum_kernel_library_matmul_family_tuned_lowering_pass", &.{ lhs_ty, rhs_ty }, &.{out_ty});
1085 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "mk,kn->mn");
1086 try fb.return_(&.{out});
1087 try fb.finish();
1088 const module = try builder.finish();
1089 defer module.deinit();
1090
1091 const caps = familyTuningTestCapabilities();
1092 const device = kernel_library.tuning.deviceFingerprint(caps);
1093 const probe = kernel_library.linalg.MatrixProduct{ .m = 5, .n = 7, .k = 3 };
1094 const thread_candidates = kernel_library.linalg.matrixProductThreadCandidatesForExtents(probe.m, probe.n);
1095 try testing.expect(thread_candidates.slice().len >= 2);
1096 var winner = probe;
1097 winner.threads = thread_candidates.slice()[0];
1098 const winner_target = try kernel_library.linalg.matrixProductFamilyTarget(allocator, winner);
1099 defer allocator.free(winner_target);
1100
1101 const records = [_]kernel_library.tuning.FamilyTuningRecord{.{
1102 .key = try kernel_library.linalg.matrixProductFamilyTuningKey(allocator, device, probe),
1103 .target = winner_target,
1104 .winner_median_ns = 800,
1105 .runner_up_median_ns = 1200,
1106 .sample_count = 30,
1107 }};
1108 const reader = kernel_library.tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
1109
1110 var options = Options{
1111 .kernel_library = .enabled,
1112 .family_tuning = &reader,
1113 };
1114 var pm = passes.PassManager.init(allocator);
1115 defer pm.deinit();
1116 try pm.addPass(einsumLoweringPassWithOptions(&options));
1117
1118 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1119 try module.verify();
1120 const kernel_call = findOpNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
1121 return error.TestExpectedKernelCall;
1122 };
1123 const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
1124 const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
1125 try testing.expectEqualStrings(winner_target, target.payload);
1126 }
1127
1128 test "einsum lowering pass keeps heuristic schedule on family tuning miss" {
1129 const allocator = testing.allocator;
1130
1131 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1132 defer builder.deinit();
1133 const lhs_ty = try builder.tensor(.f32, &.{ 5, 3 });
1134 const rhs_ty = try builder.tensor(.f32, &.{ 3, 7 });
1135 const out_ty = try builder.tensor(.f32, &.{ 5, 7 });
1136 var fb = try builder.beginFunction("einsum_kernel_library_matmul_family_tuning_miss_lowering_pass", &.{ lhs_ty, rhs_ty }, &.{out_ty});
1137 const out = try fb.einsum(&.{ fb.parameter(0), fb.parameter(1) }, out_ty, "mk,kn->mn");
1138 try fb.return_(&.{out});
1139 try fb.finish();
1140 const module = try builder.finish();
1141 defer module.deinit();
1142
1143 const reader = kernel_library.tuning.FamilyTuningReader.init(familyTuningTestCapabilities(), .{});
1144
1145 var options = Options{
1146 .kernel_library = .enabled,
1147 .family_tuning = &reader,
1148 };
1149 var pm = passes.PassManager.init(allocator);
1150 defer pm.deinit();
1151 try pm.addPass(einsumLoweringPassWithOptions(&options));
1152
1153 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1154 try module.verify();
1155 const kernel_call = findOpNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
1156 return error.TestExpectedKernelCall;
1157 };
1158 const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
1159 const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
1160 try testing.expectEqualStrings("accy.kernel.linalg.matmul_family_7x5_f32", target.payload);
1161 }
1162
1163 fn dotGeneralTunedModule(allocator: std.mem.Allocator, name: []const u8) !*semantic.SemanticModule {
1164 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1165 errdefer builder.deinit();
1166 const lhs_ty = try builder.tensor(.f32, &.{ 5, 3 });
1167 const rhs_ty = try builder.tensor(.f32, &.{ 3, 7 });
1168 const out_ty = try builder.tensor(.f32, &.{ 5, 7 });
1169 var fb = try builder.beginFunction(name, &.{ lhs_ty, rhs_ty }, &.{out_ty});
1170 const out = try fb.dotGeneral(fb.parameter(0), fb.parameter(1), out_ty, &.{1}, &.{0}, &.{}, &.{});
1171 try fb.return_(&.{out});
1172 try fb.finish();
1173 return try builder.finish();
1174 }
1175
1176 test "einsum lowering pass consults family tuning for dot general on hit" {
1177 const allocator = testing.allocator;
1178
1179 const module = try dotGeneralTunedModule(allocator, "dot_general_family_tuning_hit_lowering_pass");
1180 defer module.deinit();
1181
1182 const caps = familyTuningTestCapabilities();
1183 const device = kernel_library.tuning.deviceFingerprint(caps);
1184 const probe = kernel_library.linalg.MatrixProduct{ .m = 5, .n = 7, .k = 3 };
1185 const thread_candidates = kernel_library.linalg.matrixProductThreadCandidatesForExtents(probe.m, probe.n);
1186 try testing.expect(thread_candidates.slice().len >= 1);
1187 var winner = probe;
1188 winner.threads = thread_candidates.slice()[0];
1189 const winner_target = try kernel_library.linalg.matrixProductFamilyTarget(allocator, winner);
1190 defer allocator.free(winner_target);
1191
1192 const records = [_]kernel_library.tuning.FamilyTuningRecord{.{
1193 .key = try kernel_library.linalg.matrixProductFamilyTuningKey(allocator, device, probe),
1194 .target = winner_target,
1195 .winner_median_ns = 700,
1196 .runner_up_median_ns = 1100,
1197 .sample_count = 30,
1198 }};
1199 const reader = kernel_library.tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
1200
1201 var options = Options{
1202 .kernel_library = .enabled,
1203 .family_tuning = &reader,
1204 };
1205 var pm = passes.PassManager.init(allocator);
1206 defer pm.deinit();
1207 try pm.addPass(einsumLoweringPassWithOptions(&options));
1208
1209 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1210 try module.verify();
1211 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
1212 const kernel_call = findOpNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
1213 return error.TestExpectedKernelCall;
1214 };
1215 const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
1216 const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
1217 try testing.expectEqualStrings(winner_target, target.payload);
1218 }
1219
1220 test "einsum lowering pass keeps dot general semantic on family tuning miss" {
1221 const allocator = testing.allocator;
1222
1223 const module = try dotGeneralTunedModule(allocator, "dot_general_family_tuning_miss_lowering_pass");
1224 defer module.deinit();
1225
1226 const reader = kernel_library.tuning.FamilyTuningReader.init(familyTuningTestCapabilities(), .{});
1227
1228 var options = Options{
1229 .kernel_library = .enabled,
1230 .family_tuning = &reader,
1231 };
1232 var pm = passes.PassManager.init(allocator);
1233 defer pm.deinit();
1234 try pm.addPass(einsumLoweringPassWithOptions(&options));
1235
1236 try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1237 try module.verify();
1238 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
1239 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1240 }
1241
1242 const EinsumCase = struct {
1243 expression: []const u8 = "mk,kn->mn",
1244 shapes: []const []const i64 = &.{ &.{ 5, 3 }, &.{ 3, 7 } },
1245 output: []const i64 = &.{ 5, 7 },
1246 dtype: choir_abi.DType = .f32,
1247 copies: u32 = 1,
1248 dot: bool = false,
1249 outcome: passes.PassResult = .success,
1250 target: ?[]const u8 = null,
1251
1252 fn module(self: EinsumCase) !*semantic.SemanticModule {
1253 const allocator = testing.allocator;
1254 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.testing);
1255 defer builder.deinit();
1256 const types = try allocator.alloc(ir.Type, self.shapes.len);
1257 defer allocator.free(types);
1258 for (self.shapes, types) |shape, *typ| typ.* = try builder.tensor(self.dtype, shape);
1259 const result_type = try builder.tensor(self.dtype, self.output);
1260 var fb = try builder.beginFunction("einsum_work", types, &.{result_type});
1261 const operands = try allocator.alloc(*ir.Value, types.len);
1262 defer allocator.free(operands);
1263 for (operands, 0..) |*operand, index| operand.* = fb.parameter(index);
1264 var value: *ir.Value = undefined;
1265 for (0..self.copies) |_| {
1266 value = if (self.dot)
1267 try fb.dotGeneral(operands[0], operands[1], result_type, &.{1}, &.{0}, &.{}, &.{})
1268 else if (self.outcome == .success)
1269 try fb.einsum(operands, result_type, self.expression)
1270 else blk: {
1271 const op = try dialect_mod.AccyDialect.EinsumOp.create(
1272 fb.ctx,
1273 fb.location,
1274 operands,
1275 result_type,
1276 self.expression,
1277 );
1278 try fb.entry.addOperation(op.op);
1279 break :blk op.getResult();
1280 };
1281 }
1282 try fb.return_(&.{value});
1283 try fb.finish();
1284 if (self.outcome == .failure) {
1285 const unfinished = builder.module.?;
1286 builder.module = null;
1287 return unfinished;
1288 }
1289 return builder.finish();
1290 }
1291 };
1292
1293 fn einsumStorageImage(fixture: EinsumCase, options: Options, bounded: bool) ![]u8 {
1294 const allocator = testing.allocator;
1295 const module = try fixture.module();
1296 defer module.deinit();
1297 const root = module.choir_module;
1298 const before = try choir.bytecode.encodeModule(allocator, root);
1299 defer allocator.free(before);
1300 const bounds = try einsumWork(.{ .operation = root, .state = &options });
1301 const bytes = try allocator.alloc(u8, @intCast(if (bounded) bounds.workspace else 0));
1302 defer allocator.free(bytes);
1303 var buffer = std.heap.FixedBufferAllocator.init(bytes);
1304 const base = buffer.allocator();
1305 const vtable: std.mem.Allocator.VTable = .{
1306 .alloc = base.vtable.alloc,
1307 .resize = std.mem.Allocator.noResize,
1308 .remap = std.mem.Allocator.noRemap,
1309 .free = std.mem.Allocator.noFree,
1310 };
1311 var cache = passes.AnalysisCache.init(allocator, null);
1312 defer cache.deinit();
1313 var context = passes.PassContext.init(root, module.context(), allocator, &cache);
1314 defer context.deinit();
1315 context.allocator = if (bounded) .{ .ptr = base.ptr, .vtable = &vtable } else allocator;
1316 defer context.allocator = allocator;
1317 try testing.expectEqual(fixture.outcome, runEinsumLoweringWithOptions(&context, options));
1318 if (bounded) try testing.expect(buffer.end_index <= bounds.workspace);
1319 try testing.expectEqual(null, module.context().exhaustedSegment());
1320 if (fixture.outcome == .success) try module.verify();
1321 if (fixture.target) |target| {
1322 const name = dialect_mod.AccyDialect.KernelCallOp.operation_name;
1323 try testing.expectEqual(fixture.copies, ir.inspection.countOperationsNamed(root, name));
1324 const call = findOpNamed(root, name).?;
1325 const actual = call.getAttr("target").?.cast(ir.Attribute.DialectAttr).?.payload;
1326 try testing.expectEqualStrings(target, actual);
1327 }
1328 const image = try choir.bytecode.encodeModule(allocator, root);
1329 errdefer allocator.free(image);
1330 try testing.expect(image.len <= before.len + bounds.work.output_bytes);
1331 return image;
1332 }
1333
1334 fn checkEinsumStorage(fixture: EinsumCase, options: Options) !void {
1335 const expected = try einsumStorageImage(fixture, options, false);
1336 defer testing.allocator.free(expected);
1337 const actual = try einsumStorageImage(fixture, options, true);
1338 defer testing.allocator.free(actual);
1339 try testing.expectEqualSlices(u8, expected, actual);
1340 }
1341
1342 test "einsum lowering work covers every strategy and generic reduction path" {
1343 const chain: EinsumCase = .{
1344 .expression = "ab,bc,cd,de->ea",
1345 .shapes = &.{ &.{ 2, 3 }, &.{ 3, 4 }, &.{ 4, 5 }, &.{ 5, 6 } },
1346 .output = &.{ 6, 2 },
1347 };
1348 for (comptime std.meta.tags(Strategy)) |strategy| {
1349 try checkEinsumStorage(chain, .{ .strategy = strategy, .beam_width = 3 });
1350 }
1351 try checkEinsumStorage(chain, .{ .exact_state_limit = 0, .beam_width = 3 });
1352 try checkEinsumStorage(.{
1353 .expression = "abcd,defg->ga",
1354 .shapes = &.{ &.{ 2, 3, 4, 5 }, &.{ 5, 6, 7, 8 } },
1355 .output = &.{ 8, 2 },
1356 .copies = 17,
1357 }, .{ .strategy = .left_to_right, .kernel_library = .enabled });
1358 try checkEinsumStorage(.{ .expression = "->", .shapes = &.{&.{}}, .output = &.{} }, .{});
1359 try checkEinsumStorage(.{
1360 .expression = "abc->ca",
1361 .shapes = &.{&.{ 2, 3, 4 }},
1362 .output = &.{ 4, 2 },
1363 }, .{});
1364 try checkEinsumStorage(.{ .dot = true }, .{});
1365 }
1366
1367 test "einsum lowering work covers static and owned library families" {
1368 const cases = [_]EinsumCase{
1369 .{},
1370 .{ .dtype = .f16 },
1371 .{ .shapes = &.{ &.{ 8, 8 }, &.{ 8, 8 } }, .output = &.{ 8, 8 } },
1372 .{
1373 .expression = "bmk,bkn->bmn",
1374 .shapes = &.{ &.{ 3, 5, 7 }, &.{ 3, 7, 9 } },
1375 .output = &.{ 3, 5, 9 },
1376 },
1377 .{ .expression = "mk,k->m", .shapes = &.{ &.{ 5, 3 }, &.{3} }, .output = &.{5} },
1378 .{ .expression = "m,n->mn", .shapes = &.{ &.{5}, &.{7} } },
1379 .{ .expression = "ij->ji", .shapes = &.{&.{ 8, 8 }}, .output = &.{ 8, 8 } },
1380 .{ .expression = "i->", .shapes = &.{&.{8}}, .output = &.{} },
1381 .{ .expression = "i,i->", .shapes = &.{ &.{8}, &.{8} }, .output = &.{} },
1382 };
1383 for (cases) |fixture| {
1384 try checkEinsumStorage(fixture, .{ .kernel_library = .enabled });
1385 var repeated = fixture;
1386 repeated.copies = 17;
1387 try checkEinsumStorage(repeated, .{ .kernel_library = .enabled });
1388 }
1389 const scheduled: Options = .{
1390 .kernel_library = .enabled,
1391 .matrix_product_schedule = .{ .thread_blocks = .{ .x = 4, .y = 2 } },
1392 };
1393 try checkEinsumStorage(.{ .copies = 17 }, scheduled);
1394 try checkEinsumStorage(.{ .copies = 17, .dot = true }, scheduled);
1395 }
1396
1397 test "einsum lowering work covers maximum label and input cardinalities" {
1398 const labels = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
1399 const dims: [labels.len]i64 = @splat(1);
1400 const shapes: [63][]const i64 = @splat(&dims);
1401 var text: std.ArrayList(u8) = .empty;
1402 defer text.deinit(testing.allocator);
1403 for (shapes, 0..) |_, index| {
1404 if (index != 0) try text.append(testing.allocator, ',');
1405 try text.appendSlice(testing.allocator, labels);
1406 }
1407 try text.appendSlice(testing.allocator, "->" ++ labels);
1408 try checkEinsumStorage(.{ .expression = text.items, .shapes = &shapes, .output = &dims }, .{
1409 .strategy = .left_to_right,
1410 });
1411 }
1412
1413 fn checkEinsumAccounting(constructor: u32, denied: ?u32) !void {
1414 const allocator = testing.allocator;
1415 const revision = choir.product.revision;
1416 const module = try (EinsumCase{}).module();
1417 defer module.deinit();
1418 const root = module.choir_module;
1419 const before = try choir.bytecode.encodeModule(allocator, root);
1420 defer allocator.free(before);
1421 const options: Options = .{ .kernel_library = if (constructor == 0) .disabled else .enabled };
1422 const bounds = try einsumWork(.{ .operation = root, .state = &options });
1423 var allowance = revision.WorkVector.uniform(1 << 60);
1424 var workspace: u64 = 1 << 30;
1425 if (denied) |dimension| switch (dimension) {
1426 0 => allowance.structural_visits = bounds.work.structural_visits - 1,
1427 1 => allowance.allocation_capacity = bounds.work.allocation_capacity - 1,
1428 2 => allowance.output_bytes = bounds.work.output_bytes - 1,
1429 3 => workspace = bounds.workspace - 1,
1430 else => unreachable,
1431 };
1432 const ledger = try revision.AccountingV1.create(allocator, .{
1433 .allowance = allowance,
1434 .workspace = workspace,
1435 .events = 4,
1436 }, &.{.{ .name = einsum_lowering_pass_name, .version = 1 }});
1437 defer ledger.destroy();
1438 var cache = try passes.AnalysisCache.initAccounted(
1439 allocator,
1440 null,
1441 ledger,
1442 .{ .context = module.context() },
1443 0,
1444 );
1445 defer cache.deinit();
1446 var manager = passes.PassManager.init(allocator);
1447 defer manager.deinit();
1448 try manager.addPass(switch (constructor) {
1449 0 => einsumLoweringPass(),
1450 1 => einsumLoweringPassWithOptions(&options),
1451 2 => try einsumLoweringPassFromOptions(allocator, .{ .assignments = &.{
1452 .{ .name = "kernel-library", .value = "enabled" },
1453 } }),
1454 else => unreachable,
1455 });
1456 const result = manager.runWithAnalysisCache(root, module.context(), &cache, .{});
1457 try testing.expectEqual(if (denied == null) passes.PassResult.success else .failure, result);
1458 if (denied != null) {
1459 try testing.expectEqual(.exhausted, ledger.view().outcome);
1460 try testing.expectEqual(@as(u64, 0), manager.stats.pass_runs);
1461 const after = try choir.bytecode.encodeModule(allocator, root);
1462 defer allocator.free(after);
1463 try testing.expectEqualSlices(u8, before, after);
1464 } else {
1465 try ledger.producersComplete();
1466 try testing.expect(!ledger.view().missing_work_contract);
1467 try testing.expectEqual(@as(u64, 1), ledger.view().executed.counters.pass_runs);
1468 try module.verify();
1469 }
1470 }
1471
1472 test "einsum lowering work admits all constructors and refuses before mutation" {
1473 for (0..3) |constructor| {
1474 try checkEinsumAccounting(@intCast(constructor), null);
1475 for (0..4) |dimension| {
1476 try checkEinsumAccounting(@intCast(constructor), @intCast(dimension));
1477 }
1478 }
1479 }
1480
1481 test "einsum lowering work preserves failed parse and lowering prefixes" {
1482 const cases = [_]EinsumCase{
1483 .{ .expression = "mk,kn", .outcome = .failure },
1484 .{ .expression = "mk,kn->mn->", .outcome = .failure },
1485 .{ .expression = "mm,kn->mn", .outcome = .failure },
1486 .{ .expression = "mk,kn->x", .outcome = .failure },
1487 .{ .shapes = &.{ &.{ 5, 3 }, &.{ 4, 7 } }, .outcome = .failure },
1488 .{ .shapes = &.{ &.{ 5, -1 }, &.{ 3, 7 } }, .outcome = .failure },
1489 .{
1490 .expression = "ab,cd->",
1491 .outcome = .failure,
1492 .output = &.{},
1493 .shapes = &.{ &.{ 4_000_000_000, 4_000_000_000 }, &.{ 2, 3 } },
1494 },
1495 };
1496 for (cases) |fixture| try checkEinsumStorage(fixture, .{});
1497 const too_many: [64][]const i64 = @splat(&.{1});
1498 try checkEinsumStorage(.{
1499 .expression = "i->i",
1500 .shapes = &too_many,
1501 .output = &.{1},
1502 .outcome = .failure,
1503 }, .{});
1504 }
1505
1506 const EinsumTuningCase = struct {
1507 const tuning = kernel_library.tuning;
1508 const Reader = kernel_library.linalg.MatrixProductScheduleReader;
1509 const Record = tuning.MatrixProductFamilyScheduleTuningRecord;
1510
1511 fn instance() kernel_library.linalg.MatrixProduct {
1512 var value: kernel_library.linalg.MatrixProduct = .{ .m = 5, .n = 7, .k = 3 };
1513 const candidates = kernel_library.linalg.matrixProductThreadCandidatesForExtents(5, 7);
1514 value.threads = candidates.slice()[candidates.slice().len - 1];
1515 return value;
1516 }
1517
1518 fn matrixRecord() !Record {
1519 const value = instance();
1520 const candidates = kernel_library.linalg.matrixProductThreadCandidatesForExtents(5, 7);
1521 var threads: [8]tuning.MatrixProductFamilyScheduleThreads = undefined;
1522 for (candidates.slice(), 0..) |candidate, index| {
1523 threads[index] = .{ .x = candidate.x, .y = candidate.y };
1524 }
1525 return .{
1526 .key = try tuning.MatrixProductFamilyScheduleTuningKey.init(
1527 familyTuningTestCapabilities().identity,
1528 .{
1529 .format = .cuda_ptx,
1530 .m = 5,
1531 .n = 7,
1532 .k = 3,
1533 .dtype = value.dtype,
1534 .accumulation_dtype = value.accumulation_dtype,
1535 .family_version = kernel_library.linalg.matrix_product_family_version,
1536 .candidates = threads[0..candidates.slice().len],
1537 },
1538 ),
1539 .selection = .{
1540 .threads = .{ .x = value.threads.x, .y = value.threads.y },
1541 .winner_median_ns = 1,
1542 .runner_up_median_ns = 2,
1543 .sample_count = 3,
1544 },
1545 };
1546 }
1547
1548 fn matrixReader(records: []const Record) Reader {
1549 return .{
1550 .device = familyTuningTestCapabilities().identity,
1551 .format = .cuda_ptx,
1552 .records = records,
1553 };
1554 }
1555
1556 fn familyRecord(target: []const u8) !tuning.FamilyTuningRecord {
1557 return .{
1558 .key = try kernel_library.linalg.matrixProductFamilyTuningKey(
1559 testing.allocator,
1560 tuning.deviceFingerprint(familyTuningTestCapabilities()),
1561 instance(),
1562 ),
1563 .target = target,
1564 .winner_median_ns = 1,
1565 .runner_up_median_ns = 2,
1566 .sample_count = 3,
1567 };
1568 }
1569 };
1570
1571 test "einsum lowering work covers matrix and family tuning hits misses and precedence" {
1572 const allocator = testing.allocator;
1573 const target = try kernel_library.linalg.matrixProductFamilyTarget(
1574 allocator,
1575 EinsumTuningCase.instance(),
1576 );
1577 defer allocator.free(target);
1578 var matrix = [_]EinsumTuningCase.Record{try EinsumTuningCase.matrixRecord()};
1579 var family = [_]kernel_library.tuning.FamilyTuningRecord{
1580 try EinsumTuningCase.familyRecord(target),
1581 };
1582 const family_reader = kernel_library.tuning.FamilyTuningReader.init(
1583 familyTuningTestCapabilities(),
1584 .{ .records = &family },
1585 );
1586 for ([_]bool{ false, true }) |dot| {
1587 const fixture: EinsumCase = .{ .dot = dot, .copies = 17, .target = target };
1588 var options: Options = .{
1589 .kernel_library = .enabled,
1590 .matrix_product_tuning = EinsumTuningCase.matrixReader(&matrix),
1591 .family_tuning = &family_reader,
1592 };
1593 try checkEinsumStorage(fixture, options);
1594 matrix[0].key.k = 99;
1595 try checkEinsumStorage(fixture, options);
1596 family[0].key.device_fingerprint = 0;
1597 var miss = fixture;
1598 miss.target = null;
1599 try checkEinsumStorage(miss, options);
1600 matrix[0] = try EinsumTuningCase.matrixRecord();
1601 matrix[0].selection.threads = .{ .x = 99, .y = 99 };
1602 miss.outcome = .failure;
1603 try checkEinsumStorage(miss, options);
1604 options.matrix_product_schedule = .{ .thread_blocks = .{
1605 .x = EinsumTuningCase.instance().threads.x,
1606 .y = EinsumTuningCase.instance().threads.y,
1607 } };
1608 try checkEinsumStorage(fixture, options);
1609 matrix[0] = try EinsumTuningCase.matrixRecord();
1610 family[0] = try EinsumTuningCase.familyRecord(target);
1611 }
1612 }
1613
1614 test "einsum lowering work charges complete tuning tables and device strings" {
1615 const allocator = testing.allocator;
1616 const target = try kernel_library.linalg.matrixProductFamilyTarget(
1617 allocator,
1618 EinsumTuningCase.instance(),
1619 );
1620 defer allocator.free(target);
1621 var records: [257]EinsumTuningCase.Record = @splat(try EinsumTuningCase.matrixRecord());
1622 for (records[0..256]) |*record| record.key.k = 99;
1623 var family: [257]kernel_library.tuning.FamilyTuningRecord =
1624 @splat(try EinsumTuningCase.familyRecord(target));
1625 for (family[0..256]) |*record| record.key.device_fingerprint = 0;
1626 var family_reader = kernel_library.tuning.FamilyTuningReader.init(
1627 familyTuningTestCapabilities(),
1628 .{ .records = &family },
1629 );
1630 var options: Options = .{
1631 .kernel_library = .enabled,
1632 .matrix_product_tuning = EinsumTuningCase.matrixReader(&records),
1633 .family_tuning = &family_reader,
1634 };
1635 try checkEinsumStorage(.{ .copies = 17, .target = target }, options);
1636 const small = try einsumTuningWork(.{
1637 .matrix_product_tuning = EinsumTuningCase.matrixReader(records[256..]),
1638 }, 17);
1639 const large = try einsumTuningWork(options, 17);
1640 try testing.expect(large.input_bytes > small.input_bytes);
1641 try testing.expect(large.visits > small.visits);
1642 options.matrix_product_tuning.?.records = records[0..256];
1643 try checkEinsumStorage(.{ .copies = 17, .target = target }, options);
1644 const long_name: [4096]u8 = @splat('n');
1645 options.matrix_product_tuning.?.device.name = &long_name;
1646 options.matrix_product_tuning.?.device.driver_version = &long_name;
1647 family[0].target = &long_name;
1648 const longer = try einsumTuningWork(options, 17);
1649 try testing.expect(longer.input_bytes > large.input_bytes);
1650 try testing.expect(longer.visits > large.visits);
1651 try checkEinsumStorage(.{ .copies = 17, .target = target }, options);
1652 options.matrix_product_schedule = .{ .thread_blocks = .{ .x = 4, .y = 2 } };
1653 try testing.expectEqualDeep(EinsumTuningWork{}, try einsumTuningWork(options, 17));
1654 try testing.expectEqualDeep(EinsumTuningWork{}, try einsumTuningWork(options, 0));
1655 }
1656
1657 test "einsum lowering work bounds each owned descriptor family directly" {
1658 const requests = [_]kernel_selection.EinsumSelectionRequest{
1659 .{ .dtype = .f32, .inputs = &.{
1660 .{ .indices = "mk", .dims = &.{ 5, 3 } },
1661 .{ .indices = "kn", .dims = &.{ 3, 7 } },
1662 }, .output_indices = "mn", .output_dims = &.{ 5, 7 } },
1663 .{ .dtype = .f32, .inputs = &.{
1664 .{ .indices = "bmk", .dims = &.{ 3, 5, 7 } },
1665 .{ .indices = "bkn", .dims = &.{ 3, 7, 9 } },
1666 }, .output_indices = "bmn", .output_dims = &.{ 3, 5, 9 } },
1667 .{ .dtype = .f32, .inputs = &.{
1668 .{ .indices = "mk", .dims = &.{ 5, 3 } },
1669 .{ .indices = "k", .dims = &.{3} },
1670 }, .output_indices = "m", .output_dims = &.{5} },
1671 .{ .dtype = .f32, .inputs = &.{
1672 .{ .indices = "m", .dims = &.{5} },
1673 .{ .indices = "n", .dims = &.{7} },
1674 }, .output_indices = "mn", .output_dims = &.{ 5, 7 } },
1675 };
1676 for (requests) |request| {
1677 const bytes = try testing.allocator.alloc(u8, @intCast(einsumDescriptorStorage()));
1678 defer testing.allocator.free(bytes);
1679 var buffer = std.heap.FixedBufferAllocator.init(bytes);
1680 const base = buffer.allocator();
1681 const vtable: std.mem.Allocator.VTable = .{
1682 .alloc = base.vtable.alloc,
1683 .resize = std.mem.Allocator.noResize,
1684 .remap = std.mem.Allocator.noRemap,
1685 .free = std.mem.Allocator.noFree,
1686 };
1687 var selected = (try kernel_selection.selectOwnedEinsumCatalog(
1688 .{ .ptr = base.ptr, .vtable = &vtable },
1689 request,
1690 )).?;
1691 defer selected.deinit();
1692 try testing.expect(selected.descriptor.specialization != null);
1693 try testing.expect(buffer.end_index > 0);
1694 try testing.expect(buffer.end_index <= bytes.len);
1695 }
1696 }
1697
1698 fn checkEinsumTuningAdmission(options: Options) !void {
1699 const allocator = testing.allocator;
1700 const revision = choir.product.revision;
1701 const module = try (EinsumCase{}).module();
1702 defer module.deinit();
1703 const root = module.choir_module;
1704 const before = try choir.bytecode.encodeModule(allocator, root);
1705 defer allocator.free(before);
1706 var baseline = options;
1707 if (baseline.matrix_product_tuning) |*reader| reader.records = reader.records[0..1];
1708 var family: kernel_library.tuning.FamilyTuningReader = undefined;
1709 if (baseline.family_tuning) |reader| {
1710 family = reader.*;
1711 family.table.records = family.table.records[0..1];
1712 baseline.family_tuning = &family;
1713 }
1714 const small = try einsumWork(.{ .operation = root, .state = &baseline });
1715 var allowance = revision.WorkVector.uniform(1 << 60);
1716 allowance.input_bytes = small.work.input_bytes;
1717 const ledger = try revision.AccountingV1.create(allocator, .{
1718 .allowance = allowance,
1719 .workspace = 1 << 30,
1720 .events = 4,
1721 }, &.{.{ .name = einsum_lowering_pass_name, .version = 1 }});
1722 defer ledger.destroy();
1723 var cache = try passes.AnalysisCache.initAccounted(
1724 allocator,
1725 null,
1726 ledger,
1727 .{ .context = module.context() },
1728 0,
1729 );
1730 defer cache.deinit();
1731 var manager = passes.PassManager.init(allocator);
1732 defer manager.deinit();
1733 try manager.addPass(einsumLoweringPassWithOptions(&options));
1734 try testing.expectEqual(
1735 .failure,
1736 manager.runWithAnalysisCache(root, module.context(), &cache, .{}),
1737 );
1738 try testing.expectEqual(.exhausted, ledger.view().outcome);
1739 try testing.expectEqual(@as(u64, 0), manager.stats.pass_runs);
1740 const after = try choir.bytecode.encodeModule(allocator, root);
1741 defer allocator.free(after);
1742 try testing.expectEqualSlices(u8, before, after);
1743 }
1744
1745 test "einsum lowering work rejects enlarged tuning inputs before mutation" {
1746 const matrix: [257]EinsumTuningCase.Record = @splat(try EinsumTuningCase.matrixRecord());
1747 try checkEinsumTuningAdmission(.{
1748 .kernel_library = .enabled,
1749 .matrix_product_tuning = EinsumTuningCase.matrixReader(&matrix),
1750 });
1751 const records: [257]kernel_library.tuning.FamilyTuningRecord =
1752 @splat(try EinsumTuningCase.familyRecord("stale-target"));
1753 const reader = kernel_library.tuning.FamilyTuningReader.init(
1754 familyTuningTestCapabilities(),
1755 .{ .records = &records },
1756 );
1757 try checkEinsumTuningAdmission(.{ .kernel_library = .enabled, .family_tuning = &reader });
1758 }