lib/accy/src/preparation/saturation.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3 const choir = @import("choir");
4 const accy_root = @import("../root.zig");
5 const accy_choir = @import("../choir/root.zig");
6 const dialect_mod = accy_choir.dialect;
7 const semantic = accy_choir.semantic;
8
9 const ir = choir.ir;
10 const passes = choir.passes;
11 const work = passes.pass.work;
12 const saturation_options = passes.saturation.OptimizationOptions{
13 .candidate = isSaturationCandidate,
14 };
15
16 pub const saturation_pass_name = "accy-choir-saturate";
17 pub const saturation_pass_description =
18 "Saturate block-local Accy Choir tensor operations with dedup and algebraic identities";
19
20 /// The tensor stage adds this pass to simplify tensor operations with algebraic identities and
21 /// duplicate removal inside each block. The function returns the pass, with its name, description
22 /// and a work contract, a pass's declared name, version and cost estimate, at version 1. Every
23 /// rewrite keeps finite values, the sign of zero, infinities and which values are NaN unchanged.
24 /// The bits inside a NaN and the floating-point status flags are outside what the pass promises.
25 /// Removing a multiplication by an identity matrix applies to integer types alone, because a
26 /// floating multiplication by zero can give NaN and a floating sum can change the sign of zero. A
27 /// backend's precision tier gives no license for looser algebraic rewrites.
28 pub fn tensorSaturationPass() passes.Pass {
29 return .{
30 .name = saturation_pass_name,
31 .description = saturation_pass_description,
32 .run_fn = runTensorSaturationPass,
33 .work_contract = .{
34 .identity = .{ .name = saturation_pass_name, .version = 1 },
35 .estimate = saturationWork,
36 },
37 };
38 }
39
40 /// The compile chain calls this before the pass runs, to charge its cost against the caller's
41 /// limits. Every rule merges a value with one of its existing operands or with an ancestor through
42 /// single-input operations, and removing duplicates keeps each operation's cost. So a cheaper form
43 /// for the earliest value in a class, the set of forms the pass knows to be equal to one value, can
44 /// only be a cycle of single-input identities, and a dot product costs the most and is never the
45 /// cheaper form. Picking the cheapest form may walk that cycle, and it adds no new operation or
46 /// class. The bound comes from the shared elimination bound for this pass's nine rules and its
47 /// iteration limit.
48 fn saturationWork(input: work.Input) !work.Bounds {
49 return passes.saturation.eliminationWorkBound(input, tensor_rules.len, saturation_options.max_iterations);
50 }
51
52 test "saturation accounting covers elimination storage and rejects before mutation" {
53 for ([_]u32{ 0, 1, 8, 32 }) |depth| {
54 try checkSaturationAccounting(depth, null, false);
55 try checkSaturationAccounting(depth, null, true);
56 }
57 try checkSaturationAccounting(8, .workspace, false);
58 try checkSaturationAccounting(8, .rewrites, false);
59 }
60
61 const AccountingRefusal = enum { workspace, rewrites };
62
63 fn checkSaturationAccounting(depth: u32, refusal: ?AccountingRefusal, computed: bool) !void {
64 const allocator = std.testing.allocator;
65 const revision = choir.product.revision;
66 var builder = try semantic.Builder.init(allocator, .testing);
67 defer builder.deinit();
68 const typ = try builder.tensor(.f32, &.{ 2, 3 });
69 var function = try builder.beginFunction("saturation_accounting", &.{typ}, &.{typ});
70 var value = if (computed)
71 try function.add(function.parameter(0), function.parameter(0))
72 else
73 function.parameter(0);
74 for (0..depth) |_| value = try function.reshape(value, typ, &.{ 2, 3 });
75 try function.return_(&.{value});
76 try function.finish();
77 const module = try builder.finish();
78 defer module.deinit();
79 const selected = tensorSaturationPass();
80 const declared = try selected.work_contract.?.estimate(.{ .operation = module.choir_module });
81 var allowance = revision.WorkVector.uniform(std.math.maxInt(u64));
82 if (refusal == .rewrites) allowance.rewrite_attempts = declared.work.rewrite_attempts - 1;
83 const ledger = try revision.AccountingV1.create(allocator, .{
84 .allowance = allowance,
85 .workspace = declared.workspace - @intFromBool(refusal == .workspace),
86 .events = 16,
87 }, &.{.{ .name = saturation_pass_name, .version = 1 }});
88 defer ledger.destroy();
89 var observed = std.testing.FailingAllocator.init(allocator, .{});
90 var cache = try passes.AnalysisCache.initAccounted(observed.allocator(), null, ledger, .{}, 8);
91 defer cache.deinit();
92 var manager = passes.PassManager.init(observed.allocator());
93 defer manager.deinit();
94 try manager.addPass(selected);
95 const before = observed.allocated_bytes;
96 const result = manager.runWithAnalysisCache(module.choir_module, module.context(), &cache, .{ .max_threads = 1 });
97 const executed = observed.allocated_bytes - before;
98 const receipt = ledger.view();
99 const remaining = ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ReshapeOp.operation_name);
100 if (refusal != null) {
101 try std.testing.expectEqual(.failure, result);
102 try std.testing.expectEqual(.exhausted, receipt.outcome);
103 try std.testing.expectEqual(0, receipt.executed.counters.pass_runs);
104 try std.testing.expectEqual(depth, remaining);
105 try std.testing.expectError(error.WorkExhausted, ledger.producersComplete());
106 } else {
107 try std.testing.expectEqual(.success, result);
108 try ledger.producersComplete();
109 try std.testing.expectEqual(1, receipt.executed.counters.pass_runs);
110 try std.testing.expectEqual(0, remaining);
111 try std.testing.expect(executed <= declared.workspace);
112 try std.testing.expect(executed <= receipt.charged.allocation_capacity);
113 }
114 try module.verify();
115 }
116
117 pub fn saturateTensorOperations(
118 allocator: std.mem.Allocator,
119 choir_module: *ir.Operation,
120 ctx: *ir.Context,
121 ) !passes.saturation.OptimizationStats {
122 var rules = choir.egraph.RewriteSet.init(allocator);
123 defer rules.deinit();
124 try populateTensorRules(&rules);
125
126 const stats = try passes.runEGraphOptimization(allocator, ctx, choir_module, &rules, saturation_options);
127 return stats;
128 }
129
130 pub fn populateTensorRules(rules: *choir.egraph.RewriteSet) !void {
131 rules.setCostModel(tensorCostModel());
132 for (tensor_rules) |rule| {
133 try rules.add(.{
134 .name = rule.name,
135 .benefit = rule.benefit,
136 .apply = rule.apply,
137 });
138 }
139 }
140
141 fn runTensorSaturationPass(pass_ctx: *passes.PassContext) passes.PassResult {
142 const stats = saturateTensorOperations(
143 pass_ctx.allocator,
144 pass_ctx.op,
145 pass_ctx.ir_ctx,
146 ) catch return .failure;
147 if (stats.replacements == 0) {
148 pass_ctx.preserveAllAnalyses();
149 } else {
150 pass_ctx.markModified();
151 }
152 if (!passes.saturation.observeOptimization(pass_ctx, stats)) return .failure;
153 return .success;
154 }
155
156 fn isSaturationCandidate(_: ?*anyopaque, op: *ir.Operation) anyerror!bool {
157 if (op.getNumResults() != 1) return false;
158 return std.mem.eql(u8, op.name.getDialectNamespace(), "accy");
159 }
160
161 const TensorRule = struct {
162 name: []const u8,
163 benefit: u32,
164 apply: choir.egraph.RewriteFn,
165 };
166
167 const tensor_rules = [_]TensorRule{
168 .{ .name = "accy-transpose-identity", .benefit = 20, .apply = transposeIdentity },
169 .{ .name = "accy-reshape-identity", .benefit = 20, .apply = reshapeIdentity },
170 .{ .name = "accy-broadcast-identity", .benefit = 20, .apply = broadcastIdentity },
171 .{ .name = "accy-broadcast-in-dim-identity", .benefit = 20, .apply = broadcastInDimIdentity },
172 .{ .name = "accy-convert-identity", .benefit = 15, .apply = convertIdentity },
173 .{ .name = "accy-dot-right-identity", .benefit = 14, .apply = dotRightIdentity },
174 .{ .name = "accy-dot-left-identity", .benefit = 14, .apply = dotLeftIdentity },
175 .{ .name = "accy-convert-round-trip", .benefit = 12, .apply = convertRoundTrip },
176 .{ .name = "accy-transpose-involution", .benefit = 10, .apply = transposeInvolution },
177 };
178
179 fn tensorCostModel() choir.egraph.CostModel {
180 const AccyDialect = dialect_mod.AccyDialect;
181 return .{
182 .operation = 8,
183 .constant = 1,
184 .constant_op_name = AccyDialect.ConstantOp.operation_name,
185 .overrides = &.{
186 .{ .name = AccyDialect.ReshapeOp.operation_name, .cost = 1 },
187 .{ .name = AccyDialect.BroadcastOp.operation_name, .cost = 1 },
188 .{ .name = AccyDialect.BroadcastInDimOp.operation_name, .cost = 1 },
189 .{ .name = AccyDialect.TransposeOp.operation_name, .cost = 1 },
190 .{ .name = AccyDialect.ConvertOp.operation_name, .cost = 2 },
191 .{ .name = AccyDialect.ReduceOp.operation_name, .cost = 12 },
192 .{ .name = AccyDialect.DotGeneralOp.operation_name, .cost = 32 },
193 },
194 };
195 }
196
197 fn transposeIdentity(
198 ctx: *choir.egraph.RewriteContext,
199 class: choir.egraph.ClassId,
200 entry_node: *const choir.egraph.Node,
201 ) anyerror!bool {
202 const transpose_name = dialect_mod.AccyDialect.TransposeOp.operation_name;
203 if (!isOperation(entry_node, transpose_name)) return false;
204 const permutation = i64ListAttr(entry_node, "permutation") orelse return false;
205 if (entry_node.result_types.len != 1) return false;
206 if ((tensorRank(entry_node.result_types[0]) orelse return false) != permutation.len) return false;
207 if (!isIdentityIndexList(permutation)) return false;
208 return try mergeUnaryWithOperand(ctx, class, entry_node);
209 }
210
211 fn reshapeIdentity(
212 ctx: *choir.egraph.RewriteContext,
213 class: choir.egraph.ClassId,
214 entry_node: *const choir.egraph.Node,
215 ) anyerror!bool {
216 const reshape_name = dialect_mod.AccyDialect.ReshapeOp.operation_name;
217 if (!isOperation(entry_node, reshape_name)) return false;
218 const new_shape = i64ListAttr(entry_node, "new_shape") orelse return false;
219 if (!resultTypeShapeMatches(entry_node, new_shape)) return false;
220 return try mergeUnaryWithOperand(ctx, class, entry_node);
221 }
222
223 fn broadcastIdentity(
224 ctx: *choir.egraph.RewriteContext,
225 class: choir.egraph.ClassId,
226 entry_node: *const choir.egraph.Node,
227 ) anyerror!bool {
228 const broadcast_name = dialect_mod.AccyDialect.BroadcastOp.operation_name;
229 if (!isOperation(entry_node, broadcast_name)) return false;
230 const sizes = i64ListAttr(entry_node, "sizes") orelse return false;
231 if (sizes.len != 0) return false;
232 return try mergeUnaryWithOperand(ctx, class, entry_node);
233 }
234
235 fn broadcastInDimIdentity(
236 ctx: *choir.egraph.RewriteContext,
237 class: choir.egraph.ClassId,
238 entry_node: *const choir.egraph.Node,
239 ) anyerror!bool {
240 const broadcast_name = dialect_mod.AccyDialect.BroadcastInDimOp.operation_name;
241 if (!isOperation(entry_node, broadcast_name)) return false;
242 const dims = i64ListAttr(entry_node, "broadcast_dims") orelse return false;
243 const result_shape = i64ListAttr(entry_node, "result_shape") orelse return false;
244 if (dims.len != result_shape.len) return false;
245 if (!resultTypeShapeMatches(entry_node, result_shape)) return false;
246 if (!isIdentityIndexList(dims)) return false;
247 return try mergeUnaryWithOperand(ctx, class, entry_node);
248 }
249
250 fn convertIdentity(
251 ctx: *choir.egraph.RewriteContext,
252 class: choir.egraph.ClassId,
253 entry_node: *const choir.egraph.Node,
254 ) anyerror!bool {
255 const convert_name = dialect_mod.AccyDialect.ConvertOp.operation_name;
256 if (!isOperation(entry_node, convert_name)) return false;
257 _ = entry_node.getAttr("convert_to") orelse return false;
258 return try mergeUnaryWithOperand(ctx, class, entry_node);
259 }
260
261 fn convertRoundTrip(
262 ctx: *choir.egraph.RewriteContext,
263 class: choir.egraph.ClassId,
264 entry_node: *const choir.egraph.Node,
265 ) anyerror!bool {
266 const convert_name = dialect_mod.AccyDialect.ConvertOp.operation_name;
267 if (!isOperation(entry_node, convert_name)) return false;
268 if (entry_node.operands.len != 1) return false;
269 if (entry_node.result_types.len != 1) return false;
270 const source_dtype = tensorDType(entry_node.result_types[0]) orelse return false;
271 for (ctx.nodes(entry_node.operands[0])) |*inner_node| {
272 if (!isOperation(inner_node, convert_name)) continue;
273 if (inner_node.operands.len != 1) continue;
274 if (inner_node.result_types.len != 1) continue;
275 const intermediate_dtype = tensorDType(inner_node.result_types[0]) orelse continue;
276 if (!losslessRoundTripVia(source_dtype, intermediate_dtype)) continue;
277 if (!classHasType(ctx, inner_node.operands[0], entry_node.result_types[0])) continue;
278 return try ctx.merge(class, inner_node.operands[0]);
279 }
280 return false;
281 }
282
283 fn dotRightIdentity(
284 ctx: *choir.egraph.RewriteContext,
285 class: choir.egraph.ClassId,
286 entry_node: *const choir.egraph.Node,
287 ) anyerror!bool {
288 if (!standardMatrixProductDot(entry_node)) return false;
289 if (entry_node.operands.len != 2) return false;
290 if (entry_node.result_types.len != 1) return false;
291 const result_shape = tensorMatrixShape(entry_node.result_types[0]) orelse return false;
292 const dtype = tensorDType(entry_node.result_types[0]) orelse return false;
293 if (!classHasType(ctx, entry_node.operands[0], entry_node.result_types[0])) return false;
294 for (ctx.nodes(entry_node.operands[1])) |*rhs_node| {
295 if (!isIdentityMatrixConstant(rhs_node, dtype, result_shape[1])) continue;
296 return try ctx.merge(class, entry_node.operands[0]);
297 }
298 return false;
299 }
300
301 fn dotLeftIdentity(
302 ctx: *choir.egraph.RewriteContext,
303 class: choir.egraph.ClassId,
304 entry_node: *const choir.egraph.Node,
305 ) anyerror!bool {
306 if (!standardMatrixProductDot(entry_node)) return false;
307 if (entry_node.operands.len != 2) return false;
308 if (entry_node.result_types.len != 1) return false;
309 const result_shape = tensorMatrixShape(entry_node.result_types[0]) orelse return false;
310 const dtype = tensorDType(entry_node.result_types[0]) orelse return false;
311 if (!classHasType(ctx, entry_node.operands[1], entry_node.result_types[0])) return false;
312 for (ctx.nodes(entry_node.operands[0])) |*lhs_node| {
313 if (!isIdentityMatrixConstant(lhs_node, dtype, result_shape[0])) continue;
314 return try ctx.merge(class, entry_node.operands[1]);
315 }
316 return false;
317 }
318
319 fn transposeInvolution(
320 ctx: *choir.egraph.RewriteContext,
321 class: choir.egraph.ClassId,
322 entry_node: *const choir.egraph.Node,
323 ) anyerror!bool {
324 const transpose_name = dialect_mod.AccyDialect.TransposeOp.operation_name;
325 if (!isOperation(entry_node, transpose_name)) return false;
326 if (entry_node.operands.len != 1) return false;
327 if (entry_node.result_types.len != 1) return false;
328 const outer = i64ListAttr(entry_node, "permutation") orelse return false;
329
330 for (ctx.nodes(entry_node.operands[0])) |*inner_node| {
331 if (!isOperation(inner_node, transpose_name)) continue;
332 if (inner_node.operands.len != 1) continue;
333 const inner = i64ListAttr(inner_node, "permutation") orelse continue;
334 if (!composesToIdentity(inner, outer)) continue;
335 if (!classHasType(ctx, inner_node.operands[0], entry_node.result_types[0])) continue;
336 return try ctx.merge(class, inner_node.operands[0]);
337 }
338 return false;
339 }
340
341 fn mergeUnaryWithOperand(
342 ctx: *choir.egraph.RewriteContext,
343 class: choir.egraph.ClassId,
344 entry_node: *const choir.egraph.Node,
345 ) anyerror!bool {
346 if (entry_node.operands.len != 1) return false;
347 if (entry_node.result_types.len != 1) return false;
348 if (!classHasType(ctx, entry_node.operands[0], entry_node.result_types[0])) return false;
349 return try ctx.merge(class, entry_node.operands[0]);
350 }
351
352 fn isOperation(node: *const choir.egraph.Node, name: []const u8) bool {
353 return node.kind == .operation and std.mem.eql(u8, node.op_name, name);
354 }
355
356 fn i64ListAttr(node: *const choir.egraph.Node, name: []const u8) ?[]align(1) const i64 {
357 const attr = node.getAttr(name) orelse return null;
358 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return null;
359 if (dialect_attr.payload.len % @sizeOf(i64) != 0) return null;
360 return std.mem.bytesAsSlice(i64, dialect_attr.payload);
361 }
362
363 fn isIdentityIndexList(indices: []align(1) const i64) bool {
364 for (indices, 0..) |index, position| {
365 if (index != @as(i64, @intCast(position))) return false;
366 }
367 return true;
368 }
369
370 fn resultTypeShapeMatches(entry_node: *const choir.egraph.Node, shape: []align(1) const i64) bool {
371 if (entry_node.result_types.len != 1) return false;
372 return tensorShapeMatches(entry_node.result_types[0], shape);
373 }
374
375 fn tensorShapeMatches(typ: ir.Type, expected: []align(1) const i64) bool {
376 const type_name = typ.getDialectTypeName() orelse return false;
377 if (!std.mem.eql(u8, type_name, dialect_mod.tensor_type_name)) return false;
378 const dims_text = tensorDimsText(typ) orelse return false;
379 if (expected.len == 0) return dims_text.len == 0;
380 var iter = std.mem.splitScalar(u8, dims_text, 'x');
381 var index: usize = 0;
382 while (iter.next()) |part| {
383 if (part.len == 0) return false;
384 if (index >= expected.len) return false;
385 const dim = std.fmt.parseInt(i64, part, 10) catch return false;
386 if (dim != expected[index]) return false;
387 index += 1;
388 }
389 return index == expected.len;
390 }
391
392 fn tensorRank(typ: ir.Type) ?usize {
393 const type_name = typ.getDialectTypeName() orelse return null;
394 if (!std.mem.eql(u8, type_name, dialect_mod.tensor_type_name)) return null;
395 const dims_text = tensorDimsText(typ) orelse return null;
396 if (dims_text.len == 0) return 0;
397 var iter = std.mem.splitScalar(u8, dims_text, 'x');
398 var rank: usize = 0;
399 while (iter.next()) |part| {
400 if (part.len == 0) return null;
401 _ = std.fmt.parseInt(i64, part, 10) catch return null;
402 rank += 1;
403 }
404 return rank;
405 }
406
407 fn tensorDimsText(typ: ir.Type) ?[]const u8 {
408 const key = typ.getDialectParamKey() orelse return null;
409 const comma = std.mem.indexOfScalar(u8, key, ',') orelse return null;
410 return key[comma + 1 ..];
411 }
412
413 fn tensorDType(typ: ir.Type) ?choir_abi.DType {
414 const type_name = typ.getDialectTypeName() orelse return null;
415 if (!std.mem.eql(u8, type_name, dialect_mod.tensor_type_name)) return null;
416 const key = typ.getDialectParamKey() orelse return null;
417 const comma = std.mem.indexOfScalar(u8, key, ',') orelse return null;
418 return choir_abi.DType.fromName(key[0..comma]);
419 }
420
421 fn losslessRoundTripVia(source: choir_abi.DType, intermediate: choir_abi.DType) bool {
422 if (source == intermediate) return true;
423 if (source.isFloat() or intermediate.isFloat()) return losslessFloatRoundTripVia(source, intermediate);
424 if (source.isSignedInt()) {
425 return intermediate.isSignedInt() and intBits(intermediate) >= intBits(source);
426 }
427 if (source.isUnsignedInt()) {
428 if (intermediate.isUnsignedInt()) return intBits(intermediate) >= intBits(source);
429 if (intermediate.isSignedInt()) return intBits(intermediate) > intBits(source);
430 }
431 return false;
432 }
433
434 fn losslessFloatRoundTripVia(source: choir_abi.DType, intermediate: choir_abi.DType) bool {
435 return switch (source) {
436 .f16, .bf16 => intermediate == .f32 or intermediate == .f64,
437 .f32 => intermediate == .f64,
438 else => false,
439 };
440 }
441
442 fn intBits(dtype: choir_abi.DType) u16 {
443 return switch (dtype) {
444 .i1 => 1,
445 .i8, .u8 => 8,
446 .i16, .u16 => 16,
447 .i32, .u32 => 32,
448 .i64, .u64 => 64,
449 else => 0,
450 };
451 }
452
453 fn standardMatrixProductDot(node: *const choir.egraph.Node) bool {
454 const dot_name = dialect_mod.AccyDialect.DotGeneralOp.operation_name;
455 if (!isOperation(node, dot_name)) return false;
456 if (node.operands.len != 2) return false;
457 const lhs_batch = i64ListAttr(node, "lhs_batch") orelse return false;
458 const rhs_batch = i64ListAttr(node, "rhs_batch") orelse return false;
459 const lhs_contract = i64ListAttr(node, "lhs_contract") orelse return false;
460 const rhs_contract = i64ListAttr(node, "rhs_contract") orelse return false;
461 return lhs_batch.len == 0 and
462 rhs_batch.len == 0 and
463 i64ListEquals(lhs_contract, &[_]i64{1}) and
464 i64ListEquals(rhs_contract, &[_]i64{0});
465 }
466
467 fn i64ListEquals(lhs: []align(1) const i64, rhs: []const i64) bool {
468 if (lhs.len != rhs.len) return false;
469 for (lhs, rhs) |left, right| {
470 if (left != right) return false;
471 }
472 return true;
473 }
474
475 fn tensorMatrixShape(typ: ir.Type) ?[2]usize {
476 const type_name = typ.getDialectTypeName() orelse return null;
477 if (!std.mem.eql(u8, type_name, dialect_mod.tensor_type_name)) return null;
478 const dims_text = tensorDimsText(typ) orelse return null;
479 var iter = std.mem.splitScalar(u8, dims_text, 'x');
480 const rows_text = iter.next() orelse return null;
481 const cols_text = iter.next() orelse return null;
482 if (iter.next() != null) return null;
483 if (rows_text.len == 0 or cols_text.len == 0) return null;
484 const rows = std.fmt.parseInt(i64, rows_text, 10) catch return null;
485 const cols = std.fmt.parseInt(i64, cols_text, 10) catch return null;
486 if (rows <= 0 or cols <= 0) return null;
487 return .{ @intCast(rows), @intCast(cols) };
488 }
489
490 fn isIdentityMatrixConstant(node: *const choir.egraph.Node, dtype: choir_abi.DType, n: usize) bool {
491 const constant_name = dialect_mod.AccyDialect.ConstantOp.operation_name;
492 if (!isOperation(node, constant_name)) return false;
493 if (n == 0) return false;
494 if (node.result_types.len != 1) return false;
495 const node_dtype = tensorDType(node.result_types[0]) orelse return false;
496 if (node_dtype != dtype) return false;
497 const shape = tensorMatrixShape(node.result_types[0]) orelse return false;
498 if (shape[0] != n or shape[1] != n) return false;
499 const payload = constantPayload(node) orelse return false;
500 return identityPayloadMatchesDType(dtype, payload, n);
501 }
502
503 fn constantPayload(node: *const choir.egraph.Node) ?[]const u8 {
504 const attr = node.getAttr("payload") orelse return null;
505 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return null;
506 return dialect_attr.payload;
507 }
508
509 fn identityPayloadMatchesDType(dtype: choir_abi.DType, payload: []const u8, n: usize) bool {
510 return switch (dtype) {
511 .i1 => false,
512 .i8 => identityPayloadMatches(i8, payload, n, 0, 1),
513 .i16 => identityPayloadMatches(i16, payload, n, 0, 1),
514 .i32 => identityPayloadMatches(i32, payload, n, 0, 1),
515 .i64 => identityPayloadMatches(i64, payload, n, 0, 1),
516 .u8 => identityPayloadMatches(u8, payload, n, 0, 1),
517 .u16 => identityPayloadMatches(u16, payload, n, 0, 1),
518 .u32 => identityPayloadMatches(u32, payload, n, 0, 1),
519 .u64 => identityPayloadMatches(u64, payload, n, 0, 1),
520 .f16, .bf16, .f32, .f64, .key => false,
521 };
522 }
523
524 fn identityPayloadMatches(comptime T: type, payload: []const u8, n: usize, zero: T, one: T) bool {
525 const value_count = std.math.mul(usize, n, n) catch return false;
526 const expected_len = std.math.mul(usize, value_count, @sizeOf(T)) catch return false;
527 if (payload.len != expected_len) return false;
528 const values = std.mem.bytesAsSlice(T, payload);
529 for (0..n) |row| {
530 for (0..n) |col| {
531 const expected = if (row == col) one else zero;
532 if (values[row * n + col] != expected) return false;
533 }
534 }
535 return true;
536 }
537
538 fn composesToIdentity(inner: []align(1) const i64, outer: []align(1) const i64) bool {
539 if (inner.len != outer.len) return false;
540 for (outer, 0..) |outer_index, position| {
541 if (outer_index < 0) return false;
542 const index: usize = @intCast(outer_index);
543 if (index >= inner.len) return false;
544 if (inner[index] != @as(i64, @intCast(position))) return false;
545 }
546 return true;
547 }
548
549 fn classHasType(ctx: *choir.egraph.RewriteContext, class: choir.egraph.ClassId, expected: ir.Type) bool {
550 for (ctx.nodes(class)) |*candidate| {
551 switch (candidate.kind) {
552 .value => if (candidate.value.?.type.eql(expected)) return true,
553 .operation => {
554 if (candidate.result_types.len != 1) continue;
555 if (candidate.result_types[0].eql(expected)) return true;
556 },
557 }
558 }
559 return false;
560 }
561
562 const testing = std.testing;
563
564 fn readSymbolName(func: *ir.Operation) ?[]const u8 {
565 return ir.SymbolTable.getSymbolName(func);
566 }
567
568 fn findOpNamedInBlock(block: *ir.Block, name: []const u8) ?*ir.Operation {
569 var iter = block.operations.head;
570 while (iter) |op_ptr| {
571 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
572 if (std.mem.eql(u8, op.name.name, name)) return op;
573 iter = op.next_op;
574 }
575 return null;
576 }
577
578 fn writeIdentityI32(values: []i32, n: usize) void {
579 for (0..n) |row| {
580 for (0..n) |col| {
581 values[row * n + col] = if (row == col) 1 else 0;
582 }
583 }
584 }
585
586 test "tensor saturation rule table publishes cost model" {
587 const allocator = testing.allocator;
588 const AccyDialect = dialect_mod.AccyDialect;
589
590 var rules = choir.egraph.RewriteSet.init(allocator);
591 defer rules.deinit();
592 try populateTensorRules(&rules);
593 const model = rules.resolvedCostModel();
594
595 try testing.expectEqual(tensor_rules.len, rules.rules.items.len);
596 try testing.expectEqual(@as(u32, 8), model.operation);
597 try testing.expectEqual(@as(u32, 1), model.operationCost(AccyDialect.ConstantOp.operation_name));
598 try testing.expectEqual(@as(u32, 1), model.operationCost(AccyDialect.ReshapeOp.operation_name));
599 try testing.expectEqual(@as(u32, 1), model.operationCost(AccyDialect.BroadcastInDimOp.operation_name));
600 try testing.expectEqual(@as(u32, 2), model.operationCost(AccyDialect.ConvertOp.operation_name));
601 try testing.expectEqual(@as(u32, 32), model.operationCost(AccyDialect.DotGeneralOp.operation_name));
602 }
603
604 test "saturation eliminates duplicate accy add in one block" {
605 const allocator = testing.allocator;
606
607 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
608 defer builder.deinit();
609 const f32_4 = try builder.tensor(.f32, &.{4});
610 var fb = try builder.beginFunction("cse_add4", &.{ f32_4, f32_4 }, &.{ f32_4, f32_4 });
611 const first = try fb.add(fb.parameter(0), fb.parameter(1));
612 const second = try fb.add(fb.parameter(0), fb.parameter(1));
613 try fb.return_(&.{ first, second });
614 try fb.finish();
615 const module = try builder.finish();
616 defer module.deinit();
617
618 const choir_mod = module.choir_module;
619 const ctx = module.context();
620 var pm = passes.PassManager.init(allocator);
621 defer pm.deinit();
622 try pm.addPass(tensorSaturationPass());
623
624 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
625 try testing.expectEqual(@as(u64, 1), pm.stats.pass_runs);
626 try testing.expectEqual(@as(u64, 1), pm.stats.passes_modified);
627 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.AddOp.operation_name));
628
629 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
630 const func = ir.inspection.functionByNameInBlock(module_body, "cse_add4") orelse return error.TestExpectedFunc;
631 const entry = func.getRegion(0).?.getEntryBlock().?;
632 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
633 try testing.expect(ret.getOperand(0).? == ret.getOperand(1).?);
634 }
635
636 test "saturation keeps constants with different payload bytes distinct" {
637 const allocator = testing.allocator;
638
639 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
640 defer builder.deinit();
641 const f32_1 = try builder.tensor(.f32, &.{1});
642 var one: [4]u8 = undefined;
643 var two: [4]u8 = undefined;
644 @as(*align(1) f32, @ptrCast(&one[0])).* = 1.0;
645 @as(*align(1) f32, @ptrCast(&two[0])).* = 2.0;
646
647 var fb = try builder.beginFunction("cse_distinct_constants", &.{}, &.{ f32_1, f32_1 });
648 const c1 = try fb.constant(f32_1, &one);
649 const c2 = try fb.constant(f32_1, &two);
650 try fb.return_(&.{ c1, c2 });
651 try fb.finish();
652 const module = try builder.finish();
653 defer module.deinit();
654
655 const choir_mod = module.choir_module;
656 const ctx = module.context();
657 var pm = passes.PassManager.init(allocator);
658 defer pm.deinit();
659 try pm.addPass(tensorSaturationPass());
660
661 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
662 try testing.expectEqual(@as(u64, 1), pm.stats.pass_runs);
663 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
664 try testing.expectEqual(@as(usize, 2), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.ConstantOp.operation_name));
665 }
666
667 test "saturation pass preserves analyses when no duplicate accy ops exist" {
668 const allocator = testing.allocator;
669
670 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
671 defer builder.deinit();
672 const f32_4 = try builder.tensor(.f32, &.{4});
673 var fb = try builder.beginFunction("cse_noop_add4", &.{ f32_4, f32_4 }, &.{f32_4});
674 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
675 try fb.return_(&.{sum});
676 try fb.finish();
677 const module = try builder.finish();
678 defer module.deinit();
679
680 const choir_mod = module.choir_module;
681 const ctx = module.context();
682 var pm = passes.PassManager.init(allocator);
683 defer pm.deinit();
684 try pm.addPass(tensorSaturationPass());
685
686 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
687 try testing.expectEqual(@as(u64, 1), pm.stats.pass_runs);
688 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
689 }
690
691 test "saturation collapses identity tensor shape operations" {
692 const allocator = testing.allocator;
693
694 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
695 defer builder.deinit();
696 const f32_2x3 = try builder.tensor(.f32, &.{ 2, 3 });
697 var fb = try builder.beginFunction("shape_identity_ops", &.{f32_2x3}, &.{ f32_2x3, f32_2x3, f32_2x3, f32_2x3 });
698 const input = fb.parameter(0);
699 const transposed = try fb.transpose(input, f32_2x3, &.{ 0, 1 });
700 const reshaped = try fb.reshape(input, f32_2x3, &.{ 2, 3 });
701 const broadcasted = try fb.broadcast(input, f32_2x3, &.{});
702 const broadcast_in_dim = try fb.broadcastInDim(input, f32_2x3, &.{ 2, 3 }, &.{ 0, 1 });
703 try fb.return_(&.{ transposed, reshaped, broadcasted, broadcast_in_dim });
704 try fb.finish();
705 const module = try builder.finish();
706 defer module.deinit();
707
708 const choir_mod = module.choir_module;
709 const ctx = module.context();
710 var pm = passes.PassManager.init(allocator);
711 defer pm.deinit();
712 try pm.addPass(tensorSaturationPass());
713
714 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
715 try testing.expectEqual(@as(u64, 1), pm.stats.passes_modified);
716
717 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
718 const func = ir.inspection.functionByNameInBlock(module_body, "shape_identity_ops") orelse return error.TestExpectedFunc;
719 const entry = func.getRegion(0).?.getEntryBlock().?;
720 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
721 for (0..4) |index| {
722 try testing.expect(ret.getOperand(index).? == entry.getArgument(0).?);
723 }
724 try ir.verifyOperation(choir_mod, ir.verify.default_options);
725 }
726
727 test "saturation collapses identity convert" {
728 const allocator = testing.allocator;
729
730 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
731 defer builder.deinit();
732 const f32_4 = try builder.tensor(.f32, &.{4});
733 var fb = try builder.beginFunction("convert_identity", &.{f32_4}, &.{f32_4});
734 const converted = try fb.convert(fb.parameter(0), f32_4, .f32);
735 try fb.return_(&.{converted});
736 try fb.finish();
737 const module = try builder.finish();
738 defer module.deinit();
739
740 const choir_mod = module.choir_module;
741 const ctx = module.context();
742 var pm = passes.PassManager.init(allocator);
743 defer pm.deinit();
744 try pm.addPass(tensorSaturationPass());
745
746 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
747 try testing.expectEqual(@as(u64, 1), pm.stats.passes_modified);
748
749 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
750 const func = ir.inspection.functionByNameInBlock(module_body, "convert_identity") orelse return error.TestExpectedFunc;
751 const entry = func.getRegion(0).?.getEntryBlock().?;
752 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
753 try testing.expect(ret.getOperand(0).? == entry.getArgument(0).?);
754 try ir.verifyOperation(choir_mod, ir.verify.default_options);
755 }
756
757 test "saturation collapses lossless convert round trip" {
758 const allocator = testing.allocator;
759
760 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
761 defer builder.deinit();
762 const f32_4 = try builder.tensor(.f32, &.{4});
763 const f64_4 = try builder.tensor(.f64, &.{4});
764 var fb = try builder.beginFunction("convert_lossless_round_trip", &.{f32_4}, &.{f32_4});
765 const widened = try fb.convert(fb.parameter(0), f64_4, .f64);
766 const restored = try fb.convert(widened, f32_4, .f32);
767 try fb.return_(&.{restored});
768 try fb.finish();
769 const module = try builder.finish();
770 defer module.deinit();
771
772 const choir_mod = module.choir_module;
773 const ctx = module.context();
774 var pm = passes.PassManager.init(allocator);
775 defer pm.deinit();
776 try pm.addPass(tensorSaturationPass());
777
778 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
779 try testing.expectEqual(@as(u64, 1), pm.stats.passes_modified);
780
781 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
782 const func = ir.inspection.functionByNameInBlock(module_body, "convert_lossless_round_trip") orelse return error.TestExpectedFunc;
783 const entry = func.getRegion(0).?.getEntryBlock().?;
784 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
785 try testing.expect(ret.getOperand(0).? == entry.getArgument(0).?);
786 try ir.verifyOperation(choir_mod, ir.verify.default_options);
787 }
788
789 test "saturation preserves lossy convert round trip" {
790 const allocator = testing.allocator;
791
792 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
793 defer builder.deinit();
794 const f16_4 = try builder.tensor(.f16, &.{4});
795 const f32_4 = try builder.tensor(.f32, &.{4});
796 var fb = try builder.beginFunction("convert_lossy_round_trip", &.{f32_4}, &.{f32_4});
797 const narrowed = try fb.convert(fb.parameter(0), f16_4, .f16);
798 const restored = try fb.convert(narrowed, f32_4, .f32);
799 try fb.return_(&.{restored});
800 try fb.finish();
801 const module = try builder.finish();
802 defer module.deinit();
803
804 const choir_mod = module.choir_module;
805 const ctx = module.context();
806 var pm = passes.PassManager.init(allocator);
807 defer pm.deinit();
808 try pm.addPass(tensorSaturationPass());
809
810 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
811 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
812
813 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
814 const func = ir.inspection.functionByNameInBlock(module_body, "convert_lossy_round_trip") orelse return error.TestExpectedFunc;
815 const entry = func.getRegion(0).?.getEntryBlock().?;
816 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
817 try testing.expect(ret.getOperand(0).? != entry.getArgument(0).?);
818 try testing.expectEqual(@as(usize, 2), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.ConvertOp.operation_name));
819 try ir.verifyOperation(choir_mod, ir.verify.default_options);
820 }
821
822 test "saturation collapses right integer identity matrix product" {
823 const allocator = testing.allocator;
824
825 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
826 defer builder.deinit();
827 const i32_2x3 = try builder.tensor(.i32, &.{ 2, 3 });
828 const i32_3x3 = try builder.tensor(.i32, &.{ 3, 3 });
829 var identity: [9]i32 = undefined;
830 writeIdentityI32(identity[0..], 3);
831
832 var fb = try builder.beginFunction("dot_right_identity", &.{i32_2x3}, &.{i32_2x3});
833 const rhs = try fb.constant(i32_3x3, std.mem.sliceAsBytes(identity[0..]));
834 const product = try fb.dotGeneral(fb.parameter(0), rhs, i32_2x3, &.{1}, &.{0}, &.{}, &.{});
835 try fb.return_(&.{product});
836 try fb.finish();
837 const module = try builder.finish();
838 defer module.deinit();
839
840 const choir_mod = module.choir_module;
841 const ctx = module.context();
842 var pm = passes.PassManager.init(allocator);
843 defer pm.deinit();
844 try pm.addPass(tensorSaturationPass());
845
846 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
847 try testing.expectEqual(@as(u64, 1), pm.stats.passes_modified);
848
849 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
850 const func = ir.inspection.functionByNameInBlock(module_body, "dot_right_identity") orelse return error.TestExpectedFunc;
851 const entry = func.getRegion(0).?.getEntryBlock().?;
852 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
853 try testing.expect(ret.getOperand(0).? == entry.getArgument(0).?);
854 try ir.verifyOperation(choir_mod, ir.verify.default_options);
855 }
856
857 test "saturation collapses left integer identity matrix product" {
858 const allocator = testing.allocator;
859
860 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
861 defer builder.deinit();
862 const i32_2x2 = try builder.tensor(.i32, &.{ 2, 2 });
863 const i32_2x3 = try builder.tensor(.i32, &.{ 2, 3 });
864 var identity: [4]i32 = undefined;
865 writeIdentityI32(identity[0..], 2);
866
867 var fb = try builder.beginFunction("dot_left_identity", &.{i32_2x3}, &.{i32_2x3});
868 const lhs = try fb.constant(i32_2x2, std.mem.sliceAsBytes(identity[0..]));
869 const product = try fb.dotGeneral(lhs, fb.parameter(0), i32_2x3, &.{1}, &.{0}, &.{}, &.{});
870 try fb.return_(&.{product});
871 try fb.finish();
872 const module = try builder.finish();
873 defer module.deinit();
874
875 const choir_mod = module.choir_module;
876 const ctx = module.context();
877 var pm = passes.PassManager.init(allocator);
878 defer pm.deinit();
879 try pm.addPass(tensorSaturationPass());
880
881 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
882 try testing.expectEqual(@as(u64, 1), pm.stats.passes_modified);
883
884 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
885 const func = ir.inspection.functionByNameInBlock(module_body, "dot_left_identity") orelse return error.TestExpectedFunc;
886 const entry = func.getRegion(0).?.getEntryBlock().?;
887 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
888 try testing.expect(ret.getOperand(0).? == entry.getArgument(0).?);
889 try ir.verifyOperation(choir_mod, ir.verify.default_options);
890 }
891
892 fn expectFloatingIdentityRetained(comptime T: type, dtype: choir_abi.DType, one: T) !void {
893 const allocator = testing.allocator;
894 const payload = [_]T{ one, 0, 0, one };
895 for ([_]bool{ false, true }) |identity_on_left| {
896 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
897 defer builder.deinit();
898 const ty = try builder.tensor(dtype, &.{ 2, 2 });
899 var fb = try builder.beginFunction("floating_identity", &.{ty}, &.{ty});
900 const identity = try fb.constant(ty, std.mem.sliceAsBytes(&payload));
901 const lhs = if (identity_on_left) identity else fb.parameter(0);
902 const rhs = if (identity_on_left) fb.parameter(0) else identity;
903 const product = try fb.dotGeneral(lhs, rhs, ty, &.{1}, &.{0}, &.{}, &.{});
904 try fb.return_(&.{product});
905 try fb.finish();
906 const module = try builder.finish();
907 defer module.deinit();
908 _ = try saturateTensorOperations(allocator, module.choir_module, module.context());
909 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(
910 module.choir_module,
911 dialect_mod.AccyDialect.DotGeneralOp.operation_name,
912 ));
913 try module.verify();
914 }
915 }
916
917 test "saturation preserves floating matrix identity operations" {
918 try expectFloatingIdentityRetained(f16, .f16, 1);
919 try expectFloatingIdentityRetained(u16, .bf16, choir_abi.Bf16.fromF32(1).bits);
920 try expectFloatingIdentityRetained(f32, .f32, 1);
921 try expectFloatingIdentityRetained(f64, .f64, 1);
922 }
923
924 test "saturation preserves non-identity matrix product constant" {
925 const allocator = testing.allocator;
926
927 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
928 defer builder.deinit();
929 const f32_2x3 = try builder.tensor(.f32, &.{ 2, 3 });
930 const f32_3x3 = try builder.tensor(.f32, &.{ 3, 3 });
931 var diagonal = [_]f32{
932 1.0, 0.0, 0.0,
933 0.0, 2.0, 0.0,
934 0.0, 0.0, 1.0,
935 };
936
937 var fb = try builder.beginFunction("dot_non_identity_constant", &.{f32_2x3}, &.{f32_2x3});
938 const rhs = try fb.constant(f32_3x3, std.mem.sliceAsBytes(diagonal[0..]));
939 const product = try fb.dotGeneral(fb.parameter(0), rhs, f32_2x3, &.{1}, &.{0}, &.{}, &.{});
940 try fb.return_(&.{product});
941 try fb.finish();
942 const module = try builder.finish();
943 defer module.deinit();
944
945 const choir_mod = module.choir_module;
946 const ctx = module.context();
947 var pm = passes.PassManager.init(allocator);
948 defer pm.deinit();
949 try pm.addPass(tensorSaturationPass());
950
951 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
952 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
953
954 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
955 const func = ir.inspection.functionByNameInBlock(module_body, "dot_non_identity_constant") orelse return error.TestExpectedFunc;
956 const entry = func.getRegion(0).?.getEntryBlock().?;
957 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
958 try testing.expect(ret.getOperand(0).? != entry.getArgument(0).?);
959 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
960 try ir.verifyOperation(choir_mod, ir.verify.default_options);
961 }
962
963 test "transpose involution collapses inverse permutation pairs" {
964 const allocator = testing.allocator;
965
966 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
967 defer builder.deinit();
968 const f32_2x3 = try builder.tensor(.f32, &.{ 2, 3 });
969 const f32_3x2 = try builder.tensor(.f32, &.{ 3, 2 });
970 var fb = try builder.beginFunction("transpose_round_trip", &.{f32_2x3}, &.{f32_2x3});
971 const flipped = try fb.transpose(fb.parameter(0), f32_3x2, &.{ 1, 0 });
972 const restored = try fb.transpose(flipped, f32_2x3, &.{ 1, 0 });
973 try fb.return_(&.{restored});
974 try fb.finish();
975 const module = try builder.finish();
976 defer module.deinit();
977
978 const choir_mod = module.choir_module;
979 const ctx = module.context();
980 var pm = passes.PassManager.init(allocator);
981 defer pm.deinit();
982 try pm.addPass(tensorSaturationPass());
983
984 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
985 try testing.expectEqual(@as(u64, 1), pm.stats.passes_modified);
986
987 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
988 const func = ir.inspection.functionByNameInBlock(module_body, "transpose_round_trip") orelse return error.TestExpectedFunc;
989 const entry = func.getRegion(0).?.getEntryBlock().?;
990 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
991 try testing.expect(ret.getOperand(0).? == entry.getArgument(0).?);
992 try ir.verifyOperation(choir_mod, ir.verify.default_options);
993 }
994
995 test "transpose involution leaves non-inverse permutations alone" {
996 const allocator = testing.allocator;
997
998 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
999 defer builder.deinit();
1000 const f32_2x3x4 = try builder.tensor(.f32, &.{ 2, 3, 4 });
1001 const f32_3x4x2 = try builder.tensor(.f32, &.{ 3, 4, 2 });
1002 const f32_4x2x3 = try builder.tensor(.f32, &.{ 4, 2, 3 });
1003 var fb = try builder.beginFunction("transpose_rotation", &.{f32_2x3x4}, &.{f32_4x2x3});
1004 const rotated = try fb.transpose(fb.parameter(0), f32_3x4x2, &.{ 1, 2, 0 });
1005 const rotated_again = try fb.transpose(rotated, f32_4x2x3, &.{ 1, 2, 0 });
1006 try fb.return_(&.{rotated_again});
1007 try fb.finish();
1008 const module = try builder.finish();
1009 defer module.deinit();
1010
1011 const choir_mod = module.choir_module;
1012 const ctx = module.context();
1013 var pm = passes.PassManager.init(allocator);
1014 defer pm.deinit();
1015 try pm.addPass(tensorSaturationPass());
1016
1017 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
1018 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
1019 try testing.expectEqual(@as(usize, 2), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.TransposeOp.operation_name));
1020 }
1021
1022 test "transpose involution collapses inverse three-dimensional rotations" {
1023 const allocator = testing.allocator;
1024
1025 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1026 defer builder.deinit();
1027 const f32_2x3x4 = try builder.tensor(.f32, &.{ 2, 3, 4 });
1028 const f32_3x4x2 = try builder.tensor(.f32, &.{ 3, 4, 2 });
1029 var fb = try builder.beginFunction("transpose_inverse_rotation", &.{f32_2x3x4}, &.{f32_2x3x4});
1030 const rotated = try fb.transpose(fb.parameter(0), f32_3x4x2, &.{ 1, 2, 0 });
1031 const restored = try fb.transpose(rotated, f32_2x3x4, &.{ 2, 0, 1 });
1032 try fb.return_(&.{restored});
1033 try fb.finish();
1034 const module = try builder.finish();
1035 defer module.deinit();
1036
1037 const choir_mod = module.choir_module;
1038 const ctx = module.context();
1039 var pm = passes.PassManager.init(allocator);
1040 defer pm.deinit();
1041 try pm.addPass(tensorSaturationPass());
1042
1043 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
1044 try testing.expectEqual(@as(u64, 1), pm.stats.passes_modified);
1045
1046 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
1047 const func = ir.inspection.functionByNameInBlock(module_body, "transpose_inverse_rotation") orelse return error.TestExpectedFunc;
1048 const entry = func.getRegion(0).?.getEntryBlock().?;
1049 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
1050 try testing.expect(ret.getOperand(0).? == entry.getArgument(0).?);
1051 }