lib/accy/src/preparation/canonicalization.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_fixed = @import("alloc_fixed");
3 const choir = @import("choir");
4 const accy_choir = @import("../choir/root.zig");
5 const dialect_mod = accy_choir.dialect;
6
7 const ir = choir.ir;
8 const rewrite = ir.rewrite;
9 const passes = choir.passes;
10 const work = passes.pass.work;
11 const AccyDialect = dialect_mod.AccyDialect;
12
13 pub const canonicalization_pass_name = "accy-choir-canonicalize";
14 pub const canonicalization_pass_description =
15 "Canonicalize identity Accy Choir tensor operations";
16 const RewriteFn = *const fn (*ir.Operation, *rewrite.PatternRewriter) rewrite.PatternResult;
17
18 pub const CanonicalizationPatternEntry = struct {
19 spec: rewrite.RewritePatternSpec,
20 rewrite: RewriteFn,
21 };
22
23 pub const canonicalization_pattern_entries = [_]CanonicalizationPatternEntry{
24 .{ .spec = .{ .name = "accy-reshape-identity", .root_op_name = AccyDialect.ReshapeOp.operation_name, .products = .none }, .rewrite = rewriteReshape },
25 .{ .spec = .{ .name = "accy-transpose-identity", .root_op_name = AccyDialect.TransposeOp.operation_name, .products = .none }, .rewrite = rewriteTranspose },
26 .{ .spec = .{ .name = "accy-broadcast-identity", .root_op_name = AccyDialect.BroadcastOp.operation_name, .products = .none }, .rewrite = rewriteBroadcast },
27 .{ .spec = .{ .name = "accy-broadcast-in-dim-identity", .root_op_name = AccyDialect.BroadcastInDimOp.operation_name, .products = .none }, .rewrite = rewriteBroadcastInDim },
28 .{ .spec = .{ .name = "accy-slice-identity", .root_op_name = AccyDialect.SliceOp.operation_name, .products = .none }, .rewrite = rewriteSlice },
29 .{ .spec = .{ .name = "accy-concatenate-single-input", .root_op_name = AccyDialect.ConcatenateOp.operation_name, .products = .none }, .rewrite = rewriteConcatenate },
30 .{ .spec = .{ .name = "accy-dot-flatten-shared-rhs", .root_op_name = AccyDialect.DotGeneralOp.operation_name, .products = .none }, .rewrite = rewriteDotGeneral },
31 .{ .spec = .{ .name = "accy-cumsum-attach-scratch", .root_op_name = AccyDialect.CumsumOp.operation_name, .products = .none }, .rewrite = rewriteCumsum },
32 };
33
34 pub const cumsum_block_tile = 8192;
35
36 const greedy_config: passes.GreedyRewriteConfig = .{};
37
38 const AccyCanonicalizationPass = passes.CanonicalizationPass(.{
39 .name = canonicalization_pass_name,
40 .description = canonicalization_pass_description,
41 .populate_patterns = populateCanonicalizationPatterns,
42 .greedy_config = greedy_config,
43 });
44
45 pub fn canonicalizationPass() passes.Pass {
46 var result = AccyCanonicalizationPass.init();
47 result.work_contract = .{
48 .identity = .{ .name = canonicalization_pass_name, .version = 1 },
49 .estimate = canonicalizationWork,
50 };
51 return result;
52 }
53
54 const CanonicalizationWorkFacts = struct {
55 type_key_bytes: u64 = 0,
56 list_bytes: u64 = 0,
57 broadcasts: u64 = 0,
58 dots: u64 = 0,
59 scans: u64 = 0,
60
61 fn visit(self: *CanonicalizationWorkFacts, op: *ir.Operation) !ir.WalkResult {
62 const namespace = op.name.getDialectNamespace();
63 if (!std.mem.eql(u8, namespace, "accy") and !std.mem.eql(u8, namespace, "func") and
64 !std.mem.eql(u8, namespace, "builtin")) return error.UnsupportedDialect;
65 if (op.hasInterface(ir.interfaces.FoldOpInterface)) return error.MissingWorkContract;
66 const name = op.name.name;
67 if (std.mem.eql(u8, name, AccyDialect.BroadcastOp.operation_name) or
68 std.mem.eql(u8, name, AccyDialect.BroadcastInDimOp.operation_name))
69 {
70 self.broadcasts = try work.add(self.broadcasts, 1);
71 }
72 if (std.mem.eql(u8, name, AccyDialect.DotGeneralOp.operation_name)) {
73 self.dots = try work.add(self.dots, 1);
74 }
75 if (std.mem.eql(u8, name, AccyDialect.CumsumOp.operation_name)) {
76 self.scans = try work.add(self.scans, 1);
77 }
78 for (op.results.items) |*value| self.typeKey(value.type);
79 for (op.getOperandValues()) |value| self.typeKey(value.type);
80 var attributes = op.getAttrs();
81 while (attributes.next()) |entry| {
82 if (entry.value.cast(ir.Attribute.DialectAttr)) |attr| {
83 self.list_bytes = @max(self.list_bytes, attr.payload.len);
84 }
85 }
86 return .advance;
87 }
88
89 fn typeKey(self: *CanonicalizationWorkFacts, typ: ir.Type) void {
90 if (typ.getDialectParamKey()) |key| self.type_key_bytes = @max(self.type_key_bytes, key.len);
91 }
92 };
93
94 /// The compile chain calls this before the pass runs, to charge the costs a pass declares before it
95 /// runs against the caller's limits as the work bound, so compilation can refuse the pass on them.
96 /// The rules this pass admits remove operations that change nothing (reshape, transpose, broadcast,
97 /// `broadcast_in_dim`, slice, and a concatenate with one input), rewrite a plain broadcast as one
98 /// `broadcast_in_dim`, merge a chain of `broadcast_in_dim` into one, split one batched dot into a
99 /// reshape, a dot and a reshape, and attach one scratch value to a cumulative sum. No rule makes a
100 /// chain longer, so the operations the pass can create are bounded by counts taken from the input:
101 /// at most three per dot, two per cumulative sum, and a quadratic term in the number of broadcasts.
102 /// The bound multiplies those counts by the rewriter's iteration limit.
103 fn canonicalizationWork(input: work.Input) !work.Bounds {
104 const counts = try work.Census.inspect(input.operation);
105 var facts: CanonicalizationWorkFacts = .{};
106 _ = try input.operation.walk(.{ .order = .pre_order }, &facts, CanonicalizationWorkFacts.visit);
107 const patterns = try canonicalizationPatternStorage(input.operation.context);
108 const broadcasts = try work.multiply(facts.broadcasts, try work.add(facts.broadcasts, 1));
109 const dot_and_scan = try work.add(try work.multiply(3, facts.dots), try work.multiply(2, facts.scans));
110 const generated = try work.add(broadcasts, dot_and_scan);
111 const traversed = try work.add(counts.operations, try work.multiply(2, generated));
112 const attempts = try work.multiply(greedy_config.max_iterations, traversed);
113 const dimensions = try work.multiply(try work.add(facts.type_key_bytes, 128), @sizeOf(i64));
114 const scratch = try work.multiply(12, try work.add(
115 try work.add(dimensions, facts.list_bytes),
116 @alignOf(i64),
117 ));
118 const queues = try work.multiply(2, try work.arrayListGrowth(*ir.Operation, traversed));
119 const iterations = try work.multiply(greedy_config.max_iterations, queues);
120 const workspace = try work.add(patterns.bytes, try work.add(
121 iterations,
122 try work.multiply(attempts, scratch),
123 ));
124 const units = try work.add(try work.add(counts.atoms, counts.input_bytes), generated);
125 const scans = try work.multiply(try work.add(patterns.patterns, 64), try work.multiply(attempts, try work.add(units, 1)));
126 return .{
127 .work = .{
128 .input_bytes = try work.add(counts.input_bytes, patterns.name_bytes),
129 .output_bytes = input.operation.context.capacity.storage_bytes,
130 .structural_visits = try work.add(scans, patterns.visits),
131 .rewrite_attempts = try work.multiply(attempts, canonicalization_pattern_entries.len),
132 .allocation_capacity = workspace,
133 },
134 .workspace = workspace,
135 };
136 }
137
138 fn canonicalizationPatternStorage(
139 context: *ir.Context,
140 ) !passes.canonicalization.PatternPopulationBounds {
141 const Interface = rewrite.DialectCanonicalizationInterface;
142 const known = choir.dialects.arith.canonicalization_patterns[0..];
143 var interfaces = context.dialect_registry.interfaces.iterator();
144 while (interfaces.next()) |entry| {
145 for (entry.value_ptr.items) |interface| {
146 if (interface.id != Interface.id) continue;
147 const patterns = Interface.fromOpaque(interface.vtable).patterns;
148 if (patterns.len != 0 and (patterns.len != known.len or patterns.ptr != known.ptr)) {
149 return error.MissingWorkContract;
150 }
151 }
152 }
153 const extra = comptime blk: {
154 var specs: [canonicalization_pattern_entries.len]rewrite.RewritePatternSpec = undefined;
155 for (canonicalization_pattern_entries, 0..) |entry, index| specs[index] = entry.spec;
156 break :blk specs;
157 };
158 return passes.canonicalization.patternPopulationBounds(context, &extra);
159 }
160
161 pub fn canonicalizeModule(
162 allocator: std.mem.Allocator,
163 choir_module: *ir.Operation,
164 ctx: *ir.Context,
165 ) !usize {
166 var pm = passes.PassManager.init(allocator);
167 defer pm.deinit();
168 pm.enableVerifier();
169 try pm.addPass(canonicalizationPass());
170 if (pm.run(choir_module, ctx) == .failure) return error.CanonicalizationFailed;
171 return pm.stats.passes_modified;
172 }
173
174 pub fn populateCanonicalizationPatterns(patterns: *rewrite.RewritePatternSet) !void {
175 for (canonicalization_pattern_entries) |entry| {
176 try patterns.add(rewrite.RewritePattern.init(entry.spec, entry.rewrite));
177 }
178 }
179
180 fn rewriteResult(result: anyerror!bool) rewrite.PatternResult {
181 return if (result catch return .failure) .success else .failure;
182 }
183
184 fn rewriteReshape(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {
185 return rewriteResult(canonicalizeReshape(rewriter.allocator, op, rewriter));
186 }
187
188 fn rewriteTranspose(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {
189 return rewriteResult(canonicalizeTranspose(rewriter.allocator, op, rewriter));
190 }
191
192 fn rewriteBroadcast(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {
193 return rewriteResult(canonicalizeBroadcast(rewriter.allocator, op, rewriter));
194 }
195
196 fn rewriteBroadcastInDim(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {
197 return rewriteResult(canonicalizeBroadcastInDim(rewriter.allocator, op, rewriter));
198 }
199
200 fn rewriteSlice(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {
201 return rewriteResult(canonicalizeSlice(rewriter.allocator, op, rewriter));
202 }
203
204 fn rewriteConcatenate(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {
205 return rewriteResult(canonicalizeConcatenate(rewriter.allocator, op, rewriter));
206 }
207
208 fn broadcastBatchedSource(rhs: *ir.Value) ?*ir.Value {
209 const def_any = rhs.getDefiningOp() orelse return null;
210 const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));
211 if (!std.mem.eql(u8, def_op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) return null;
212 if (def_op.getNumOperands() != 1) return null;
213 var dims_buffer: [8]i64 = undefined;
214 const attr = def_op.getAttrAs(ir.Attribute.DialectAttr, "broadcast_dims") orelse return null;
215 if (attr.payload.len != 2 * @sizeOf(i64)) return null;
216 @memcpy(std.mem.sliceAsBytes(dims_buffer[0..2]), attr.payload[0 .. 2 * @sizeOf(i64)]);
217 if (dims_buffer[0] != 1 or dims_buffer[1] != 2) return null;
218 return def_op.getOperand(0);
219 }
220
221 fn rewriteDotGeneral(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {
222 return rewriteResult(canonicalizeDotGeneral(rewriter.allocator, op, rewriter));
223 }
224
225 fn rewriteCumsum(op: *ir.Operation, rewriter: *rewrite.PatternRewriter) rewrite.PatternResult {
226 return rewriteResult(canonicalizeCumsum(rewriter.allocator, op, rewriter));
227 }
228
229 fn canonicalizeCumsum(
230 allocator: std.mem.Allocator,
231 op: *ir.Operation,
232 rewriter: *rewrite.PatternRewriter,
233 ) !bool {
234 if (op.getNumOperands() != 1 or op.getNumResults() != 1) return false;
235 const input = op.getOperand(0) orelse return false;
236 const result = op.getResult(0) orelse return false;
237
238 const shape = try tensorShapeAlloc(allocator, input.type) orelse return false;
239 defer allocator.free(shape);
240 if (shape.len != 1) return false;
241 const total = shape[0];
242 if (total <= cumsum_block_tile) return false;
243
244 const blocks = @divTrunc(total + cumsum_block_tile - 1, cumsum_block_tile);
245 const words = 1 + 3 * blocks;
246
247 const ctx = rewriter.ir_ctx;
248 const scratch_type = try dialect_mod.accyTensorType(ctx, .u32, &.{words});
249 const scratch = try AccyDialect.ScratchOp.create(ctx, op.location, scratch_type, words);
250 _ = try rewriter.insert(scratch.op);
251
252 const axis_attr = op.getAttrAs(ir.Attribute.IntegerAttr, "axis") orelse return false;
253 const replacement = try AccyDialect.CumsumOp.createWithScratch(ctx, op.location, input, scratch.getResult(), result.type, axis_attr.getValue());
254 _ = try rewriter.insert(replacement.op);
255 try rewriter.replaceOpWithOperation(op, replacement.op);
256 return true;
257 }
258
259 fn canonicalizeDotGeneral(
260 allocator: std.mem.Allocator,
261 op: *ir.Operation,
262 rewriter: *rewrite.PatternRewriter,
263 ) !bool {
264 if (op.getNumOperands() != 2 or op.getNumResults() != 1) return false;
265 const lhs = op.getOperand(0) orelse return false;
266 var rhs = op.getOperand(1) orelse return false;
267 const result = op.getResult(0) orelse return false;
268
269 const lhs_shape = try tensorShapeAlloc(allocator, lhs.type) orelse return false;
270 defer allocator.free(lhs_shape);
271 const result_shape = try tensorShapeAlloc(allocator, result.type) orelse return false;
272 defer allocator.free(result_shape);
273 if (lhs_shape.len != 3 or result_shape.len != 3) return false;
274
275 const lhs_contract = try readI64ListAttrAlloc(allocator, op, "lhs_contract", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("lhs_contract")) orelse return false;
276 defer allocator.free(lhs_contract);
277 const rhs_contract = try readI64ListAttrAlloc(allocator, op, "rhs_contract", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("rhs_contract")) orelse return false;
278 defer allocator.free(rhs_contract);
279 const lhs_batch = try readI64ListAttrAlloc(allocator, op, "lhs_batch", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("lhs_batch")) orelse return false;
280 defer allocator.free(lhs_batch);
281 const rhs_batch = try readI64ListAttrAlloc(allocator, op, "rhs_batch", dialect_mod.AccyDialect.DotGeneralOp.dialectAttrName("rhs_batch")) orelse return false;
282 defer allocator.free(rhs_batch);
283
284 if (lhs_contract.len != 1 or rhs_contract.len != 1) return false;
285 if (lhs_contract[0] != 2) return false;
286
287 if (rhs_batch.len == 0) {
288 if (rhs_contract[0] != 0) return false;
289 if (lhs_batch.len > 1) return false;
290 if (lhs_batch.len == 1 and lhs_batch[0] != 0) return false;
291 } else {
292 if (lhs_batch.len != 1 or rhs_batch.len != 1) return false;
293 if (lhs_batch[0] != 0 or rhs_batch[0] != 0) return false;
294 if (rhs_contract[0] != 1) return false;
295 rhs = broadcastBatchedSource(rhs) orelse return false;
296 }
297
298 const rhs_shape = try tensorShapeAlloc(allocator, rhs.type) orelse return false;
299 defer allocator.free(rhs_shape);
300 if (rhs_shape.len != 2) return false;
301
302 const batch = lhs_shape[0];
303 const m = lhs_shape[1];
304 const k = lhs_shape[2];
305 const n = rhs_shape[1];
306 if (rhs_shape[0] != k) return false;
307 if (result_shape[0] != batch or result_shape[1] != m or result_shape[2] != n) return false;
308 const flat_m = std.math.mul(i64, batch, m) catch return false;
309
310 var dtype_arena_buffer: [256]u8 = undefined;
311 var dtype_arena = alloc_fixed.FixedBuffer.init(dtype_arena_buffer[0..]);
312 const lhs_decoded = dialect_mod.decodeTensorType(dtype_arena.allocator(), lhs.type) catch return false;
313
314 const ctx = rewriter.ir_ctx;
315 const flat_lhs_type = try dialect_mod.accyTensorType(ctx, lhs_decoded.dtype, &.{ flat_m, k });
316 const flat_out_type = try dialect_mod.accyTensorType(ctx, lhs_decoded.dtype, &.{ flat_m, n });
317
318 const flat_lhs = try dialect_mod.AccyDialect.ReshapeOp.create(ctx, op.location, lhs, flat_lhs_type, &.{ flat_m, k });
319 _ = try rewriter.insert(flat_lhs.op);
320 const flat_dot = try dialect_mod.AccyDialect.DotGeneralOp.create(ctx, op.location, flat_lhs.getResult(), rhs, flat_out_type, &.{}, &.{}, &.{1}, &.{0});
321 _ = try rewriter.insert(flat_dot.op);
322 const restored = try dialect_mod.AccyDialect.ReshapeOp.create(ctx, op.location, flat_dot.getResult(), result.type, result_shape);
323 _ = try rewriter.insert(restored.op);
324 try rewriter.replaceOpWithOperation(op, restored.op);
325 return true;
326 }
327
328 fn canonicalizeReshape(
329 allocator: std.mem.Allocator,
330 op: *ir.Operation,
331 rewriter: *rewrite.PatternRewriter,
332 ) !bool {
333 const input = identityUnaryInput(op) orelse return false;
334 const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;
335 defer allocator.free(input_shape);
336 const new_shape = try readI64ListAttrAlloc(allocator, op, "new_shape", dialect_mod.AccyDialect.ReshapeOp.dialectAttrName("new_shape")) orelse return false;
337 defer allocator.free(new_shape);
338 if (!std.mem.eql(i64, input_shape, new_shape)) return false;
339 return replaceWithInput(op, input, rewriter);
340 }
341
342 fn canonicalizeTranspose(
343 allocator: std.mem.Allocator,
344 op: *ir.Operation,
345 rewriter: *rewrite.PatternRewriter,
346 ) !bool {
347 const input = identityUnaryInput(op) orelse return false;
348 const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;
349 defer allocator.free(input_shape);
350 const permutation = try readI64ListAttrAlloc(allocator, op, "permutation", dialect_mod.AccyDialect.TransposeOp.dialectAttrName("permutation")) orelse return false;
351 defer allocator.free(permutation);
352 if (permutation.len != input_shape.len) return false;
353 if (!isIdentityPermutation(permutation)) return false;
354 return replaceWithInput(op, input, rewriter);
355 }
356
357 fn canonicalizeBroadcast(
358 allocator: std.mem.Allocator,
359 op: *ir.Operation,
360 rewriter: *rewrite.PatternRewriter,
361 ) !bool {
362 if (op.getNumOperands() != 1 or op.getNumResults() != 1) return false;
363 const input = op.getOperand(0) orelse return false;
364 const result = op.getResult(0) orelse return false;
365 const sizes = try readI64ListAttrAlloc(allocator, op, "sizes", dialect_mod.AccyDialect.BroadcastOp.dialectAttrName("sizes")) orelse return false;
366 defer allocator.free(sizes);
367 if (sizes.len == 0) {
368 if (!result.type.eql(input.type)) return false;
369 return replaceWithInput(op, input, rewriter);
370 }
371
372 const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;
373 defer allocator.free(input_shape);
374
375 const result_shape = try allocator.alloc(i64, sizes.len + input_shape.len);
376 defer allocator.free(result_shape);
377 @memcpy(result_shape[0..sizes.len], sizes);
378 @memcpy(result_shape[sizes.len..], input_shape);
379
380 const broadcast_dims = try allocator.alloc(i64, input_shape.len);
381 defer allocator.free(broadcast_dims);
382 for (broadcast_dims, 0..) |*dim, index| dim.* = @intCast(sizes.len + index);
383
384 const replacement = try dialect_mod.AccyDialect.BroadcastInDimOp.create(
385 rewriter.ir_ctx,
386 op.location,
387 input,
388 result.type,
389 broadcast_dims,
390 result_shape,
391 );
392 _ = try rewriter.insert(replacement.op);
393 try rewriter.replaceOpWithOperation(op, replacement.op);
394 return true;
395 }
396
397 fn canonicalizeBroadcastInDim(
398 allocator: std.mem.Allocator,
399 op: *ir.Operation,
400 rewriter: *rewrite.PatternRewriter,
401 ) !bool {
402 if (try collapseChainedBroadcastInDim(allocator, op, rewriter)) return true;
403 const input = identityUnaryInput(op) orelse return false;
404 const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;
405 defer allocator.free(input_shape);
406 const broadcast_dims = try readI64ListAttrAlloc(allocator, op, "broadcast_dims", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("broadcast_dims")) orelse return false;
407 defer allocator.free(broadcast_dims);
408 const result_shape = try readI64ListAttrAlloc(allocator, op, "result_shape", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("result_shape")) orelse return false;
409 defer allocator.free(result_shape);
410 if (broadcast_dims.len != input_shape.len) return false;
411 if (!isIdentityPermutation(broadcast_dims)) return false;
412 if (!std.mem.eql(i64, input_shape, result_shape)) return false;
413 return replaceWithInput(op, input, rewriter);
414 }
415
416 fn collapseChainedBroadcastInDim(
417 allocator: std.mem.Allocator,
418 op: *ir.Operation,
419 rewriter: *rewrite.PatternRewriter,
420 ) !bool {
421 if (op.getNumOperands() != 1 or op.getNumResults() != 1) return false;
422 const mid = op.getOperand(0) orelse return false;
423 const result = op.getResult(0) orelse return false;
424 const inner_any = mid.getDefiningOp() orelse return false;
425 const inner: *ir.Operation = @ptrCast(@alignCast(inner_any));
426 if (!std.mem.eql(u8, inner.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) return false;
427 if (inner.getNumOperands() != 1) return false;
428 const source = inner.getOperand(0) orelse return false;
429
430 const outer_dims = try readI64ListAttrAlloc(allocator, op, "broadcast_dims", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("broadcast_dims")) orelse return false;
431 defer allocator.free(outer_dims);
432 const outer_shape = try readI64ListAttrAlloc(allocator, op, "result_shape", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("result_shape")) orelse return false;
433 defer allocator.free(outer_shape);
434 const inner_dims = try readI64ListAttrAlloc(allocator, inner, "broadcast_dims", dialect_mod.AccyDialect.BroadcastInDimOp.dialectAttrName("broadcast_dims")) orelse return false;
435 defer allocator.free(inner_dims);
436
437 const composed = try allocator.alloc(i64, inner_dims.len);
438 defer allocator.free(composed);
439 for (inner_dims, 0..) |mid_dim, index| {
440 if (mid_dim < 0 or mid_dim >= outer_dims.len) return false;
441 composed[index] = outer_dims[@intCast(mid_dim)];
442 }
443
444 const replacement = try dialect_mod.AccyDialect.BroadcastInDimOp.create(
445 rewriter.ir_ctx,
446 op.location,
447 source,
448 result.type,
449 composed,
450 outer_shape,
451 );
452 _ = try rewriter.insert(replacement.op);
453 try rewriter.replaceOpWithOperation(op, replacement.op);
454 return true;
455 }
456
457 fn canonicalizeSlice(
458 allocator: std.mem.Allocator,
459 op: *ir.Operation,
460 rewriter: *rewrite.PatternRewriter,
461 ) !bool {
462 const input = identityUnaryInput(op) orelse return false;
463 const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;
464 defer allocator.free(input_shape);
465 const starts = try readI64ListAttrAlloc(allocator, op, "starts", dialect_mod.AccyDialect.SliceOp.dialectAttrName("starts")) orelse return false;
466 defer allocator.free(starts);
467 const limits = try readI64ListAttrAlloc(allocator, op, "limits", dialect_mod.AccyDialect.SliceOp.dialectAttrName("limits")) orelse return false;
468 defer allocator.free(limits);
469 const strides = try readI64ListAttrAlloc(allocator, op, "strides", dialect_mod.AccyDialect.SliceOp.dialectAttrName("strides")) orelse return false;
470 defer allocator.free(strides);
471
472 if (starts.len != input_shape.len or
473 limits.len != input_shape.len or
474 strides.len != input_shape.len)
475 {
476 return false;
477 }
478 for (starts, limits, strides, input_shape) |start, limit, stride, dim| {
479 if (start != 0 or limit != dim or stride != 1) return false;
480 }
481 return replaceWithInput(op, input, rewriter);
482 }
483
484 fn canonicalizeConcatenate(
485 allocator: std.mem.Allocator,
486 op: *ir.Operation,
487 rewriter: *rewrite.PatternRewriter,
488 ) !bool {
489 if (op.getNumOperands() != 1 or op.getNumResults() != 1) return false;
490 const input = op.getOperand(0) orelse return false;
491 const result = op.getResult(0) orelse return false;
492 if (!result.type.eql(input.type)) return false;
493 const input_shape = try tensorShapeAlloc(allocator, input.type) orelse return false;
494 defer allocator.free(input_shape);
495 const dimension = readIntegerAttr(op, "dimension") orelse return false;
496 if (dimension < 0 or dimension >= @as(i64, @intCast(input_shape.len))) return false;
497 try rewriter.replaceOpWithValue(op, input);
498 return true;
499 }
500
501 fn identityUnaryInput(op: *ir.Operation) ?*ir.Value {
502 if (op.getNumOperands() != 1 or op.getNumResults() != 1) return null;
503 const input = op.getOperand(0) orelse return null;
504 const result = op.getResult(0) orelse return null;
505 if (!result.type.eql(input.type)) return null;
506 return input;
507 }
508
509 fn replaceWithInput(
510 op: *ir.Operation,
511 input: *ir.Value,
512 rewriter: *rewrite.PatternRewriter,
513 ) !bool {
514 try rewriter.replaceOpWithValue(op, input);
515 return true;
516 }
517
518 pub fn readDialectPayload(op: *const ir.Operation, attr_name: []const u8, dialect_attr_name: []const u8) ?[]const u8 {
519 const attr = op.getAttr(attr_name) orelse return null;
520 if (!std.mem.eql(u8, attr.abstract.name, dialect_attr_name)) return null;
521 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return null;
522 return dialect_attr.payload;
523 }
524
525 pub fn readI64ListAttrAlloc(
526 allocator: std.mem.Allocator,
527 op: *const ir.Operation,
528 attr_name: []const u8,
529 dialect_attr_name: []const u8,
530 ) !?[]i64 {
531 const payload = readDialectPayload(op, attr_name, dialect_attr_name) orelse return null;
532 if (payload.len % @sizeOf(i64) != 0) return null;
533 const values = try allocator.alloc(i64, payload.len / @sizeOf(i64));
534 for (values, 0..) |*value, i| {
535 const start = i * @sizeOf(i64);
536 @memcpy(std.mem.asBytes(value), payload[start..][0..@sizeOf(i64)]);
537 }
538 return values;
539 }
540
541 fn readIntegerAttr(op: *const ir.Operation, attr_name: []const u8) ?i64 {
542 const attr = op.getAttr(attr_name) orelse return null;
543 if (!std.mem.eql(u8, attr.abstract.name, "builtin.integer")) return null;
544 const int_attr = attr.cast(ir.Attribute.IntegerAttr) orelse return null;
545 return int_attr.getValue();
546 }
547
548 fn tensorShapeAlloc(allocator: std.mem.Allocator, typ: ir.Type) !?[]i64 {
549 const type_name = typ.getDialectTypeName() orelse return null;
550 if (!std.mem.eql(u8, type_name, dialect_mod.tensor_type_name)) return null;
551 const key = typ.getDialectParamKey() orelse return null;
552 const comma = std.mem.indexOfScalar(u8, key, ',') orelse return null;
553 const dims_text = key[comma + 1 ..];
554 if (dims_text.len == 0) return try allocator.alloc(i64, 0);
555
556 var dim_count: usize = 1;
557 for (dims_text) |ch| {
558 if (ch == 'x') dim_count += 1;
559 }
560
561 const dims = try allocator.alloc(i64, dim_count);
562 errdefer allocator.free(dims);
563 var iter = std.mem.splitScalar(u8, dims_text, 'x');
564 var index: usize = 0;
565 while (iter.next()) |part| {
566 if (part.len == 0 or index >= dim_count) {
567 allocator.free(dims);
568 return null;
569 }
570 const dim = std.fmt.parseInt(i64, part, 10) catch {
571 allocator.free(dims);
572 return null;
573 };
574 if (dim < 0) {
575 allocator.free(dims);
576 return null;
577 }
578 dims[index] = dim;
579 index += 1;
580 }
581 if (index != dim_count) {
582 allocator.free(dims);
583 return null;
584 }
585 return dims;
586 }
587
588 fn isIdentityPermutation(values: []const i64) bool {
589 for (values, 0..) |value, i| {
590 if (value != @as(i64, @intCast(i))) return false;
591 }
592 return true;
593 }
594
595 const testing = std.testing;
596 const semantic = accy_choir.semantic;
597
598 fn readSymbolName(func: *ir.Operation) ?[]const u8 {
599 return ir.SymbolTable.getSymbolName(func);
600 }
601
602 fn findOpNamedInBlock(block: *ir.Block, name: []const u8) ?*ir.Operation {
603 var iter = block.operations.head;
604 while (iter) |op_ptr| {
605 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
606 if (std.mem.eql(u8, op.name.name, name)) return op;
607 iter = op.next_op;
608 }
609 return null;
610 }
611
612 test "canonicalization patterns publish Choir rewrite specs" {
613 const allocator = testing.allocator;
614
615 var patterns = rewrite.RewritePatternSet.init(allocator);
616 defer patterns.deinit();
617
618 try populateCanonicalizationPatterns(&patterns);
619
620 try testing.expectEqual(canonicalization_pattern_entries.len, patterns.patterns.items.len);
621 for (canonicalization_pattern_entries, patterns.patterns.items) |entry, pattern| {
622 try testing.expectEqualStrings(entry.spec.name, pattern.spec.name);
623 try testing.expectEqualStrings(entry.spec.root_op_name, pattern.spec.root_op_name);
624 try testing.expectEqual(entry.spec.benefit, pattern.spec.benefit);
625 try testing.expectEqual(entry.spec.kind, pattern.spec.kind);
626 switch (entry.spec.products) {
627 .none => {},
628 else => return error.TestExpectedNoProducts,
629 }
630 }
631 try testing.expectEqualStrings(AccyDialect.ReshapeOp.operation_name, canonicalization_pattern_entries[0].spec.root_op_name);
632 }
633
634 test "canonicalizeModule removes identity shape operations" {
635 const allocator = testing.allocator;
636
637 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
638 defer builder.deinit();
639 const f32_2x3 = try builder.tensor(.f32, &.{ 2, 3 });
640 var fb = try builder.beginFunction("identity_shapes", &.{f32_2x3}, &.{f32_2x3});
641 const r0 = try fb.reshape(fb.parameter(0), f32_2x3, &.{ 2, 3 });
642 const r1 = try fb.transpose(r0, f32_2x3, &.{ 0, 1 });
643 const r2 = try fb.broadcast(r1, f32_2x3, &.{});
644 const r3 = try fb.broadcastInDim(r2, f32_2x3, &.{ 2, 3 }, &.{ 0, 1 });
645 const r4 = try fb.slice(r3, f32_2x3, &.{ 0, 0 }, &.{ 2, 3 }, &.{ 1, 1 });
646 const r5 = try fb.concatenate(&.{r4}, f32_2x3, 0);
647 try fb.return_(&.{r5});
648 try fb.finish();
649 const module = try builder.finish();
650 defer module.deinit();
651
652 const choir_mod = module.choir_module;
653 const ctx = module.context();
654 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.ReshapeOp.operation_name));
655 try testing.expectEqual(@as(usize, 1), try canonicalizeModule(allocator, choir_mod, ctx));
656 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.ReshapeOp.operation_name));
657 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.TransposeOp.operation_name));
658 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.BroadcastOp.operation_name));
659 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name));
660 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.SliceOp.operation_name));
661 try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.ConcatenateOp.operation_name));
662
663 const module_body = choir_mod.getRegion(0).?.getEntryBlock().?;
664 const func = ir.inspection.functionByNameInBlock(module_body, "identity_shapes") orelse return error.TestExpectedFunc;
665 const entry = func.getRegion(0).?.getEntryBlock().?;
666 const ret = findOpNamedInBlock(entry, "func.return") orelse return error.TestExpectedReturn;
667 try testing.expectEqual(entry.getArgument(0).?, ret.getOperand(0).?);
668 }
669
670 test "canonicalization pass preserves analyses when it makes no changes" {
671 const allocator = testing.allocator;
672
673 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
674 defer builder.deinit();
675 const f32_4 = try builder.tensor(.f32, &.{4});
676 var fb = try builder.beginFunction("canonicalize_add4", &.{ f32_4, f32_4 }, &.{f32_4});
677 const sum = try fb.add(fb.parameter(0), fb.parameter(1));
678 try fb.return_(&.{sum});
679 try fb.finish();
680 const module = try builder.finish();
681 defer module.deinit();
682
683 const choir_mod = module.choir_module;
684 const ctx = module.context();
685 var pm = passes.PassManager.init(allocator);
686 defer pm.deinit();
687 try pm.addPass(canonicalizationPass());
688
689 try testing.expectEqual(passes.PassResult.success, pm.run(choir_mod, ctx));
690 try testing.expectEqual(@as(u64, 1), pm.stats.pass_runs);
691 try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
692 }
693
694 test "canonicalizeModule flattens shared-rhs batched dots" {
695 const allocator = testing.allocator;
696
697 var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
698 defer builder.deinit();
699 const lhs_ty = try builder.tensor(.f32, &.{ 8, 2, 4 });
700 const rhs_ty = try builder.tensor(.f32, &.{ 4, 3 });
701 const out_ty = try builder.tensor(.f32, &.{ 8, 2, 3 });
702 var fb = try builder.beginFunction("flatten_dot", &.{ lhs_ty, rhs_ty }, &.{out_ty});
703 const product = try fb.dotGeneral(fb.parameter(0), fb.parameter(1), out_ty, &.{2}, &.{0}, &.{}, &.{});
704 try fb.return_(&.{product});
705 try fb.finish();
706 const module = try builder.finish();
707 defer module.deinit();
708
709 const choir_mod = module.choir_module;
710 const ctx = module.context();
711 _ = try canonicalizeModule(allocator, choir_mod, ctx);
712 try testing.expectEqual(@as(usize, 2), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.ReshapeOp.operation_name));
713 try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(choir_mod, dialect_mod.AccyDialect.DotGeneralOp.operation_name));
714 }
715
716 test "canonicalization accounting reserves worker storage and closes actual identity rewrites" {
717 const module = try accountingIdentityModule(testing.allocator);
718 defer module.deinit();
719 try checkCanonicalizationAccounting(module, 6);
720 }
721
722 fn checkCanonicalizationAccounting(module: *semantic.SemanticModule, rewrites: u64) !void {
723 const revision = choir.product.revision;
724 const allocator = testing.allocator;
725 const selected = canonicalizationPass();
726 const declared = try selected.work_contract.?.estimate(.{ .operation = module.choir_module });
727 const ledger = try revision.AccountingV1.create(allocator, .{
728 .allowance = .uniform(std.math.maxInt(u64)),
729 .workspace = declared.workspace,
730 .events = 16,
731 }, &.{.{ .name = canonicalization_pass_name, .version = 1 }});
732 defer ledger.destroy();
733 var observed = testing.FailingAllocator.init(allocator, .{});
734 var cache = try passes.AnalysisCache.initAccounted(observed.allocator(), null, ledger, .{}, 8);
735 defer cache.deinit();
736 var manager = passes.PassManager.init(observed.allocator());
737 defer manager.deinit();
738 try manager.addPass(selected);
739 const before = observed.allocated_bytes;
740 try testing.expectEqual(.success, manager.runWithAnalysisCache(
741 module.choir_module,
742 module.context(),
743 &cache,
744 .{ .max_threads = 1 },
745 ));
746 const worker_bytes = observed.allocated_bytes - before;
747 try ledger.producersComplete();
748 const receipt = ledger.view();
749 try testing.expectEqual(1, receipt.executed.counters.pass_runs);
750 try testing.expectEqual(@intFromBool(rewrites != 0), receipt.executed.counters.passes_modified);
751 try testing.expectEqual(rewrites, receipt.executed.counters.successful_rewrites);
752 const iterations: u64 = if (rewrites == 0) 1 else 2;
753 try testing.expectEqual(iterations, receipt.executed.counters.rewrite_iterations);
754 var pass_events: u32 = 0;
755 for (receipt.events) |event| {
756 if (event.phase != .pass) continue;
757 pass_events += 1;
758 try testing.expectEqual(declared.workspace, event.workspace);
759 try testing.expect(event.workspace >= worker_bytes);
760 try testing.expect(event.charged.allocation_capacity >= worker_bytes);
761 try testing.expect(event.charged.rewrite_attempts >= rewrites);
762 }
763 try testing.expectEqual(1, pass_events);
764 try module.verify();
765 }
766
767 fn accountingIdentityModule(allocator: std.mem.Allocator) !*semantic.SemanticModule {
768 var builder = try semantic.Builder.init(allocator, .testing);
769 defer builder.deinit();
770 const typ = try builder.tensor(.f32, &.{ 2, 3 });
771 var function = try builder.beginFunction("accounted_identity", &.{typ}, &.{typ});
772 const reshape = try function.reshape(function.parameter(0), typ, &.{ 2, 3 });
773 const transpose = try function.transpose(reshape, typ, &.{ 0, 1 });
774 const broadcast = try function.broadcast(transpose, typ, &.{});
775 const in_dim = try function.broadcastInDim(broadcast, typ, &.{ 2, 3 }, &.{ 0, 1 });
776 const slice = try function.slice(in_dim, typ, &.{ 0, 0 }, &.{ 2, 3 }, &.{ 1, 1 });
777 const concat = try function.concatenate(&.{slice}, typ, 0);
778 try function.return_(&.{concat});
779 try function.finish();
780 const module = try builder.finish();
781 return module;
782 }
783
784 test "canonicalization accounting refuses workspace and rewrite budgets before mutation" {
785 inline for (.{ false, true }) |exhaust_workspace| {
786 try checkCanonicalizationRefusal(exhaust_workspace);
787 }
788 }
789
790 fn checkCanonicalizationRefusal(exhaust_workspace: bool) !void {
791 const revision = choir.product.revision;
792 const allocator = testing.allocator;
793 const module = try accountingIdentityModule(allocator);
794 defer module.deinit();
795 const selected = canonicalizationPass();
796 const declared = try selected.work_contract.?.estimate(.{ .operation = module.choir_module });
797 try testing.expect(declared.workspace > 0);
798 try testing.expect(declared.work.rewrite_attempts > 0);
799 var allowance = revision.WorkVector.uniform(std.math.maxInt(u64));
800 if (!exhaust_workspace) allowance.rewrite_attempts = declared.work.rewrite_attempts - 1;
801 const ledger = try revision.AccountingV1.create(allocator, .{
802 .allowance = allowance,
803 .workspace = declared.workspace - @intFromBool(exhaust_workspace),
804 .events = 16,
805 }, &.{.{ .name = canonicalization_pass_name, .version = 1 }});
806 defer ledger.destroy();
807 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
808 defer cache.deinit();
809 var manager = passes.PassManager.init(allocator);
810 defer manager.deinit();
811 try manager.addPass(selected);
812 try testing.expectEqual(.failure, manager.runWithAnalysisCache(
813 module.choir_module,
814 module.context(),
815 &cache,
816 .{ .max_threads = 1 },
817 ));
818 const receipt = ledger.view();
819 try testing.expectEqual(.exhausted, receipt.outcome);
820 try testing.expectEqual(0, receipt.executed.counters.pass_runs);
821 try testing.expectEqual(0, receipt.executed.counters.passes_modified);
822 try testing.expectEqual(0, receipt.executed.counters.successful_rewrites);
823 try testing.expectEqual(0, receipt.executed.counters.rewrite_iterations);
824 try testing.expectEqual(1, ir.inspection.countOperationsNamed(
825 module.choir_module,
826 AccyDialect.ReshapeOp.operation_name,
827 ));
828 try testing.expectError(error.WorkExhausted, ledger.producersComplete());
829 try module.verify();
830 }
831
832 test "canonicalization accounting covers generated dots scans and broadcast chains" {
833 const allocator = testing.allocator;
834 const dot = try accountingDotModule(allocator);
835 defer dot.deinit();
836 try checkCanonicalizationAccounting(dot, 1);
837 const scan = try accountingScanModule(allocator);
838 defer scan.deinit();
839 try checkCanonicalizationAccounting(scan, 1);
840 for ([_]u32{ 0, 1, 4, 16 }) |depth| {
841 const chain = try accountingBroadcastModule(allocator, depth);
842 defer chain.deinit();
843 try checkCanonicalizationAccounting(chain, if (depth == 0) 0 else 2 * depth - 1);
844 try testing.expectEqual(@intFromBool(depth > 0), ir.inspection.countOperationsNamed(
845 chain.choir_module,
846 AccyDialect.BroadcastInDimOp.operation_name,
847 ));
848 }
849 }
850
851 fn accountingDotModule(allocator: std.mem.Allocator) !*semantic.SemanticModule {
852 var builder = try semantic.Builder.init(allocator, .testing);
853 defer builder.deinit();
854 const lhs = try builder.tensor(.f32, &.{ 8, 2, 4 });
855 const rhs = try builder.tensor(.f32, &.{ 4, 3 });
856 const output = try builder.tensor(.f32, &.{ 8, 2, 3 });
857 var function = try builder.beginFunction("accounted_dot", &.{ lhs, rhs }, &.{output});
858 const value = try function.dotGeneral(
859 function.parameter(0),
860 function.parameter(1),
861 output,
862 &.{2},
863 &.{0},
864 &.{},
865 &.{},
866 );
867 try function.return_(&.{value});
868 try function.finish();
869 return builder.finish();
870 }
871
872 fn accountingScanModule(allocator: std.mem.Allocator) !*semantic.SemanticModule {
873 var builder = try semantic.Builder.init(allocator, .testing);
874 defer builder.deinit();
875 const typ = try builder.tensor(.f32, &.{16384});
876 var function = try builder.beginFunction("accounted_scan", &.{typ}, &.{typ});
877 const value = try function.cumsum(function.parameter(0), typ, 0);
878 try function.return_(&.{value});
879 try function.finish();
880 return builder.finish();
881 }
882
883 fn accountingBroadcastModule(allocator: std.mem.Allocator, depth: u32) !*semantic.SemanticModule {
884 std.debug.assert(depth <= 16);
885 var builder = try semantic.Builder.init(allocator, .testing);
886 defer builder.deinit();
887 var dimensions: [17]i64 = @splat(2);
888 dimensions[depth] = 1;
889 const typ = try builder.tensor(.f32, &.{1});
890 const output = try builder.tensor(.f32, dimensions[0 .. depth + 1]);
891 var function = try builder.beginFunction("accounted_broadcast", &.{typ}, &.{output});
892 var value = function.parameter(0);
893 for (0..depth) |index| {
894 const result = try builder.tensor(.f32, dimensions[depth - index - 1 .. depth + 1]);
895 value = try function.broadcast(value, result, &.{2});
896 }
897 try function.return_(&.{value});
898 try function.finish();
899 return builder.finish();
900 }
901
902 test "canonicalization accounting refuses an unmodeled registered callback" {
903 const allocator = testing.allocator;
904 const module = try accountingIdentityModule(allocator);
905 defer module.deinit();
906 const extra = comptime [_]rewrite.RewritePattern{rewrite.RewritePattern.init(.{
907 .name = "unmodeled-reshape",
908 .root_op_name = AccyDialect.ReshapeOp.operation_name,
909 .benefit = 20,
910 }, unmodeledCanonicalization)};
911 const registered = comptime rewrite.DialectCanonicalizationInterface.entryFor("accy", &extra);
912 try module.context().registerDialectInterface("accy", registered);
913 const ledger = try choir.product.revision.AccountingV1.create(allocator, .{
914 .allowance = .uniform(std.math.maxInt(u64)),
915 .workspace = 128 * 1024 * 1024,
916 .events = 16,
917 }, &.{.{ .name = canonicalization_pass_name, .version = 1 }});
918 defer ledger.destroy();
919 var cache = try passes.AnalysisCache.initAccounted(allocator, null, ledger, .{}, 8);
920 defer cache.deinit();
921 var manager = passes.PassManager.init(allocator);
922 defer manager.deinit();
923 try manager.addPass(canonicalizationPass());
924 try testing.expectEqual(.failure, manager.runWithAnalysisCache(
925 module.choir_module,
926 module.context(),
927 &cache,
928 .{ .max_threads = 1 },
929 ));
930 try testing.expectEqual(.rejected, ledger.view().outcome);
931 try testing.expect(ledger.view().missing_work_contract);
932 try testing.expectEqual(0, ledger.view().executed.counters.pass_runs);
933 try testing.expectEqual(0, ledger.view().executed.counters.successful_rewrites);
934 try testing.expectEqual(1, ir.inspection.countOperationsNamed(
935 module.choir_module,
936 AccyDialect.ReshapeOp.operation_name,
937 ));
938 }
939
940 fn unmodeledCanonicalization(_: *ir.Operation, _: *rewrite.PatternRewriter) rewrite.PatternResult {
941 @panic("an unmodeled callback reached accounted work");
942 }