lib/accy/src/choir/einsum/lowering.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3 const testing = std.testing;
4
5 const choir = @import("choir");
6 const choir_root = @import("../root.zig");
7 const planner = @import("planner.zig");
8 const spec = @import("spec.zig");
9
10 const ir = choir.ir;
11 const rewrite = ir.rewrite;
12 const DType = choir_abi.DType;
13 const Equation = spec.Equation;
14 const IndexSet = spec.IndexSet;
15 const Plan = planner.Plan;
16 const Step = planner.Step;
17 const AccyDialect = choir_root.AccyDialect;
18
19 pub const LowerError = error{
20 InputCountMismatch,
21 InvalidPlan,
22 InvalidDimension,
23 } || std.mem.Allocator.Error;
24
25 const State = struct {
26 value: ?*ir.Value = null,
27 order: []const u8 = &.{},
28 owned: bool = false,
29 };
30
31 /// Four figures for one call of `lowerPlan`: bytes of local scratch, structural
32 /// visits, operations requested from the IR sink, and tensor types requested
33 /// from it, so a compiler pass can charge one lowering against its work budget
34 /// before it runs it. The figures assume an equation from `parse` and a plan
35 /// from `createPlan`, so at most 62 labels, at most 63 inputs, and exactly one
36 /// step fewer than inputs. Growth of the caller's arena, the internals of the
37 /// context and of the sink, rewriter queues, parsing and planning are left out,
38 /// and their owners add their own charges, as the einsum pass does.
39 pub const LoweringBounds = struct {
40 scratch_bytes: u64,
41 structural_visits: u64,
42 operation_requests: u64,
43 tensor_type_requests: u64,
44
45 /// Bytes charged for one list of `T`: four times its grown capacity for 62
46 /// items, alignment slack for each allocation, and one more slice for a
47 /// final owned copy, so the bound functions can charge one list of at most
48 /// 62 items.
49 fn listBytes(comptime T: type) u64 {
50 return 4 * std.ArrayList(T).growCapacity(62) * @sizeOf(T) +
51 4 * @alignOf(T) + sliceBytes(T);
52 }
53
54 fn sliceBytes(comptime T: type) u64 {
55 return 62 * @sizeOf(T) + @alignOf(T) - 1;
56 }
57
58 /// Scratch charged for one contraction step: the reduction of each operand,
59 /// the batch and contracted axis lists, the broadcast-and-reduce path, and
60 /// the final transpose, so `loweringBounds` can multiply it by the step
61 /// count to charge the scratch of a whole plan. The charge is 11 byte
62 /// lists, 6 axis lists, 14 slices of dimensions or permutations, and 1 kept
63 /// order.
64 fn stepScratch() u64 {
65 return 11 * listBytes(u8) + 6 * listBytes(i64) +
66 14 * sliceBytes(i64) + sliceBytes(u8);
67 }
68 };
69
70 /// Returns the figures for an equation with `input_count` inputs, or
71 /// `error.EmptyEquation` for none and `error.TooManyInputs` above 63, and that
72 /// check keeps every product in range. The einsum pass calls this with the
73 /// operand count to add the lowering's share to its work estimate. Each step is
74 /// charged for both operand reductions and for the general broadcast-and-reduce
75 /// path, even when the matrix-product path skips them. One reduction may
76 /// transpose, make a zero, reshape, reduce, and reshape back. The operation
77 /// figure is 19 per step and the tensor-type figure 17 per step, and a single
78 /// input is charged 3 of each for its own reduction and the final transpose.
79 /// Visits count searches over labels and axes, which grow with the square of
80 /// the rank, together with local copies and cleanup. Visits measure neither
81 /// machine instructions nor tensor execution.
82 pub fn loweringBounds(input_count: usize) !LoweringBounds {
83 if (input_count == 0) return error.EmptyEquation;
84 if (input_count > 63) return error.TooManyInputs;
85 const steps: u64 = input_count - 1;
86 const states: u64 = input_count + steps;
87 const scratch = states * @sizeOf(State) + @alignOf(State) - 1 + if (steps == 0)
88 LoweringBounds.listBytes(u8) + LoweringBounds.listBytes(i64) +
89 3 * LoweringBounds.sliceBytes(i64)
90 else
91 steps * LoweringBounds.stepScratch();
92 return .{
93 .scratch_bytes = scratch,
94 .structural_visits = 8 * scratch + 64 * 63 * 63 * @max(steps, 1) + 8 * states,
95 .operation_requests = if (steps == 0) 3 else 19 * steps,
96 .tensor_type_requests = if (steps == 0) 3 else 17 * steps,
97 };
98 }
99
100 pub fn lowerPlan(
101 allocator: std.mem.Allocator,
102 fb: *choir_root.semantic.FunctionBuilder,
103 equation: *const Equation,
104 plan: *const Plan,
105 inputs: []const *ir.Value,
106 dtype: DType,
107 ) !*ir.Value {
108 var sink = FunctionSink{ .fb = fb };
109 return try lowerPlanWithSink(allocator, &sink, equation, plan, inputs, dtype);
110 }
111
112 pub fn lowerPlanWithRewriter(
113 allocator: std.mem.Allocator,
114 rewriter: *rewrite.PatternRewriter,
115 equation: *const Equation,
116 plan: *const Plan,
117 inputs: []const *ir.Value,
118 dtype: DType,
119 ) !*ir.Value {
120 var sink = RewriterSink{ .rewriter = rewriter };
121 return try lowerPlanWithSink(allocator, &sink, equation, plan, inputs, dtype);
122 }
123
124 const FunctionSink = struct {
125 fb: *choir_root.semantic.FunctionBuilder,
126
127 fn context(self: *FunctionSink) *ir.Context {
128 return self.fb.ctx;
129 }
130
131 fn constant(self: *FunctionSink, result_type: ir.Type, bytes: []const u8) !*ir.Value {
132 return try self.fb.constant(result_type, bytes);
133 }
134
135 fn dotGeneral(
136 self: *FunctionSink,
137 lhs: *ir.Value,
138 rhs: *ir.Value,
139 result_type: ir.Type,
140 contracting_lhs: []const i64,
141 contracting_rhs: []const i64,
142 batching_lhs: []const i64,
143 batching_rhs: []const i64,
144 ) !*ir.Value {
145 return try self.fb.dotGeneral(
146 lhs,
147 rhs,
148 result_type,
149 contracting_lhs,
150 contracting_rhs,
151 batching_lhs,
152 batching_rhs,
153 );
154 }
155
156 fn reduce(
157 self: *FunctionSink,
158 input: *ir.Value,
159 init: *ir.Value,
160 result_type: ir.Type,
161 reducer_kind: []const u8,
162 dimensions: []const i64,
163 ) !*ir.Value {
164 return try self.fb.reduce(input, init, result_type, reducer_kind, dimensions);
165 }
166
167 fn reshape(self: *FunctionSink, input: *ir.Value, result_type: ir.Type, new_shape: []const i64) !*ir.Value {
168 return try self.fb.reshape(input, result_type, new_shape);
169 }
170
171 fn broadcastInDim(
172 self: *FunctionSink,
173 input: *ir.Value,
174 result_type: ir.Type,
175 result_shape: []const i64,
176 broadcast_dims: []const i64,
177 ) !*ir.Value {
178 return try self.fb.broadcastInDim(input, result_type, result_shape, broadcast_dims);
179 }
180
181 fn mul(self: *FunctionSink, lhs: *ir.Value, rhs: *ir.Value) !*ir.Value {
182 return try self.fb.mul(lhs, rhs);
183 }
184
185 fn transpose(self: *FunctionSink, input: *ir.Value, result_type: ir.Type, permutation: []const i64) !*ir.Value {
186 return try self.fb.transpose(input, result_type, permutation);
187 }
188 };
189
190 const RewriterSink = struct {
191 rewriter: *rewrite.PatternRewriter,
192
193 fn context(self: *RewriterSink) *ir.Context {
194 return self.rewriter.ir_ctx;
195 }
196
197 fn constant(self: *RewriterSink, result_type: ir.Type, bytes: []const u8) !*ir.Value {
198 const op = try AccyDialect.ConstantOp.create(self.rewriter.ir_ctx, ir.Location.getUnknown(), bytes, result_type);
199 _ = try self.rewriter.insert(op.op);
200 return op.getResult();
201 }
202
203 fn dotGeneral(
204 self: *RewriterSink,
205 lhs: *ir.Value,
206 rhs: *ir.Value,
207 result_type: ir.Type,
208 contracting_lhs: []const i64,
209 contracting_rhs: []const i64,
210 batching_lhs: []const i64,
211 batching_rhs: []const i64,
212 ) !*ir.Value {
213 const op = try AccyDialect.DotGeneralOp.create(
214 self.rewriter.ir_ctx,
215 ir.Location.getUnknown(),
216 lhs,
217 rhs,
218 result_type,
219 batching_lhs,
220 batching_rhs,
221 contracting_lhs,
222 contracting_rhs,
223 );
224 _ = try self.rewriter.insert(op.op);
225 return op.getResult();
226 }
227
228 fn reduce(
229 self: *RewriterSink,
230 input: *ir.Value,
231 init: *ir.Value,
232 result_type: ir.Type,
233 reducer_kind: []const u8,
234 dimensions: []const i64,
235 ) !*ir.Value {
236 const op = try AccyDialect.ReduceOp.create(
237 self.rewriter.ir_ctx,
238 ir.Location.getUnknown(),
239 input,
240 init,
241 result_type,
242 reducer_kind,
243 dimensions,
244 );
245 _ = try self.rewriter.insert(op.op);
246 return op.getResult();
247 }
248
249 fn reshape(self: *RewriterSink, input: *ir.Value, result_type: ir.Type, new_shape: []const i64) !*ir.Value {
250 const op = try AccyDialect.ReshapeOp.create(
251 self.rewriter.ir_ctx,
252 ir.Location.getUnknown(),
253 input,
254 result_type,
255 new_shape,
256 );
257 _ = try self.rewriter.insert(op.op);
258 return op.getResult();
259 }
260
261 fn broadcastInDim(
262 self: *RewriterSink,
263 input: *ir.Value,
264 result_type: ir.Type,
265 result_shape: []const i64,
266 broadcast_dims: []const i64,
267 ) !*ir.Value {
268 const op = try AccyDialect.BroadcastInDimOp.create(
269 self.rewriter.ir_ctx,
270 ir.Location.getUnknown(),
271 input,
272 result_type,
273 broadcast_dims,
274 result_shape,
275 );
276 _ = try self.rewriter.insert(op.op);
277 return op.getResult();
278 }
279
280 fn mul(self: *RewriterSink, lhs: *ir.Value, rhs: *ir.Value) !*ir.Value {
281 const op = try AccyDialect.MulOp.create(self.rewriter.ir_ctx, ir.Location.getUnknown(), lhs, rhs);
282 _ = try self.rewriter.insert(op.op);
283 return op.getResult();
284 }
285
286 fn transpose(self: *RewriterSink, input: *ir.Value, result_type: ir.Type, permutation: []const i64) !*ir.Value {
287 const op = try AccyDialect.TransposeOp.create(
288 self.rewriter.ir_ctx,
289 ir.Location.getUnknown(),
290 input,
291 result_type,
292 permutation,
293 );
294 _ = try self.rewriter.insert(op.op);
295 return op.getResult();
296 }
297 };
298
299 fn lowerPlanWithSink(
300 allocator: std.mem.Allocator,
301 sink: anytype,
302 equation: *const Equation,
303 plan: *const Plan,
304 inputs: []const *ir.Value,
305 dtype: DType,
306 ) !*ir.Value {
307 if (inputs.len != equation.inputs.len) return error.InputCountMismatch;
308 const state_count = equation.inputs.len + plan.steps.len;
309 const states = try allocator.alloc(State, state_count);
310 defer allocator.free(states);
311 @memset(states, .{});
312 defer {
313 for (states) |state| {
314 if (state.owned) allocator.free(state.order);
315 }
316 }
317
318 for (inputs, 0..) |input, index| {
319 states[index] = .{
320 .value = input,
321 .order = equation.inputs[index].indices,
322 };
323 }
324
325 if (plan.steps.len == 0) {
326 if (equation.inputs.len != 1) return error.InvalidPlan;
327 return try lowerSingleInput(allocator, sink, equation, states[0], dtype);
328 }
329
330 for (plan.steps) |step| {
331 try lowerStep(allocator, sink, equation, step, states, dtype);
332 }
333
334 const final_state = states[states.len - 1];
335 return final_state.value orelse error.InvalidPlan;
336 }
337
338 fn lowerSingleInput(
339 allocator: std.mem.Allocator,
340 sink: anytype,
341 equation: *const Equation,
342 input: State,
343 dtype: DType,
344 ) !*ir.Value {
345 var value = input.value orelse return error.InvalidPlan;
346 var order = input.order;
347
348 var reduce_axes = std.ArrayListUnmanaged(i64).empty;
349 defer reduce_axes.deinit(allocator);
350 var kept = std.ArrayListUnmanaged(u8).empty;
351 defer kept.deinit(allocator);
352 for (order, 0..) |index, axis| {
353 if (equation.output_set.contains(index)) {
354 try kept.append(allocator, index);
355 } else {
356 try reduce_axes.append(allocator, @intCast(axis));
357 }
358 }
359
360 if (reduce_axes.items.len != 0) {
361 const result_dims = try dimsForOrder(allocator, equation, kept.items);
362 defer allocator.free(result_dims);
363 const result_type = try choir_root.accyTensorType(sink.context(), dtype, result_dims);
364 const scalar_type = try choir_root.accyTensorType(sink.context(), dtype, &.{});
365 const zero = try zeroConstant(sink, scalar_type, dtype);
366 value = try sink.reduce(value, zero, result_type, "sum", reduce_axes.items);
367 order = kept.items;
368 }
369
370 return try transposeToOrder(allocator, sink, equation, value, order, equation.output, dtype);
371 }
372
373 fn lowerStep(
374 allocator: std.mem.Allocator,
375 sink: anytype,
376 equation: *const Equation,
377 step: Step,
378 states: []State,
379 dtype: DType,
380 ) !void {
381 if (step.lhs >= states.len or step.rhs >= states.len or step.result >= states.len) return error.InvalidPlan;
382 const lhs = states[step.lhs];
383 const rhs = states[step.rhs];
384 const initial_lhs_value = lhs.value orelse return error.InvalidPlan;
385 const initial_rhs_value = rhs.value orelse return error.InvalidPlan;
386
387 const initial_lhs_set = setFromOrder(lhs.order);
388 const initial_rhs_set = setFromOrder(rhs.order);
389 const reduced_lhs = try reduceOperandLocalSummedIndices(allocator, sink, equation, initial_lhs_value, lhs.order, step.summed_indices, initial_rhs_set, dtype);
390 defer reduced_lhs.deinit(allocator);
391 const reduced_rhs = try reduceOperandLocalSummedIndices(allocator, sink, equation, initial_rhs_value, rhs.order, step.summed_indices, initial_lhs_set, dtype);
392 defer reduced_rhs.deinit(allocator);
393
394 var lhs_batch = std.ArrayListUnmanaged(i64).empty;
395 defer lhs_batch.deinit(allocator);
396 var rhs_batch = std.ArrayListUnmanaged(i64).empty;
397 defer rhs_batch.deinit(allocator);
398 var lhs_contract = std.ArrayListUnmanaged(i64).empty;
399 defer lhs_contract.deinit(allocator);
400 var rhs_contract = std.ArrayListUnmanaged(i64).empty;
401 defer rhs_contract.deinit(allocator);
402 var contracted_order = std.ArrayListUnmanaged(u8).empty;
403 defer contracted_order.deinit(allocator);
404 var raw_order = std.ArrayListUnmanaged(u8).empty;
405 defer raw_order.deinit(allocator);
406
407 const rhs_set = setFromOrder(reduced_rhs.order);
408 for (reduced_lhs.order, 0..) |index, axis| {
409 if (!rhs_set.contains(index)) continue;
410 if (step.result_indices.contains(index)) {
411 try lhs_batch.append(allocator, @intCast(axis));
412 try rhs_batch.append(allocator, try axisOf(reduced_rhs.order, index));
413 try raw_order.append(allocator, index);
414 }
415 }
416
417 for (reduced_lhs.order, 0..) |index, axis| {
418 if (!step.summed_indices.contains(index)) continue;
419 if (!rhs_set.contains(index)) return error.InvalidPlan;
420 try lhs_contract.append(allocator, @intCast(axis));
421 try rhs_contract.append(allocator, try axisOf(reduced_rhs.order, index));
422 try contracted_order.append(allocator, index);
423 }
424
425 for (reduced_lhs.order) |index| {
426 if (step.summed_indices.contains(index)) continue;
427 if (contains(raw_order.items, index)) continue;
428 try raw_order.append(allocator, index);
429 }
430
431 for (reduced_rhs.order) |index| {
432 if (step.summed_indices.contains(index)) continue;
433 if (contains(raw_order.items, index)) continue;
434 try raw_order.append(allocator, index);
435 }
436
437 if (!setFromOrder(raw_order.items).eql(step.result_indices)) return error.InvalidPlan;
438
439 const raw_dims = try dimsForOrder(allocator, equation, raw_order.items);
440 defer allocator.free(raw_dims);
441 const raw_type = try choir_root.accyTensorType(sink.context(), dtype, raw_dims);
442 const raw_value = if (useDotGeneral(
443 reduced_lhs.order,
444 reduced_rhs.order,
445 lhs_batch.items,
446 rhs_batch.items,
447 lhs_contract.items,
448 rhs_contract.items,
449 raw_order.items,
450 ))
451 try sink.dotGeneral(
452 reduced_lhs.value,
453 reduced_rhs.value,
454 raw_type,
455 lhs_contract.items,
456 rhs_contract.items,
457 lhs_batch.items,
458 rhs_batch.items,
459 )
460 else
461 try multiplyBroadcastReduce(
462 allocator,
463 sink,
464 equation,
465 raw_type,
466 raw_dims,
467 reduced_lhs,
468 reduced_rhs,
469 raw_order.items,
470 contracted_order.items,
471 dtype,
472 );
473
474 const final_mask = allInputsMask(equation.inputs.len);
475 const desired = if (step.result_inputs == final_mask) equation.output else raw_order.items;
476 const value = try transposeToOrder(allocator, sink, equation, raw_value, raw_order.items, desired, dtype);
477 states[step.result] = .{
478 .value = value,
479 .order = try allocator.dupe(u8, desired),
480 .owned = true,
481 };
482 }
483
484 fn useDotGeneral(
485 lhs_order: []const u8,
486 rhs_order: []const u8,
487 lhs_batch: []const i64,
488 rhs_batch: []const i64,
489 lhs_contract: []const i64,
490 rhs_contract: []const i64,
491 result_order: []const u8,
492 ) bool {
493 return lhs_batch.len == 0 and
494 rhs_batch.len == 0 and
495 lhs_contract.len == 1 and
496 rhs_contract.len == 1 and
497 lhs_order.len == 2 and
498 rhs_order.len == 2 and
499 result_order.len == 2;
500 }
501
502 fn multiplyBroadcastReduce(
503 allocator: std.mem.Allocator,
504 sink: anytype,
505 equation: *const Equation,
506 result_type: ir.Type,
507 result_shape: []const i64,
508 lhs: ReducedState,
509 rhs: ReducedState,
510 result_order: []const u8,
511 contracted_order: []const u8,
512 dtype: DType,
513 ) !*ir.Value {
514 var product_order = std.ArrayListUnmanaged(u8).empty;
515 defer product_order.deinit(allocator);
516 try product_order.appendSlice(allocator, result_order);
517 try product_order.appendSlice(allocator, contracted_order);
518
519 const product_shape = try dimsForOrder(allocator, equation, product_order.items);
520 defer allocator.free(product_shape);
521 const product_type = try choir_root.accyTensorType(sink.context(), dtype, product_shape);
522 const lhs_dims = try broadcastDimsForOrder(allocator, product_order.items, lhs.order);
523 defer allocator.free(lhs_dims);
524 const rhs_dims = try broadcastDimsForOrder(allocator, product_order.items, rhs.order);
525 defer allocator.free(rhs_dims);
526 const lhs_value = try sink.broadcastInDim(lhs.value, product_type, product_shape, lhs_dims);
527 const rhs_value = try sink.broadcastInDim(rhs.value, product_type, product_shape, rhs_dims);
528 const product = try sink.mul(lhs_value, rhs_value);
529 return try reduceToOrder(
530 allocator,
531 sink,
532 equation,
533 product,
534 product_order.items,
535 result_order,
536 dtype,
537 result_type,
538 result_shape,
539 );
540 }
541
542 const ReducedState = struct {
543 value: *ir.Value,
544 order: []const u8,
545 owned: bool = false,
546
547 fn deinit(self: ReducedState, allocator: std.mem.Allocator) void {
548 if (self.owned) allocator.free(self.order);
549 }
550 };
551
552 fn reduceOperandLocalSummedIndices(
553 allocator: std.mem.Allocator,
554 sink: anytype,
555 equation: *const Equation,
556 value: *ir.Value,
557 order: []const u8,
558 summed: IndexSet,
559 other: IndexSet,
560 dtype: DType,
561 ) !ReducedState {
562 var reduce_axes = std.ArrayListUnmanaged(i64).empty;
563 defer reduce_axes.deinit(allocator);
564 var kept = std.ArrayListUnmanaged(u8).empty;
565 defer kept.deinit(allocator);
566
567 for (order, 0..) |index, axis| {
568 if (summed.contains(index) and !other.contains(index)) {
569 try reduce_axes.append(allocator, @intCast(axis));
570 } else {
571 try kept.append(allocator, index);
572 }
573 }
574
575 if (reduce_axes.items.len == 0) {
576 return .{
577 .value = value,
578 .order = order,
579 };
580 }
581
582 const owned_order = try kept.toOwnedSlice(allocator);
583 errdefer allocator.free(owned_order);
584 const result_dims = try dimsForOrder(allocator, equation, owned_order);
585 defer allocator.free(result_dims);
586 const result_type = try choir_root.accyTensorType(sink.context(), dtype, result_dims);
587 const reduced = try reduceToOrder(allocator, sink, equation, value, order, owned_order, dtype, result_type, result_dims);
588 return .{
589 .value = reduced,
590 .order = owned_order,
591 .owned = true,
592 };
593 }
594
595 fn reduceToOrder(
596 allocator: std.mem.Allocator,
597 sink: anytype,
598 equation: *const Equation,
599 value: *ir.Value,
600 source: []const u8,
601 target: []const u8,
602 dtype: DType,
603 target_type: ir.Type,
604 target_shape: []const i64,
605 ) !*ir.Value {
606 if (std.mem.eql(u8, source, target)) return value;
607
608 var reduced_order = std.ArrayListUnmanaged(u8).empty;
609 defer reduced_order.deinit(allocator);
610 for (source) |index| {
611 if (!contains(target, index)) try reduced_order.append(allocator, index);
612 }
613 if (reduced_order.items.len == 0) return try transposeToOrder(allocator, sink, equation, value, source, target, dtype);
614
615 var canonical_order = std.ArrayListUnmanaged(u8).empty;
616 defer canonical_order.deinit(allocator);
617 try canonical_order.appendSlice(allocator, target);
618 try canonical_order.appendSlice(allocator, reduced_order.items);
619 const canonical = try transposeToOrder(allocator, sink, equation, value, source, canonical_order.items, dtype);
620
621 const reduced_elements = try elementCountI64(equation, reduced_order.items);
622 const scalar_type = try choir_root.accyTensorType(sink.context(), dtype, &.{});
623 const zero = try zeroConstant(sink, scalar_type, dtype);
624 if (target.len == 0 and reduced_order.items.len == 1 and canonical_order.items.len == 1) {
625 return try sink.reduce(canonical, zero, target_type, "sum", &.{0});
626 }
627 if (target.len == 1 and reduced_order.items.len == 1 and canonical_order.items.len == 2) {
628 return try sink.reduce(canonical, zero, target_type, "sum", &.{1});
629 }
630 if (target.len == 0) {
631 const flat_shape = [_]i64{reduced_elements};
632 const flat_type = try choir_root.accyTensorType(sink.context(), dtype, &flat_shape);
633 const flat = try sink.reshape(canonical, flat_type, &flat_shape);
634 return try sink.reduce(flat, zero, target_type, "sum", &.{0});
635 }
636
637 const kept_elements = try elementCountI64(equation, target);
638 const flat_shape = [_]i64{ kept_elements, reduced_elements };
639 const flat_type = try choir_root.accyTensorType(sink.context(), dtype, &flat_shape);
640 const flat = try sink.reshape(canonical, flat_type, &flat_shape);
641 const reduced_shape = [_]i64{kept_elements};
642 const reduced_type = try choir_root.accyTensorType(sink.context(), dtype, &reduced_shape);
643 const reduced = try sink.reduce(flat, zero, reduced_type, "sum", &.{1});
644 return try sink.reshape(reduced, target_type, target_shape);
645 }
646
647 fn transposeToOrder(
648 allocator: std.mem.Allocator,
649 sink: anytype,
650 equation: *const Equation,
651 value: *ir.Value,
652 source: []const u8,
653 target: []const u8,
654 dtype: DType,
655 ) !*ir.Value {
656 if (std.mem.eql(u8, source, target)) return value;
657 const permutation = try permutationForOrder(allocator, source, target);
658 defer allocator.free(permutation);
659 const dims = try dimsForOrder(allocator, equation, target);
660 defer allocator.free(dims);
661 const result_type = try choir_root.accyTensorType(sink.context(), dtype, dims);
662 return try sink.transpose(value, result_type, permutation);
663 }
664
665 fn zeroConstant(sink: anytype, scalar_type: ir.Type, dtype: DType) !*ir.Value {
666 var bytes = @as([8]u8, @splat(0));
667 return try sink.constant(scalar_type, bytes[0..dtype.sizeOf()]);
668 }
669
670 fn dimsForOrder(allocator: std.mem.Allocator, equation: *const Equation, order: []const u8) ![]i64 {
671 const dims = try allocator.alloc(i64, order.len);
672 errdefer allocator.free(dims);
673 for (order, 0..) |index, axis| {
674 const dim = equation.dimension(index);
675 if (dim > @as(u64, @intCast(std.math.maxInt(i64)))) return error.InvalidDimension;
676 dims[axis] = @intCast(dim);
677 }
678 return dims;
679 }
680
681 fn elementCountI64(equation: *const Equation, order: []const u8) !i64 {
682 const count = equation.elementCount(setFromOrder(order));
683 if (count > @as(u128, @intCast(std.math.maxInt(i64)))) return error.InvalidDimension;
684 return @intCast(count);
685 }
686
687 fn permutationForOrder(allocator: std.mem.Allocator, source: []const u8, target: []const u8) ![]i64 {
688 if (source.len != target.len) return error.InvalidPlan;
689 const permutation = try allocator.alloc(i64, target.len);
690 errdefer allocator.free(permutation);
691 for (target, 0..) |index, axis| {
692 permutation[axis] = try axisOf(source, index);
693 }
694 return permutation;
695 }
696
697 fn broadcastDimsForOrder(allocator: std.mem.Allocator, result: []const u8, source: []const u8) ![]i64 {
698 const dims = try allocator.alloc(i64, source.len);
699 errdefer allocator.free(dims);
700 for (source, 0..) |index, axis| {
701 dims[axis] = try axisOf(result, index);
702 }
703 return dims;
704 }
705
706 fn setFromOrder(order: []const u8) IndexSet {
707 var set: IndexSet = .{};
708 for (order) |index| set.add(index);
709 return set;
710 }
711
712 fn axisOf(order: []const u8, index: u8) !i64 {
713 for (order, 0..) |candidate, axis| {
714 if (candidate == index) return @intCast(axis);
715 }
716 return error.InvalidPlan;
717 }
718
719 fn contains(order: []const u8, index: u8) bool {
720 for (order) |candidate| {
721 if (candidate == index) return true;
722 }
723 return false;
724 }
725
726 fn allInputsMask(input_count: usize) u64 {
727 return (@as(u64, 1) << @intCast(input_count)) - 1;
728 }
729
730 test "einsum lowering emits planned dot_general chain" {
731 const shapes = [_][]const u64{ &.{ 1000, 2 }, &.{ 2, 100 }, &.{ 100, 10 } };
732 var equation = try spec.parse(testing.allocator, "ik,kl,lj->ij", &shapes);
733 defer equation.deinit();
734 var plan = try planner.createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
735 defer plan.deinit();
736
737 var builder = try choir_root.SemanticBuilder.init(testing.allocator, choir_root.SemanticBuilder.ContextLimits.testing);
738 defer builder.deinit();
739 const a_ty = try builder.tensor(.f32, &.{ 1000, 2 });
740 const b_ty = try builder.tensor(.f32, &.{ 2, 100 });
741 const c_ty = try builder.tensor(.f32, &.{ 100, 10 });
742 const out_ty = try builder.tensor(.f32, &.{ 1000, 10 });
743 var fb = try builder.beginFunction("einsum_lower_chain", &.{ a_ty, b_ty, c_ty }, &.{out_ty});
744
745 const value = try lowerPlan(
746 testing.allocator,
747 &fb,
748 &equation,
749 &plan,
750 &.{ fb.parameter(0), fb.parameter(1), fb.parameter(2) },
751 .f32,
752 );
753 try fb.return_(&.{value});
754 try fb.finish();
755
756 const module = try builder.finish();
757 defer module.deinit();
758 try module.verify();
759 try testing.expectEqual(@as(usize, 2), ir.inspection.countOperationsNamed(module.choir_module, choir_root.AccyDialect.DotGeneralOp.operation_name));
760 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, choir_root.AccyDialect.TransposeOp.operation_name));
761 }
762
763 test "einsum lowering transposes final output order" {
764 const shapes = [_][]const u64{ &.{ 2, 3 }, &.{ 3, 5 } };
765 var equation = try spec.parse(testing.allocator, "ab,bc->ca", &shapes);
766 defer equation.deinit();
767 var plan = try planner.createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
768 defer plan.deinit();
769
770 var builder = try choir_root.SemanticBuilder.init(testing.allocator, choir_root.SemanticBuilder.ContextLimits.testing);
771 defer builder.deinit();
772 const a_ty = try builder.tensor(.f32, &.{ 2, 3 });
773 const b_ty = try builder.tensor(.f32, &.{ 3, 5 });
774 const out_ty = try builder.tensor(.f32, &.{ 5, 2 });
775 var fb = try builder.beginFunction("einsum_lower_transpose", &.{ a_ty, b_ty }, &.{out_ty});
776
777 const value = try lowerPlan(testing.allocator, &fb, &equation, &plan, &.{ fb.parameter(0), fb.parameter(1) }, .f32);
778 try fb.return_(&.{value});
779 try fb.finish();
780
781 const module = try builder.finish();
782 defer module.deinit();
783 try module.verify();
784 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, choir_root.AccyDialect.DotGeneralOp.operation_name));
785 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, choir_root.AccyDialect.TransposeOp.operation_name));
786 }
787
788 test "einsum lowering emits single input reduction" {
789 const shapes = [_][]const u64{&.{ 4, 7 }};
790 var equation = try spec.parse(testing.allocator, "ij->i", &shapes);
791 defer equation.deinit();
792 var plan = try planner.createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
793 defer plan.deinit();
794
795 var builder = try choir_root.SemanticBuilder.init(testing.allocator, choir_root.SemanticBuilder.ContextLimits.testing);
796 defer builder.deinit();
797 const input_ty = try builder.tensor(.f32, &.{ 4, 7 });
798 const out_ty = try builder.tensor(.f32, &.{4});
799 var fb = try builder.beginFunction("einsum_lower_reduce", &.{input_ty}, &.{out_ty});
800
801 const value = try lowerPlan(testing.allocator, &fb, &equation, &plan, &.{fb.parameter(0)}, .f32);
802 try fb.return_(&.{value});
803 try fb.finish();
804
805 const module = try builder.finish();
806 defer module.deinit();
807 try module.verify();
808 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, choir_root.AccyDialect.ReduceOp.operation_name));
809 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, choir_root.AccyDialect.DotGeneralOp.operation_name));
810 }
811
812 test "einsum lowering reduces operand-local summed axes before dot" {
813 const shapes = [_][]const u64{ &.{ 2, 3 }, &.{ 5, 7 } };
814 var equation = try spec.parse(testing.allocator, "ab,cd->ac", &shapes);
815 defer equation.deinit();
816 var plan = try planner.createPlan(testing.allocator, &equation, .{ .strategy = .optimal });
817 defer plan.deinit();
818
819 var builder = try choir_root.SemanticBuilder.init(testing.allocator, choir_root.SemanticBuilder.ContextLimits.testing);
820 defer builder.deinit();
821 const lhs_ty = try builder.tensor(.f32, &.{ 2, 3 });
822 const rhs_ty = try builder.tensor(.f32, &.{ 5, 7 });
823 const out_ty = try builder.tensor(.f32, &.{ 2, 5 });
824 var fb = try builder.beginFunction("einsum_lower_local_reductions", &.{ lhs_ty, rhs_ty }, &.{out_ty});
825
826 const value = try lowerPlan(testing.allocator, &fb, &equation, &plan, &.{ fb.parameter(0), fb.parameter(1) }, .f32);
827 try fb.return_(&.{value});
828 try fb.finish();
829
830 const module = try builder.finish();
831 defer module.deinit();
832 try module.verify();
833 try testing.expectEqual(@as(usize, 2), ir.inspection.countOperationsNamed(module.choir_module, choir_root.AccyDialect.ReduceOp.operation_name));
834 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, choir_root.AccyDialect.DotGeneralOp.operation_name));
835 try testing.expectEqual(@as(usize, 2), ir.inspection.countOperationsNamed(module.choir_module, choir_root.AccyDialect.BroadcastInDimOp.operation_name));
836 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, choir_root.AccyDialect.MulOp.operation_name));
837 }
838
839 const BoundSink = enum { function, rewriter };
840
841 fn boundTensorType(
842 builder: *choir_root.SemanticBuilder,
843 dimensions: []const u64,
844 ) !ir.Type {
845 var dims: [62]i64 = undefined;
846 for (dimensions, 0..) |dimension, index| dims[index] = @intCast(dimension);
847 return builder.tensor(.f32, dims[0..dimensions.len]);
848 }
849
850 fn boundLoweredImage(
851 scratch: std.mem.Allocator,
852 equation: *const Equation,
853 plan: *const Plan,
854 mode: BoundSink,
855 ) ![]u8 {
856 var builder = try choir_root.SemanticBuilder.init(
857 testing.allocator,
858 choir_root.SemanticBuilder.ContextLimits.testing,
859 );
860 defer builder.deinit();
861 var types: [63]ir.Type = undefined;
862 for (equation.inputs, 0..) |input, index| {
863 types[index] = try boundTensorType(&builder, input.dims);
864 }
865 const dims = try dimsForOrder(testing.allocator, equation, equation.output);
866 defer testing.allocator.free(dims);
867 const output_type = try builder.tensor(.f32, dims);
868 var fb = try builder.beginFunction(
869 "lower_bound",
870 types[0..equation.inputs.len],
871 &.{output_type},
872 );
873 var inputs: [63]*ir.Value = undefined;
874 for (0..equation.inputs.len) |index| inputs[index] = fb.parameter(index);
875 const operands = inputs[0..equation.inputs.len];
876 const before = fb.ctx.operationCount();
877 var rewriter = rewrite.PatternRewriter.init(testing.allocator, fb.ctx);
878 defer rewriter.deinit();
879 rewriter.setInsertionPoint(fb.entry);
880 const result = switch (mode) {
881 .function => try lowerPlan(scratch, &fb, equation, plan, operands, .f32),
882 .rewriter => try lowerPlanWithRewriter(scratch, &rewriter, equation, plan, operands, .f32),
883 };
884 const bound = try loweringBounds(equation.inputs.len);
885 try testing.expect(fb.ctx.operationCount() - before <= bound.operation_requests);
886 try fb.return_(&.{result});
887 try fb.finish();
888 const module = try builder.finish();
889 defer module.deinit();
890 rewriter.finalize(module.choir_module);
891 try module.verify();
892 return choir.bytecode.encodeModule(testing.allocator, module.choir_module);
893 }
894
895 fn loweringStorageWitness(text: []const u8, shapes: []const []const u64) !void {
896 var equation = try spec.parse(testing.allocator, text, shapes);
897 defer equation.deinit();
898 var plan = try planner.createPlan(testing.allocator, &equation, .{
899 .strategy = .left_to_right,
900 });
901 defer plan.deinit();
902 const expected = try boundLoweredImage(testing.allocator, &equation, &plan, .function);
903 defer testing.allocator.free(expected);
904 const bound = try loweringBounds(shapes.len);
905 const bytes = try testing.allocator.alloc(u8, @intCast(bound.scratch_bytes));
906 defer testing.allocator.free(bytes);
907 for (comptime std.meta.tags(BoundSink)) |mode| {
908 var buffer = std.heap.FixedBufferAllocator.init(bytes);
909 const base = buffer.allocator();
910 const vtable: std.mem.Allocator.VTable = .{
911 .alloc = base.vtable.alloc,
912 .resize = std.mem.Allocator.noResize,
913 .remap = std.mem.Allocator.noRemap,
914 .free = std.mem.Allocator.noFree,
915 };
916 const allocator: std.mem.Allocator = .{ .ptr = base.ptr, .vtable = &vtable };
917 const actual = try boundLoweredImage(allocator, &equation, &plan, mode);
918 defer testing.allocator.free(actual);
919 try testing.expectEqualSlices(u8, expected, actual);
920 try testing.expect(buffer.end_index <= bound.scratch_bytes);
921 }
922 }
923
924 test "einsum lowering bounds preserve both sinks across contractions and reductions" {
925 try loweringStorageWitness("->", &.{&.{}});
926 try loweringStorageWitness("a->a", &.{&.{3}});
927 try loweringStorageWitness("ab->ba", &.{&.{ 2, 3 }});
928 try loweringStorageWitness("abc->ca", &.{&.{ 2, 3, 4 }});
929 try loweringStorageWitness("ab,bc->ca", &.{ &.{ 2, 3 }, &.{ 3, 5 } });
930 try loweringStorageWitness("abc,acd->abd", &.{ &.{ 2, 3, 4 }, &.{ 2, 4, 5 } });
931 try loweringStorageWitness("ab,cd->ac", &.{ &.{ 2, 3 }, &.{ 5, 7 } });
932 try loweringStorageWitness("abcd,defg->ga", &.{ &.{ 2, 3, 4, 5 }, &.{ 5, 6, 7, 8 } });
933 try loweringStorageWitness("ab,bc,cd->da", &.{ &.{ 2, 3 }, &.{ 3, 4 }, &.{ 4, 5 } });
934 }
935
936 test "einsum lowering bounds cover full alphabet and maximum plan arity" {
937 const labels = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
938 const dims: [labels.len]u64 = @splat(1);
939 const shapes: [63][]const u64 = @splat(&dims);
940 var text: std.ArrayList(u8) = .empty;
941 defer text.deinit(testing.allocator);
942 for (shapes, 0..) |_, index| {
943 if (index != 0) try text.append(testing.allocator, ',');
944 try text.appendSlice(testing.allocator, labels);
945 }
946 try text.appendSlice(testing.allocator, "->");
947 var reversed: [labels.len]u8 = undefined;
948 for (labels, 0..) |label, index| reversed[labels.len - index - 1] = label;
949 try text.appendSlice(testing.allocator, &reversed);
950 try loweringStorageWitness(text.items, &shapes);
951 }
952
953 test "einsum lowering bounds reject invalid arity and preserve dimension failure" {
954 try testing.expectError(error.EmptyEquation, loweringBounds(0));
955 try testing.expectError(error.TooManyInputs, loweringBounds(64));
956 try testing.expectError(error.TooManyInputs, loweringBounds(std.math.maxInt(usize)));
957 const shapes = [_][]const u64{ &.{ 4_000_000_000, 4_000_000_000 }, &.{ 2, 3 } };
958 var equation = try spec.parse(testing.allocator, "ab,cd->", &shapes);
959 defer equation.deinit();
960 var plan = try planner.createPlan(testing.allocator, &equation, .{
961 .strategy = .left_to_right,
962 });
963 defer plan.deinit();
964 const bound = try loweringBounds(shapes.len);
965 const bytes = try testing.allocator.alloc(u8, @intCast(bound.scratch_bytes));
966 defer testing.allocator.free(bytes);
967 for (comptime std.meta.tags(BoundSink)) |mode| {
968 var buffer = std.heap.FixedBufferAllocator.init(bytes);
969 try testing.expectError(error.InvalidDimension, boundLoweredImage(
970 buffer.allocator(),
971 &equation,
972 &plan,
973 mode,
974 ));
975 }
976 }
977
978 fn loweringAllocationWitness(allocator: std.mem.Allocator, mode: BoundSink) !void {
979 const shapes = [_][]const u64{ &.{ 2, 3, 4, 5 }, &.{ 5, 6, 7, 8 } };
980 var equation = try spec.parse(testing.allocator, "abcd,defg->ga", &shapes);
981 defer equation.deinit();
982 var plan = try planner.createPlan(testing.allocator, &equation, .{
983 .strategy = .left_to_right,
984 });
985 defer plan.deinit();
986 const image = try boundLoweredImage(allocator, &equation, &plan, mode);
987 defer testing.allocator.free(image);
988 }
989
990 test "einsum lowering releases scratch after every partial allocation" {
991 for (comptime std.meta.tags(BoundSink)) |mode| {
992 try testing.checkAllAllocationFailures(
993 testing.allocator,
994 loweringAllocationWitness,
995 .{mode},
996 );
997 }
998 }