lib/choir/src/dialects/arith/effects.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const ir = @import("../../core/root.zig");
3 const arith = @import("root.zig");
4 const types = arith.types;
5 const scalar = arith.scalar;
6 const facts = ir.interfaces.effects;
7
8 /// Each registered operation selects one declaration, independently of pass policy.
9 pub const Kind = enum {
10 constant,
11 add,
12 sub,
13 mul,
14 div,
15 rem,
16 neg,
17 abs,
18 min,
19 max,
20 umulhi,
21 @"and",
22 @"or",
23 xor,
24 not,
25 popcount,
26 shl,
27 shr,
28 ushr,
29 fma,
30 sqrt,
31 exp,
32 log,
33 tanh,
34 sin,
35 cos,
36 tan,
37 pow,
38 floor,
39 round,
40 trunc,
41 cmp,
42 select,
43 cast,
44 bitcast,
45 splat,
46 extract,
47 insert,
48 vec_cmp,
49 vec_shuffle,
50 vec_constant,
51 addo,
52 subo,
53 mulo,
54 };
55
56 pub fn entries(comptime kind: Kind) []const ir.interfaces.InterfaceEntry {
57 return &.{facts.EffectOpInterface.entryFor(.{
58 .capacity = .{ .entries = 6, .per_result = 1 },
59 .enumerate = struct {
60 fn enumerate(op: *const ir.Operation, collector: *facts.Collector) void {
61 declare(op, collector, kind);
62 }
63 }.enumerate,
64 })};
65 }
66
67 fn declare(op: *const ir.Operation, collector: *facts.Collector, kind: Kind) void {
68 collector.valueResults(op);
69 if (kind == .addo or kind == .subo or kind == .mulo) {
70 if (op.getNumOperands() != 2 or op.getNumResults() != 2) return;
71 if (op.regions.items.len != 0) return;
72 for (op.getOperandValues()) |operand| {
73 if (types.scalarKindFromType(operand.type) != .i64) return;
74 }
75 if (types.scalarKindFromType(op.results.items[0].type) != .i64) return;
76 if (types.scalarKindFromType(op.results.items[1].type) != .bool) return;
77 collector.complete = true;
78 return;
79 }
80 if (op.getNumResults() != 1 or op.regions.items.len != 0) return;
81 const typ = op.results.items[0].type;
82 if (kind == .constant) {
83 if (op.getNumOperands() != 0) return;
84 collector.complete = constantMatches(op, types.scalarKindFromType(typ) orelse return);
85 return;
86 }
87 if (vectorOperation(kind)) return declareVector(op, collector, kind);
88 const result_kind = types.scalarKindFromType(typ) orelse return;
89 if (kind == .cast or kind == .bitcast) return declareCast(op, collector, kind, result_kind);
90 if (kind == .cmp) return declareComparison(op, collector, result_kind);
91 if (kind == .select) return declareSelect(op, collector, result_kind);
92 if (op.getNumOperands() != arity(kind)) return;
93 for (op.getOperandValues()) |operand| if (!operand.type.eql(typ)) return;
94 if (!supports(kind, result_kind)) return;
95 collector.complete = true;
96 if (types.scalarKindIsFloat(result_kind)) return floatingEnvironment(op, collector);
97 switch (kind) {
98 .div, .rem => division(op, collector, result_kind),
99 .shl, .shr, .ushr => shift(op, collector, result_kind),
100 else => {},
101 }
102 }
103
104 fn arity(kind: Kind) usize {
105 return switch (kind) {
106 .neg,
107 .abs,
108 .not,
109 .popcount,
110 .sqrt,
111 .exp,
112 .log,
113 .tanh,
114 .sin,
115 .cos,
116 .tan,
117 .floor,
118 .round,
119 .trunc,
120 => 1,
121 .fma => 3,
122 else => 2,
123 };
124 }
125
126 fn supports(kind: Kind, typ: types.ScalarKind) bool {
127 const integer = types.scalarKindIsInteger(typ);
128 const floating = types.scalarKindIsFloat(typ);
129 return switch (kind) {
130 .add, .sub, .mul, .div, .neg, .abs, .min, .max => integer or floating,
131 .@"and", .@"or", .xor, .not => integer or typ == .bool,
132 .rem, .popcount, .umulhi, .shl, .shr, .ushr => integer,
133 .fma, .sqrt, .exp, .log, .tanh, .sin, .cos, .tan, .pow, .floor, .round, .trunc => floating,
134 else => false,
135 };
136 }
137
138 /// Ruling 36: every floating declaration uses this one Context premise.
139 fn floatingEnvironment(op: *const ir.Operation, collector: *facts.Collector) void {
140 if (op.getContext().arithmetic_policy.permitsFloatingValues()) {
141 collector.append(.{ .premise = .floating_environment });
142 return;
143 }
144 collector.append(.{ .event = .{
145 .kind = .state_observe,
146 .resource = .{
147 .subject = .{ .global = "arithmetic.environment" },
148 .state_key = "floating_environment",
149 },
150 } });
151 collector.append(.{ .requirement = .{ .kind = .execution_context, .subject = .operation } });
152 }
153
154 fn failure(
155 collector: *facts.Collector,
156 requirement: facts.RequirementKind,
157 operand: usize,
158 name: []const u8,
159 ) void {
160 collector.append(.{ .requirement = .{
161 .kind = requirement,
162 .subject = .{ .operand = operand },
163 } });
164 collector.append(.{ .event = .{ .kind = .failure, .failure_name = name } });
165 }
166
167 fn literal(value: *ir.Value) ?ir.Attribute {
168 const raw = value.getDefiningOp() orelse return null;
169 const definition: *ir.Operation = @ptrCast(@alignCast(raw));
170 if (!std.mem.eql(u8, definition.name.name, "arith.constant")) return null;
171 return definition.getAttr("value");
172 }
173
174 fn literalInt(op: *const ir.Operation, index: usize) ?i64 {
175 const attr = literal(op.operands.items[index].value) orelse return null;
176 return (attr.cast(ir.Attribute.IntegerAttr) orelse return null).getValue();
177 }
178
179 fn division(op: *const ir.Operation, collector: *facts.Collector, typ: types.ScalarKind) void {
180 const bits = types.scalarBitWidth(typ);
181 const rhs = literalInt(op, 1);
182 const nonzero = if (rhs) |value| scalar.maskToBits(value, bits) != 0 else false;
183 if (!nonzero) failure(collector, .nonzero, 1, "DivisionByZero");
184 if (types.scalarKindIsUnsignedInteger(typ) or typ == .index) return;
185 const representable = if (rhs) |right|
186 scalar.truncate(right, bits) != -1 or if (literalInt(op, 0)) |left|
187 @as(i128, scalar.truncate(left, bits)) != scalar.intLimits(bits).min
188 else
189 false
190 else
191 false;
192 if (!representable) failure(collector, .quotient_representable, 0, "SignedDivisionOverflow");
193 }
194
195 fn shift(op: *const ir.Operation, collector: *facts.Collector, typ: types.ScalarKind) void {
196 if (literalInt(op, 1)) |count| {
197 if (scalar.shiftCount(count, types.scalarBitWidth(typ)) != null) return;
198 }
199 failure(collector, .in_bounds, 1, "InvalidShiftAmount");
200 }
201
202 fn constantMatches(op: *const ir.Operation, kind: types.ScalarKind) bool {
203 const attr = op.getAttr("value") orelse return false;
204 return switch (types.scalarDescriptor(kind).class) {
205 .boolean => attr.cast(ir.Attribute.BoolAttr) != null or
206 attr.cast(ir.Attribute.IntegerAttr) != null,
207 .signed_integer, .unsigned_integer, .index => attr.cast(ir.Attribute.IntegerAttr) != null,
208 .float => attr.cast(ir.Attribute.FloatAttr) != null,
209 };
210 }
211
212 fn declareCast(
213 op: *const ir.Operation,
214 collector: *facts.Collector,
215 kind: Kind,
216 result: types.ScalarKind,
217 ) void {
218 if (op.getNumOperands() != 1) return;
219 const source = types.scalarKindFromType(op.operands.items[0].value.type) orelse return;
220 if (kind == .bitcast) {
221 if (types.scalarBitWidth(source) != types.scalarBitWidth(result)) return;
222 } else {
223 if (source == .bool or result == .bool) {
224 if (source != result) return;
225 }
226 }
227 collector.complete = true;
228 if (kind == .cast and types.scalarKindIsFloat(source) and types.scalarKindIsInteger(result)) {
229 floatConversion(op, collector, source, result);
230 }
231 if (types.scalarKindIsFloat(source) or types.scalarKindIsFloat(result)) {
232 floatingEnvironment(op, collector);
233 }
234 }
235
236 fn floatConversion(
237 op: *const ir.Operation,
238 collector: *facts.Collector,
239 source: types.ScalarKind,
240 result: types.ScalarKind,
241 ) void {
242 const admitted = admitted: {
243 const attr = literal(op.operands.items[0].value) orelse break :admitted false;
244 const value = attr.cast(ir.Attribute.FloatAttr) orelse break :admitted false;
245 const rounded = arith.eval.roundedFloat(value.getValue(), source) catch
246 break :admitted false;
247 break :admitted scalar.floatToInt(
248 rounded,
249 types.scalarBitWidth(result),
250 types.scalarKindIsSignedInteger(result) and result != .index,
251 ) != null;
252 };
253 if (!admitted) failure(collector, .conversion_representable, 0, "InvalidFloatToInteger");
254 }
255
256 fn validPredicate(op: *const ir.Operation, kind: types.ScalarKind) bool {
257 const attr = op.getAttrAs(ir.Attribute.DialectAttr, "predicate") orelse return false;
258 const predicate = std.meta.stringToEnum(arith.CmpPredicate, attr.payload) orelse return false;
259 if (kind == .bool) return predicate == .eq or predicate == .ne;
260 if (!types.scalarKindIsFloat(kind)) return true;
261 return switch (predicate) {
262 .eq, .ne, .lt, .le, .gt, .ge => true,
263 else => false,
264 };
265 }
266
267 fn declareComparison(
268 op: *const ir.Operation,
269 collector: *facts.Collector,
270 result: types.ScalarKind,
271 ) void {
272 if (result != .bool or op.getNumOperands() != 2) return;
273 const lhs = op.operands.items[0].value.type;
274 if (!lhs.eql(op.operands.items[1].value.type)) return;
275 const kind = types.scalarKindFromType(lhs) orelse return;
276 if (!validPredicate(op, kind)) return;
277 collector.complete = true;
278 if (types.scalarKindIsFloat(kind)) floatingEnvironment(op, collector);
279 }
280
281 fn declareSelect(
282 op: *const ir.Operation,
283 collector: *facts.Collector,
284 kind: types.ScalarKind,
285 ) void {
286 if (op.getNumOperands() != 3) return;
287 if (types.scalarKindFromType(op.operands.items[0].value.type) != .bool) return;
288 for (op.getOperandValues()[1..]) |operand| {
289 if (!operand.type.eql(op.results.items[0].type)) return;
290 }
291 collector.complete = true;
292 if (types.scalarKindIsFloat(kind)) floatingEnvironment(op, collector);
293 }
294
295 fn vectorOperation(kind: Kind) bool {
296 return switch (kind) {
297 .splat, .extract, .insert, .vec_cmp, .vec_shuffle, .vec_constant => true,
298 else => false,
299 };
300 }
301
302 fn vector(typ: ir.Type) ?types.VectorTypeInfo {
303 return types.parseVectorTypeName(typ.getDialectTypeName() orelse return null);
304 }
305
306 fn declareVector(op: *const ir.Operation, collector: *facts.Collector, kind: Kind) void {
307 const result = op.results.items[0].type;
308 const shape = if (kind == .extract or kind == .vec_cmp) shape: {
309 if (op.getNumOperands() != (if (kind == .extract) @as(usize, 1) else 2)) return;
310 break :shape vector(op.operands.items[0].value.type) orelse return;
311 } else vector(result) orelse return;
312 const element = types.scalarKindFromTypeName(shape.elem_type_name).?;
313 if (!validVector(op, kind, shape, element)) return;
314 collector.complete = true;
315 if (types.scalarKindIsFloat(element)) floatingEnvironment(op, collector);
316 }
317
318 fn validVector(
319 op: *const ir.Operation,
320 kind: Kind,
321 shape: types.VectorTypeInfo,
322 element: types.ScalarKind,
323 ) bool {
324 const result = op.results.items[0].type;
325 switch (kind) {
326 .vec_constant => return op.getNumOperands() == 0 and constantMatches(op, element),
327 .splat => return op.getNumOperands() == 1 and
328 types.scalarKindFromType(op.operands.items[0].value.type) == element,
329 .extract => {
330 if (types.scalarKindFromType(result) != element) return false;
331 return validLane(op, shape.width);
332 },
333 .insert => {
334 if (op.getNumOperands() != 2) return false;
335 if (!op.operands.items[0].value.type.eql(result)) return false;
336 if (types.scalarKindFromType(op.operands.items[1].value.type) != element) return false;
337 return validLane(op, shape.width);
338 },
339 .vec_cmp => {
340 if (op.getNumOperands() != 2) return false;
341 if (!op.operands.items[0].value.type.eql(op.operands.items[1].value.type)) return false;
342 const mask = vector(result) orelse return false;
343 if (mask.width != shape.width) return false;
344 const mask_kind = types.scalarKindFromTypeName(mask.elem_type_name) orelse return false;
345 if (types.scalarBitWidth(mask_kind) != types.scalarBitWidth(element)) return false;
346 if (mask_kind != element and !types.scalarKindIsInteger(mask_kind)) return false;
347 return validPredicate(op, element);
348 },
349 .vec_shuffle => return validShuffle(op, shape, element),
350 else => unreachable,
351 }
352 }
353
354 fn validLane(op: *const ir.Operation, width: u32) bool {
355 const attr = op.getAttrAs(ir.Attribute.IntegerAttr, "index") orelse return false;
356 const index = attr.getValue();
357 return index >= 0 and index < width;
358 }
359
360 fn validShuffle(
361 op: *const ir.Operation,
362 result: types.VectorTypeInfo,
363 element: types.ScalarKind,
364 ) bool {
365 if (op.getNumOperands() != 1) return false;
366 const source = vector(op.operands.items[0].value.type) orelse return false;
367 if (types.scalarKindFromTypeName(source.elem_type_name) != element) return false;
368 const attr = op.getAttrAs(ir.Attribute.StringAttr, "indices") orelse return false;
369 var chunks = std.mem.splitScalar(u8, attr.value, ',');
370 for (0..result.width) |_| {
371 const chunk = chunks.next() orelse return false;
372 const index = std.fmt.parseInt(i64, chunk, 10) catch return false;
373 if (index < 0 or index >= source.width) return false;
374 }
375 return chunks.next() == null;
376 }