lib/choir/src/dialects/registry.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_arena = @import("alloc_arena");
3 const ir = @import("../core/root.zig");
4 const core_dialects = @import("../core/root.zig").dialects;
5 const interfaces = @import("../core/root.zig").interfaces;
6
7 const builtin = @import("builtin.zig");
8 const arith = @import("arith/root.zig");
9 const memref = @import("memref.zig");
10 const rc = @import("rc.zig");
11 const scf = @import("scf.zig");
12 const func = @import("func.zig");
13 const tile = @import("tile.zig");
14 const aarch64 = @import("aarch64.zig");
15
16 const BuiltinDialect = builtin.BuiltinDialect;
17 const MemrefDialect = memref.MemrefDialect;
18 const RcDialect = rc.RcDialect;
19 const ScfDialect = scf.ScfDialect;
20 const FuncDialect = func.FuncDialect;
21 const TileDialect = tile.TileDialect;
22 const AArch64Dialect = aarch64.AArch64Dialect;
23
24 fn loadBuiltinDialect(ctx: *ir.Context) !void {
25 try ir.dialects.loadDialectSpec(ctx, BuiltinDialect.spec);
26 }
27
28 fn loadTileDialect(ctx: *ir.Context) !void {
29 try ir.dialects.loadDialectSpec(ctx, TileDialect.spec);
30 }
31
32 fn loadAArch64Dialect(ctx: *ir.Context) !void {
33 try ir.dialects.loadDialectSpec(ctx, AArch64Dialect.spec);
34 }
35
36 fn loadArithDialect(ctx: *ir.Context) !void {
37 try ir.dialects.loadDialectSpec(ctx, arith.spec);
38 }
39
40 fn loadMemrefDialect(ctx: *ir.Context) !void {
41 try ir.dialects.loadDialectSpec(ctx, MemrefDialect.spec);
42 }
43
44 fn loadRcDialect(ctx: *ir.Context) !void {
45 try ir.dialects.loadDialectSpec(ctx, RcDialect.spec);
46 }
47
48 fn loadScfDialect(ctx: *ir.Context) !void {
49 try ir.dialects.loadDialectSpec(ctx, ScfDialect.spec);
50 }
51
52 fn loadFuncDialect(ctx: *ir.Context) !void {
53 try ir.dialects.loadDialectSpec(ctx, FuncDialect.spec);
54 }
55
56 pub const choir_registry = core_dialects.DialectRegistrySpec{
57 .entries = &[_]core_dialects.DialectRegistryEntry{
58 .{ .name = "builtin", .load = loadBuiltinDialect },
59 .{ .name = "tile", .load = loadTileDialect },
60 .{ .name = "aarch64", .load = loadAArch64Dialect },
61 .{ .name = "arith", .load = loadArithDialect },
62 .{ .name = "memref", .load = loadMemrefDialect },
63 .{ .name = "rc", .load = loadRcDialect },
64 .{ .name = "scf", .load = loadScfDialect },
65 .{ .name = "func", .load = loadFuncDialect },
66 },
67 };
68
69 fn registerChoirDialectLoader(ctx: *ir.Context, entry: core_dialects.DialectRegistryEntry) !void {
70 ctx.registerDialectLoader(entry.name, entry.load) catch |err| switch (err) {
71 error.DuplicateDialectLoader => {
72 const existing = ctx.dialect_registry.loaders.get(entry.name) orelse return err;
73 if (existing != entry.load) return err;
74 },
75 else => return err,
76 };
77 }
78
79 pub fn registerChoirDialect(ctx: *ir.Context) !void {
80 inline for (choir_registry.entries) |entry| {
81 try registerChoirDialectLoader(ctx, entry);
82 }
83 }
84
85 pub fn registerAllDialects(ctx: *ir.Context) !void {
86 try registerChoirDialect(ctx);
87 inline for (choir_registry.entries) |entry| {
88 _ = try ctx.getOrLoadDialect(entry.name);
89 }
90 }
91
92 test {
93 std.testing.refAllDecls(@This());
94 }
95
96 test "DialectSpec typeNames skips non-string entries without aborting" {
97 const testing = std.testing;
98
99 var arena = alloc_arena.Arena.init(std.testing.allocator);
100 defer arena.deinit();
101 const allocator = arena.allocator();
102
103 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
104 defer ctx.deinit(allocator);
105
106 const Names = struct {
107 pub const first = "choir.test.first";
108 pub const bad: *const [2]u16 = &[_]u16{ 1, 2 };
109 pub const second = "choir.test.second";
110 };
111
112 try ir.dialects.loadDialectSpec(&ctx, .{
113 .name = "choir.test",
114 .types = ir.dialects.typeNames(Names),
115 });
116
117 try testing.expect(ctx.lookupType(Names.first) != null);
118 try testing.expect(ctx.lookupType(Names.second) != null);
119 }
120
121 test "registerChoirDialect lazily loads builtin, arith, memref, rc, scf, and func for strict contexts" {
122 const testing = std.testing;
123
124 var arena = alloc_arena.Arena.init(std.testing.allocator);
125 defer arena.deinit();
126 const allocator = arena.allocator();
127
128 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
129 defer ctx.deinit(allocator);
130
131 try testing.expectError(error.UnknownDialect, ctx.getOrLoadDialect("builtin"));
132
133 try registerChoirDialect(&ctx);
134 _ = try ctx.getOrLoadDialect("builtin");
135
136 try testing.expect(ctx.lookupOperation("arith.add") == null);
137 try testing.expect(ctx.lookupOperation("memref.load") == null);
138 try testing.expect(ctx.lookupOperation("rc.retain") == null);
139 try testing.expect(ctx.lookupOperation("scf.for") == null);
140 try testing.expect(ctx.lookupOperation("func.func") == null);
141 try testing.expect(ctx.lookupOperation("func.call") == null);
142 try testing.expect(ctx.lookupOperation("func.return") == null);
143 try testing.expect(ctx.lookupType(arith.type_names.int32) == null);
144 try testing.expect(ctx.lookupType(MemrefDialect.name) == null);
145
146 _ = try ctx.getOrLoadDialect("arith");
147 try testing.expect(ctx.lookupOperation("arith.add") != null);
148 try testing.expect(ctx.lookupType(arith.type_names.int32) != null);
149 try testing.expect(ctx.lookupType(arith.type_names.vec4xf32) != null);
150
151 const loc = ir.Location.getUnknown();
152 const arith_state = ir.Operation.State.init("arith.add", loc);
153 _ = try ctx.createOperation(arith_state);
154
155 _ = try ctx.getOrLoadDialect("memref");
156 try testing.expect(ctx.lookupOperation("memref.load") != null);
157 try testing.expect(ctx.lookupType(MemrefDialect.name) != null);
158 const memref_type = try ctx.getDialectTypeFromNameWithKey("memref", "64,arith.f32,device");
159 try testing.expect((try ctx.getTypeParamPayload(memref_type, MemrefDialect.MemrefTypePayload)) != null);
160 try testing.expect(ctx.typeInterface(memref_type, interfaces.ShapedTypeInterface) != null);
161 const memref_state = ir.Operation.State.init("memref.load", loc);
162 _ = try ctx.createOperation(memref_state);
163
164 _ = try ctx.getOrLoadDialect("rc");
165 try testing.expect(ctx.lookupOperation("rc.retain") != null);
166 try testing.expect(ctx.lookupOperation("rc.release") != null);
167 const rc_state = ir.Operation.State.init("rc.retain", loc);
168 _ = try ctx.createOperation(rc_state);
169
170 _ = try ctx.getOrLoadDialect("scf");
171 try testing.expect(ctx.lookupOperation("scf.for") != null);
172 const scf_state = ir.Operation.State.init("scf.for", loc);
173 _ = try ctx.createOperation(scf_state);
174
175 _ = try ctx.getOrLoadDialect("func");
176 const func_info = ctx.lookupOperation("func.func") orelse return error.TestExpectedOperation;
177 try testing.expect(func_info.hasInterface(interfaces.SymbolOpInterface.id));
178 const call_info = ctx.lookupOperation("func.call") orelse return error.TestExpectedOperation;
179 try testing.expect(call_info.hasInterface(interfaces.CallOpInterface.id));
180 try testing.expect(call_info.hasInterface(interfaces.SymbolUserOpInterface.id));
181 try testing.expect(ctx.lookupOperation("func.return") != null);
182 const func_state = ir.Operation.State.init("func.func", loc);
183 _ = try ctx.createOperation(func_state);
184
185 const bad_state = ir.Operation.State.init("builtin.unknown", loc);
186 try testing.expectError(error.UnknownOperation, ctx.createOperation(bad_state));
187 }
188
189 test "registerAllDialects is repeatable after lazy registry registration" {
190 const testing = std.testing;
191
192 var arena = alloc_arena.Arena.init(std.testing.allocator);
193 defer arena.deinit();
194 const allocator = arena.allocator();
195
196 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
197 defer ctx.deinit(allocator);
198
199 try registerChoirDialect(&ctx);
200 try registerAllDialects(&ctx);
201 try registerAllDialects(&ctx);
202
203 inline for (choir_registry.entries) |entry| {
204 try testing.expect(ctx.isDialectLoaded(entry.name));
205 }
206 }
207
208 test "Choir dialect operation metadata matches co-owned capacity witnesses" {
209 const testing = std.testing;
210
211 var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);
212 defer ctx.deinit(testing.allocator);
213 try registerAllDialects(&ctx);
214
215 var maximum_traits: usize = 0;
216 var maximum_interfaces: usize = 0;
217 var maximum_inherent_attribute_names: usize = 0;
218 var maximum_required_attribute_names: usize = 0;
219 var total_inherent_attribute_names: usize = 0;
220 var total_required_attribute_names: usize = 0;
221 var nonempty_inherent_attribute_names: usize = 0;
222 var nonempty_required_attribute_names: usize = 0;
223 var operation_count: usize = 0;
224 var operation_iter = ctx.dialect_registry.operation_registry.ops.valueIterator();
225 while (operation_iter.next()) |info_ptr| {
226 operation_count += 1;
227 maximum_traits = @max(maximum_traits, info_ptr.*.getDynamicTraitIds().len);
228 maximum_interfaces = @max(maximum_interfaces, info_ptr.*.getNumInterfaces());
229 const inherent_attribute_names = info_ptr.*.getInherentAttributeNames().len;
230 const required_attribute_names = info_ptr.*.getRequiredAttributeNames().len;
231 maximum_inherent_attribute_names = @max(
232 maximum_inherent_attribute_names,
233 inherent_attribute_names,
234 );
235 maximum_required_attribute_names = @max(
236 maximum_required_attribute_names,
237 required_attribute_names,
238 );
239 total_inherent_attribute_names += inherent_attribute_names;
240 total_required_attribute_names += required_attribute_names;
241 if (inherent_attribute_names != 0) nonempty_inherent_attribute_names += 1;
242 if (required_attribute_names != 0) nonempty_required_attribute_names += 1;
243 }
244 try testing.expectEqual(
245 interfaces.OperationInfo.dynamic_trait_inline_capacity,
246 maximum_traits,
247 );
248 try testing.expectEqual(
249 @as(usize, 3),
250 maximum_interfaces,
251 );
252 try testing.expectEqual(@as(usize, 6), maximum_inherent_attribute_names);
253 try testing.expectEqual(@as(usize, 4), maximum_required_attribute_names);
254 try testing.expectEqual(@as(usize, 87), total_inherent_attribute_names);
255 try testing.expectEqual(@as(usize, 71), total_required_attribute_names);
256 try testing.expectEqual(@as(usize, 58), nonempty_inherent_attribute_names);
257 try testing.expectEqual(@as(usize, 51), nonempty_required_attribute_names);
258 try testing.expectEqual(@as(usize, 121), operation_count);
259 }
260
261 test "standard compiler dialect specs expose exact rollback record capacities" {
262 const testing = std.testing;
263 const specs = .{
264 BuiltinDialect.spec,
265 arith.spec,
266 FuncDialect.spec,
267 MemrefDialect.spec,
268 ScfDialect.spec,
269 };
270
271 var operation_count: usize = 0;
272 var type_count: usize = 0;
273 var region_count: usize = 0;
274 inline for (specs) |spec| {
275 operation_count += spec.operations.len;
276 type_count += spec.types.len;
277 if (spec.operations.len != 0 or spec.types.len != 0) region_count += 1;
278 }
279
280 try testing.expectEqual(@as(usize, 72), operation_count);
281 try testing.expectEqual(@as(usize, 64), type_count);
282 try testing.expectEqual(@as(usize, 5), region_count);
283 }
284
285 test "registerChoirDialect lazily loads first-class builtin dialect (tick 41 rename)" {
286 const testing = std.testing;
287
288 var arena = alloc_arena.Arena.init(std.testing.allocator);
289 defer arena.deinit();
290 const allocator = arena.allocator();
291
292 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
293 defer ctx.deinit(allocator);
294
295 try registerChoirDialect(&ctx);
296
297 try testing.expect(ctx.lookupOperation("builtin.module") == null);
298
299 _ = try ctx.getOrLoadDialect("builtin");
300
301 try testing.expect(ctx.lookupOperation("builtin.module") != null);
302
303 try testing.expectError(error.UnknownDialect, ctx.getOrLoadDialect("choir"));
304 }
305
306 test "registerChoirDialect leaves accelerator dialects to package extensions" {
307 const testing = std.testing;
308
309 var arena = alloc_arena.Arena.init(std.testing.allocator);
310 defer arena.deinit();
311 const allocator = arena.allocator();
312
313 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
314 defer ctx.deinit(allocator);
315
316 try registerChoirDialect(&ctx);
317
318 inline for (.{ "gpu", "spirv", "nvptx" }) |dialect_name| {
319 try testing.expectError(error.UnknownDialect, ctx.getOrLoadDialect(dialect_name));
320 }
321 }
322
323 test "registerChoirDialect lazily loads first-class tile dialect (tick 40 rename)" {
324 const testing = std.testing;
325
326 var arena = alloc_arena.Arena.init(std.testing.allocator);
327 defer arena.deinit();
328 const allocator = arena.allocator();
329
330 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
331 defer ctx.deinit(allocator);
332
333 try registerChoirDialect(&ctx);
334
335 try testing.expect(ctx.lookupOperation("tile.load") == null);
336 try testing.expect(ctx.lookupType(tile.type_names.tile) == null);
337
338 _ = try ctx.getOrLoadDialect("tile");
339
340 try testing.expect(ctx.lookupOperation("tile.load") != null);
341 try testing.expect(ctx.lookupOperation("tile.store") != null);
342 try testing.expect(ctx.lookupOperation("tile.mma") != null);
343 try testing.expect(ctx.lookupOperation("tile.copy") != null);
344 const tile_type_info = ctx.lookupType(tile.type_names.tile) orelse return error.TestExpectedType;
345 try testing.expect(tile_type_info.hasInterface(interfaces.TypeParamInterface.id));
346 try testing.expect(ctx.lookupType(tile.type_names.barrier) != null);
347 }
348
349 test "registerChoirDialect lazily loads first-class aarch64 dialect (tick 40 rename)" {
350 const testing = std.testing;
351
352 var arena = alloc_arena.Arena.init(std.testing.allocator);
353 defer arena.deinit();
354 const allocator = arena.allocator();
355
356 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
357 defer ctx.deinit(allocator);
358
359 try registerChoirDialect(&ctx);
360
361 try testing.expect(ctx.lookupOperation("aarch64.add") == null);
362 try testing.expect(ctx.lookupType(aarch64.type_names.gpr) == null);
363 try testing.expect(ctx.lookupType(aarch64.type_names.fpr) == null);
364
365 _ = try ctx.getOrLoadDialect("aarch64");
366
367 try testing.expect(ctx.lookupOperation("aarch64.add") != null);
368 try testing.expect(ctx.lookupOperation("aarch64.add_imm") != null);
369 try testing.expect(ctx.lookupOperation("aarch64.get_reg") != null);
370 try testing.expect(ctx.lookupType(aarch64.type_names.gpr) != null);
371 try testing.expect(ctx.lookupType(aarch64.type_names.fpr) != null);
372 }
373
374 test "memref type reconstruction lazy-loads interfaces in strict contexts" {
375 const testing = std.testing;
376
377 var arena = alloc_arena.Arena.init(std.testing.allocator);
378 defer arena.deinit();
379 const allocator = arena.allocator();
380
381 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
382 defer ctx.deinit(allocator);
383
384 try registerChoirDialect(&ctx);
385 const memref_type = try ctx.getDialectTypeFromNameWithKey("memref", "64,arith.f32,device");
386
387 try testing.expect(ctx.lookupType(MemrefDialect.name) != null);
388 try testing.expect((try ctx.getTypeParamPayload(memref_type, MemrefDialect.MemrefTypePayload)) != null);
389 try testing.expect(ctx.typeInterface(memref_type, interfaces.ShapedTypeInterface) != null);
390 }