lib/choir/src/dialects/arith/ops.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const ir = @import("../../core/root.zig");
3 const format = ir.format;
4 const types = @import("types.zig");
5 const fold_mod = @import("folds.zig");
6 const predicate_mod = @import("predicate.zig");
7
8 const effects = @import("root.zig").effects;
9
10 const commutative_op_traits = ir.OperationTraits{ .is_commutative = true };
11 const CmpPredicate = predicate_mod.CmpPredicate;
12
13 pub const ArithDialect = struct {
14 pub const name = "arith";
15 const op_specs = ir.dialects.opSpec.dialect(@This());
16 const op_templates = ir.dialects.operationTemplate.dialect(@This());
17 const folds = fold_mod.Folds(@This());
18 const spec = ir.dialects.dialectSpec(@This(), .{
19 .types = types.registeredTypeSpecs(),
20 });
21
22 pub const ScalarTypeKind: type = types.ScalarKind;
23
24 pub const ConstantOp = struct {
25 op: *ir.Operation,
26
27 pub const operation_spec = op_specs.leaf(.{
28 .mnemonic = "constant",
29 .operands = 0,
30 .results = 1,
31 .required_attrs = &.{"value"},
32 .properties = ir.singleAttributePropertiesModel("arith.constant.properties", "value"),
33 .interfaces = effects.entries(.constant),
34 });
35 pub const operation_name = operation_spec.name;
36
37 pub fn createInt(ctx: *ir.Context, loc: ir.Location, result_type: ir.Type, value: i64) !ConstantOp {
38 var builder = ir.OperationBuilder.init(ctx);
39 var state = op_specs.state(@This(), loc);
40 state.addTypes(&.{result_type});
41 const value_attr = try getIntAttr(ctx, value);
42 const uses_properties = try state.setPropertiesAttrIfRegistered(ctx, value_attr);
43
44 const op = try builder.create(state);
45 errdefer op.erase();
46 if (!uses_properties) try op.setAttr("value", value_attr);
47
48 return .{ .op = op };
49 }
50
51 pub fn createFloat(ctx: *ir.Context, loc: ir.Location, result_type: ir.Type, value: f64) !ConstantOp {
52 var builder = ir.OperationBuilder.init(ctx);
53 var state = op_specs.state(@This(), loc);
54 state.addTypes(&.{result_type});
55 const value_attr = try getFloatAttr(ctx, value);
56 const uses_properties = try state.setPropertiesAttrIfRegistered(ctx, value_attr);
57
58 const op = try builder.create(state);
59 errdefer op.erase();
60 if (!uses_properties) try op.setAttr("value", value_attr);
61
62 return .{ .op = op };
63 }
64
65 pub fn createBool(ctx: *ir.Context, loc: ir.Location, value: bool) !ConstantOp {
66 var builder = ir.OperationBuilder.init(ctx);
67 const bool_type = try getScalarType(ctx, .bool);
68 var state = op_specs.state(@This(), loc);
69 state.addTypes(&.{bool_type});
70 const value_attr = try getBoolAttr(ctx, value);
71 const uses_properties = try state.setPropertiesAttrIfRegistered(ctx, value_attr);
72
73 const op = try builder.create(state);
74 errdefer op.erase();
75 if (!uses_properties) try op.setAttr("value", value_attr);
76
77 return .{ .op = op };
78 }
79
80 pub fn getResult(self: *const ConstantOp) *ir.Value {
81 return self.op.getResult(0).?;
82 }
83
84 pub fn getIntValue(self: ConstantOp) ?i64 {
85 const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "value") orelse return null;
86 return int_attr.getValue();
87 }
88
89 pub fn getFloatValue(self: ConstantOp) ?f64 {
90 const float_attr = self.op.getAttrAs(ir.Attribute.FloatAttr, "value") orelse return null;
91 return float_attr.getValue();
92 }
93 };
94
95 pub const AddOp: type = op_templates.binarySameTypeFold("add", effectOptions(.add, commutativeSameTypeOptions()), folds.foldAdd);
96
97 pub const SubOp: type = op_templates.binarySameTypeFold(
98 "sub",
99 effectOptions(.sub, sameTypeOptions(.{})),
100 folds.foldSub,
101 );
102
103 pub const MulOp: type = op_templates.binarySameTypeFold("mul", effectOptions(.mul, commutativeSameTypeOptions()), folds.foldMul);
104
105 pub const UmulhiOp: type = op_templates.binarySameType("umulhi", effectOptions(.umulhi, commutativeSameTypeOptions()));
106
107 /// Overflow arithmetic returns the exact result modulo 2^64, then whether the exact
108 /// mathematical result is unrepresentable in the operand type. Only i64 is admitted today.
109 /// These operations neither trap nor inherit the single-result operations' folds.
110 pub const AddoOp = overflowOp("addo", .addo);
111 pub const SuboOp = overflowOp("subo", .subo);
112 pub const MuloOp = overflowOp("mulo", .mulo);
113
114 pub const DivOp: type = op_templates.binarySameTypeFold(
115 "div",
116 effectOptions(.div, sameTypeOptions(.{})),
117 folds.foldDiv,
118 );
119
120 pub const MaxOp: type = op_templates.binarySameType("max", effectOptions(.max, sameTypeOptions(.{})));
121
122 pub const MinOp: type = op_templates.binarySameType("min", effectOptions(.min, sameTypeOptions(.{})));
123
124 pub const CmpOp = struct {
125 op: *ir.Operation,
126
127 pub const operation_spec = op_specs.leaf(.{
128 .mnemonic = "cmp",
129 .interfaces = effects.entries(.cmp),
130 .operands = 2,
131 .results = 1,
132 .required_attrs = &.{"predicate"},
133 .result_types = &.{ir.dialects.typeConstraint.exact(0, types.type_names.boolean)},
134 .dynamic_traits = .{ir.traits.SameTypeOperands},
135 });
136 pub const operation_name = operation_spec.name;
137 pub const fold = folds.foldCmp;
138
139 pub fn create(ctx: *ir.Context, loc: ir.Location, predicate: CmpPredicate, lhs: *ir.Value, rhs: *ir.Value) !CmpOp {
140 var builder = ir.OperationBuilder.init(ctx);
141 const bool_type = try getScalarType(ctx, .bool);
142 var state = op_specs.state(@This(), loc);
143 state.addOperands(&.{ lhs, rhs });
144 state.addTypes(&.{bool_type});
145
146 const op = try builder.create(state);
147 errdefer op.erase();
148
149 var buf: [8]u8 = undefined;
150 const pred_str = try std.fmt.bufPrint(&buf, "{s}", .{predicate.toString()});
151 const pred_attr = try ctx.getDialectAttr("arith.predicate", pred_str);
152 try op.setAttr("predicate", pred_attr);
153
154 return .{ .op = op };
155 }
156
157 pub fn getResult(self: *const CmpOp) *ir.Value {
158 return self.op.getResult(0).?;
159 }
160
161 pub fn getPredicate(self: CmpOp) ?CmpPredicate {
162 if (self.op.getAttrAs(ir.Attribute.DialectAttr, "predicate")) |dialect_attr| {
163 inline for (@typeInfo(CmpPredicate).@"enum".field_names, std.meta.tags(CmpPredicate)) |field_name, value| {
164 if (std.mem.eql(u8, dialect_attr.payload, field_name)) {
165 return value;
166 }
167 }
168 }
169 return null;
170 }
171 };
172
173 pub const CastOp: type = op_templates.unaryExplicitTypeFold(
174 "cast",
175 effectOptions(.cast, .{}),
176 folds.foldSameTypeUnary,
177 );
178
179 pub const SelectOp: type = op_templates.selectSameTypeFold("select", .{
180 .interfaces = effects.entries(.select),
181 .traits = .{},
182 .dynamic_traits = ir.dialects.opSpec.dynamicTraits(.{
183 ir.traits.TypesMatchWith(.{
184 .label = "select_result_eq_true",
185 .target = .{ .result = 0 },
186 .source = .{ .operand = 1 },
187 }),
188 ir.traits.TypesMatchWith(.{
189 .label = "select_true_eq_false",
190 .target = .{ .operand = 1 },
191 .source = .{ .operand = 2 },
192 }),
193 }),
194 .operand_types = &.{ir.dialects.typeConstraint.exact(0, types.type_names.boolean)},
195 }, folds.foldSelect);
196
197 pub const RemOp: type = op_templates.binarySameType("rem", effectOptions(.rem, sameTypeOptions(.{})));
198
199 pub const FmaOp: type = op_templates.ternarySameType("fma", effectOptions(.fma, sameTypeOptions(.{})));
200
201 pub const NegOp: type = op_templates.unarySameType("neg", effectOptions(.neg, sameTypeOptions(.{})));
202
203 pub const AbsOp: type = op_templates.unarySameType("abs", effectOptions(.abs, sameTypeOptions(.{})));
204
205 pub const SqrtOp: type = op_templates.unarySameType("sqrt", effectOptions(.sqrt, sameTypeOptions(.{})));
206
207 pub const ExpOp: type = op_templates.unarySameType("exp", effectOptions(.exp, sameTypeOptions(.{})));
208
209 pub const LogOp: type = op_templates.unarySameType("log", effectOptions(.log, sameTypeOptions(.{})));
210
211 pub const TanhOp: type = op_templates.unarySameType("tanh", effectOptions(.tanh, sameTypeOptions(.{})));
212
213 pub const SinOp: type = op_templates.unarySameType("sin", effectOptions(.sin, sameTypeOptions(.{})));
214
215 pub const CosOp: type = op_templates.unarySameType("cos", effectOptions(.cos, sameTypeOptions(.{})));
216
217 pub const TanOp: type = op_templates.unarySameType("tan", effectOptions(.tan, sameTypeOptions(.{})));
218
219 pub const FloorOp: type = op_templates.unarySameType("floor", effectOptions(.floor, sameTypeOptions(.{})));
220
221 pub const RoundOp: type = op_templates.unarySameType("round", effectOptions(.round, sameTypeOptions(.{})));
222
223 pub const TruncOp: type = op_templates.unarySameType("trunc", effectOptions(.trunc, sameTypeOptions(.{})));
224
225 pub const Tf32RoundOp: type = op_templates.unarySameType("tf32_round", sameTypeOptions(.{}));
226
227 pub const PowOp: type = op_templates.binarySameType("pow", effectOptions(.pow, sameTypeOptions(.{})));
228
229 pub const Atan2Op: type = op_templates.binarySameType("atan2", sameTypeOptions(.{}));
230
231 pub const AndOp: type = op_templates.binarySameTypeFold(
232 "and",
233 effectOptions(.@"and", commutativeSameTypeOptions()),
234 folds.foldAnd,
235 );
236
237 pub const OrOp: type = op_templates.binarySameTypeFold(
238 "or",
239 effectOptions(.@"or", commutativeSameTypeOptions()),
240 folds.foldOr,
241 );
242
243 pub const XorOp: type = op_templates.binarySameTypeFold(
244 "xor",
245 effectOptions(.xor, commutativeSameTypeOptions()),
246 folds.foldXor,
247 );
248
249 pub const NotOp: type = op_templates.unarySameTypeFold(
250 "not",
251 effectOptions(.not, sameTypeOptions(.{})),
252 folds.foldNot,
253 );
254
255 pub const PopCountOp: type = op_templates.unarySameType(
256 "popcount",
257 effectOptions(.popcount, sameTypeOptions(.{})),
258 );
259
260 pub const ShlOp: type = op_templates.binarySameTypeFold(
261 "shl",
262 effectOptions(.shl, sameTypeOptions(.{})),
263 folds.foldShift,
264 );
265
266 pub const ShrOp: type = op_templates.binarySameTypeFold(
267 "shr",
268 effectOptions(.shr, sameTypeOptions(.{})),
269 folds.foldShift,
270 );
271
272 pub const UshrOp: type = op_templates.binarySameTypeFold(
273 "ushr",
274 effectOptions(.ushr, sameTypeOptions(.{})),
275 folds.foldShift,
276 );
277
278 pub const BitcastOp: type = op_templates.unaryExplicitTypeFold(
279 "bitcast",
280 effectOptions(.bitcast, .{}),
281 folds.foldSameTypeUnary,
282 );
283
284 pub const SplatOp: type = op_templates.unaryExplicitType("splat", effectOptions(.splat, .{}));
285
286 pub const ExtractOp = struct {
287 op: *ir.Operation,
288
289 pub const operation_spec = op_specs.leaf(.{
290 .mnemonic = "extract",
291 .interfaces = effects.entries(.extract),
292 .operands = 1,
293 .results = 1,
294 .required_attrs = &.{"index"},
295 });
296 pub const operation_name = operation_spec.name;
297
298 pub fn create(
299 ctx: *ir.Context,
300 loc: ir.Location,
301 vector: *ir.Value,
302 index: i64,
303 result_type: ir.Type,
304 ) !ExtractOp {
305 var builder = ir.OperationBuilder.init(ctx);
306 var state = op_specs.state(@This(), loc);
307 state.addOperands(&.{vector});
308 state.addTypes(&.{result_type});
309
310 const op = try builder.create(state);
311 errdefer op.erase();
312 const idx_attr = try getIntAttr(ctx, index);
313 try op.setAttr("index", idx_attr);
314
315 return .{ .op = op };
316 }
317
318 pub fn getResult(self: *const ExtractOp) *ir.Value {
319 return self.op.getResult(0).?;
320 }
321
322 pub fn getVector(self: ExtractOp) *ir.Value {
323 return self.op.operands.items[0].value;
324 }
325
326 pub fn getIndex(self: ExtractOp) ?i64 {
327 const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "index") orelse return null;
328 return int_attr.getValue();
329 }
330 };
331
332 pub const InsertOp = struct {
333 op: *ir.Operation,
334
335 pub const operation_spec = op_specs.leaf(.{
336 .mnemonic = "insert",
337 .interfaces = effects.entries(.insert),
338 .operands = 2,
339 .results = 1,
340 .required_attrs = &.{"index"},
341 });
342 pub const operation_name = operation_spec.name;
343
344 pub fn create(
345 ctx: *ir.Context,
346 loc: ir.Location,
347 vector: *ir.Value,
348 scalar: *ir.Value,
349 index: i64,
350 ) !InsertOp {
351 var builder = ir.OperationBuilder.init(ctx);
352 var state = op_specs.state(@This(), loc);
353 state.addOperands(&.{ vector, scalar });
354 state.addTypes(&.{vector.type});
355
356 const op = try builder.create(state);
357 errdefer op.erase();
358 const idx_attr = try getIntAttr(ctx, index);
359 try op.setAttr("index", idx_attr);
360
361 return .{ .op = op };
362 }
363
364 pub fn getResult(self: *const InsertOp) *ir.Value {
365 return self.op.getResult(0).?;
366 }
367
368 pub fn getVector(self: InsertOp) *ir.Value {
369 return self.op.operands.items[0].value;
370 }
371
372 pub fn getScalar(self: InsertOp) *ir.Value {
373 return self.op.operands.items[1].value;
374 }
375
376 pub fn getIndex(self: InsertOp) ?i64 {
377 const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "index") orelse return null;
378 return int_attr.getValue();
379 }
380 };
381
382 pub const VecCmpOp = struct {
383 op: *ir.Operation,
384
385 pub const operation_spec = op_specs.leaf(.{
386 .mnemonic = "vec_cmp",
387 .interfaces = effects.entries(.vec_cmp),
388 .operands = 2,
389 .results = 1,
390 .required_attrs = &.{"predicate"},
391 .dynamic_traits = .{ir.traits.SameTypeOperands},
392 });
393 pub const operation_name = operation_spec.name;
394
395 pub fn create(
396 ctx: *ir.Context,
397 loc: ir.Location,
398 predicate: CmpPredicate,
399 lhs: *ir.Value,
400 rhs: *ir.Value,
401 result_type: ir.Type,
402 ) !VecCmpOp {
403 var builder = ir.OperationBuilder.init(ctx);
404 var state = op_specs.state(@This(), loc);
405 state.addOperands(&.{ lhs, rhs });
406 state.addTypes(&.{result_type});
407
408 const op = try builder.create(state);
409 errdefer op.erase();
410
411 var buf: [8]u8 = undefined;
412 const pred_str = try std.fmt.bufPrint(&buf, "{s}", .{predicate.toString()});
413 const pred_attr = try ctx.getDialectAttr(ArithDialect.name ++ ".predicate", pred_str);
414 try op.setAttr("predicate", pred_attr);
415
416 return .{ .op = op };
417 }
418
419 pub fn getResult(self: *const VecCmpOp) *ir.Value {
420 return self.op.getResult(0).?;
421 }
422
423 pub fn getPredicate(self: VecCmpOp) ?CmpPredicate {
424 if (self.op.getAttrAs(ir.Attribute.DialectAttr, "predicate")) |dialect_attr| {
425 inline for (
426 @typeInfo(CmpPredicate).@"enum".field_names,
427 std.meta.tags(CmpPredicate),
428 ) |field_name, value| {
429 if (std.mem.eql(u8, dialect_attr.payload, field_name)) {
430 return value;
431 }
432 }
433 }
434 return null;
435 }
436 };
437
438 pub const VecShuffleOp = struct {
439 op: *ir.Operation,
440
441 pub const operation_spec = op_specs.leaf(.{
442 .mnemonic = "vec_shuffle",
443 .interfaces = effects.entries(.vec_shuffle),
444 .operands = 1,
445 .results = 1,
446 .required_attrs = &.{"indices"},
447 });
448 pub const operation_name = operation_spec.name;
449
450 pub const IndexParseError = error{
451 MissingIndices,
452 InvalidIndices,
453 TooManyIndices,
454 };
455
456 pub fn create(
457 ctx: *ir.Context,
458 loc: ir.Location,
459 vector: *ir.Value,
460 result_type: ir.Type,
461 indices: []const i64,
462 ) !VecShuffleOp {
463 var builder = ir.OperationBuilder.init(ctx);
464 var state = op_specs.state(@This(), loc);
465 state.addOperands(&.{vector});
466 state.addTypes(&.{result_type});
467
468 const op = try builder.create(state);
469 errdefer op.erase();
470
471 var buf: [256]u8 = undefined;
472 var pos: usize = 0;
473 for (indices, 0..) |idx, i| {
474 if (i != 0) {
475 if (pos >= buf.len) return error.OutOfMemory;
476 buf[pos] = ',';
477 pos += 1;
478 }
479 pos = format.appendFmt(buf[0..], pos, "{d}", .{idx}) catch return error.OutOfMemory;
480 }
481
482 const indices_attr = try ctx.getStringAttr(buf[0..pos]);
483 try op.setAttr("indices", indices_attr);
484
485 return .{ .op = op };
486 }
487
488 pub fn getResult(self: *const VecShuffleOp) *ir.Value {
489 return self.op.getResult(0).?;
490 }
491
492 pub fn getVector(self: VecShuffleOp) *ir.Value {
493 return self.op.operands.items[0].value;
494 }
495
496 pub fn getIndices(self: VecShuffleOp, out: []i64) IndexParseError![]i64 {
497 const str_attr = self.op.getAttrAs(ir.Attribute.StringAttr, "indices") orelse {
498 if (self.op.getAttr("indices") == null) return error.MissingIndices;
499 return error.InvalidIndices;
500 };
501
502 var it = std.mem.splitScalar(u8, str_attr.value, ',');
503 var count: usize = 0;
504 while (it.next()) |chunk| {
505 if (chunk.len == 0) return error.InvalidIndices;
506 if (count >= out.len) return error.TooManyIndices;
507 const value = std.fmt.parseInt(i64, chunk, 10) catch return error.InvalidIndices;
508 out[count] = value;
509 count += 1;
510 }
511 return out[0..count];
512 }
513 };
514
515 pub const VecConstantOp = struct {
516 op: *ir.Operation,
517
518 pub const operation_spec = op_specs.leaf(.{
519 .mnemonic = "vec_constant",
520 .interfaces = effects.entries(.vec_constant),
521 .operands = 0,
522 .results = 1,
523 .required_attrs = &.{"value"},
524 .properties = ir.singleAttributePropertiesModel("arith.vec_constant.properties", "value"),
525 });
526 pub const operation_name = operation_spec.name;
527
528 pub fn createInt(
529 ctx: *ir.Context,
530 loc: ir.Location,
531 result_type: ir.Type,
532 value: i64,
533 ) !VecConstantOp {
534 var builder = ir.OperationBuilder.init(ctx);
535 var state = op_specs.state(@This(), loc);
536 state.addTypes(&.{result_type});
537 const value_attr = try getIntAttr(ctx, value);
538 const uses_properties = try state.setPropertiesAttrIfRegistered(ctx, value_attr);
539
540 const op = try builder.create(state);
541 errdefer op.erase();
542 if (!uses_properties) try op.setAttr("value", value_attr);
543
544 return .{ .op = op };
545 }
546
547 pub fn createFloat(
548 ctx: *ir.Context,
549 loc: ir.Location,
550 result_type: ir.Type,
551 value: f64,
552 ) !VecConstantOp {
553 var builder = ir.OperationBuilder.init(ctx);
554 var state = op_specs.state(@This(), loc);
555 state.addTypes(&.{result_type});
556 const value_attr = try getFloatAttr(ctx, value);
557 const uses_properties = try state.setPropertiesAttrIfRegistered(ctx, value_attr);
558
559 const op = try builder.create(state);
560 errdefer op.erase();
561 if (!uses_properties) try op.setAttr("value", value_attr);
562
563 return .{ .op = op };
564 }
565
566 pub fn getResult(self: *const VecConstantOp) *ir.Value {
567 return self.op.getResult(0).?;
568 }
569
570 pub fn getIntValue(self: VecConstantOp) ?i64 {
571 const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "value") orelse return null;
572 return int_attr.getValue();
573 }
574
575 pub fn getFloatValue(self: VecConstantOp) ?f64 {
576 const float_attr = self.op.getAttrAs(ir.Attribute.FloatAttr, "value") orelse return null;
577 return float_attr.getValue();
578 }
579 };
580
581 fn overflowOp(comptime mnemonic: []const u8, comptime kind: effects.Kind) type {
582 return struct {
583 op: *ir.Operation,
584
585 pub const operation_spec = op_specs.leaf(.{
586 .mnemonic = mnemonic,
587 .operands = 2,
588 .results = 2,
589 .operand_types = &.{
590 ir.dialects.typeConstraint.exact(0, types.type_names.int64),
591 ir.dialects.typeConstraint.exact(1, types.type_names.int64),
592 },
593 .result_types = &.{
594 ir.dialects.typeConstraint.exact(0, types.type_names.int64),
595 ir.dialects.typeConstraint.exact(1, types.type_names.boolean),
596 },
597 .traits = ir.OperationTraits{ .is_commutative = kind != .subo },
598 .interfaces = effects.entries(kind),
599 });
600 pub const operation_name = operation_spec.name;
601
602 pub fn create(
603 ctx: *ir.Context,
604 loc: ir.Location,
605 lhs: *ir.Value,
606 rhs: *ir.Value,
607 ) !@This() {
608 var builder = ir.OperationBuilder.init(ctx);
609 var state = op_specs.state(@This(), loc);
610 state.addOperands(&.{ lhs, rhs });
611 state.addTypes(&.{ lhs.type, try getScalarType(ctx, .bool) });
612 return .{ .op = try builder.create(state) };
613 }
614
615 pub fn getResult(self: @This()) *ir.Value {
616 return self.op.getResult(0).?;
617 }
618
619 pub fn getOverflow(self: @This()) *ir.Value {
620 return self.op.getResult(1).?;
621 }
622 };
623 }
624
625 fn loadSpec(ctx: *ir.Context) !void {
626 try ir.dialects.loadDialectSpec(ctx, spec);
627 }
628
629 fn sameTypeOptions(comptime traits: ir.OperationTraits) ir.dialects.opSpec.Options {
630 return .{
631 .traits = traits,
632 .dynamic_traits = ir.dialects.opSpec.dynamicTraits(.{ir.traits.SameOperandsAndResultType}),
633 };
634 }
635
636 fn effectOptions(
637 comptime kind: effects.Kind,
638 base: ir.dialects.opSpec.Options,
639 ) ir.dialects.opSpec.Options {
640 var result = base;
641 result.interfaces = effects.entries(kind);
642 return result;
643 }
644
645 fn commutativeSameTypeOptions() ir.dialects.opSpec.Options {
646 return sameTypeOptions(commutative_op_traits);
647 }
648
649 pub fn getScalarType(ctx: *ir.Context, kind: ScalarTypeKind) !ir.Type {
650 try loadSpec(ctx);
651 return ctx.getDialectTypeFromName(types.scalarTypeName(kind));
652 }
653
654 pub fn getI8Type(ctx: *ir.Context) !ir.Type {
655 return getScalarType(ctx, .i8);
656 }
657
658 pub fn getI16Type(ctx: *ir.Context) !ir.Type {
659 return getScalarType(ctx, .i16);
660 }
661
662 pub fn getU8Type(ctx: *ir.Context) !ir.Type {
663 return getScalarType(ctx, .u8);
664 }
665
666 pub fn getU16Type(ctx: *ir.Context) !ir.Type {
667 return getScalarType(ctx, .u16);
668 }
669
670 pub fn getI32Type(ctx: *ir.Context) !ir.Type {
671 return getScalarType(ctx, .i32);
672 }
673
674 pub fn getU32Type(ctx: *ir.Context) !ir.Type {
675 return getScalarType(ctx, .u32);
676 }
677
678 pub fn getU64Type(ctx: *ir.Context) !ir.Type {
679 return getScalarType(ctx, .u64);
680 }
681
682 pub fn getF16Type(ctx: *ir.Context) !ir.Type {
683 return getScalarType(ctx, .f16);
684 }
685
686 pub fn getBf16Type(ctx: *ir.Context) !ir.Type {
687 return getScalarType(ctx, .bf16);
688 }
689
690 pub fn getF64Type(ctx: *ir.Context) !ir.Type {
691 return getScalarType(ctx, .f64);
692 }
693
694 pub fn getIndexType(ctx: *ir.Context) !ir.Type {
695 return getScalarType(ctx, .index);
696 }
697
698 pub fn getVecType(ctx: *ir.Context, width: u32, element_type_name: []const u8) !?ir.Type {
699 try loadSpec(ctx);
700 const vec_name = types.vectorTypeNameForElement(width, element_type_name) orelse return null;
701 return try ctx.getDialectTypeFromName(vec_name);
702 }
703
704 pub fn getVec4xF32Type(ctx: *ir.Context) !ir.Type {
705 try loadSpec(ctx);
706 return ctx.getDialectTypeFromName(types.type_names.vec4xf32);
707 }
708
709 pub fn getVec4xI32Type(ctx: *ir.Context) !ir.Type {
710 try loadSpec(ctx);
711 return ctx.getDialectTypeFromName(types.type_names.vec4xi32);
712 }
713
714 pub fn getVec8xF32Type(ctx: *ir.Context) !ir.Type {
715 try loadSpec(ctx);
716 return ctx.getDialectTypeFromName(types.type_names.vec8xf32);
717 }
718
719 pub fn getVec8xI32Type(ctx: *ir.Context) !ir.Type {
720 try loadSpec(ctx);
721 return ctx.getDialectTypeFromName(types.type_names.vec8xi32);
722 }
723
724 pub fn getIntAttr(ctx: *ir.Context, value: i64) !ir.Attribute {
725 return ctx.getI64Attr(value);
726 }
727
728 pub fn getIntValue(attr: ir.Attribute) ?i64 {
729 const int_attr = attr.cast(ir.Attribute.IntegerAttr) orelse return null;
730 return int_attr.getValue();
731 }
732
733 pub fn getFloatAttr(ctx: *ir.Context, value: f64) !ir.Attribute {
734 return ctx.getF64Attr(value);
735 }
736
737 pub fn getFloatValue(attr: ir.Attribute) ?f64 {
738 const float_attr = attr.cast(ir.Attribute.FloatAttr) orelse return null;
739 return float_attr.getValue();
740 }
741
742 pub fn getBoolAttr(ctx: *ir.Context, value: bool) !ir.Attribute {
743 return ctx.getBoolAttr(value);
744 }
745
746 pub fn getBoolValue(attr: ir.Attribute) ?bool {
747 const bool_attr = attr.cast(ir.Attribute.BoolAttr) orelse return null;
748 return bool_attr.getValue();
749 }
750 };
751
752 const ConstructorResourceCounts = struct {
753 operations: usize,
754
755 fn capture(ctx: *const ir.Context) ConstructorResourceCounts {
756 return .{
757 .operations = ctx.operationCount(),
758 };
759 }
760
761 fn expectEqual(self: ConstructorResourceCounts, ctx: *const ir.Context) !void {
762 try std.testing.expectEqual(self.operations, ctx.operationCount());
763 }
764 };
765
766 fn expectConstructorCleanup(
767 baseline: ConstructorResourceCounts,
768 ctx: *ir.Context,
769 constructed: anytype,
770 ) !void {
771 const value = constructed catch |err| {
772 try baseline.expectEqual(ctx);
773 return err;
774 };
775 value.op.erase();
776 try baseline.expectEqual(ctx);
777 }
778
779 fn checkArithConstructorAllocationFailures(allocator: std.mem.Allocator) !void {
780 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
781 defer ctx.deinit(allocator);
782 const loc = ir.Location.getUnknown();
783 const i32_type = try ArithDialect.getI32Type(&ctx);
784 const f32_type = try ArithDialect.getScalarType(&ctx, .f32);
785 const vec4xi32_type = try ArithDialect.getVec4xI32Type(&ctx);
786 const vec4xf32_type = try ArithDialect.getVec4xF32Type(&ctx);
787 var lhs = try ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 10);
788 var rhs = try ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 20);
789 var scalar = try ArithDialect.ConstantOp.createFloat(&ctx, loc, f32_type, 3.0);
790 var vector = try ArithDialect.VecConstantOp.createInt(&ctx, loc, vec4xi32_type, 1);
791 var vector_rhs = try ArithDialect.VecConstantOp.createInt(&ctx, loc, vec4xi32_type, 2);
792 const baseline = ConstructorResourceCounts.capture(&ctx);
793
794 try expectConstructorCleanup(
795 baseline,
796 &ctx,
797 ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 30),
798 );
799 try expectConstructorCleanup(
800 baseline,
801 &ctx,
802 ArithDialect.ConstantOp.createFloat(&ctx, loc, f32_type, 4.0),
803 );
804 try expectConstructorCleanup(baseline, &ctx, ArithDialect.ConstantOp.createBool(&ctx, loc, true));
805 try expectConstructorCleanup(
806 baseline,
807 &ctx,
808 ArithDialect.CmpOp.create(&ctx, loc, .lt, lhs.getResult(), rhs.getResult()),
809 );
810 try expectConstructorCleanup(
811 baseline,
812 &ctx,
813 ArithDialect.ExtractOp.create(&ctx, loc, vector.getResult(), 3, i32_type),
814 );
815 try expectConstructorCleanup(
816 baseline,
817 &ctx,
818 ArithDialect.InsertOp.create(&ctx, loc, vector.getResult(), lhs.getResult(), 2),
819 );
820 try expectConstructorCleanup(
821 baseline,
822 &ctx,
823 ArithDialect.VecCmpOp.create(
824 &ctx,
825 loc,
826 .gt,
827 vector.getResult(),
828 vector_rhs.getResult(),
829 vec4xi32_type,
830 ),
831 );
832 try expectConstructorCleanup(
833 baseline,
834 &ctx,
835 ArithDialect.VecShuffleOp.create(
836 &ctx,
837 loc,
838 vector.getResult(),
839 vec4xi32_type,
840 &.{ 3, 1, 2, 0 },
841 ),
842 );
843 try expectConstructorCleanup(
844 baseline,
845 &ctx,
846 ArithDialect.VecConstantOp.createInt(&ctx, loc, vec4xi32_type, 5),
847 );
848 try expectConstructorCleanup(
849 baseline,
850 &ctx,
851 ArithDialect.VecConstantOp.createFloat(&ctx, loc, vec4xf32_type, 6.0),
852 );
853 try std.testing.expectEqual(@as(f64, 3.0), scalar.getFloatValue().?);
854 }
855
856 test "ArithDialect constructors clean every allocation failure" {
857 try std.testing.checkAllAllocationFailures(
858 std.testing.allocator,
859 checkArithConstructorAllocationFailures,
860 .{},
861 );
862 }