lib/accy/src/choir/dialect.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3 const alloc_arena = @import("alloc_arena");
4 const alloc_fixed = @import("alloc_fixed");
5 const builtin = @import("builtin");
6 const choir = @import("choir");
7 const ir = choir.ir;
8 const dialects_mod = choir.dialects;
9 const semantics = @import("semantics.zig");
10 const arith_mod = dialects_mod.arith;
11
12 const native_endian = builtin.cpu.arch.endian();
13
14 const effect_facts = ir.interfaces.effects;
15
16 pub const AccyChoirVerifyError = error{
17 UnknownAccyOperation,
18 ExpectedAccyTensorType,
19 MalformedAccyTensorType,
20 UnknownAccyDType,
21 MissingAttribute,
22 AttributeKindMismatch,
23 ResultTypeMismatch,
24 UnsupportedDType,
25 ConstantPayloadLengthMismatch,
26 InvalidKernelContract,
27 };
28
29 pub const AccyDialect = struct {
30 pub const name = "accy";
31 const op_templates = ir.dialects.operationTemplate.dialect(@This());
32 const op_attr = ir.dialects.attribute;
33 pub const spec = ir.dialects.dialectSpec(@This(), .{
34 .types = &.{ir.dialects.typeName(tensor_type_name)},
35 });
36
37 pub const IotaOp = struct {
38 op: *ir.Operation,
39
40 pub const leaf = op_templates.explicitLeaf(@This(), .{
41 .mnemonic = "iota",
42 .interfaces = &.{accyEffectsEntry()},
43 .operands = 0,
44 .results = .{"result"},
45
46 .required_attrs = .{op_attr.integer("iota_dimension")},
47 });
48 pub const operation_name = leaf.operation_name;
49 pub const getResult = leaf.getResult;
50 pub const verify = verifyAccyChoirOp;
51
52 pub fn create(
53 ctx: *ir.Context,
54 loc: ir.Location,
55 result_type: ir.Type,
56 iota_dimension: i64,
57 ) !IotaOp {
58 const self = try leaf.createLeaf(ctx, loc, &.{}, &.{result_type});
59 errdefer self.op.erase();
60 try leaf.setI64Attr(self, "iota_dimension", iota_dimension);
61 return self;
62 }
63 };
64
65 pub const ConstantOp = struct {
66 op: *ir.Operation,
67
68 pub const leaf = op_templates.explicitLeaf(@This(), .{
69 .mnemonic = "constant",
70 .interfaces = &.{accyEffectsEntry()},
71 .operands = 0,
72 .results = .{"result"},
73
74 .required_attrs = .{op_attr.dialect("payload", "accy.constant_payload")},
75 });
76 pub const operation_name = leaf.operation_name;
77 pub const getResult = leaf.getResult;
78 pub const verify = verifyAccyChoirOp;
79 pub const payload_attr_name = leaf.dialectAttrName("payload");
80
81 pub fn create(
82 ctx: *ir.Context,
83 loc: ir.Location,
84 payload: []const u8,
85 result_type: ir.Type,
86 ) !ConstantOp {
87 const self = try leaf.createLeaf(ctx, loc, &.{}, &.{result_type});
88 errdefer self.op.erase();
89 try leaf.setDialectAttrPayload(self, "payload", payload);
90 return self;
91 }
92
93 pub fn getPayload(self: ConstantOp) ?[]const u8 {
94 return leaf.getDialectAttrPayload(self, "payload");
95 }
96 };
97
98 pub const AddOp: type = op_templates.binarySameType("add", verifiedOptions());
99 pub const SubOp: type = op_templates.binarySameType("sub", verifiedOptions());
100 pub const MulOp: type = op_templates.binarySameType("mul", verifiedOptions());
101 pub const DivOp: type = op_templates.binarySameType("div", verifiedOptions());
102 pub const MaxOp: type = op_templates.binarySameType("max", verifiedOptions());
103 pub const MinOp: type = op_templates.binarySameType("min", verifiedOptions());
104 pub const NegOp: type = op_templates.unarySameType("neg", verifiedOptions());
105 pub const ExpOp: type = op_templates.unarySameType("exp", verifiedOptions());
106 pub const LogOp: type = op_templates.unarySameType("log", verifiedOptions());
107 pub const TanhOp: type = op_templates.unarySameType("tanh", verifiedOptions());
108 pub const SqrtOp: type = op_templates.unarySameType("sqrt", verifiedOptions());
109
110 pub const ActivationOp = struct {
111 op: *ir.Operation,
112
113 pub const leaf = op_templates.explicitLeaf(@This(), .{
114 .mnemonic = "activation",
115 .operands = .{"input"},
116 .results = .{"result"},
117
118 .required_attrs = .{op_attr.dialect("activation_kind", "accy.activation_kind")},
119 .interfaces = &.{ ir.dialects.opSpec.verifier(verifyAccyChoirOp), accyEffectsEntry() },
120 });
121 pub const operation_name = leaf.operation_name;
122 pub const getResult = leaf.getResult;
123 pub const verify = verifyAccyChoirOp;
124 pub const dialectAttrName = leaf.dialectAttrName;
125 pub const activation_kind_attr_name = leaf.dialectAttrName("activation_kind");
126
127 pub fn create(
128 ctx: *ir.Context,
129 loc: ir.Location,
130 input: *ir.Value,
131 result_type: ir.Type,
132 kind: semantics.ActivationKind,
133 ) !ActivationOp {
134 const self = try leaf.createLeaf(ctx, loc, &.{input}, &.{result_type});
135 errdefer self.op.erase();
136 try leaf.setDialectAttrPayload(self, "activation_kind", @tagName(kind));
137 return self;
138 }
139
140 pub fn getKind(self: ActivationOp) ?[]const u8 {
141 return leaf.getDialectAttrPayload(self, "activation_kind");
142 }
143 };
144
145 pub const AbsOp: type = op_templates.unarySameType("abs", verifiedOptions());
146 pub const SinOp: type = op_templates.unarySameType("sin", verifiedOptions());
147 pub const CosOp: type = op_templates.unarySameType("cos", verifiedOptions());
148 pub const TanOp: type = op_templates.unarySameType("tan", verifiedOptions());
149 pub const FloorOp: type = op_templates.unarySameType("floor", verifiedOptions());
150 pub const RoundOp: type = op_templates.unarySameType("round", verifiedOptions());
151 pub const TruncOp: type = op_templates.unarySameType("trunc", verifiedOptions());
152 pub const PowOp: type = op_templates.binarySameType("pow", verifiedOptions());
153 pub const Atan2Op: type = op_templates.binarySameType("atan2", verifiedOptions());
154
155 pub const CompareOp = struct {
156 op: *ir.Operation,
157
158 pub const leaf = op_templates.explicitLeaf(@This(), .{
159 .mnemonic = "compare",
160 .interfaces = &.{accyEffectsEntry()},
161 .operands = .{ "lhs", "rhs" },
162 .results = .{"result"},
163
164 .required_attrs = .{op_attr.dialect("compare_direction", "accy.compare_direction")},
165 });
166 pub const operation_name = leaf.operation_name;
167 pub const getResult = leaf.getResult;
168 pub const verify = verifyAccyChoirOp;
169 pub const dialectAttrName = leaf.dialectAttrName;
170
171 pub fn create(
172 ctx: *ir.Context,
173 loc: ir.Location,
174 lhs: *ir.Value,
175 rhs: *ir.Value,
176 result_type: ir.Type,
177 direction: []const u8,
178 ) !CompareOp {
179 const self = try leaf.createLeaf(ctx, loc, &.{ lhs, rhs }, &.{result_type});
180 errdefer self.op.erase();
181 try leaf.setDialectAttrPayload(self, "compare_direction", direction);
182 return self;
183 }
184 };
185
186 pub const ConvertOp = struct {
187 op: *ir.Operation,
188
189 pub const leaf = op_templates.explicitLeaf(@This(), .{
190 .mnemonic = "convert",
191 .interfaces = &.{accyEffectsEntry()},
192 .operands = .{"input"},
193 .results = .{"result"},
194
195 .required_attrs = .{op_attr.dialect("convert_to", "accy.convert_to")},
196 });
197 pub const operation_name = leaf.operation_name;
198 pub const getResult = leaf.getResult;
199 pub const verify = verifyAccyChoirOp;
200 pub const dialectAttrName = leaf.dialectAttrName;
201
202 pub fn create(
203 ctx: *ir.Context,
204 loc: ir.Location,
205 input: *ir.Value,
206 result_type: ir.Type,
207 target_dtype: []const u8,
208 ) !ConvertOp {
209 const self = try leaf.createLeaf(ctx, loc, &.{input}, &.{result_type});
210 errdefer self.op.erase();
211 try leaf.setDialectAttrPayload(self, "convert_to", target_dtype);
212 return self;
213 }
214 };
215
216 pub const SelectOp: type = op_templates.selectSameType("select", verifiedOptions());
217
218 pub const ReduceOp = struct {
219 op: *ir.Operation,
220
221 pub const leaf = op_templates.explicitLeaf(@This(), .{
222 .mnemonic = "reduce",
223 .interfaces = &.{accyEffectsEntry()},
224 .operands = .{ "operand", "init" },
225 .results = .{"result"},
226
227 .required_attrs = .{
228 op_attr.dialect("dimensions", "accy.reduce_dimensions"),
229 op_attr.dialect("reducer_kind", "accy.reducer_kind"),
230 },
231 });
232 pub const operation_name = leaf.operation_name;
233 pub const getResult = leaf.getResult;
234 pub const verify = verifyAccyChoirOp;
235 pub const dialectAttrName = leaf.dialectAttrName;
236
237 pub fn create(
238 ctx: *ir.Context,
239 loc: ir.Location,
240 operand: *ir.Value,
241 init: *ir.Value,
242 result_type: ir.Type,
243 reducer_kind: []const u8,
244 dimensions: []const i64,
245 ) !ReduceOp {
246 const self = try leaf.createLeaf(ctx, loc, &.{ operand, init }, &.{result_type});
247 errdefer self.op.erase();
248 try leaf.setDialectAttrPayload(self, "reducer_kind", reducer_kind);
249 try leaf.setDialectAttrPayload(self, "dimensions", std.mem.sliceAsBytes(dimensions));
250 return self;
251 }
252 };
253
254 pub const DotGeneralOp = struct {
255 op: *ir.Operation,
256
257 pub const leaf = op_templates.explicitLeaf(@This(), .{
258 .mnemonic = "dot_general",
259 .interfaces = &.{accyEffectsEntry()},
260 .operands = .{ "lhs", "rhs" },
261 .results = .{"result"},
262
263 .required_attrs = .{
264 op_attr.dialect("lhs_batch", "accy.dot_lhs_batch"),
265 op_attr.dialect("lhs_contract", "accy.dot_lhs_contract"),
266 op_attr.dialect("rhs_batch", "accy.dot_rhs_batch"),
267 op_attr.dialect("rhs_contract", "accy.dot_rhs_contract"),
268 },
269 });
270 pub const operation_name = leaf.operation_name;
271 pub const getResult = leaf.getResult;
272 pub const verify = verifyAccyChoirOp;
273 pub const dialectAttrName = leaf.dialectAttrName;
274
275 pub fn create(
276 ctx: *ir.Context,
277 loc: ir.Location,
278 lhs: *ir.Value,
279 rhs: *ir.Value,
280 result_type: ir.Type,
281 lhs_batch: []const i64,
282 rhs_batch: []const i64,
283 lhs_contract: []const i64,
284 rhs_contract: []const i64,
285 ) !DotGeneralOp {
286 const self = try leaf.createLeaf(ctx, loc, &.{ lhs, rhs }, &.{result_type});
287 errdefer self.op.erase();
288 try leaf.setDialectAttrPayload(self, "lhs_batch", std.mem.sliceAsBytes(lhs_batch));
289 try leaf.setDialectAttrPayload(self, "rhs_batch", std.mem.sliceAsBytes(rhs_batch));
290 try leaf.setDialectAttrPayload(self, "lhs_contract", std.mem.sliceAsBytes(lhs_contract));
291 try leaf.setDialectAttrPayload(self, "rhs_contract", std.mem.sliceAsBytes(rhs_contract));
292 return self;
293 }
294
295 pub fn getLhsBatchPayload(self: DotGeneralOp) ?[]const u8 {
296 return leaf.getDialectAttrPayload(self, "lhs_batch");
297 }
298
299 pub fn getRhsBatchPayload(self: DotGeneralOp) ?[]const u8 {
300 return leaf.getDialectAttrPayload(self, "rhs_batch");
301 }
302
303 pub fn getLhsContractPayload(self: DotGeneralOp) ?[]const u8 {
304 return leaf.getDialectAttrPayload(self, "lhs_contract");
305 }
306
307 pub fn getRhsContractPayload(self: DotGeneralOp) ?[]const u8 {
308 return leaf.getDialectAttrPayload(self, "rhs_contract");
309 }
310 };
311
312 pub const EinsumOp = struct {
313 op: *ir.Operation,
314
315 pub const leaf = op_templates.explicitLeaf(@This(), .{
316 .mnemonic = "einsum",
317 .operands = ir.dialects.shape.atLeast(1),
318 .results = .{"result"},
319
320 .required_attrs = .{op_attr.dialect("equation", "accy.einsum_equation")},
321 .interfaces = &.{ ir.dialects.opSpec.verifier(verifyAccyChoirOp), accyEffectsEntry() },
322 });
323 pub const operation_name = leaf.operation_name;
324 pub const getResult = leaf.getResult;
325 pub const verify = verifyAccyChoirOp;
326 pub const dialectAttrName = leaf.dialectAttrName;
327 pub const equation_attr_name = leaf.dialectAttrName("equation");
328
329 pub fn create(
330 ctx: *ir.Context,
331 loc: ir.Location,
332 operands: []const *ir.Value,
333 result_type: ir.Type,
334 equation: []const u8,
335 ) !EinsumOp {
336 const self = try leaf.createLeaf(ctx, loc, operands, &.{result_type});
337 errdefer self.op.erase();
338 try leaf.setDialectAttrPayload(self, "equation", equation);
339 return self;
340 }
341
342 pub fn getEquation(self: EinsumOp) ?[]const u8 {
343 return leaf.getDialectAttrPayload(self, "equation");
344 }
345 };
346
347 pub const BroadcastOp = struct {
348 op: *ir.Operation,
349
350 pub const leaf = op_templates.explicitLeaf(@This(), .{
351 .mnemonic = "broadcast",
352 .interfaces = &.{accyEffectsEntry()},
353 .operands = .{"input"},
354 .results = .{"result"},
355
356 .required_attrs = .{op_attr.dialect("sizes", "accy.broadcast_sizes")},
357 });
358 pub const operation_name = leaf.operation_name;
359 pub const getResult = leaf.getResult;
360 pub const verify = verifyAccyChoirOp;
361 pub const dialectAttrName = leaf.dialectAttrName;
362
363 pub fn create(
364 ctx: *ir.Context,
365 loc: ir.Location,
366 input: *ir.Value,
367 result_type: ir.Type,
368 sizes: []const i64,
369 ) !BroadcastOp {
370 const self = try leaf.createLeaf(ctx, loc, &.{input}, &.{result_type});
371 errdefer self.op.erase();
372 try leaf.setDialectAttrPayload(self, "sizes", std.mem.sliceAsBytes(sizes));
373 return self;
374 }
375 };
376
377 pub const BroadcastInDimOp = struct {
378 op: *ir.Operation,
379
380 pub const leaf = op_templates.explicitLeaf(@This(), .{
381 .mnemonic = "broadcast_in_dim",
382 .interfaces = &.{accyEffectsEntry()},
383 .operands = .{"input"},
384 .results = .{"result"},
385
386 .required_attrs = .{
387 op_attr.dialect("broadcast_dims", "accy.broadcast_dims"),
388 op_attr.dialect("result_shape", "accy.broadcast_result_shape"),
389 },
390 });
391 pub const operation_name = leaf.operation_name;
392 pub const getResult = leaf.getResult;
393 pub const verify = verifyAccyChoirOp;
394 pub const dialectAttrName = leaf.dialectAttrName;
395
396 pub fn create(
397 ctx: *ir.Context,
398 loc: ir.Location,
399 input: *ir.Value,
400 result_type: ir.Type,
401 broadcast_dims: []const i64,
402 result_shape: []const i64,
403 ) !BroadcastInDimOp {
404 const self = try leaf.createLeaf(ctx, loc, &.{input}, &.{result_type});
405 errdefer self.op.erase();
406 try leaf.setDialectAttrPayload(self, "broadcast_dims", std.mem.sliceAsBytes(broadcast_dims));
407 try leaf.setDialectAttrPayload(self, "result_shape", std.mem.sliceAsBytes(result_shape));
408 return self;
409 }
410 };
411
412 pub const ReshapeOp = struct {
413 op: *ir.Operation,
414
415 pub const leaf = op_templates.explicitLeaf(@This(), .{
416 .mnemonic = "reshape",
417 .interfaces = &.{accyEffectsEntry()},
418 .operands = .{"input"},
419 .results = .{"result"},
420
421 .required_attrs = .{op_attr.dialect("new_shape", "accy.reshape_new_shape")},
422 });
423 pub const operation_name = leaf.operation_name;
424 pub const getResult = leaf.getResult;
425 pub const verify = verifyAccyChoirOp;
426 pub const dialectAttrName = leaf.dialectAttrName;
427
428 pub fn create(
429 ctx: *ir.Context,
430 loc: ir.Location,
431 input: *ir.Value,
432 result_type: ir.Type,
433 new_shape: []const i64,
434 ) !ReshapeOp {
435 const self = try leaf.createLeaf(ctx, loc, &.{input}, &.{result_type});
436 errdefer self.op.erase();
437 try leaf.setDialectAttrPayload(self, "new_shape", std.mem.sliceAsBytes(new_shape));
438 return self;
439 }
440 };
441
442 pub const TransposeOp = struct {
443 op: *ir.Operation,
444
445 pub const leaf = op_templates.explicitLeaf(@This(), .{
446 .mnemonic = "transpose",
447 .interfaces = &.{accyEffectsEntry()},
448 .operands = .{"input"},
449 .results = .{"result"},
450
451 .required_attrs = .{op_attr.dialect("permutation", "accy.transpose_permutation")},
452 });
453 pub const operation_name = leaf.operation_name;
454 pub const getResult = leaf.getResult;
455 pub const verify = verifyAccyChoirOp;
456 pub const dialectAttrName = leaf.dialectAttrName;
457
458 pub fn create(
459 ctx: *ir.Context,
460 loc: ir.Location,
461 input: *ir.Value,
462 result_type: ir.Type,
463 permutation: []const i64,
464 ) !TransposeOp {
465 const self = try leaf.createLeaf(ctx, loc, &.{input}, &.{result_type});
466 errdefer self.op.erase();
467 try leaf.setDialectAttrPayload(self, "permutation", std.mem.sliceAsBytes(permutation));
468 return self;
469 }
470 };
471
472 pub const SliceOp = struct {
473 op: *ir.Operation,
474
475 pub const leaf = op_templates.explicitLeaf(@This(), .{
476 .mnemonic = "slice",
477 .interfaces = &.{accyEffectsEntry()},
478 .operands = .{"input"},
479 .results = .{"result"},
480
481 .required_attrs = .{
482 op_attr.dialect("limits", "accy.slice_limits"),
483 op_attr.dialect("starts", "accy.slice_starts"),
484 op_attr.dialect("strides", "accy.slice_strides"),
485 },
486 });
487 pub const operation_name = leaf.operation_name;
488 pub const getResult = leaf.getResult;
489 pub const verify = verifyAccyChoirOp;
490 pub const dialectAttrName = leaf.dialectAttrName;
491
492 pub fn create(
493 ctx: *ir.Context,
494 loc: ir.Location,
495 input: *ir.Value,
496 result_type: ir.Type,
497 starts: []const i64,
498 limits: []const i64,
499 strides: []const i64,
500 ) !SliceOp {
501 const self = try leaf.createLeaf(ctx, loc, &.{input}, &.{result_type});
502 errdefer self.op.erase();
503 try leaf.setDialectAttrPayload(self, "starts", std.mem.sliceAsBytes(starts));
504 try leaf.setDialectAttrPayload(self, "limits", std.mem.sliceAsBytes(limits));
505 try leaf.setDialectAttrPayload(self, "strides", std.mem.sliceAsBytes(strides));
506 return self;
507 }
508 };
509
510 pub const GatherOp = struct {
511 op: *ir.Operation,
512
513 pub const leaf = op_templates.explicitLeaf(@This(), .{
514 .mnemonic = "gather",
515 .interfaces = &.{accyEffectsEntry()},
516 .operands = .{ "input", "indices" },
517 .results = .{"result"},
518
519 .required_attrs = .{op_attr.integer("axis")},
520 });
521 pub const operation_name = leaf.operation_name;
522 pub const getResult = leaf.getResult;
523 pub const verify = verifyAccyChoirOp;
524
525 pub fn create(
526 ctx: *ir.Context,
527 loc: ir.Location,
528 input: *ir.Value,
529 indices: *ir.Value,
530 result_type: ir.Type,
531 axis: i64,
532 ) !GatherOp {
533 const self = try leaf.createLeaf(ctx, loc, &.{ input, indices }, &.{result_type});
534 errdefer self.op.erase();
535 try leaf.setI64Attr(self, "axis", axis);
536 return self;
537 }
538
539 pub fn getAxis(self: GatherOp) ?i64 {
540 const attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "axis") orelse return null;
541 return attr.getValue();
542 }
543 };
544
545 pub const ScatterOp = struct {
546 op: *ir.Operation,
547
548 pub const leaf = op_templates.explicitLeaf(@This(), .{
549 .mnemonic = "scatter",
550 .interfaces = &.{accyEffectsEntry()},
551 .operands = .{ "input", "indices", "updates" },
552 .results = .{"result"},
553
554 .required_attrs = .{op_attr.integer("axis")},
555 });
556 pub const operation_name = leaf.operation_name;
557 pub const getResult = leaf.getResult;
558 pub const verify = verifyAccyChoirOp;
559
560 pub fn create(
561 ctx: *ir.Context,
562 loc: ir.Location,
563 input: *ir.Value,
564 indices: *ir.Value,
565 updates: *ir.Value,
566 result_type: ir.Type,
567 axis: i64,
568 ) !ScatterOp {
569 const self = try leaf.createLeaf(ctx, loc, &.{ input, indices, updates }, &.{result_type});
570 errdefer self.op.erase();
571 try leaf.setI64Attr(self, "axis", axis);
572 return self;
573 }
574
575 pub fn getAxis(self: ScatterOp) ?i64 {
576 const attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "axis") orelse return null;
577 return attr.getValue();
578 }
579 };
580
581 pub const ScatterAddOp = struct {
582 op: *ir.Operation,
583
584 pub const leaf = op_templates.explicitLeaf(@This(), .{
585 .mnemonic = "scatter_add",
586 .interfaces = &.{accyEffectsEntry()},
587 .operands = .{ "input", "indices", "updates" },
588 .results = .{"result"},
589
590 .required_attrs = .{op_attr.integer("axis")},
591 });
592 pub const operation_name = leaf.operation_name;
593 pub const getResult = leaf.getResult;
594 pub const verify = verifyAccyChoirOp;
595
596 pub fn create(
597 ctx: *ir.Context,
598 loc: ir.Location,
599 input: *ir.Value,
600 indices: *ir.Value,
601 updates: *ir.Value,
602 result_type: ir.Type,
603 axis: i64,
604 ) !ScatterAddOp {
605 const self = try leaf.createLeaf(ctx, loc, &.{ input, indices, updates }, &.{result_type});
606 errdefer self.op.erase();
607 try leaf.setI64Attr(self, "axis", axis);
608 return self;
609 }
610
611 pub fn getAxis(self: ScatterAddOp) ?i64 {
612 const attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "axis") orelse return null;
613 return attr.getValue();
614 }
615 };
616
617 pub const SparseCrossEntropyOp = struct {
618 op: *ir.Operation,
619
620 pub const leaf = op_templates.explicitLeaf(@This(), .{
621 .mnemonic = "sparse_cross_entropy",
622 .interfaces = &.{accyEffectsEntry()},
623 .operands = .{ "logits", "targets" },
624 .results = .{"result"},
625 });
626 pub const operation_name = leaf.operation_name;
627 pub const getResult = leaf.getResult;
628 pub const verify = verifyAccyChoirOp;
629
630 pub fn create(
631 ctx: *ir.Context,
632 loc: ir.Location,
633 logits: *ir.Value,
634 targets: *ir.Value,
635 result_type: ir.Type,
636 ) !SparseCrossEntropyOp {
637 return try leaf.createLeaf(ctx, loc, &.{ logits, targets }, &.{result_type});
638 }
639 };
640
641 pub const PadOp = struct {
642 op: *ir.Operation,
643
644 pub const leaf = op_templates.explicitLeaf(@This(), .{
645 .mnemonic = "pad",
646 .interfaces = &.{accyEffectsEntry()},
647 .operands = .{ "input", "padding_value" },
648 .results = .{"result"},
649
650 .required_attrs = .{
651 op_attr.dialect("edge_high", "accy.pad_edge_high"),
652 op_attr.dialect("edge_low", "accy.pad_edge_low"),
653 op_attr.dialect("interior", "accy.pad_interior"),
654 },
655 });
656 pub const operation_name = leaf.operation_name;
657 pub const getResult = leaf.getResult;
658 pub const verify = verifyAccyChoirOp;
659 pub const dialectAttrName = leaf.dialectAttrName;
660
661 pub fn create(
662 ctx: *ir.Context,
663 loc: ir.Location,
664 input: *ir.Value,
665 padding_value: *ir.Value,
666 result_type: ir.Type,
667 edge_low: []const i64,
668 edge_high: []const i64,
669 interior: []const i64,
670 ) !PadOp {
671 const self = try leaf.createLeaf(ctx, loc, &.{ input, padding_value }, &.{result_type});
672 errdefer self.op.erase();
673 try leaf.setDialectAttrPayload(self, "edge_low", std.mem.sliceAsBytes(edge_low));
674 try leaf.setDialectAttrPayload(self, "edge_high", std.mem.sliceAsBytes(edge_high));
675 try leaf.setDialectAttrPayload(self, "interior", std.mem.sliceAsBytes(interior));
676 return self;
677 }
678 };
679
680 pub const ConcatenateOp = struct {
681 op: *ir.Operation,
682
683 pub const leaf = op_templates.explicitLeaf(@This(), .{
684 .mnemonic = "concatenate",
685 .interfaces = &.{accyEffectsEntry()},
686 .operands = ir.dialects.shape.atLeast(1),
687 .results = .{"result"},
688
689 .required_attrs = .{op_attr.integer("dimension")},
690 });
691 pub const operation_name = leaf.operation_name;
692 pub const getResult = leaf.getResult;
693 pub const verify = verifyAccyChoirOp;
694
695 pub fn create(
696 ctx: *ir.Context,
697 loc: ir.Location,
698 operands: []const *ir.Value,
699 result_type: ir.Type,
700 dimension: i64,
701 ) !ConcatenateOp {
702 const self = try leaf.createLeaf(ctx, loc, operands, &.{result_type});
703 errdefer self.op.erase();
704 try leaf.setI64Attr(self, "dimension", dimension);
705 return self;
706 }
707 };
708
709 pub const KernelCallOp = struct {
710 op: *ir.Operation,
711
712 pub const leaf = op_templates.explicitLeaf(@This(), .{
713 .mnemonic = "kernel_call",
714 .results = ir.dialects.shape.atLeast(1),
715 .required_attrs = .{
716 op_attr.dialect("target", "accy.kernel_call_target"),
717 op_attr.integer("version"),
718 op_attr.boolean("has_side_effects"),
719 op_attr.dialect("operand_effects", "accy.kernel_call_operand_effects"),
720 op_attr.dialect("result_aliases", "accy.kernel_call_result_aliases"),
721 },
722 .interfaces = &.{
723 ir.dialects.opSpec.verifier(verifyAccyChoirOp),
724 effect_facts.EffectOpInterface.entryFor(.{
725 .capacity = .{ .entries = 2, .per_operand = 2, .per_result = 1 },
726 .enumerate = kernelCallEffects,
727 }),
728 },
729 });
730 pub const operation_name = leaf.operation_name;
731 pub const verify = verifyAccyChoirOp;
732 pub const target_attr_name = leaf.dialectAttrName("target");
733 pub const operand_effects_attr_name = leaf.dialectAttrName("operand_effects");
734 pub const result_aliases_attr_name = leaf.dialectAttrName("result_aliases");
735 pub const runtime_scalars_attr_name = "accy.kernel_call_runtime_scalars";
736 pub const dialectAttrName = leaf.dialectAttrName;
737
738 pub fn create(
739 ctx: *ir.Context,
740 loc: ir.Location,
741 operands: []const *ir.Value,
742 result_types: []const ir.Type,
743 target: []const u8,
744 version: u32,
745 has_side_effects: bool,
746 operand_effects: []const semantics.KernelOperandEffect,
747 result_aliases: []const ?usize,
748 ) !KernelCallOp {
749 const kernel_call = try leaf.createLeaf(ctx, loc, operands, result_types);
750 errdefer kernel_call.op.erase();
751 try leaf.setDialectAttrPayload(kernel_call, "target", target);
752 try leaf.setI64Attr(kernel_call, "version", @intCast(version));
753 try leaf.setBoolAttr(kernel_call, "has_side_effects", has_side_effects);
754 try leaf.setDialectAttrPayload(kernel_call, "operand_effects", std.mem.sliceAsBytes(operand_effects));
755 try setKernelCallResultAliasesAttr(ctx, kernel_call.op, result_aliases);
756 return kernel_call;
757 }
758
759 pub fn getResult(self: KernelCallOp, index: usize) ?*ir.Value {
760 return self.op.getResult(index);
761 }
762
763 pub fn getFirstResult(self: KernelCallOp) *ir.Value {
764 return self.op.getResult(0).?;
765 }
766 };
767
768 pub const KernelCallScalarKind = enum(u64) {
769 i32,
770 u32,
771 i64,
772 u64,
773 f32,
774 f64,
775 };
776
777 pub const KernelCallScalar = extern struct {
778 kind: KernelCallScalarKind,
779 bits: u64,
780 };
781
782 pub const max_kernel_call_runtime_scalars = 16;
783
784 pub const KernelCallRuntimeScalars = struct {
785 count: usize = 0,
786 items: [max_kernel_call_runtime_scalars]KernelCallScalar = undefined,
787
788 pub fn slice(self: *const KernelCallRuntimeScalars) []const KernelCallScalar {
789 return self.items[0..self.count];
790 }
791 };
792
793 pub const ScratchOp = struct {
794 op: *ir.Operation,
795
796 pub const leaf = op_templates.explicitLeaf(@This(), .{
797 .mnemonic = "scratch",
798 .interfaces = &.{accyEffectsEntry()},
799 .results = .{"result"},
800 .required_attrs = .{op_attr.integer("words")},
801 });
802 pub const operation_name = leaf.operation_name;
803 pub const getResult = leaf.getResult;
804 pub const verify = verifyAccyChoirOp;
805
806 pub fn create(
807 ctx: *ir.Context,
808 loc: ir.Location,
809 result_type: ir.Type,
810 words: i64,
811 ) !ScratchOp {
812 const self = try leaf.createLeaf(ctx, loc, &.{}, &.{result_type});
813 errdefer self.op.erase();
814 try leaf.setI64Attr(self, "words", words);
815 return self;
816 }
817 };
818
819 pub const CumsumOp = struct {
820 op: *ir.Operation,
821
822 pub const leaf = op_templates.explicitLeaf(@This(), .{
823 .mnemonic = "cumsum",
824 .interfaces = &.{accyEffectsEntry()},
825 .operands = ir.dialects.shape.between(1, 2),
826 .results = .{"result"},
827
828 .required_attrs = .{op_attr.integer("axis")},
829 });
830 pub const operation_name = leaf.operation_name;
831 pub const getResult = leaf.getResult;
832 pub const verify = verifyAccyChoirOp;
833
834 pub fn create(
835 ctx: *ir.Context,
836 loc: ir.Location,
837 operand: *ir.Value,
838 result_type: ir.Type,
839 axis: i64,
840 ) !CumsumOp {
841 const self = try leaf.createLeaf(ctx, loc, &.{operand}, &.{result_type});
842 errdefer self.op.erase();
843 try leaf.setI64Attr(self, "axis", axis);
844 return self;
845 }
846
847 pub fn createWithScratch(
848 ctx: *ir.Context,
849 loc: ir.Location,
850 operand: *ir.Value,
851 scratch: *ir.Value,
852 result_type: ir.Type,
853 axis: i64,
854 ) !CumsumOp {
855 const self = try leaf.createLeaf(ctx, loc, &.{ operand, scratch }, &.{result_type});
856 errdefer self.op.erase();
857 try leaf.setI64Attr(self, "axis", axis);
858 return self;
859 }
860 };
861
862 pub const IterateOp = struct {
863 op: *ir.Operation,
864
865 pub const template = op_templates.explicit(@This(), .{
866 .mnemonic = "iterate",
867 .interfaces = &.{accyEffectsEntry()},
868 .operands = ir.dialects.shape.atLeast(1),
869 .results = ir.dialects.shape.atLeast(1),
870 .regions = ir.dialects.shape.exactly(1),
871 .region_names = .{"body"},
872
873 .required_attrs = .{op_attr.integer("max_iters")},
874 });
875 pub const operation_name = template.operation_name;
876 pub const verify = verifyAccyChoirOp;
877 pub const getRegion = template.getRegion;
878
879 pub fn create(
880 ctx: *ir.Context,
881 loc: ir.Location,
882 carries: []const *ir.Value,
883 max_iters: i64,
884 ) !IterateOp {
885 var body = ir.context.initRegion(ctx);
886 defer body.deinit();
887 var body_builder = ir.OperationBuilder.init(ctx);
888 var carry_types_buffer: [8]ir.Type = undefined;
889 if (carries.len > carry_types_buffer.len) return error.OutOfMemory;
890 for (carries, 0..) |carry, index| {
891 carry_types_buffer[index] = carry.type;
892 }
893 _ = try body_builder.createBlockWithLoc(&body, carry_types_buffer[0..carries.len], loc);
894 var regions = [_]*ir.Region{&body};
895 const self = try template.createOperation(ctx, loc, carries, carry_types_buffer[0..carries.len], ®ions, &.{});
896 errdefer self.op.erase();
897 try template.setI64Attr(self, "max_iters", max_iters);
898 return self;
899 }
900
901 pub fn bodyBlock(self: IterateOp) ?*ir.Block {
902 const region = self.op.getRegion(0) orelse return null;
903 return region.getEntryBlock();
904 }
905
906 pub fn getMaxIters(self: IterateOp) ?i64 {
907 return template.getI64Attr(self, "max_iters");
908 }
909
910 pub fn getResult(self: IterateOp, index: usize) ?*ir.Value {
911 return self.op.getResult(index);
912 }
913 };
914
915 pub const IterateYieldOp = struct {
916 op: *ir.Operation,
917
918 pub const term = op_templates.explicitTerminator(@This(), .{
919 .mnemonic = "iterate_yield",
920 .interfaces = &.{accyEffectsEntry()},
921 });
922 pub const operation_name = term.operation_name;
923 pub const verify = verifyAccyChoirOp;
924
925 pub fn create(
926 ctx: *ir.Context,
927 loc: ir.Location,
928 predicate: *ir.Value,
929 carries: []const *ir.Value,
930 ) !IterateYieldOp {
931 var operands_buffer: [9]*ir.Value = undefined;
932 if (carries.len + 1 > operands_buffer.len) return error.OutOfMemory;
933 operands_buffer[0] = predicate;
934 for (carries, 0..) |carry, index| {
935 operands_buffer[index + 1] = carry;
936 }
937 return try term.createTerminator(ctx, loc, operands_buffer[0 .. carries.len + 1], &.{});
938 }
939 };
940
941 pub const ReturnOp = struct {
942 op: *ir.Operation,
943
944 pub const term = op_templates.explicitTerminator(@This(), .{
945 .mnemonic = "return",
946 });
947 pub const operation_name = term.operation_name;
948 pub const verify = verifyAccyChoirOp;
949
950 pub fn create(
951 ctx: *ir.Context,
952 loc: ir.Location,
953 operands: []const *ir.Value,
954 ) !ReturnOp {
955 return try term.createTerminator(ctx, loc, operands, &.{});
956 }
957 };
958
959 fn verifiedOptions() ir.dialects.opSpec.Options {
960 return .{
961 .interfaces = &.{ ir.dialects.opSpec.verifier(verifyAccyChoirOp), accyEffectsEntry() },
962 };
963 }
964
965 pub fn setKernelCallRuntimeScalars(
966 ctx: *ir.Context,
967 op: *ir.Operation,
968 scalars: []const KernelCallScalar,
969 ) !void {
970 if (scalars.len > max_kernel_call_runtime_scalars) return error.InvalidKernelCallContract;
971 try op.setAttr(
972 "runtime_scalars",
973 try ctx.getDialectAttr(KernelCallOp.runtime_scalars_attr_name, std.mem.sliceAsBytes(scalars)),
974 );
975 }
976
977 pub fn kernelCallRuntimeScalars(op: *const ir.Operation) !?KernelCallRuntimeScalars {
978 const attr = op.getAttr("runtime_scalars") orelse return null;
979 if (!std.mem.eql(u8, attr.abstract.name, KernelCallOp.runtime_scalars_attr_name)) {
980 return error.InvalidKernelCallContract;
981 }
982 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return error.InvalidKernelCallContract;
983 const payload = dialect_attr.payload;
984 const record_size = @sizeOf(KernelCallScalar);
985 if (payload.len % record_size != 0) return error.InvalidKernelCallContract;
986 const count = payload.len / record_size;
987 if (count > max_kernel_call_runtime_scalars) return error.InvalidKernelCallContract;
988
989 var decoded = KernelCallRuntimeScalars{ .count = count };
990 for (0..count) |index| {
991 const record = payload[index * record_size ..][0..record_size];
992 const kind_bits = std.mem.readInt(u64, record[0..8], native_endian);
993 const kind = std.enums.fromInt(KernelCallScalarKind, kind_bits) orelse {
994 return error.InvalidKernelCallContract;
995 };
996 decoded.items[index] = .{
997 .kind = kind,
998 .bits = std.mem.readInt(u64, record[8..16], native_endian),
999 };
1000 }
1001 return decoded;
1002 }
1003 };
1004
1005 fn setKernelCallResultAliasesAttr(ctx: *ir.Context, op: *ir.Operation, result_aliases: []const ?usize) !void {
1006 const allocator = ir.context.transientAllocator(ctx);
1007 const aliases = try allocator.alloc(i64, result_aliases.len);
1008 defer allocator.free(aliases);
1009 for (result_aliases, 0..) |alias, i| {
1010 aliases[i] = if (alias) |operand_index| @intCast(operand_index) else -1;
1011 }
1012 try op.setAttr(
1013 "result_aliases",
1014 try ctx.getDialectAttr(AccyDialect.KernelCallOp.result_aliases_attr_name, std.mem.sliceAsBytes(aliases)),
1015 );
1016 }
1017
1018 fn verifyAccyChoirOp(op_ptr: *const anyopaque) anyerror!void {
1019 const op: *ir.Operation = @ptrCast(@alignCast(@constCast(op_ptr)));
1020 const kind = accyKindFromChoirName(op.name.name) orelse return AccyChoirVerifyError.UnknownAccyOperation;
1021
1022 var stack_buffer: [8192]u8 = undefined;
1023 var stack_fallback = alloc_fixed.Fallback.init(
1024 &stack_buffer,
1025 ir.context.transientAllocator(op.context),
1026 );
1027 var arena_state = alloc_arena.Arena.init(stack_fallback.allocator());
1028 defer arena_state.deinit();
1029 const arena = arena_state.allocator();
1030
1031 const input_types = try arena.alloc(semantics.Type, op.getNumOperands());
1032 for (op.getOperandValues(), 0..) |operand, i| {
1033 input_types[i] = try decodeTensorType(arena, operand.type);
1034 }
1035
1036 const result_types = try arena.alloc(semantics.Type, op.getNumResults());
1037 for (op.getResultTypes(), 0..) |typ, i| {
1038 result_types[i] = try decodeTensorType(arena, typ);
1039 }
1040
1041 const attrs = try decodeAttrsForAccyOp(arena, op, kind, result_types);
1042 const inferred = try semantics.inferShape(kind, arena, input_types, attrs);
1043 try verifyDTypeLegality(kind, input_types, inferred);
1044
1045 if (inferred.len != result_types.len) return AccyChoirVerifyError.ResultTypeMismatch;
1046 for (inferred, result_types) |expected, actual| {
1047 if (!expected.eql(actual)) return AccyChoirVerifyError.ResultTypeMismatch;
1048 }
1049
1050 if (kind == .constant and result_types.len == 1 and attrs.len >= 1) {
1051 try verifyConstantPayloadLength(result_types[0], attrs[0].bytes.len);
1052 }
1053
1054 if (kind == .iterate) {
1055 try verifyIterateRegion(op);
1056 }
1057 }
1058
1059 fn verifyIterateRegion(op: *ir.Operation) anyerror!void {
1060 if (op.regions.items.len != 1) return AccyChoirVerifyError.InvalidKernelContract;
1061 const iterate = AccyDialect.IterateOp{ .op = op };
1062 const block = iterate.bodyBlock() orelse return AccyChoirVerifyError.InvalidKernelContract;
1063 const carry_count = op.getNumOperands();
1064 if (block.arguments.items.len != carry_count) return AccyChoirVerifyError.InvalidKernelContract;
1065 for (block.arguments.items, op.getOperandValues()) |arg, operand| {
1066 if (!arg.type.eql(operand.type)) return AccyChoirVerifyError.ResultTypeMismatch;
1067 }
1068 const terminator_any = block.operations.tail orelse return AccyChoirVerifyError.InvalidKernelContract;
1069 const terminator: *ir.Operation = @ptrCast(@alignCast(terminator_any));
1070 if (!std.mem.eql(u8, terminator.name.name, AccyDialect.IterateYieldOp.operation_name)) {
1071 return AccyChoirVerifyError.InvalidKernelContract;
1072 }
1073 if (terminator.getNumOperands() != carry_count + 1) return AccyChoirVerifyError.InvalidKernelContract;
1074 const yields = terminator.getOperandValues();
1075 for (yields[1..], op.getOperandValues()) |yielded, operand| {
1076 if (!yielded.type.eql(operand.type)) return AccyChoirVerifyError.ResultTypeMismatch;
1077 }
1078 }
1079
1080 fn accyKindFromChoirName(choir_name: []const u8) ?semantics.OpKind {
1081 const prefix = "accy.";
1082 if (!std.mem.startsWith(u8, choir_name, prefix)) return null;
1083 const local_name = choir_name[prefix.len..];
1084 inline for (
1085 @typeInfo(semantics.OpKind).@"enum".field_names,
1086 @typeInfo(semantics.OpKind).@"enum".field_values,
1087 ) |field_name, field_name_value| {
1088 const field = .{ .name = field_name, .value = field_name_value };
1089 const kind: semantics.OpKind = @fromBackingInt(@intCast(field.value));
1090 if (std.mem.eql(u8, local_name, semantics.info(kind).name)) return kind;
1091 }
1092 return null;
1093 }
1094
1095 pub fn decodeTensorType(arena: std.mem.Allocator, typ: ir.Type) !semantics.Type {
1096 const type_name = typ.getDialectTypeName() orelse return AccyChoirVerifyError.ExpectedAccyTensorType;
1097 if (!std.mem.eql(u8, type_name, tensor_type_name)) return AccyChoirVerifyError.ExpectedAccyTensorType;
1098 const key = typ.getDialectParamKey() orelse return AccyChoirVerifyError.MalformedAccyTensorType;
1099 const comma = std.mem.indexOfScalar(u8, key, ',') orelse return AccyChoirVerifyError.MalformedAccyTensorType;
1100 const dtype_name = key[0..comma];
1101 const dtype = choir_abi.DType.fromName(dtype_name) orelse return AccyChoirVerifyError.UnknownAccyDType;
1102 const dims_text = key[comma + 1 ..];
1103 if (dims_text.len == 0) {
1104 return .{ .dtype = dtype, .dims = try arena.alloc(i64, 0) };
1105 }
1106
1107 var dim_count: usize = 1;
1108 for (dims_text) |ch| {
1109 if (ch == 'x') dim_count += 1;
1110 }
1111 const dims = try arena.alloc(i64, dim_count);
1112 var iter = std.mem.splitScalar(u8, dims_text, 'x');
1113 var index: usize = 0;
1114 while (iter.next()) |part| {
1115 if (part.len == 0) return AccyChoirVerifyError.MalformedAccyTensorType;
1116 const dim = std.fmt.parseInt(i64, part, 10) catch return AccyChoirVerifyError.MalformedAccyTensorType;
1117 if (dim < 0) return AccyChoirVerifyError.MalformedAccyTensorType;
1118 dims[index] = dim;
1119 index += 1;
1120 }
1121 if (index != dim_count) return AccyChoirVerifyError.MalformedAccyTensorType;
1122 return .{ .dtype = dtype, .dims = dims };
1123 }
1124
1125 fn decodeAttrsForAccyOp(
1126 arena: std.mem.Allocator,
1127 op: *ir.Operation,
1128 kind: semantics.OpKind,
1129 result_types: []const semantics.Type,
1130 ) ![]const semantics.Attribute {
1131 switch (kind) {
1132 .constant => {
1133 const result_type = try requireSingleResultType(result_types);
1134 const attrs = try arena.alloc(semantics.Attribute, 3);
1135 attrs[0] = .{ .bytes = try requireDialectPayload(op, "payload", AccyDialect.ConstantOp.payload_attr_name) };
1136 attrs[1] = .{ .dtype = result_type.dtype };
1137 attrs[2] = .{ .i64_list = result_type.dims };
1138 return attrs;
1139 },
1140 .iota => {
1141 const result_type = try requireSingleResultType(result_types);
1142 const attrs = try arena.alloc(semantics.Attribute, 3);
1143 attrs[0] = .{ .i64 = try requireIntegerAttr(op, "iota_dimension") };
1144 attrs[1] = .{ .dtype = result_type.dtype };
1145 attrs[2] = .{ .i64_list = result_type.dims };
1146 return attrs;
1147 },
1148 .compare => {
1149 const attrs = try arena.alloc(semantics.Attribute, 1);
1150 const payload = try requireDialectPayload(op, "compare_direction", AccyDialect.CompareOp.dialectAttrName("compare_direction"));
1151 attrs[0] = .{ .compare_direction = try parseEnumTag(semantics.CompareDirection, payload) };
1152 return attrs;
1153 },
1154 .activation => {
1155 const attrs = try arena.alloc(semantics.Attribute, 1);
1156 const payload = try requireDialectPayload(op, "activation_kind", AccyDialect.ActivationOp.activation_kind_attr_name);
1157 attrs[0] = .{ .activation_kind = try parseEnumTag(semantics.ActivationKind, payload) };
1158 return attrs;
1159 },
1160 .convert => {
1161 const attrs = try arena.alloc(semantics.Attribute, 1);
1162 const payload = try requireDialectPayload(op, "convert_to", AccyDialect.ConvertOp.dialectAttrName("convert_to"));
1163 attrs[0] = .{ .dtype = choir_abi.DType.fromName(payload) orelse return AccyChoirVerifyError.UnknownAccyDType };
1164 return attrs;
1165 },
1166 .reduce => {
1167 const attrs = try arena.alloc(semantics.Attribute, 2);
1168 const reducer_payload = try requireDialectPayload(op, "reducer_kind", AccyDialect.ReduceOp.dialectAttrName("reducer_kind"));
1169 attrs[0] = .{ .reducer_kind = try parseEnumTag(semantics.ReducerKind, reducer_payload) };
1170 attrs[1] = .{ .i64_list = try requireI64ListAttr(arena, op, "dimensions", AccyDialect.ReduceOp.dialectAttrName("dimensions")) };
1171 return attrs;
1172 },
1173 .dot_general => {
1174 const result_type = try requireSingleResultType(result_types);
1175 const attrs = try arena.alloc(semantics.Attribute, 5);
1176 attrs[0] = .{ .i64_list = try requireI64ListAttr(arena, op, "lhs_batch", AccyDialect.DotGeneralOp.dialectAttrName("lhs_batch")) };
1177 attrs[1] = .{ .i64_list = try requireI64ListAttr(arena, op, "rhs_batch", AccyDialect.DotGeneralOp.dialectAttrName("rhs_batch")) };
1178 attrs[2] = .{ .i64_list = try requireI64ListAttr(arena, op, "lhs_contract", AccyDialect.DotGeneralOp.dialectAttrName("lhs_contract")) };
1179 attrs[3] = .{ .i64_list = try requireI64ListAttr(arena, op, "rhs_contract", AccyDialect.DotGeneralOp.dialectAttrName("rhs_contract")) };
1180 attrs[4] = .{ .dtype = result_type.dtype };
1181 return attrs;
1182 },
1183 .einsum => {
1184 const attrs = try arena.alloc(semantics.Attribute, 1);
1185 attrs[0] = .{ .einsum = try requireDialectPayload(op, "equation", AccyDialect.EinsumOp.equation_attr_name) };
1186 return attrs;
1187 },
1188 .broadcast => {
1189 const attrs = try arena.alloc(semantics.Attribute, 1);
1190 attrs[0] = .{ .i64_list = try requireI64ListAttr(arena, op, "sizes", AccyDialect.BroadcastOp.dialectAttrName("sizes")) };
1191 return attrs;
1192 },
1193 .broadcast_in_dim => {
1194 const attrs = try arena.alloc(semantics.Attribute, 2);
1195 attrs[0] = .{ .i64_list = try requireI64ListAttr(arena, op, "broadcast_dims", AccyDialect.BroadcastInDimOp.dialectAttrName("broadcast_dims")) };
1196 attrs[1] = .{ .i64_list = try requireI64ListAttr(arena, op, "result_shape", AccyDialect.BroadcastInDimOp.dialectAttrName("result_shape")) };
1197 return attrs;
1198 },
1199 .reshape => {
1200 const attrs = try arena.alloc(semantics.Attribute, 1);
1201 attrs[0] = .{ .i64_list = try requireI64ListAttr(arena, op, "new_shape", AccyDialect.ReshapeOp.dialectAttrName("new_shape")) };
1202 return attrs;
1203 },
1204 .transpose => {
1205 const attrs = try arena.alloc(semantics.Attribute, 1);
1206 attrs[0] = .{ .i64_list = try requireI64ListAttr(arena, op, "permutation", AccyDialect.TransposeOp.dialectAttrName("permutation")) };
1207 return attrs;
1208 },
1209 .slice => {
1210 const attrs = try arena.alloc(semantics.Attribute, 3);
1211 attrs[0] = .{ .i64_list = try requireI64ListAttr(arena, op, "starts", AccyDialect.SliceOp.dialectAttrName("starts")) };
1212 attrs[1] = .{ .i64_list = try requireI64ListAttr(arena, op, "limits", AccyDialect.SliceOp.dialectAttrName("limits")) };
1213 attrs[2] = .{ .i64_list = try requireI64ListAttr(arena, op, "strides", AccyDialect.SliceOp.dialectAttrName("strides")) };
1214 return attrs;
1215 },
1216 .gather => {
1217 const attrs = try arena.alloc(semantics.Attribute, 1);
1218 attrs[0] = .{ .i64 = try requireIntegerAttr(op, "axis") };
1219 return attrs;
1220 },
1221 .iterate => {
1222 const attrs = try arena.alloc(semantics.Attribute, 1);
1223 attrs[0] = .{ .i64 = try requireIntegerAttr(op, "max_iters") };
1224 return attrs;
1225 },
1226 .cumsum => {
1227 const attrs = try arena.alloc(semantics.Attribute, 1);
1228 attrs[0] = .{ .i64 = try requireIntegerAttr(op, "axis") };
1229 return attrs;
1230 },
1231 .scratch => {
1232 const attrs = try arena.alloc(semantics.Attribute, 1);
1233 attrs[0] = .{ .i64 = try requireIntegerAttr(op, "words") };
1234 return attrs;
1235 },
1236 .scatter, .scatter_add => {
1237 const attrs = try arena.alloc(semantics.Attribute, 1);
1238 attrs[0] = .{ .i64 = try requireIntegerAttr(op, "axis") };
1239 return attrs;
1240 },
1241 .pad => {
1242 const attrs = try arena.alloc(semantics.Attribute, 3);
1243 attrs[0] = .{ .i64_list = try requireI64ListAttr(arena, op, "edge_low", AccyDialect.PadOp.dialectAttrName("edge_low")) };
1244 attrs[1] = .{ .i64_list = try requireI64ListAttr(arena, op, "edge_high", AccyDialect.PadOp.dialectAttrName("edge_high")) };
1245 attrs[2] = .{ .i64_list = try requireI64ListAttr(arena, op, "interior", AccyDialect.PadOp.dialectAttrName("interior")) };
1246 return attrs;
1247 },
1248 .concatenate => {
1249 const attrs = try arena.alloc(semantics.Attribute, 1);
1250 attrs[0] = .{ .i64 = try requireIntegerAttr(op, "dimension") };
1251 return attrs;
1252 },
1253 .kernel_call => {
1254 const version_i64 = try requireIntegerAttr(op, "version");
1255 if (version_i64 < 0 or version_i64 > std.math.maxInt(u32)) return AccyChoirVerifyError.AttributeKindMismatch;
1256 const attrs = try arena.alloc(semantics.Attribute, 1);
1257 attrs[0] = .{ .kernel_call = .{
1258 .target = try requireDialectPayload(op, "target", AccyDialect.KernelCallOp.target_attr_name),
1259 .version = @intCast(version_i64),
1260 .has_side_effects = try requireBoolAttr(op, "has_side_effects"),
1261 .operand_effects = try requireKernelOperandEffectsAttr(arena, op, op.getNumOperands()),
1262 .result_aliases = try requireKernelResultAliasesAttr(arena, op, op.getNumResults()),
1263 .results = result_types,
1264 } };
1265 return attrs;
1266 },
1267 .parameter,
1268 => return AccyChoirVerifyError.UnknownAccyOperation,
1269 .add,
1270 .sub,
1271 .mul,
1272 .div,
1273 .max,
1274 .min,
1275 .pow,
1276 .atan2,
1277 .neg,
1278 .exp,
1279 .log,
1280 .tanh,
1281 .sqrt,
1282 .abs,
1283 .sin,
1284 .cos,
1285 .tan,
1286 .floor,
1287 .round,
1288 .trunc,
1289 .select,
1290 .sparse_cross_entropy,
1291 .iterate_yield,
1292 .@"return",
1293 => return arena.alloc(semantics.Attribute, 0),
1294 }
1295 }
1296
1297 fn requireSingleResultType(result_types: []const semantics.Type) AccyChoirVerifyError!semantics.Type {
1298 if (result_types.len != 1) return AccyChoirVerifyError.ResultTypeMismatch;
1299 return result_types[0];
1300 }
1301
1302 fn requireDialectPayload(op: *const ir.Operation, attr_name: []const u8, dialect_attr_name: []const u8) AccyChoirVerifyError![]const u8 {
1303 const attr = op.getAttr(attr_name) orelse return AccyChoirVerifyError.MissingAttribute;
1304 if (!std.mem.eql(u8, attr.abstract.name, dialect_attr_name)) return AccyChoirVerifyError.AttributeKindMismatch;
1305 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return AccyChoirVerifyError.AttributeKindMismatch;
1306 return dialect_attr.payload;
1307 }
1308
1309 fn requireIntegerAttr(op: *const ir.Operation, attr_name: []const u8) AccyChoirVerifyError!i64 {
1310 const attr = op.getAttr(attr_name) orelse return AccyChoirVerifyError.MissingAttribute;
1311 if (!std.mem.eql(u8, attr.abstract.name, "builtin.integer")) return AccyChoirVerifyError.AttributeKindMismatch;
1312 const int_attr = attr.cast(ir.Attribute.IntegerAttr) orelse return AccyChoirVerifyError.AttributeKindMismatch;
1313 return int_attr.getValue();
1314 }
1315
1316 fn requireBoolAttr(op: *const ir.Operation, attr_name: []const u8) AccyChoirVerifyError!bool {
1317 const attr = op.getAttr(attr_name) orelse return AccyChoirVerifyError.MissingAttribute;
1318 if (!std.mem.eql(u8, attr.abstract.name, "builtin.bool")) return AccyChoirVerifyError.AttributeKindMismatch;
1319 const bool_attr = attr.cast(ir.Attribute.BoolAttr) orelse return AccyChoirVerifyError.AttributeKindMismatch;
1320 return bool_attr.getValue();
1321 }
1322
1323 fn requireI64ListAttr(
1324 arena: std.mem.Allocator,
1325 op: *const ir.Operation,
1326 attr_name: []const u8,
1327 dialect_attr_name: []const u8,
1328 ) ![]const i64 {
1329 const payload = try requireDialectPayload(op, attr_name, dialect_attr_name);
1330 if (payload.len % @sizeOf(i64) != 0) return AccyChoirVerifyError.AttributeKindMismatch;
1331 const values = try arena.alloc(i64, payload.len / @sizeOf(i64));
1332 for (values, 0..) |*value, i| {
1333 const start = i * @sizeOf(i64);
1334 @memcpy(std.mem.asBytes(value), payload[start..][0..@sizeOf(i64)]);
1335 }
1336 return values;
1337 }
1338
1339 fn requireKernelOperandEffectsAttr(
1340 arena: std.mem.Allocator,
1341 op: *const ir.Operation,
1342 expected_len: usize,
1343 ) ![]const semantics.KernelOperandEffect {
1344 const payload = try requireDialectPayload(op, "operand_effects", AccyDialect.KernelCallOp.operand_effects_attr_name);
1345 if (payload.len != expected_len) return AccyChoirVerifyError.InvalidKernelContract;
1346 const effects = try arena.alloc(semantics.KernelOperandEffect, payload.len);
1347 for (payload, 0..) |byte, i| {
1348 effects[i] = semantics.KernelOperandEffect.fromByte(byte) orelse return AccyChoirVerifyError.InvalidKernelContract;
1349 }
1350 return effects;
1351 }
1352
1353 fn requireKernelResultAliasesAttr(
1354 arena: std.mem.Allocator,
1355 op: *const ir.Operation,
1356 expected_len: usize,
1357 ) ![]const ?usize {
1358 const values = try requireI64ListAttr(arena, op, "result_aliases", AccyDialect.KernelCallOp.result_aliases_attr_name);
1359 if (values.len != expected_len) return AccyChoirVerifyError.InvalidKernelContract;
1360 const aliases = try arena.alloc(?usize, values.len);
1361 for (values, 0..) |value, i| {
1362 aliases[i] = if (value == -1) null else blk: {
1363 if (value < 0) return AccyChoirVerifyError.InvalidKernelContract;
1364 break :blk @intCast(value);
1365 };
1366 }
1367 return aliases;
1368 }
1369
1370 fn parseEnumTag(comptime E: type, payload: []const u8) AccyChoirVerifyError!E {
1371 inline for (
1372 @typeInfo(E).@"enum".field_names,
1373 @typeInfo(E).@"enum".field_values,
1374 ) |field_name, field_name_value| {
1375 const field = .{ .name = field_name, .value = field_name_value };
1376 if (std.mem.eql(u8, payload, field.name)) return @fromBackingInt(@intCast(field.value));
1377 }
1378 return AccyChoirVerifyError.AttributeKindMismatch;
1379 }
1380
1381 fn verifyDTypeLegality(
1382 kind: semantics.OpKind,
1383 input_types: []const semantics.Type,
1384 inferred: []const semantics.Type,
1385 ) AccyChoirVerifyError!void {
1386 const op_info = semantics.info(kind);
1387 if (input_types.len > 0) {
1388 for (input_types) |typ| {
1389 if (!op_info.supported_dtypes.allows(typ.dtype)) return AccyChoirVerifyError.UnsupportedDType;
1390 }
1391 return;
1392 }
1393 for (inferred) |typ| {
1394 if (!op_info.supported_dtypes.allows(typ.dtype)) return AccyChoirVerifyError.UnsupportedDType;
1395 }
1396 }
1397
1398 fn verifyConstantPayloadLength(result_type: semantics.Type, payload_len: usize) AccyChoirVerifyError!void {
1399 var elements: usize = 1;
1400 for (result_type.dims) |dim| {
1401 if (dim < 0) return AccyChoirVerifyError.MalformedAccyTensorType;
1402 const dim_usize: usize = @intCast(dim);
1403 elements = std.math.mul(usize, elements, dim_usize) catch return AccyChoirVerifyError.ConstantPayloadLengthMismatch;
1404 }
1405 const expected = std.math.mul(usize, elements, result_type.dtype.sizeOf()) catch return AccyChoirVerifyError.ConstantPayloadLengthMismatch;
1406 if (payload_len != expected) return AccyChoirVerifyError.ConstantPayloadLengthMismatch;
1407 }
1408
1409 fn loadAccyDialect(ctx: *ir.Context) !void {
1410 try ir.dialects.loadDialectSpec(ctx, AccyDialect.spec);
1411 }
1412
1413 pub const tensor_type_name: []const u8 = "accy.tensor";
1414
1415 pub fn accyTensorType(
1416 ctx: *ir.Context,
1417 dt: choir_abi.DType,
1418 dims: []const i64,
1419 ) !ir.Type {
1420 var buf: std.ArrayListUnmanaged(u8) = .empty;
1421 const allocator = ir.context.transientAllocator(ctx);
1422 defer buf.deinit(allocator);
1423 try buf.appendSlice(allocator, dt.name());
1424 try buf.append(allocator, ',');
1425 var num_buf: [24]u8 = undefined;
1426 for (dims, 0..) |d, i| {
1427 if (i > 0) try buf.append(allocator, 'x');
1428 const num_str = try std.fmt.bufPrint(num_buf[0..], "{d}", .{d});
1429 try buf.appendSlice(allocator, num_str);
1430 }
1431 return ctx.getDialectTypeFromNameWithKey(tensor_type_name, buf.items);
1432 }
1433
1434 pub const accy_package_extension = choir.extensions.PackageExtension{
1435 .name = "accy-choir-dialect",
1436 .dialects = &.{.{
1437 .name = "accy",
1438 .load = loadAccyDialect,
1439 }},
1440 };
1441
1442 pub fn registerAccyDialect(ctx: *ir.Context) !void {
1443 try accy_package_extension.registerContext(ctx);
1444 }
1445
1446 const testing = std.testing;
1447
1448 const FactoryResourceCounts = struct {
1449 operations: usize,
1450
1451 fn capture(ctx: *const ir.Context) FactoryResourceCounts {
1452 return .{
1453 .operations = ctx.operationCount(),
1454 };
1455 }
1456
1457 fn expectEqual(self: FactoryResourceCounts, ctx: *const ir.Context) !void {
1458 try testing.expectEqual(self.operations, ctx.operationCount());
1459 }
1460 };
1461
1462 fn exerciseAttributedFactories(ctx: *ir.Context) !void {
1463 const loc = ir.Location.getUnknown();
1464 const tensor_type = try accyTensorType(ctx, .f32, &.{4});
1465 const scalar_type = try accyTensorType(ctx, .f32, &.{});
1466 const indices_type = try accyTensorType(ctx, .i32, &.{4});
1467 const predicate_type = try accyTensorType(ctx, .i1, &.{4});
1468
1469 const lhs = try AccyDialect.IotaOp.create(ctx, loc, tensor_type, 0);
1470 defer lhs.op.erase();
1471 const rhs = try AccyDialect.IotaOp.create(ctx, loc, tensor_type, 0);
1472 defer rhs.op.erase();
1473 const indices = try AccyDialect.IotaOp.create(ctx, loc, indices_type, 0);
1474 defer indices.op.erase();
1475 const predicate = try AccyDialect.IotaOp.create(ctx, loc, predicate_type, 0);
1476 defer predicate.op.erase();
1477
1478 const tensor_payload = [_]f32{ 0, 1, 2, 3 };
1479 const constant = try AccyDialect.ConstantOp.create(ctx, loc, std.mem.sliceAsBytes(&tensor_payload), tensor_type);
1480 defer constant.op.erase();
1481 const scalar_payload = [_]f32{0};
1482 const padding_value = try AccyDialect.ConstantOp.create(ctx, loc, std.mem.sliceAsBytes(&scalar_payload), scalar_type);
1483 defer padding_value.op.erase();
1484
1485 const activation = try AccyDialect.ActivationOp.create(ctx, loc, lhs.getResult(), tensor_type, .gelu);
1486 defer activation.op.erase();
1487 const compare = try AccyDialect.CompareOp.create(ctx, loc, lhs.getResult(), rhs.getResult(), predicate_type, "lt");
1488 defer compare.op.erase();
1489 const convert = try AccyDialect.ConvertOp.create(ctx, loc, lhs.getResult(), indices_type, "i32");
1490 defer convert.op.erase();
1491 const reduce = try AccyDialect.ReduceOp.create(ctx, loc, lhs.getResult(), padding_value.getResult(), scalar_type, "sum", &.{0});
1492 defer reduce.op.erase();
1493 const dot_general = try AccyDialect.DotGeneralOp.create(
1494 ctx,
1495 loc,
1496 lhs.getResult(),
1497 rhs.getResult(),
1498 tensor_type,
1499 &.{},
1500 &.{},
1501 &.{0},
1502 &.{0},
1503 );
1504 defer dot_general.op.erase();
1505 const einsum = try AccyDialect.EinsumOp.create(ctx, loc, &.{ lhs.getResult(), rhs.getResult() }, tensor_type, "i,i->i");
1506 defer einsum.op.erase();
1507 const broadcast = try AccyDialect.BroadcastOp.create(ctx, loc, padding_value.getResult(), tensor_type, &.{4});
1508 defer broadcast.op.erase();
1509 const broadcast_in_dim = try AccyDialect.BroadcastInDimOp.create(ctx, loc, lhs.getResult(), tensor_type, &.{0}, &.{4});
1510 defer broadcast_in_dim.op.erase();
1511 const reshape = try AccyDialect.ReshapeOp.create(ctx, loc, lhs.getResult(), tensor_type, &.{4});
1512 defer reshape.op.erase();
1513 const transpose = try AccyDialect.TransposeOp.create(ctx, loc, lhs.getResult(), tensor_type, &.{0});
1514 defer transpose.op.erase();
1515 const slice = try AccyDialect.SliceOp.create(ctx, loc, lhs.getResult(), tensor_type, &.{0}, &.{4}, &.{1});
1516 defer slice.op.erase();
1517 const gather = try AccyDialect.GatherOp.create(ctx, loc, lhs.getResult(), indices.getResult(), tensor_type, 0);
1518 defer gather.op.erase();
1519 const scatter = try AccyDialect.ScatterOp.create(ctx, loc, lhs.getResult(), indices.getResult(), rhs.getResult(), tensor_type, 0);
1520 defer scatter.op.erase();
1521 const scatter_add = try AccyDialect.ScatterAddOp.create(ctx, loc, lhs.getResult(), indices.getResult(), rhs.getResult(), tensor_type, 0);
1522 defer scatter_add.op.erase();
1523 const pad = try AccyDialect.PadOp.create(ctx, loc, lhs.getResult(), padding_value.getResult(), tensor_type, &.{1}, &.{1}, &.{0});
1524 defer pad.op.erase();
1525 const concatenate = try AccyDialect.ConcatenateOp.create(ctx, loc, &.{ lhs.getResult(), rhs.getResult() }, tensor_type, 0);
1526 defer concatenate.op.erase();
1527 const kernel_call = try AccyDialect.KernelCallOp.create(
1528 ctx,
1529 loc,
1530 &.{lhs.getResult()},
1531 &.{tensor_type},
1532 "allocation_failure_test",
1533 1,
1534 false,
1535 &.{.read},
1536 &.{null},
1537 );
1538 defer kernel_call.op.erase();
1539 const scratch = try AccyDialect.ScratchOp.create(ctx, loc, indices_type, 4);
1540 defer scratch.op.erase();
1541 const cumsum = try AccyDialect.CumsumOp.create(ctx, loc, lhs.getResult(), tensor_type, 0);
1542 defer cumsum.op.erase();
1543 const cumsum_with_scratch = try AccyDialect.CumsumOp.createWithScratch(
1544 ctx,
1545 loc,
1546 lhs.getResult(),
1547 scratch.getResult(),
1548 tensor_type,
1549 0,
1550 );
1551 defer cumsum_with_scratch.op.erase();
1552 const iterate = try AccyDialect.IterateOp.create(ctx, loc, &.{predicate.getResult()}, 4);
1553 defer iterate.op.erase();
1554 }
1555
1556 fn checkAttributedFactoryAllocationFailures(allocator: std.mem.Allocator) !void {
1557 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1558 defer ctx.deinit(allocator);
1559 try registerAccyDialect(&ctx);
1560 _ = try ctx.getOrLoadDialect("accy");
1561 const baseline = FactoryResourceCounts.capture(&ctx);
1562 exerciseAttributedFactories(&ctx) catch |err| {
1563 try baseline.expectEqual(&ctx);
1564 return err;
1565 };
1566 try baseline.expectEqual(&ctx);
1567 }
1568
1569 test "accy attributed operation factories clean every allocation failure" {
1570 try @import("../fixture/root.zig").checkAllAllocationFailures(
1571 checkAttributedFactoryAllocationFailures,
1572 .{},
1573 );
1574 }
1575
1576 test "accy dialect spec owns verifier interfaces" {
1577 try testing.expect(AccyDialect.spec.operations.len > 0);
1578 var add_shape = false;
1579 var neg_shape = false;
1580 var select_shape = false;
1581 var activation_shape = false;
1582 var einsum_shape = false;
1583 var kernel_call_shape = false;
1584 var kernel_call_memory_effects = false;
1585 for (AccyDialect.spec.operations) |op_spec| {
1586 var has_verify = false;
1587 for (op_spec.interfaces) |entry| {
1588 if (entry.id == ir.VerifyOpInterface.id) has_verify = true;
1589 }
1590 try testing.expect(has_verify);
1591
1592 if (std.mem.eql(u8, op_spec.name, AccyDialect.AddOp.operation_name)) {
1593 add_shape = op_spec.shape.operands.allows(2) and
1594 !op_spec.shape.operands.allows(1) and
1595 op_spec.shape.results.allows(1) and
1596 !op_spec.shape.results.allows(0) and
1597 op_spec.shape.regions.allows(0) and
1598 !op_spec.shape.regions.allows(1);
1599 }
1600 if (std.mem.eql(u8, op_spec.name, AccyDialect.NegOp.operation_name)) {
1601 neg_shape = op_spec.shape.operands.allows(1) and
1602 !op_spec.shape.operands.allows(2) and
1603 op_spec.shape.results.allows(1) and
1604 !op_spec.shape.results.allows(0) and
1605 op_spec.shape.successors.allows(0) and
1606 !op_spec.shape.successors.allows(1);
1607 }
1608 if (std.mem.eql(u8, op_spec.name, AccyDialect.SelectOp.operation_name)) {
1609 select_shape = op_spec.shape.operands.allows(3) and
1610 !op_spec.shape.operands.allows(2) and
1611 op_spec.shape.results.allows(1) and
1612 !op_spec.shape.results.allows(2);
1613 }
1614 if (std.mem.eql(u8, op_spec.name, AccyDialect.ActivationOp.operation_name)) {
1615 activation_shape = op_spec.shape.operands.allows(1) and
1616 !op_spec.shape.operands.allows(2) and
1617 op_spec.shape.results.allows(1) and
1618 !op_spec.shape.results.allows(0);
1619 }
1620 if (std.mem.eql(u8, op_spec.name, AccyDialect.EinsumOp.operation_name)) {
1621 einsum_shape = op_spec.shape.operands.allows(1) and
1622 op_spec.shape.operands.allows(3) and
1623 !op_spec.shape.operands.allows(0) and
1624 op_spec.shape.results.allows(1) and
1625 !op_spec.shape.results.allows(2);
1626 }
1627 if (std.mem.eql(u8, op_spec.name, AccyDialect.KernelCallOp.operation_name)) {
1628 kernel_call_shape = op_spec.shape.operands.allows(0) and
1629 op_spec.shape.operands.allows(3) and
1630 op_spec.shape.results.allows(1) and
1631 op_spec.shape.results.allows(2) and
1632 !op_spec.shape.results.allows(0) and
1633 op_spec.shape.regions.allows(0) and
1634 !op_spec.shape.regions.allows(1);
1635 for (op_spec.interfaces) |entry| {
1636 if (entry.id == ir.interfaces.EffectOpInterface.id) {
1637 kernel_call_memory_effects = true;
1638 }
1639 }
1640 }
1641 }
1642 try testing.expect(add_shape);
1643 try testing.expect(neg_shape);
1644 try testing.expect(select_shape);
1645 try testing.expect(activation_shape);
1646 try testing.expect(einsum_shape);
1647 try testing.expect(kernel_call_shape);
1648 try testing.expect(kernel_call_memory_effects);
1649 }
1650
1651 test "accy dialect lazy-loads in a strict Choir context" {
1652 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1653 defer arena.deinit();
1654 const allocator = arena.allocator();
1655
1656 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1657 defer ctx.deinit(allocator);
1658
1659 try ctx.requireRegistered();
1660
1661 try registerAccyDialect(&ctx);
1662
1663 try testing.expect(ctx.lookupOperation("accy.iota") == null);
1664 try testing.expect(ctx.lookupOperation("accy.add") == null);
1665 try testing.expect(ctx.lookupOperation("accy.tanh") == null);
1666 try testing.expect(ctx.lookupOperation("accy.activation") == null);
1667 try testing.expect(ctx.lookupOperation("accy.compare") == null);
1668 try testing.expect(ctx.lookupOperation("accy.return") == null);
1669
1670 _ = try ctx.getOrLoadDialect("accy");
1671
1672 for (AccyDialect.spec.operations) |op_spec| {
1673 const info = ctx.lookupOperation(op_spec.name) orelse return error.TestExpectedOp;
1674 try testing.expect(info.hasInterface(ir.VerifyOpInterface.id));
1675 }
1676 try testing.expect(ctx.lookupType(tensor_type_name) != null);
1677 }
1678
1679 test "accy Choir verifier accepts valid elementwise ops" {
1680 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1681 defer arena.deinit();
1682 const allocator = arena.allocator();
1683
1684 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1685 defer ctx.deinit(allocator);
1686 try registerAccyDialect(&ctx);
1687 _ = try ctx.getOrLoadDialect("accy");
1688
1689 const tensor_type = try accyTensorType(&ctx, .f32, &.{ 2, 3 });
1690 const loc = ir.Location.getUnknown();
1691
1692 const lhs = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 0);
1693 const rhs = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 1);
1694 const add = try AccyDialect.AddOp.create(&ctx, loc, lhs.getResult(), rhs.getResult());
1695
1696 try ir.verifyOperation(add.op, .{ .recursive = false });
1697 }
1698
1699 test "accy Choir verifier accepts einsum shape inference" {
1700 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1701 defer arena.deinit();
1702 const allocator = arena.allocator();
1703
1704 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1705 defer ctx.deinit(allocator);
1706 try registerAccyDialect(&ctx);
1707 _ = try ctx.getOrLoadDialect("accy");
1708
1709 const lhs_type = try accyTensorType(&ctx, .f32, &.{ 4, 8 });
1710 const rhs_type = try accyTensorType(&ctx, .f32, &.{ 8, 16 });
1711 const result_type = try accyTensorType(&ctx, .f32, &.{ 4, 16 });
1712 const loc = ir.Location.getUnknown();
1713
1714 const lhs = try AccyDialect.IotaOp.create(&ctx, loc, lhs_type, 0);
1715 const rhs = try AccyDialect.IotaOp.create(&ctx, loc, rhs_type, 0);
1716 const op = try AccyDialect.EinsumOp.create(&ctx, loc, &.{ lhs.getResult(), rhs.getResult() }, result_type, "ik,kj->ij");
1717
1718 try testing.expectEqualStrings("accy.einsum", op.op.name.name);
1719 try testing.expectEqualStrings("ik,kj->ij", op.getEquation().?);
1720 try testing.expect(op.getResult().type.eql(result_type));
1721 try ir.verifyOperation(op.op, .{ .recursive = false });
1722 }
1723
1724 test "accy Choir verifier accepts activation shape inference" {
1725 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1726 defer arena.deinit();
1727 const allocator = arena.allocator();
1728
1729 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1730 defer ctx.deinit(allocator);
1731 try registerAccyDialect(&ctx);
1732 _ = try ctx.getOrLoadDialect("accy");
1733
1734 const tensor_type = try accyTensorType(&ctx, .f32, &.{8});
1735 const loc = ir.Location.getUnknown();
1736
1737 const input = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 0);
1738 const op = try AccyDialect.ActivationOp.create(&ctx, loc, input.getResult(), tensor_type, .gelu);
1739
1740 try testing.expectEqualStrings("accy.activation", op.op.name.name);
1741 try testing.expectEqualStrings("gelu", op.getKind().?);
1742 try testing.expect(op.getResult().type.eql(tensor_type));
1743 try ir.verifyOperation(op.op, .{ .recursive = false });
1744 }
1745
1746 test "accy Choir verifier accepts kernel_call contracts" {
1747 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1748 defer arena.deinit();
1749 const allocator = arena.allocator();
1750
1751 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1752 defer ctx.deinit(allocator);
1753 try registerAccyDialect(&ctx);
1754 _ = try ctx.getOrLoadDialect("accy");
1755
1756 const tensor_type = try accyTensorType(&ctx, .f32, &.{ 2, 3 });
1757 const loc = ir.Location.getUnknown();
1758
1759 const input = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 0);
1760 const call = try AccyDialect.KernelCallOp.create(
1761 &ctx,
1762 loc,
1763 &.{input.getResult()},
1764 &.{tensor_type},
1765 "scale_f32",
1766 1,
1767 false,
1768 &.{.none},
1769 &.{null},
1770 );
1771
1772 try testing.expectEqualStrings("accy.kernel_call", call.op.name.name);
1773 try testing.expect(call.getResult(0).?.type.eql(tensor_type));
1774 try testing.expect(call.getFirstResult().type.eql(tensor_type));
1775 try ir.verifyOperation(call.op, .{ .recursive = false });
1776
1777 var summary = try choir.passes.effects.EffectSummary.init(allocator, call.op);
1778 defer summary.deinit();
1779 try testing.expect(!summary.discard());
1780 try testing.expect(!summary.repeatableExpression());
1781
1782 const inplace = try AccyDialect.KernelCallOp.create(
1783 &ctx,
1784 loc,
1785 &.{input.getResult()},
1786 &.{tensor_type},
1787 "update_f32",
1788 1,
1789 false,
1790 &.{.read_write},
1791 &.{0},
1792 );
1793 var inplace_summary = try choir.passes.effects.EffectSummary.init(allocator, inplace.op);
1794 defer inplace_summary.deinit();
1795 try testing.expect(inplace_summary.invalidatesStores());
1796
1797 const effectful = try AccyDialect.KernelCallOp.create(
1798 &ctx,
1799 loc,
1800 &.{input.getResult()},
1801 &.{tensor_type},
1802 "stateful_f32",
1803 1,
1804 true,
1805 &.{.none},
1806 &.{null},
1807 );
1808 var effectful_summary = try choir.passes.effects.EffectSummary.init(allocator, effectful.op);
1809 defer effectful_summary.deinit();
1810 try testing.expect(!effectful_summary.discard());
1811 }
1812
1813 test "accy Choir verifier rejects malformed kernel_call contracts" {
1814 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1815 defer arena.deinit();
1816 const allocator = arena.allocator();
1817
1818 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1819 defer ctx.deinit(allocator);
1820 try registerAccyDialect(&ctx);
1821 _ = try ctx.getOrLoadDialect("accy");
1822
1823 const tensor_type = try accyTensorType(&ctx, .f32, &.{4});
1824 const loc = ir.Location.getUnknown();
1825
1826 const input = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 0);
1827 const empty_target = try AccyDialect.KernelCallOp.create(
1828 &ctx,
1829 loc,
1830 &.{input.getResult()},
1831 &.{tensor_type},
1832 "",
1833 1,
1834 false,
1835 &.{.none},
1836 &.{null},
1837 );
1838 try testing.expectError(error.InvalidKernelContract, ir.verifyOperation(empty_target.op, .{ .recursive = false }));
1839
1840 const zero_version = try AccyDialect.KernelCallOp.create(
1841 &ctx,
1842 loc,
1843 &.{input.getResult()},
1844 &.{tensor_type},
1845 "scale_f32",
1846 0,
1847 false,
1848 &.{.none},
1849 &.{null},
1850 );
1851 try testing.expectError(error.InvalidKernelContract, ir.verifyOperation(zero_version.op, .{ .recursive = false }));
1852
1853 const missing_effect = try AccyDialect.KernelCallOp.create(
1854 &ctx,
1855 loc,
1856 &.{input.getResult()},
1857 &.{tensor_type},
1858 "scale_f32",
1859 1,
1860 false,
1861 &.{},
1862 &.{null},
1863 );
1864 try testing.expectError(error.InvalidKernelContract, ir.verifyOperation(missing_effect.op, .{ .recursive = false }));
1865
1866 const read_alias = try AccyDialect.KernelCallOp.create(
1867 &ctx,
1868 loc,
1869 &.{input.getResult()},
1870 &.{tensor_type},
1871 "update_f32",
1872 1,
1873 false,
1874 &.{.read},
1875 &.{0},
1876 );
1877 try testing.expectError(error.InvalidKernelContract, ir.verifyOperation(read_alias.op, .{ .recursive = false }));
1878
1879 const duplicate_alias = try AccyDialect.KernelCallOp.create(
1880 &ctx,
1881 loc,
1882 &.{input.getResult()},
1883 &.{ tensor_type, tensor_type },
1884 "update_f32",
1885 1,
1886 false,
1887 &.{.read_write},
1888 &.{ 0, 0 },
1889 );
1890 try testing.expectError(error.InvalidKernelContract, ir.verifyOperation(duplicate_alias.op, .{ .recursive = false }));
1891 }
1892
1893 test "accy Choir verifier rejects shape mismatch" {
1894 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1895 defer arena.deinit();
1896 const allocator = arena.allocator();
1897
1898 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1899 defer ctx.deinit(allocator);
1900 try registerAccyDialect(&ctx);
1901 _ = try ctx.getOrLoadDialect("accy");
1902
1903 const lhs_type = try accyTensorType(&ctx, .f32, &.{2});
1904 const rhs_type = try accyTensorType(&ctx, .f32, &.{3});
1905 const loc = ir.Location.getUnknown();
1906
1907 const lhs = try AccyDialect.IotaOp.create(&ctx, loc, lhs_type, 0);
1908 const rhs = try AccyDialect.IotaOp.create(&ctx, loc, rhs_type, 0);
1909 const add = try AccyDialect.AddOp.create(&ctx, loc, lhs.getResult(), rhs.getResult());
1910
1911 try testing.expectError(error.ShapeMismatch, ir.verifyOperation(add.op, .{ .recursive = false }));
1912 }
1913
1914 test "accy Choir verifier rejects result type mismatch" {
1915 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1916 defer arena.deinit();
1917 const allocator = arena.allocator();
1918
1919 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1920 defer ctx.deinit(allocator);
1921 try registerAccyDialect(&ctx);
1922 _ = try ctx.getOrLoadDialect("accy");
1923
1924 const operand_type = try accyTensorType(&ctx, .f32, &.{2});
1925 const wrong_result_type = try accyTensorType(&ctx, .f32, &.{3});
1926 const loc = ir.Location.getUnknown();
1927
1928 const lhs = try AccyDialect.IotaOp.create(&ctx, loc, operand_type, 0);
1929 const rhs = try AccyDialect.IotaOp.create(&ctx, loc, operand_type, 0);
1930
1931 var builder = ir.OperationBuilder.init(&ctx);
1932 var state = ir.Operation.State.init(AccyDialect.AddOp.operation_name, loc);
1933 state.addOperands(&.{ lhs.getResult(), rhs.getResult() });
1934 state.addTypes(&.{wrong_result_type});
1935 const add = try builder.create(state);
1936
1937 try testing.expectError(AccyChoirVerifyError.ResultTypeMismatch, ir.verifyOperation(add, .{ .recursive = false }));
1938 }
1939
1940 test "accy Choir verifier rejects unsupported dtype" {
1941 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1942 defer arena.deinit();
1943 const allocator = arena.allocator();
1944
1945 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1946 defer ctx.deinit(allocator);
1947 try registerAccyDialect(&ctx);
1948 _ = try ctx.getOrLoadDialect("accy");
1949
1950 const tensor_type = try accyTensorType(&ctx, .i32, &.{4});
1951 const loc = ir.Location.getUnknown();
1952
1953 const x = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 0);
1954 const sqrt = try AccyDialect.SqrtOp.create(&ctx, loc, x.getResult());
1955
1956 try testing.expectError(AccyChoirVerifyError.UnsupportedDType, ir.verifyOperation(sqrt.op, .{ .recursive = false }));
1957 }
1958
1959 test "accy Choir verifier rejects malformed constant payload length" {
1960 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1961 defer arena.deinit();
1962 const allocator = arena.allocator();
1963
1964 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1965 defer ctx.deinit(allocator);
1966 try registerAccyDialect(&ctx);
1967 _ = try ctx.getOrLoadDialect("accy");
1968
1969 const scalar_i32 = try accyTensorType(&ctx, .i32, &.{});
1970 const one_byte = [_]u8{0};
1971 const c = try AccyDialect.ConstantOp.create(&ctx, ir.Location.getUnknown(), &one_byte, scalar_i32);
1972
1973 try testing.expectError(AccyChoirVerifyError.ConstantPayloadLengthMismatch, ir.verifyOperation(c.op, .{ .recursive = false }));
1974 }
1975
1976 test "accy Choir verifier rejects malformed tensor type keys" {
1977 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1978 defer arena.deinit();
1979 const allocator = arena.allocator();
1980
1981 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1982 defer ctx.deinit(allocator);
1983 try registerAccyDialect(&ctx);
1984 _ = try ctx.getOrLoadDialect("accy");
1985
1986 const malformed_type = try ctx.getDialectTypeFromNameWithKey(tensor_type_name, "f32,2x");
1987 const iota = try AccyDialect.IotaOp.create(&ctx, ir.Location.getUnknown(), malformed_type, 0);
1988
1989 try testing.expectError(AccyChoirVerifyError.MalformedAccyTensorType, ir.verifyOperation(iota.op, .{ .recursive = false }));
1990 }
1991
1992 test "accy.iota constructor produces a result of the requested tensor type" {
1993 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1994 defer arena.deinit();
1995 const allocator = arena.allocator();
1996
1997 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1998 defer ctx.deinit(allocator);
1999 try registerAccyDialect(&ctx);
2000 _ = try ctx.getOrLoadDialect("accy");
2001
2002 const tensor_type = try accyTensorType(&ctx, .f32, &.{ 4, 8 });
2003 const loc = ir.Location.getUnknown();
2004
2005 const iota = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 0);
2006 try testing.expectEqualStrings("accy.iota", iota.op.name.name);
2007
2008 const result = iota.getResult();
2009 try testing.expect(result.type.eql(tensor_type));
2010 }
2011
2012 test "accy elementwise binary constructors produce uniform-shape results" {
2013 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2014 defer arena.deinit();
2015 const allocator = arena.allocator();
2016
2017 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2018 defer ctx.deinit(allocator);
2019 try registerAccyDialect(&ctx);
2020 _ = try ctx.getOrLoadDialect("accy");
2021
2022 const tensor_type = try accyTensorType(&ctx, .f32, &.{ 2, 3 });
2023 const loc = ir.Location.getUnknown();
2024
2025 const lhs_op = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 0);
2026 const rhs_op = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 1);
2027
2028 const add = try AccyDialect.AddOp.create(&ctx, loc, lhs_op.getResult(), rhs_op.getResult());
2029 const sub = try AccyDialect.SubOp.create(&ctx, loc, lhs_op.getResult(), rhs_op.getResult());
2030 const mul = try AccyDialect.MulOp.create(&ctx, loc, lhs_op.getResult(), rhs_op.getResult());
2031 const div = try AccyDialect.DivOp.create(&ctx, loc, lhs_op.getResult(), rhs_op.getResult());
2032 const atan2 = try AccyDialect.Atan2Op.create(&ctx, loc, lhs_op.getResult(), rhs_op.getResult());
2033
2034 try testing.expectEqualStrings("accy.add", add.op.name.name);
2035 try testing.expectEqualStrings("accy.sub", sub.op.name.name);
2036 try testing.expectEqualStrings("accy.mul", mul.op.name.name);
2037 try testing.expectEqualStrings("accy.div", div.op.name.name);
2038 try testing.expectEqualStrings("accy.atan2", atan2.op.name.name);
2039
2040 try testing.expect(add.getResult().type.eql(tensor_type));
2041 try testing.expect(sub.getResult().type.eql(tensor_type));
2042 try testing.expect(mul.getResult().type.eql(tensor_type));
2043 try testing.expect(div.getResult().type.eql(tensor_type));
2044 try testing.expect(atan2.getResult().type.eql(tensor_type));
2045 }
2046
2047 test "accy.return accepts variadic operand counts" {
2048 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2049 defer arena.deinit();
2050 const allocator = arena.allocator();
2051
2052 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2053 defer ctx.deinit(allocator);
2054 try registerAccyDialect(&ctx);
2055 _ = try ctx.getOrLoadDialect("accy");
2056
2057 const tensor_type = try accyTensorType(&ctx, .i32, &.{4});
2058 const loc = ir.Location.getUnknown();
2059
2060 const v0 = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 0);
2061 const v1 = try AccyDialect.IotaOp.create(&ctx, loc, tensor_type, 0);
2062
2063 const ret_empty = try AccyDialect.ReturnOp.create(&ctx, loc, &.{});
2064 try testing.expectEqualStrings("accy.return", ret_empty.op.name.name);
2065 try testing.expectEqual(@as(usize, 0), ret_empty.op.operands.items.len);
2066
2067 const ret_one = try AccyDialect.ReturnOp.create(&ctx, loc, &.{v0.getResult()});
2068 try testing.expectEqual(@as(usize, 1), ret_one.op.operands.items.len);
2069
2070 const ret_two = try AccyDialect.ReturnOp.create(&ctx, loc, &.{ v0.getResult(), v1.getResult() });
2071 try testing.expectEqual(@as(usize, 2), ret_two.op.operands.items.len);
2072 }
2073
2074 comptime {
2075 _ = arith_mod;
2076 }
2077
2078 fn kernelCallEffects(op: *const ir.Operation, collector: *effect_facts.Collector) void {
2079 collector.append(.{ .event = .{ .kind = .foreign } });
2080 collector.append(.{ .requirement = .{ .kind = .callee_contract, .subject = .operation } });
2081 kernelCallAccessEffects(op, collector);
2082 const aliases = op.getAttrAs(ir.Attribute.DialectAttr, "result_aliases");
2083 for (0..op.getNumResults()) |index| {
2084 var result = effect_facts.ResultFact{ .index = index };
2085 if (aliases) |attribute| {
2086 const width = @sizeOf(i64);
2087 if (index < attribute.payload.len / width) {
2088 const bytes = attribute.payload[index * width ..][0..width];
2089 const alias = std.mem.readInt(i64, bytes, native_endian);
2090 if (std.math.cast(usize, alias)) |operand_index| {
2091 if (operand_index < op.getNumOperands()) {
2092 result.alias = .{ .operand = operand_index };
2093 }
2094 }
2095 }
2096 }
2097 collector.append(.{ .result = result });
2098 }
2099 }
2100
2101 fn kernelCallAccessEffects(op: *const ir.Operation, collector: *effect_facts.Collector) void {
2102 const raw = op.getAttr("operand_effects") orelse return;
2103 const attribute = raw.cast(ir.Attribute.DialectAttr) orelse return;
2104 if (!std.mem.eql(u8, raw.abstract.name, AccyDialect.KernelCallOp.operand_effects_attr_name)) {
2105 return;
2106 }
2107 const count = @min(op.getNumOperands(), attribute.payload.len);
2108 for (attribute.payload[0..count], 0..) |byte, index| {
2109 const access = semantics.KernelOperandEffect.fromByte(byte) orelse continue;
2110 const resource = effect_facts.Resource{ .subject = .{ .operand = index } };
2111 if (access == .read or access == .read_write) {
2112 collector.append(.{ .event = .{ .kind = .read, .resource = resource } });
2113 }
2114 if (access == .write or access == .read_write) {
2115 collector.append(.{ .event = .{ .kind = .write, .resource = resource } });
2116 }
2117 }
2118 }
2119
2120 test "accy effect declarations never certify caller supplied kernel purity" {
2121 var ctx = try ir.Context.init(std.testing.allocator, ir.Context.Limits.testing);
2122 defer ctx.deinit(std.testing.allocator);
2123 try registerAccyDialect(&ctx);
2124 const typ = try accyTensorType(&ctx, .f32, &.{2});
2125 const input = try AccyDialect.IotaOp.create(&ctx, .unknown, typ, 0);
2126 const call = try AccyDialect.KernelCallOp.create(
2127 &ctx,
2128 .unknown,
2129 &.{input.getResult()},
2130 &.{typ},
2131 "effect_fixture",
2132 1,
2133 false,
2134 &.{.read_write},
2135 &.{0},
2136 );
2137 var declaration = try effect_facts.inspect(std.testing.allocator, call.op);
2138 defer declaration.deinit(std.testing.allocator);
2139 try std.testing.expect(!declaration.facts.complete);
2140 try std.testing.expect(!effect_facts.repeatableExpression(declaration.facts));
2141 try std.testing.expectEqual(
2142 effect_facts.EventKind.foreign,
2143 declaration.facts.records[0].event.kind,
2144 );
2145 try std.testing.expectEqual(@as(usize, 0), declaration.facts.records[4].result.alias.?.operand);
2146 const zero_operand_call = try AccyDialect.KernelCallOp.create(
2147 &ctx,
2148 .unknown,
2149 &.{},
2150 &.{typ},
2151 "effect_fixture",
2152 1,
2153 false,
2154 &.{},
2155 &.{null},
2156 );
2157 var zero = try effect_facts.inspect(std.testing.allocator, zero_operand_call.op);
2158 defer zero.deinit(std.testing.allocator);
2159 try std.testing.expectEqual(effect_facts.EventKind.foreign, zero.facts.records[0].event.kind);
2160 try std.testing.expect(!effect_facts.discard(zero.facts));
2161 try std.testing.expectEqual(
2162 effect_facts.Ownership.unknown,
2163 zero.facts.records[2].result.ownership,
2164 );
2165 }
2166
2167 fn accyEffectsEntry() ir.interfaces.InterfaceEntry {
2168 return effect_facts.EffectOpInterface.entryFor(.{
2169 .capacity = .{ .entries = 6, .per_operand = 1, .per_result = 1, .per_region = 1 },
2170 .enumerate = accyEffects,
2171 });
2172 }
2173
2174 const StaticTensor = struct {
2175 dtype: choir_abi.DType,
2176 elements: usize,
2177 dimensions: []const u8,
2178
2179 fn extent(self: StaticTensor, axis: usize) ?usize {
2180 var dims = std.mem.splitScalar(u8, self.dimensions, 'x');
2181 var index: usize = 0;
2182 while (dims.next()) |part| : (index += 1) {
2183 if (index == axis) return std.fmt.parseInt(usize, part, 10) catch null;
2184 }
2185 return null;
2186 }
2187
2188 fn rank(self: StaticTensor) usize {
2189 if (self.dimensions.len == 0) return 0;
2190 return std.mem.count(u8, self.dimensions, "x") + 1;
2191 }
2192 };
2193
2194 fn staticTensor(typ: ir.Type) ?StaticTensor {
2195 if (!std.mem.eql(u8, typ.getDialectTypeName() orelse return null, tensor_type_name)) {
2196 return null;
2197 }
2198 const key = typ.getDialectParamKey() orelse return null;
2199 const comma = std.mem.indexOfScalar(u8, key, ',') orelse return null;
2200 const dtype = choir_abi.DType.fromName(key[0..comma]) orelse return null;
2201 const dimensions = key[comma + 1 ..];
2202 var elements: usize = 1;
2203 if (dimensions.len != 0) {
2204 var dims = std.mem.splitScalar(u8, dimensions, 'x');
2205 while (dims.next()) |part| {
2206 const dim = std.fmt.parseInt(usize, part, 10) catch return null;
2207 elements = std.math.mul(usize, elements, dim) catch return null;
2208 }
2209 }
2210 return .{ .dtype = dtype, .elements = elements, .dimensions = dimensions };
2211 }
2212
2213 fn accyEffects(op: *const ir.Operation, collector: *effect_facts.Collector) void {
2214 const kind = accyKindFromChoirName(op.name.name) orelse return;
2215 if (op.getAttr("has_side_effects") != null) return;
2216 switch (kind) {
2217 .kernel_call, .@"return", .parameter => return,
2218 .scratch => return scratchEffects(op, collector),
2219 .iterate, .iterate_yield => return iterationEffects(op, collector),
2220 else => {},
2221 }
2222 var floating = false;
2223 for (op.operands.items) |operand| {
2224 const typ = staticTensor(operand.value.type) orelse return;
2225 floating = floating or typ.dtype.isFloat();
2226 }
2227 for (op.results.items) |result| {
2228 const typ = staticTensor(result.type) orelse return;
2229 floating = floating or typ.dtype.isFloat();
2230 }
2231 switch (kind) {
2232 .max, .min, .neg, .abs => if (!floating) return,
2233 else => {},
2234 }
2235 collector.complete = true;
2236 collector.valueResults(op);
2237 if (floating) floatingEnvironmentEffects(op, collector);
2238 switch (kind) {
2239 .div => integerDivisionEffects(op, collector),
2240 .convert => conversionEffects(op, collector),
2241 .gather, .scatter, .scatter_add, .sparse_cross_entropy => indexEffects(op, kind, collector),
2242 .constant,
2243 .iota,
2244 .add,
2245 .sub,
2246 .mul,
2247 .max,
2248 .min,
2249 .pow,
2250 .compare,
2251 .neg,
2252 .exp,
2253 .log,
2254 .tanh,
2255 .sqrt,
2256 .activation,
2257 .abs,
2258 .sin,
2259 .cos,
2260 .tan,
2261 .floor,
2262 .round,
2263 .trunc,
2264 .atan2,
2265 .reduce,
2266 .dot_general,
2267 .einsum,
2268 .broadcast,
2269 .broadcast_in_dim,
2270 .reshape,
2271 .transpose,
2272 .slice,
2273 .pad,
2274 .concatenate,
2275 .select,
2276 .cumsum,
2277 => {},
2278 .kernel_call, .@"return", .parameter, .scratch, .iterate, .iterate_yield => unreachable,
2279 }
2280 }
2281
2282 fn floatingEnvironmentEffects(op: *const ir.Operation, collector: *effect_facts.Collector) void {
2283 if (op.getContext().arithmetic_policy.permitsFloatingValues()) {
2284 collector.append(.{ .premise = .floating_environment });
2285 return;
2286 }
2287 collector.append(.{ .event = .{
2288 .kind = .state_observe,
2289 .resource = .{
2290 .subject = .{ .global = "arithmetic.environment" },
2291 .state_key = "floating_environment",
2292 },
2293 } });
2294 collector.append(.{ .requirement = .{ .kind = .execution_context, .subject = .operation } });
2295 }
2296
2297 fn conditionedFailure(
2298 collector: *effect_facts.Collector,
2299 kind: effect_facts.RequirementKind,
2300 subject: effect_facts.Subject,
2301 name: []const u8,
2302 ) void {
2303 collector.append(.{ .requirement = .{ .kind = kind, .subject = subject } });
2304 collector.append(.{ .event = .{ .kind = .failure, .failure_name = name } });
2305 }
2306
2307 const ConstantTensor = struct {
2308 typ: StaticTensor,
2309 payload: []const u8,
2310
2311 fn number(self: ConstantTensor, index: usize) union(enum) { integer: i128, float: f64 } {
2312 std.debug.assert(index < self.typ.elements);
2313 const width = self.typ.dtype.sizeOf();
2314 const bytes = self.payload[index * width ..][0..width];
2315 return switch (self.typ.dtype) {
2316 .i1 => .{ .integer = if (bytes[0] == 0) 0 else 1 },
2317 .key => unreachable,
2318 .bf16 => .{ .float = (choir_abi.Bf16{ .bits = std.mem.readInt(
2319 u16,
2320 bytes[0..2],
2321 native_endian,
2322 ) }).toF32() },
2323 inline else => |dt| value: {
2324 const T = dt.ZigType();
2325 const value = std.mem.bytesAsValue(T, bytes[0..@sizeOf(T)]).*;
2326 if (comptime dt.isFloat()) break :value .{ .float = @floatCast(value) };
2327 break :value .{ .integer = @intCast(value) };
2328 },
2329 };
2330 }
2331 };
2332
2333 fn constantTensor(value: *ir.Value) ?ConstantTensor {
2334 const typ = staticTensor(value.type) orelse return null;
2335 if (typ.dtype == .key) return null;
2336 const raw = value.getDefiningOp() orelse return null;
2337 const op: *ir.Operation = @ptrCast(@alignCast(raw));
2338 if (!std.mem.eql(u8, op.name.name, AccyDialect.ConstantOp.operation_name)) return null;
2339 const payload = (AccyDialect.ConstantOp{ .op = op }).getPayload() orelse return null;
2340 const bytes = std.math.mul(usize, typ.elements, typ.dtype.sizeOf()) catch return null;
2341 if (payload.len != bytes) return null;
2342 return .{ .typ = typ, .payload = payload };
2343 }
2344
2345 fn integerDivisionEffects(op: *const ir.Operation, collector: *effect_facts.Collector) void {
2346 if (op.getNumOperands() != 2) return;
2347 const lhs = op.getOperand(0).?;
2348 const rhs = op.getOperand(1).?;
2349 const typ = staticTensor(rhs.type) orelse return;
2350 if (typ.dtype.isFloat()) return;
2351 const divisor = constantTensor(rhs);
2352 const dividend = constantTensor(lhs);
2353 var nonzero = divisor != null;
2354 var representable = !typ.dtype.isSignedInt() or divisor != null;
2355 if (divisor) |constant| {
2356 const bits: u7 = @intCast(typ.dtype.sizeOf() * 8);
2357 const minimum = -(@as(i128, 1) << @intCast(bits - 1));
2358 for (0..constant.typ.elements) |index| {
2359 const d = constant.number(index).integer;
2360 if (d == 0) nonzero = false;
2361 if (typ.dtype.isSignedInt() and d == -1) {
2362 if (dividend) |numerator| {
2363 if (index >= numerator.typ.elements) {
2364 representable = false;
2365 } else if (numerator.number(index).integer == minimum) {
2366 representable = false;
2367 }
2368 } else representable = false;
2369 }
2370 }
2371 }
2372 if (!nonzero) conditionedFailure(collector, .nonzero, .{ .operand = 1 }, "DivisionByZero");
2373 if (!representable) {
2374 conditionedFailure(
2375 collector,
2376 .quotient_representable,
2377 .operation,
2378 "SignedDivisionOverflow",
2379 );
2380 }
2381 }
2382
2383 fn conversionEffects(op: *const ir.Operation, collector: *effect_facts.Collector) void {
2384 if (op.getNumOperands() != 1 or op.getNumResults() != 1) return;
2385 const input = op.getOperand(0).?;
2386 const source = staticTensor(input.type) orelse return;
2387 const target = staticTensor(op.results.items[0].type) orelse return;
2388 if (!source.dtype.isFloat() or target.dtype.isFloat()) return;
2389 if (floatConversionInRange(input, target.dtype)) return;
2390 conditionedFailure(
2391 collector,
2392 .conversion_representable,
2393 .{ .operand = 0 },
2394 "InvalidFloatToInteger",
2395 );
2396 }
2397
2398 fn floatConversionInRange(input: *ir.Value, target: choir_abi.DType) bool {
2399 const constant = constantTensor(input) orelse return false;
2400 if (!target.isSignedInt() and !target.isUnsignedInt() and target != .i1) return false;
2401 const bits: u7 = if (target == .i1) 1 else @intCast(target.sizeOf() * 8);
2402 const exponent = bits - @as(u7, if (target.isSignedInt()) 1 else 0);
2403 const upper: f64 = @floatFromInt(@as(u128, 1) << exponent);
2404 const lower: f64 = if (target.isSignedInt()) -upper else 0;
2405 for (0..constant.typ.elements) |index| {
2406 const value = @trunc(constant.number(index).float);
2407 if (!std.math.isFinite(value) or value < lower or value >= upper) return false;
2408 }
2409 return true;
2410 }
2411
2412 fn indexEffects(
2413 op: *const ir.Operation,
2414 kind: semantics.OpKind,
2415 collector: *effect_facts.Collector,
2416 ) void {
2417 if (op.getNumOperands() < 2) return;
2418 const source = staticTensor(op.getOperand(0).?.type) orelse return;
2419 const axis: ?usize = if (kind == .sparse_cross_entropy) blk: {
2420 if (source.rank() == 0) break :blk 0;
2421 break :blk source.rank() - 1;
2422 } else blk: {
2423 const attr = op.getAttrAs(ir.Attribute.IntegerAttr, "axis") orelse break :blk null;
2424 break :blk std.math.cast(usize, attr.value);
2425 };
2426 const extent = if (axis) |value| source.extent(value) else null;
2427 const indices = constantTensor(op.getOperand(1).?);
2428 if (extent != null and indices != null and !indices.?.typ.dtype.isFloat()) {
2429 var in_bounds = true;
2430 for (0..indices.?.typ.elements) |index| {
2431 const value = indices.?.number(index).integer;
2432 if (value < 0 or value >= extent.?) in_bounds = false;
2433 }
2434 if (in_bounds) return;
2435 }
2436 conditionedFailure(collector, .in_bounds, .{ .operand = 1 }, "IndexOutOfBounds");
2437 }
2438
2439 fn scratchEffects(op: *const ir.Operation, collector: *effect_facts.Collector) void {
2440 collector.append(.{ .event = .{
2441 .kind = .allocate,
2442 .resource = .{ .allocator_domain = "accy.scratch" },
2443 } });
2444 collector.append(.{ .event = .{ .kind = .failure, .failure_name = "AllocationFailure" } });
2445 for (0..op.getNumResults()) |index| {
2446 collector.append(.{ .result = .{
2447 .index = index,
2448 .fresh_identity = true,
2449 .ownership = .owned,
2450 } });
2451 }
2452 }
2453
2454 fn iterationEffects(op: *const ir.Operation, collector: *effect_facts.Collector) void {
2455 for (0..op.getNumResults()) |index| collector.append(.{ .result = .{ .index = index } });
2456 for (0..op.getNumRegions()) |index| {
2457 collector.append(.{ .region = .{ .index = index, .execution = .repeated } });
2458 for (0..op.getNumOperands()) |argument| {
2459 collector.append(.{ .binding = .{
2460 .region = index,
2461 .argument = argument,
2462 .source = .{ .operand = argument },
2463 } });
2464 }
2465 }
2466 if (op.getNumRegions() == 0) {
2467 for (0..op.getNumOperands()) |index| {
2468 collector.append(.{ .event = .{
2469 .kind = .move,
2470 .resource = .{ .subject = .{ .operand = index } },
2471 } });
2472 }
2473 }
2474 }
2475
2476 fn effectScalar(ctx: *ir.Context, comptime T: type, value: T) !*ir.Value {
2477 const dtype = choir_abi.DType.fromZigType(T) orelse unreachable;
2478 const typ = try accyTensorType(ctx, dtype, &.{});
2479 const bytes = [1]T{value};
2480 return (try AccyDialect.ConstantOp.create(
2481 ctx,
2482 .unknown,
2483 std.mem.sliceAsBytes(&bytes),
2484 typ,
2485 )).getResult();
2486 }
2487
2488 fn expectAccyPermission(op: *ir.Operation, expected: bool) !void {
2489 var declaration = try effect_facts.inspect(std.testing.allocator, op);
2490 defer declaration.deinit(std.testing.allocator);
2491 try testing.expectEqual(expected, effect_facts.discard(declaration.facts));
2492 try testing.expectEqual(expected, effect_facts.duplicate(declaration.facts, .{}));
2493 try testing.expectEqual(expected, effect_facts.speculate(declaration.facts, true));
2494 }
2495
2496 test "accy effect floating policy withdraws facts and cached permissions" {
2497 var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);
2498 defer ctx.deinit(testing.allocator);
2499 try registerAccyDialect(&ctx);
2500 const one = try effectScalar(&ctx, f32, 1);
2501 const add = try AccyDialect.AddOp.create(&ctx, .unknown, one, one);
2502 try expectAccyPermission(add.op, true);
2503 var before = try choir.passes.effects.EffectSummary.init(testing.allocator, add.op);
2504 defer before.deinit();
2505 ctx.arithmetic_policy.environment_observable = true;
2506 try testing.expect(!before.isCurrent());
2507 try expectAccyPermission(add.op, false);
2508 ctx.arithmetic_policy = .{ .exceptions_masked = false };
2509 try expectAccyPermission(add.op, false);
2510 ctx.arithmetic_policy = .{ .default_rounding = false };
2511 try expectAccyPermission(add.op, false);
2512 ctx.arithmetic_policy = .{};
2513 try expectAccyPermission(add.op, true);
2514 try add.op.setAttr("has_side_effects", try ctx.getBoolAttr(false));
2515 try expectAccyPermission(add.op, false);
2516 }
2517
2518 test "accy effect declarations reject dynamic tensor shapes" {
2519 var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);
2520 defer ctx.deinit(testing.allocator);
2521 try registerAccyDialect(&ctx);
2522 _ = try ctx.getOrLoadDialect("accy");
2523 const dynamic = try ctx.getDialectTypeFromNameWithKey(tensor_type_name, "f32,?x4");
2524 const value = try AccyDialect.IotaOp.create(&ctx, .unknown, dynamic, 0);
2525 try expectAccyPermission(value.op, false);
2526 }
2527
2528 fn divisionMotionCase(divisor_value: ?i32, numerator_value: ?i32, allowed: bool) !void {
2529 var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);
2530 defer ctx.deinit(testing.allocator);
2531 try registerAccyDialect(&ctx);
2532 try dialects_mod.registerChoirDialect(&ctx);
2533 const module = try dialects_mod.BuiltinDialect.ModuleOp.create(&ctx, .unknown);
2534 const body = module.getBodyBlock();
2535 const tensor = try accyTensorType(&ctx, .i32, &.{});
2536 const numerator = if (numerator_value) |v| value: {
2537 break :value try effectScalar(&ctx, i32, v);
2538 } else try body.addArgument(
2539 tensor,
2540 .unknown,
2541 );
2542 const divisor = if (divisor_value) |v| try effectScalar(&ctx, i32, v) else try body.addArgument(
2543 tensor,
2544 .unknown,
2545 );
2546 for ([_]*ir.Value{ numerator, divisor }) |value| {
2547 if (value.getDefiningOp()) |raw| {
2548 const op: *ir.Operation = @ptrCast(@alignCast(raw));
2549 try body.addOperation(op);
2550 }
2551 }
2552 const arith = dialects_mod.ArithDialect;
2553 const index_type = try arith.getScalarType(&ctx, .i64);
2554 const zero = try arith.ConstantOp.createInt(&ctx, .unknown, index_type, 0);
2555 const one = try arith.ConstantOp.createInt(&ctx, .unknown, index_type, 1);
2556 try body.addOperation(zero.op);
2557 try body.addOperation(one.op);
2558 const scf = dialects_mod.ScfDialect;
2559 const loop = try scf.ForOp.create(
2560 &ctx,
2561 .unknown,
2562 zero.getResult(),
2563 zero.getResult(),
2564 one.getResult(),
2565 &.{},
2566 &.{},
2567 );
2568 try body.addOperation(loop.op);
2569 const loop_body = loop.op.getRegion(0).?.getEntryBlock().?;
2570 const division = try AccyDialect.DivOp.create(&ctx, .unknown, numerator, divisor);
2571 try loop_body.addOperation(division.op);
2572 const yield = try scf.YieldOp.create(&ctx, .unknown, &.{});
2573 try loop_body.addOperation(yield.op);
2574 var motion = choir.passes.PassManager.init(testing.allocator);
2575 defer motion.deinit();
2576 try motion.addPass(choir.passes.createLoopInvariantCodeMotionPass());
2577 try testing.expectEqual(choir.passes.PassResult.success, motion.run(module.op, &ctx));
2578 try testing.expectEqual(if (allowed) body else loop_body, division.op.getBlock().?);
2579 try expectAccyPermission(division.op, allowed);
2580 var dce = choir.passes.PassManager.init(testing.allocator);
2581 defer dce.deinit();
2582 try dce.addPass(choir.passes.createDeadCodeEliminationPass());
2583 try testing.expectEqual(choir.passes.PassResult.success, dce.run(module.op, &ctx));
2584 try testing.expectEqual(
2585 @as(usize, if (allowed) 0 else 1),
2586 ir.inspection.countOperationsNamed(module.op, AccyDialect.DivOp.operation_name),
2587 );
2588 }
2589
2590 test "accy effect integer division motion discharges only constant domain conditions" {
2591 try divisionMotionCase(null, null, false);
2592 try divisionMotionCase(0, null, false);
2593 try divisionMotionCase(-1, null, false);
2594 try divisionMotionCase(-1, std.math.minInt(i32), false);
2595 try divisionMotionCase(2, null, true);
2596 try divisionMotionCase(-1, 7, true);
2597 }
2598
2599 test "accy effect conversion and indexing keep unresolved failures" {
2600 var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);
2601 defer ctx.deinit(testing.allocator);
2602 try registerAccyDialect(&ctx);
2603 const int_type = try accyTensorType(&ctx, .i32, &.{});
2604 for ([_]f64{ 1.5, 2147483648, std.math.nan(f64), std.math.inf(f64) }, 0..) |value, index| {
2605 const input = try effectScalar(&ctx, f64, value);
2606 const convert = try AccyDialect.ConvertOp.create(&ctx, .unknown, input, int_type, "i32");
2607 try expectAccyPermission(convert.op, index == 0);
2608 }
2609 const tensor = try accyTensorType(&ctx, .f32, &.{4});
2610 const input = try AccyDialect.IotaOp.create(&ctx, .unknown, tensor, 0);
2611 const scalar_float = try accyTensorType(&ctx, .f32, &.{});
2612 for ([_]i32{ -1, 0, 3, 4 }) |index| {
2613 const value = try effectScalar(&ctx, i32, index);
2614 const gather = try AccyDialect.GatherOp.create(
2615 &ctx,
2616 .unknown,
2617 input.getResult(),
2618 value,
2619 scalar_float,
2620 0,
2621 );
2622 try ir.verifyOperation(gather.op, .{ .recursive = false });
2623 try expectAccyPermission(gather.op, index >= 0 and index < 4);
2624 }
2625 const indices_type = try accyTensorType(&ctx, .i32, &.{1});
2626 const gathered_type = try accyTensorType(&ctx, .f32, &.{1});
2627 const runtime = try AccyDialect.IotaOp.create(&ctx, .unknown, indices_type, 0);
2628 const gather = try AccyDialect.GatherOp.create(
2629 &ctx,
2630 .unknown,
2631 input.getResult(),
2632 runtime.getResult(),
2633 gathered_type,
2634 0,
2635 );
2636 try ir.verifyOperation(gather.op, .{ .recursive = false });
2637 try expectAccyPermission(gather.op, false);
2638 const scratch = try AccyDialect.ScratchOp.create(&ctx, .unknown, int_type, 4);
2639 try expectAccyPermission(scratch.op, false);
2640 }
2641
2642 test "accy effect foreign kernel calls survive CSE and DCE" {
2643 var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);
2644 defer ctx.deinit(testing.allocator);
2645 try registerAccyDialect(&ctx);
2646 try dialects_mod.registerChoirDialect(&ctx);
2647 const module = try dialects_mod.BuiltinDialect.ModuleOp.create(&ctx, .unknown);
2648 const body = module.getBodyBlock();
2649 const typ = try accyTensorType(&ctx, .f32, &.{2});
2650 for (0..2) |_| {
2651 const call = try AccyDialect.KernelCallOp.create(
2652 &ctx,
2653 .unknown,
2654 &.{},
2655 &.{typ},
2656 "effect_fixture",
2657 1,
2658 true,
2659 &.{},
2660 &.{null},
2661 );
2662 try body.addOperation(call.op);
2663 try ir.verifyOperation(call.op, .{ .recursive = false });
2664 }
2665 var manager = choir.passes.PassManager.init(testing.allocator);
2666 defer manager.deinit();
2667 try manager.addPass(choir.passes.createCommonSubexpressionEliminationPass());
2668 try manager.addPass(choir.passes.createDeadCodeEliminationPass());
2669 try testing.expectEqual(choir.passes.PassResult.success, manager.run(module.op, &ctx));
2670 try testing.expectEqual(
2671 @as(usize, 2),
2672 ir.inspection.countOperationsNamed(module.op, AccyDialect.KernelCallOp.operation_name),
2673 );
2674 }
2675
2676 test "accy effect integer unary domains remain outside floating qualification" {
2677 var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);
2678 defer ctx.deinit(testing.allocator);
2679 try registerAccyDialect(&ctx);
2680 for ([_]i32{ 1, std.math.minInt(i32) }) |value| {
2681 const input = try effectScalar(&ctx, i32, value);
2682 const negative = try AccyDialect.NegOp.create(&ctx, .unknown, input);
2683 const absolute = try AccyDialect.AbsOp.create(&ctx, .unknown, input);
2684 try ir.verifyOperation(negative.op, .{ .recursive = false });
2685 try ir.verifyOperation(absolute.op, .{ .recursive = false });
2686 try expectAccyPermission(negative.op, false);
2687 try expectAccyPermission(absolute.op, false);
2688 }
2689 }