lib/accy/src/choir/einsum/spec.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const testing = std.testing;
3
4 pub const IndexSet = struct {
5 words: [4]u64 = .{ 0, 0, 0, 0 },
6
7 pub fn add(self: *IndexSet, index: u8) void {
8 self.words[wordIndex(index)] |= bitMask(index);
9 }
10
11 pub fn contains(self: IndexSet, index: u8) bool {
12 return (self.words[wordIndex(index)] & bitMask(index)) != 0;
13 }
14
15 pub fn unioned(self: IndexSet, other: IndexSet) IndexSet {
16 return .{ .words = .{
17 self.words[0] | other.words[0],
18 self.words[1] | other.words[1],
19 self.words[2] | other.words[2],
20 self.words[3] | other.words[3],
21 } };
22 }
23
24 pub fn intersected(self: IndexSet, other: IndexSet) IndexSet {
25 return .{ .words = .{
26 self.words[0] & other.words[0],
27 self.words[1] & other.words[1],
28 self.words[2] & other.words[2],
29 self.words[3] & other.words[3],
30 } };
31 }
32
33 pub fn without(self: IndexSet, other: IndexSet) IndexSet {
34 return .{ .words = .{
35 self.words[0] & ~other.words[0],
36 self.words[1] & ~other.words[1],
37 self.words[2] & ~other.words[2],
38 self.words[3] & ~other.words[3],
39 } };
40 }
41
42 pub fn eql(self: IndexSet, other: IndexSet) bool {
43 return std.mem.eql(u64, &self.words, &other.words);
44 }
45
46 pub fn isEmpty(self: IndexSet) bool {
47 return self.words[0] == 0 and self.words[1] == 0 and self.words[2] == 0 and self.words[3] == 0;
48 }
49
50 pub fn count(self: IndexSet) u32 {
51 return @popCount(self.words[0]) + @popCount(self.words[1]) + @popCount(self.words[2]) + @popCount(self.words[3]);
52 }
53
54 fn wordIndex(index: u8) usize {
55 return @intCast(index >> 6);
56 }
57
58 fn bitMask(index: u8) u64 {
59 const shift: u6 = @intCast(index & 63);
60 return @as(u64, 1) << shift;
61 }
62 };
63
64 pub const Input = struct {
65 indices: []const u8,
66 index_set: IndexSet,
67 dims: []const u64,
68 };
69
70 pub const Equation = struct {
71 allocator: std.mem.Allocator,
72 inputs: []Input,
73 output: []const u8,
74 output_set: IndexSet,
75 index_set: IndexSet,
76 dimensions: [256]u64,
77
78 pub fn deinit(self: *Equation) void {
79 for (self.inputs) |input| {
80 self.allocator.free(input.indices);
81 self.allocator.free(input.dims);
82 }
83 self.allocator.free(self.inputs);
84 self.allocator.free(self.output);
85 self.* = undefined;
86 }
87
88 pub fn dimension(self: *const Equation, index: u8) u64 {
89 std.debug.assert(self.index_set.contains(index));
90 return self.dimensions[index];
91 }
92
93 pub fn elementCount(self: *const Equation, indices: IndexSet) u128 {
94 var product: u128 = 1;
95 for (indices.words, 0..) |initial_word, word_index| {
96 var word = initial_word;
97 while (word != 0) {
98 const bit: usize = @intCast(@ctz(word));
99 const index: u8 = @intCast(word_index * 64 + bit);
100 product = saturatingMul(product, self.dimension(index));
101 word &= word - 1;
102 }
103 }
104 return product;
105 }
106
107 pub fn inputSet(self: *const Equation, input_mask: u64) IndexSet {
108 var set: IndexSet = .{};
109 for (self.inputs, 0..) |input, index| {
110 if ((input_mask & (@as(u64, 1) << @intCast(index))) == 0) continue;
111 set = set.unioned(input.index_set);
112 }
113 return set;
114 }
115 };
116
117 pub const ParseError = error{
118 MissingArrow,
119 MultipleArrows,
120 EmptyInputList,
121 InvalidIndex,
122 RepeatedInputIndex,
123 RepeatedOutputIndex,
124 OutputIndexMissing,
125 ArityMismatch,
126 RankMismatch,
127 DimensionMismatch,
128 TooManyInputs,
129 } || std.mem.Allocator.Error;
130
131 /// Two figures for one parse: the bytes the parser asks its allocator for,
132 /// counting list growth it cannot give back, and the structural visits. The
133 /// einsum pass reads these two figures to charge a parse before it runs it. The
134 /// figures leave out stack storage, the internals of the allocator or arena,
135 /// and the borrowed equation text and shapes. The parser reads an input's
136 /// dimensions only after that input's rank matches its labels, which are unique
137 /// letters and digits.
138 pub const ParseBounds = struct {
139 allocation_capacity: u64 = 0,
140 structural_visits: u64 = 0,
141
142 const Error = error{WorkOverflow};
143
144 fn add(a: u64, b: u64) Error!u64 {
145 return std.math.add(u64, a, b) catch return error.WorkOverflow;
146 }
147
148 fn multiply(a: u64, b: u64) Error!u64 {
149 return std.math.mul(u64, a, b) catch return error.WorkOverflow;
150 }
151
152 fn slice(comptime T: type, count: usize) Error!u64 {
153 if (count == 0) return 0;
154 return add(try multiply(@sizeOf(T), count), @alignOf(T) - 1);
155 }
156
157 /// Charges the bytes for one growing list of `count` items of `T` for the
158 /// parse bound: every capacity the list grows through, then one more slice
159 /// for the final `toOwnedSlice` copy that happens when shrinking cannot
160 /// resize or remap in place.
161 fn list(comptime T: type, count: usize) Error!u64 {
162 std.debug.assert(count <= 63);
163 var capacity: usize = 0;
164 var bytes: u64 = 0;
165 while (capacity < count) {
166 capacity = std.ArrayList(T).growCapacity(capacity + 1);
167 bytes = try add(bytes, try slice(T, capacity));
168 }
169 return add(bytes, try slice(T, count));
170 }
171 };
172
173 /// Returns the figures for an equation of `equation_bytes` bytes and
174 /// `input_count` inputs, computed from those two numbers alone with no parsing
175 /// and no allocation, so the einsum pass can call this with the equation length
176 /// and operand count before parsing. The figures cover a successful parse and
177 /// every error the parse can stop at. The parser looks for the arrow and for a
178 /// second arrow before it checks the input limit, so above 63 inputs the
179 /// figures charge only the text scan and no storage. Label lists stop at 62
180 /// distinct letters and digits, so extra whitespace lengthens the scan but adds
181 /// no storage. Structural visits cover copies, checks of dimensions and label
182 /// sets, transfers, and cleanup after a failed prefix. A test parses inside a
183 /// fixed buffer of the charged size that never frees.
184 pub fn parseBounds(equation_bytes: usize, input_count: usize) ParseBounds.Error!ParseBounds {
185 var bounds: ParseBounds = .{
186 .structural_visits = try ParseBounds.add(1, try ParseBounds.multiply(16, equation_bytes)),
187 };
188 if (input_count > 63) return bounds;
189 const labels = @min(equation_bytes, 62);
190 const label_storage = try ParseBounds.list(u8, labels);
191 const input_storage = try ParseBounds.add(label_storage, try ParseBounds.slice(u64, labels));
192 bounds.allocation_capacity = try ParseBounds.add(
193 try ParseBounds.list(Input, input_count),
194 try ParseBounds.add(try ParseBounds.multiply(input_count, input_storage), label_storage),
195 );
196 bounds.structural_visits = try ParseBounds.add(bounds.structural_visits, try ParseBounds.add(
197 try ParseBounds.multiply(4, bounds.allocation_capacity),
198 16 * (input_count + 1),
199 ));
200 return bounds;
201 }
202
203 pub fn parse(allocator: std.mem.Allocator, equation: []const u8, input_shapes: []const []const u64) ParseError!Equation {
204 const arrow = findArrow(equation) orelse return error.MissingArrow;
205 if (findArrow(equation[arrow + 2 ..]) != null) return error.MultipleArrows;
206 if (input_shapes.len > 63) return error.TooManyInputs;
207
208 const lhs = equation[0..arrow];
209 const rhs = equation[arrow + 2 ..];
210 var dimensions: [256]u64 = undefined;
211 var global_indices: IndexSet = .{};
212 var inputs = std.ArrayListUnmanaged(Input).empty;
213 errdefer inputs.deinit(allocator);
214 errdefer freeInputs(allocator, inputs.items);
215
216 var input_index: usize = 0;
217 var term_iter = std.mem.splitScalar(u8, lhs, ',');
218 while (term_iter.next()) |term| {
219 if (input_index >= input_shapes.len) return error.ArityMismatch;
220 const input = try parseInput(allocator, term, input_shapes[input_index], &dimensions, &global_indices);
221 var input_owned = true;
222 errdefer if (input_owned) freeInput(allocator, input);
223 try inputs.append(allocator, input);
224 input_owned = false;
225 input_index += 1;
226 }
227 if (input_index == 0) return error.EmptyInputList;
228 if (input_index != input_shapes.len) return error.ArityMismatch;
229
230 const output = try parseOutput(allocator, rhs, global_indices);
231 errdefer allocator.free(output.indices);
232
233 return .{
234 .allocator = allocator,
235 .inputs = try inputs.toOwnedSlice(allocator),
236 .output = output.indices,
237 .output_set = output.index_set,
238 .index_set = global_indices,
239 .dimensions = dimensions,
240 };
241 }
242
243 fn parseInput(
244 allocator: std.mem.Allocator,
245 term: []const u8,
246 shape: []const u64,
247 dimensions: *[256]u64,
248 global_indices: *IndexSet,
249 ) ParseError!Input {
250 var local_indices = std.ArrayListUnmanaged(u8).empty;
251 errdefer local_indices.deinit(allocator);
252 var local_set: IndexSet = .{};
253 for (term) |index| {
254 if (isSpace(index)) continue;
255 if (!validIndex(index)) return error.InvalidIndex;
256 if (local_set.contains(index)) return error.RepeatedInputIndex;
257 local_set.add(index);
258 try local_indices.append(allocator, index);
259 }
260 if (local_indices.items.len != shape.len) return error.RankMismatch;
261 for (local_indices.items, shape) |index, dim| {
262 if (global_indices.contains(index)) {
263 if (dimensions[index] != dim) return error.DimensionMismatch;
264 } else {
265 global_indices.add(index);
266 dimensions[index] = dim;
267 }
268 }
269 const owned_indices = try local_indices.toOwnedSlice(allocator);
270 errdefer allocator.free(owned_indices);
271 const owned_dims = try allocator.dupe(u64, shape);
272 return .{
273 .indices = owned_indices,
274 .index_set = local_set,
275 .dims = owned_dims,
276 };
277 }
278
279 const OutputParse = struct {
280 indices: []const u8,
281 index_set: IndexSet,
282 };
283
284 fn parseOutput(allocator: std.mem.Allocator, term: []const u8, global_indices: IndexSet) ParseError!OutputParse {
285 var indices = std.ArrayListUnmanaged(u8).empty;
286 errdefer indices.deinit(allocator);
287 var output_set: IndexSet = .{};
288 for (term) |index| {
289 if (isSpace(index)) continue;
290 if (!validIndex(index)) return error.InvalidIndex;
291 if (output_set.contains(index)) return error.RepeatedOutputIndex;
292 if (!global_indices.contains(index)) return error.OutputIndexMissing;
293 output_set.add(index);
294 try indices.append(allocator, index);
295 }
296 return .{
297 .indices = try indices.toOwnedSlice(allocator),
298 .index_set = output_set,
299 };
300 }
301
302 fn freeInputs(allocator: std.mem.Allocator, inputs: []const Input) void {
303 for (inputs) |input| freeInput(allocator, input);
304 }
305
306 fn freeInput(allocator: std.mem.Allocator, input: Input) void {
307 allocator.free(input.indices);
308 allocator.free(input.dims);
309 }
310
311 fn findArrow(equation: []const u8) ?usize {
312 return std.mem.indexOf(u8, equation, "->");
313 }
314
315 fn validIndex(index: u8) bool {
316 return std.ascii.isAlphanumeric(index);
317 }
318
319 fn isSpace(index: u8) bool {
320 return switch (index) {
321 ' ', '\n', '\r', '\t' => true,
322 else => false,
323 };
324 }
325
326 fn saturatingMul(lhs: u128, rhs: u64) u128 {
327 if (lhs == 0 or rhs == 0) return 0;
328 const rhs_wide: u128 = rhs;
329 const max = std.math.maxInt(u128);
330 if (lhs > max / rhs_wide) return max;
331 return lhs * rhs_wide;
332 }
333
334 fn expectSetContains(set: IndexSet, indices: []const u8) !void {
335 for (indices) |index| try testing.expect(set.contains(index));
336 try testing.expectEqual(@as(u32, @intCast(indices.len)), set.count());
337 }
338
339 test "einsum parser validates matrix product equation" {
340 const shapes = [_][]const u64{ &.{ 2, 3 }, &.{ 3, 5 } };
341 var equation = try parse(testing.allocator, "ik, kj -> ij", &shapes);
342 defer equation.deinit();
343
344 try testing.expectEqual(@as(usize, 2), equation.inputs.len);
345 try testing.expectEqualSlices(u8, "ik", equation.inputs[0].indices);
346 try testing.expectEqualSlices(u8, "kj", equation.inputs[1].indices);
347 try testing.expectEqualSlices(u8, "ij", equation.output);
348 try expectSetContains(equation.index_set, "ijk");
349 try testing.expectEqual(@as(u64, 2), equation.dimension('i'));
350 try testing.expectEqual(@as(u64, 3), equation.dimension('k'));
351 try testing.expectEqual(@as(u64, 5), equation.dimension('j'));
352 }
353
354 test "einsum parser accepts scalar operands and scalar outputs" {
355 const dot_shapes = [_][]const u64{ &.{4}, &.{4} };
356 var dot = try parse(testing.allocator, "i,i->", &dot_shapes);
357 defer dot.deinit();
358
359 try testing.expectEqual(@as(usize, 0), dot.output.len);
360 try testing.expect(dot.output_set.isEmpty());
361
362 const scalar_shapes = [_][]const u64{&.{}};
363 var scalar = try parse(testing.allocator, "->", &scalar_shapes);
364 defer scalar.deinit();
365
366 try testing.expectEqual(@as(usize, 1), scalar.inputs.len);
367 try testing.expectEqual(@as(usize, 0), scalar.inputs[0].indices.len);
368 try testing.expectEqual(@as(usize, 0), scalar.output.len);
369 }
370
371 test "einsum parser rejects inconsistent dimensions and missing output indices" {
372 const bad_shapes = [_][]const u64{ &.{ 2, 3 }, &.{ 4, 5 } };
373 try testing.expectError(error.DimensionMismatch, parse(testing.allocator, "ik,kj->ij", &bad_shapes));
374
375 const shapes = [_][]const u64{ &.{ 2, 3 }, &.{ 3, 5 } };
376 try testing.expectError(error.OutputIndexMissing, parse(testing.allocator, "ik,kj->ix", &shapes));
377 }
378
379 test "einsum parser rejects repeated indices in one input" {
380 const shapes = [_][]const u64{&.{ 4, 4 }};
381 try testing.expectError(error.RepeatedInputIndex, parse(testing.allocator, "ii->i", &shapes));
382 }
383
384 fn parseStorageWitness(text: []const u8, shapes: []const []const u64) !void {
385 const bound = try parseBounds(text.len, shapes.len);
386 const bytes = try testing.allocator.alloc(u8, @intCast(bound.allocation_capacity));
387 defer testing.allocator.free(bytes);
388 var buffer = std.heap.FixedBufferAllocator.init(bytes);
389 const base = buffer.allocator();
390 const vtable: std.mem.Allocator.VTable = .{
391 .alloc = base.vtable.alloc,
392 .resize = std.mem.Allocator.noResize,
393 .remap = std.mem.Allocator.noRemap,
394 .free = std.mem.Allocator.noFree,
395 };
396 const allocator: std.mem.Allocator = .{ .ptr = base.ptr, .vtable = &vtable };
397 var expected = parse(testing.allocator, text, shapes) catch |err| {
398 try testing.expectError(err, parse(allocator, text, shapes));
399 try testing.expect(buffer.end_index <= bound.allocation_capacity);
400 return;
401 };
402 defer expected.deinit();
403 var actual = try parse(allocator, text, shapes);
404 defer actual.deinit();
405 try testing.expectEqualDeep(expected.inputs, actual.inputs);
406 try testing.expectEqualSlices(u8, expected.output, actual.output);
407 try testing.expectEqual(expected.index_set, actual.index_set);
408 try testing.expectEqual(expected.output_set, actual.output_set);
409 for (0..256) |label| {
410 if (actual.index_set.contains(@intCast(label))) {
411 try testing.expectEqual(expected.dimensions[label], actual.dimensions[label]);
412 }
413 }
414 try testing.expect(buffer.end_index <= bound.allocation_capacity);
415 }
416
417 test "einsum parser bounds cover unreclaimed growth and full label arity" {
418 const labels = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
419 const dims: [labels.len]u64 = @splat(2);
420 for ([_]usize{ 1, 2, 8, 9, 20, 63 }) |count| {
421 var text: std.ArrayList(u8) = .empty;
422 defer text.deinit(testing.allocator);
423 const shapes = try testing.allocator.alloc([]const u64, count);
424 defer testing.allocator.free(shapes);
425 for (shapes, 0..) |*shape, index| {
426 if (index != 0) try text.append(testing.allocator, ',');
427 try text.appendSlice(testing.allocator, labels);
428 shape.* = &dims;
429 }
430 try text.appendSlice(testing.allocator, "->" ++ labels);
431 try parseStorageWitness(text.items, shapes);
432 }
433 try parseStorageWitness("->", &.{&.{}});
434 try parseStorageWitness(",,->", &.{ &.{}, &.{}, &.{} });
435 try parseStorageWitness("i,i->", &.{ &.{4}, &.{4} });
436 var spaced: [4100]u8 = @splat(' ');
437 @memcpy(spaced[4096..], "a->a");
438 try parseStorageWitness(&spaced, &.{&.{7}});
439 }
440
441 test "einsum parser bounds preserve errors after partial allocations" {
442 const Case = struct { text: []const u8, shapes: []const []const u64, err: ParseError };
443 const cases = [_]Case{
444 .{ .text = "a", .shapes = &.{&.{2}}, .err = error.MissingArrow },
445 .{ .text = "a->a->", .shapes = &.{&.{2}}, .err = error.MultipleArrows },
446 .{ .text = "a,b?->", .shapes = &.{ &.{2}, &.{3} }, .err = error.InvalidIndex },
447 .{ .text = "a,bb->", .shapes = &.{ &.{2}, &.{ 3, 3 } }, .err = error.RepeatedInputIndex },
448 .{ .text = "a->aa", .shapes = &.{&.{2}}, .err = error.RepeatedOutputIndex },
449 .{ .text = "a->ab", .shapes = &.{&.{2}}, .err = error.OutputIndexMissing },
450 .{ .text = "a,b->", .shapes = &.{&.{2}}, .err = error.ArityMismatch },
451 .{ .text = "a->", .shapes = &.{ &.{2}, &.{3} }, .err = error.ArityMismatch },
452 .{ .text = "->", .shapes = &.{}, .err = error.ArityMismatch },
453 .{ .text = "a,b->", .shapes = &.{ &.{2}, &.{ 3, 3 } }, .err = error.RankMismatch },
454 .{ .text = "a,a->", .shapes = &.{ &.{2}, &.{3} }, .err = error.DimensionMismatch },
455 };
456 for (cases) |case| {
457 try testing.expectError(case.err, parse(testing.allocator, case.text, case.shapes));
458 try parseStorageWitness(case.text, case.shapes);
459 }
460 const too_many: [64][]const u64 = @splat(&.{});
461 try testing.expectError(error.TooManyInputs, parse(testing.allocator, "->", &too_many));
462 try parseStorageWitness("->", &too_many);
463 }
464
465 test "einsum parser bounds include rejected text scans and checked arithmetic" {
466 const short = try parseBounds(2, 64);
467 const long = try parseBounds(8192, 64);
468 try testing.expectEqual(@as(u64, 0), long.allocation_capacity);
469 try testing.expect(long.structural_visits >= 8192);
470 try testing.expect(long.structural_visits > short.structural_visits);
471 const labels = try parseBounds(62, 2);
472 const spaced = try parseBounds(8192, 2);
473 try testing.expectEqual(labels.allocation_capacity, spaced.allocation_capacity);
474 try testing.expect(spaced.structural_visits > labels.structural_visits);
475 try testing.expectError(error.WorkOverflow, parseBounds(std.math.maxInt(usize), 1));
476 try testing.expectError(error.WorkOverflow, ParseBounds.add(std.math.maxInt(u64), 1));
477 }
478
479 fn parseAllocationWitness(allocator: std.mem.Allocator) !void {
480 const labels = "abcdefghi";
481 const dims: [labels.len]u64 = @splat(2);
482 const shapes: [9][]const u64 = @splat(&dims);
483 const text = "abcdefghi,abcdefghi,abcdefghi,abcdefghi,abcdefghi," ++
484 "abcdefghi,abcdefghi,abcdefghi,abcdefghi->abcdefghi";
485 var equation = try parse(allocator, text, &shapes);
486 defer equation.deinit();
487 try testing.expectEqualSlices(u8, labels, equation.output);
488 }
489
490 test "einsum parser releases every partially allocated equation" {
491 try testing.checkAllAllocationFailures(testing.allocator, parseAllocationWitness, .{});
492 }