lib/choir/src/dialects/tile.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_arena = @import("alloc_arena");
3 const ir = @import("../core/root.zig");
4 const interfaces = @import("../core/root.zig").interfaces;
5
6 pub const type_names = struct {
7 pub const tile = "tile";
8 pub const barrier = "tile.barrier";
9 };
10
11 pub const TileMemLevel = enum {
12 register,
13 shared,
14 global,
15
16 pub fn toString(self: TileMemLevel) []const u8 {
17 return @tagName(self);
18 }
19
20 pub fn fromString(s: []const u8) ?TileMemLevel {
21 inline for (@typeInfo(TileMemLevel).@"enum".field_names, std.meta.tags(TileMemLevel)) |field_name, value| {
22 if (std.mem.eql(u8, s, field_name)) return value;
23 }
24 return null;
25 }
26 };
27
28 pub const TileDialect = struct {
29 pub const name = type_names.tile;
30 const op_specs = ir.dialects.opSpec.dialect(@This());
31
32 pub const spec = ir.dialects.dialectSpec(@This(), .{
33 .types = ir.dialects.typeNames(type_specs),
34 });
35
36 const type_specs = struct {
37 pub const tile = ir.dialects.TypeSpec{
38 .name = type_names.tile,
39 .interfaces = &.{
40 interfaces.TypeParamInterface.entry(&tile_type_param_vtable),
41 },
42 };
43 pub const barrier = type_names.barrier;
44 };
45
46 pub const TileTypePayload = struct {
47 element_type_name: []const u8,
48 element_type: ?ir.Type,
49 mem_level: TileMemLevel,
50 shape: []u32,
51 };
52
53 const tile_type_param_vtable = interfaces.TypeParamInterface.VTable{
54 .parse = parseTileTypeParams,
55 };
56
57 pub const LoadOp = struct {
58 op: *ir.Operation,
59
60 pub const operation_spec = op_specs.leaf(.{
61 .mnemonic = "load",
62 .operands = 2,
63 .results = 1,
64 });
65 pub const operation_name = operation_spec.name;
66
67 pub fn create(
68 ctx: *ir.Context,
69 loc: ir.Location,
70 buffer: *ir.Value,
71 coord: *ir.Value,
72 result_type: ir.Type,
73 ) !LoadOp {
74 var builder = ir.OperationBuilder.init(ctx);
75 var state = op_specs.state(@This(), loc);
76 state.addOperands(&.{ buffer, coord });
77 state.addTypes(&.{result_type});
78 const op = try builder.create(state);
79 return .{ .op = op };
80 }
81
82 pub fn getResult(self: *const LoadOp) *ir.Value {
83 return self.op.getResult(0).?;
84 }
85
86 pub fn getBuffer(self: LoadOp) *ir.Value {
87 return self.op.operands.items[0].value;
88 }
89
90 pub fn getCoord(self: LoadOp) *ir.Value {
91 return self.op.operands.items[1].value;
92 }
93 };
94
95 pub const StoreOp = struct {
96 op: *ir.Operation,
97
98 pub const operation_spec = op_specs.leaf(.{
99 .mnemonic = "store",
100 .operands = 3,
101 .results = 0,
102 });
103 pub const operation_name = operation_spec.name;
104
105 pub fn create(
106 ctx: *ir.Context,
107 loc: ir.Location,
108 tile: *ir.Value,
109 buffer: *ir.Value,
110 coord: *ir.Value,
111 ) !StoreOp {
112 var builder = ir.OperationBuilder.init(ctx);
113 var state = op_specs.state(@This(), loc);
114 state.addOperands(&.{ tile, buffer, coord });
115 const op = try builder.create(state);
116 return .{ .op = op };
117 }
118
119 pub fn getTile(self: StoreOp) *ir.Value {
120 return self.op.operands.items[0].value;
121 }
122
123 pub fn getBuffer(self: StoreOp) *ir.Value {
124 return self.op.operands.items[1].value;
125 }
126
127 pub fn getCoord(self: StoreOp) *ir.Value {
128 return self.op.operands.items[2].value;
129 }
130 };
131
132 pub const MmaOp = struct {
133 op: *ir.Operation,
134
135 pub const operation_spec = op_specs.leaf(.{
136 .mnemonic = "mma",
137 .operands = 3,
138 .results = 1,
139 });
140 pub const operation_name = operation_spec.name;
141
142 pub fn create(
143 ctx: *ir.Context,
144 loc: ir.Location,
145 a: *ir.Value,
146 b: *ir.Value,
147 c: *ir.Value,
148 result_type: ir.Type,
149 ) !MmaOp {
150 var builder = ir.OperationBuilder.init(ctx);
151 var state = op_specs.state(@This(), loc);
152 state.addOperands(&.{ a, b, c });
153 state.addTypes(&.{result_type});
154 const op = try builder.create(state);
155 return .{ .op = op };
156 }
157
158 pub fn getResult(self: *const MmaOp) *ir.Value {
159 return self.op.getResult(0).?;
160 }
161
162 pub fn getA(self: MmaOp) *ir.Value {
163 return self.op.operands.items[0].value;
164 }
165
166 pub fn getB(self: MmaOp) *ir.Value {
167 return self.op.operands.items[1].value;
168 }
169
170 pub fn getC(self: MmaOp) *ir.Value {
171 return self.op.operands.items[2].value;
172 }
173 };
174
175 pub const CopyOp = struct {
176 op: *ir.Operation,
177
178 pub const operation_spec = op_specs.leaf(.{
179 .mnemonic = "copy",
180 .operands = 2,
181 .results = 0,
182 });
183 pub const operation_name = operation_spec.name;
184
185 pub fn create(
186 ctx: *ir.Context,
187 loc: ir.Location,
188 src: *ir.Value,
189 dst: *ir.Value,
190 ) !CopyOp {
191 var builder = ir.OperationBuilder.init(ctx);
192 var state = op_specs.state(@This(), loc);
193 state.addOperands(&.{ src, dst });
194 const op = try builder.create(state);
195 return .{ .op = op };
196 }
197
198 pub fn getSrc(self: CopyOp) *ir.Value {
199 return self.op.operands.items[0].value;
200 }
201
202 pub fn getDst(self: CopyOp) *ir.Value {
203 return self.op.operands.items[1].value;
204 }
205 };
206
207 pub const BarrierOp = struct {
208 op: *ir.Operation,
209
210 pub const operation_spec = op_specs.leaf(.{
211 .mnemonic = "barrier",
212 .operands = 0,
213 .results = 0,
214 });
215 pub const operation_name = operation_spec.name;
216
217 pub fn create(ctx: *ir.Context, loc: ir.Location) !BarrierOp {
218 var builder = ir.OperationBuilder.init(ctx);
219 const state = op_specs.state(@This(), loc);
220 const op = try builder.create(state);
221 return .{ .op = op };
222 }
223 };
224
225 pub const ArriveOp = struct {
226 op: *ir.Operation,
227
228 pub const operation_spec = op_specs.leaf(.{
229 .mnemonic = "arrive",
230 .operands = 1,
231 .results = 0,
232 });
233 pub const operation_name = operation_spec.name;
234
235 pub fn create(ctx: *ir.Context, loc: ir.Location, barrier: *ir.Value) !ArriveOp {
236 var builder = ir.OperationBuilder.init(ctx);
237 var state = op_specs.state(@This(), loc);
238 state.addOperands(&.{barrier});
239 const op = try builder.create(state);
240 return .{ .op = op };
241 }
242
243 pub fn getBarrier(self: ArriveOp) *ir.Value {
244 return self.op.operands.items[0].value;
245 }
246 };
247
248 pub const WaitOp = struct {
249 op: *ir.Operation,
250
251 pub const operation_spec = op_specs.leaf(.{
252 .mnemonic = "wait",
253 .operands = 2,
254 .results = 0,
255 });
256 pub const operation_name = operation_spec.name;
257
258 pub fn create(ctx: *ir.Context, loc: ir.Location, barrier: *ir.Value, phase: *ir.Value) !WaitOp {
259 var builder = ir.OperationBuilder.init(ctx);
260 var state = op_specs.state(@This(), loc);
261 state.addOperands(&.{ barrier, phase });
262 const op = try builder.create(state);
263 return .{ .op = op };
264 }
265
266 pub fn getBarrier(self: WaitOp) *ir.Value {
267 return self.op.operands.items[0].value;
268 }
269
270 pub fn getPhase(self: WaitOp) *ir.Value {
271 return self.op.operands.items[1].value;
272 }
273 };
274
275 fn deinitTilePayload(allocator: std.mem.Allocator, ptr: *anyopaque) void {
276 const payload: *TileTypePayload = @ptrCast(@alignCast(ptr));
277 if (payload.shape.len > 0) {
278 allocator.free(payload.shape);
279 }
280 allocator.destroy(payload);
281 }
282
283 fn parseShape(allocator: std.mem.Allocator, shape_str: []const u8) ![]u32 {
284 if (shape_str.len == 0) return error.InvalidTileShape;
285
286 var dims = std.ArrayListUnmanaged(u32).empty;
287 errdefer dims.deinit(allocator);
288
289 var iter = std.mem.splitScalar(u8, shape_str, 'x');
290 while (iter.next()) |part| {
291 if (part.len == 0) return error.InvalidTileShape;
292 const value = std.fmt.parseInt(u32, part, 10) catch return error.InvalidTileShape;
293 try dims.append(allocator, value);
294 }
295
296 if (dims.items.len == 0) return error.InvalidTileShape;
297 return dims.toOwnedSlice(allocator);
298 }
299
300 fn parseTileTypeParams(
301 type_ptr: *const anyopaque,
302 ctx_opaque: *const interfaces.ContextOpaque,
303 ) anyerror!?interfaces.TypeParamPayload {
304 const ctx = interfaces.castContext(ir.Context, ctx_opaque);
305 const storage: *const ir.Type.DialectTypeStorage = @ptrCast(@alignCast(type_ptr));
306 if (storage.param_key.len == 0) return null;
307
308 var iter = std.mem.splitScalar(u8, storage.param_key, ',');
309 const elem_name = iter.next() orelse return null;
310 const mem_level_str = iter.next() orelse return null;
311 const shape_str = iter.next() orelse return null;
312 if (iter.next() != null) return null;
313
314 const mem_level = TileMemLevel.fromString(mem_level_str) orelse return null;
315 const allocator = ir.context.typePayloadAllocator(ctx);
316 const shape = parseShape(allocator, shape_str) catch return null;
317 errdefer allocator.free(shape);
318
319 const payload = try allocator.create(TileTypePayload);
320 payload.* = .{
321 .element_type_name = elem_name,
322 .element_type = ctx.getDialectTypeFromName(elem_name) catch null,
323 .mem_level = mem_level,
324 .shape = shape,
325 };
326 return .{ .ptr = payload, .deinit = deinitTilePayload };
327 }
328
329 fn loadSpec(ctx: *ir.Context) !void {
330 ir.dialects.loadDialectSpec(ctx, spec) catch |err| switch (err) {
331 error.ContextFrozen => {},
332 else => return err,
333 };
334 }
335
336 fn payloadFromType(ctx: *ir.Context, typ: ir.Type) ?*const TileTypePayload {
337 loadSpec(ctx) catch return null;
338 return ctx.getTypeParamPayload(typ, TileTypePayload) catch null;
339 }
340
341 pub fn getTileType(
342 ctx: *ir.Context,
343 element_type: ir.Type,
344 mem_level: TileMemLevel,
345 shape: []const u32,
346 ) !ir.Type {
347 try loadSpec(ctx);
348 if (shape.len == 0) return error.InvalidTileShape;
349
350 var buf: [128]u8 = undefined;
351 const elem_name = element_type.getDialectTypeName() orelse "unknown";
352 var pos: usize = 0;
353 pos = try ir.format.appendFmt(buf[0..], pos, "{s},{s},", .{ elem_name, mem_level.toString() });
354 for (shape, 0..) |dim, idx| {
355 if (idx > 0) {
356 pos = try ir.format.appendFmt(buf[0..], pos, "x", .{});
357 }
358 pos = try ir.format.appendFmt(buf[0..], pos, "{d}", .{dim});
359 }
360 return ctx.getDialectTypeFromNameWithKey(type_names.tile, buf[0..pos]);
361 }
362
363 pub fn getTileElementType(ctx: *ir.Context, typ: ir.Type) ?ir.Type {
364 const payload = payloadFromType(ctx, typ) orelse return null;
365 return payload.element_type;
366 }
367
368 pub fn getTileMemLevel(ctx: *ir.Context, typ: ir.Type) ?TileMemLevel {
369 const payload = payloadFromType(ctx, typ) orelse return null;
370 return payload.mem_level;
371 }
372
373 pub fn getTileShape(ctx: *ir.Context, typ: ir.Type) ?[]const u32 {
374 const payload = payloadFromType(ctx, typ) orelse return null;
375 return payload.shape;
376 }
377
378 pub fn getBarrierType(ctx: *ir.Context) !ir.Type {
379 try loadSpec(ctx);
380 return ctx.getDialectTypeFromName(type_names.barrier);
381 }
382 };
383
384 test "TileDialect.TileType roundtrip" {
385 const testing = std.testing;
386
387 var arena = alloc_arena.Arena.init(std.testing.allocator);
388 defer arena.deinit();
389 const allocator = arena.allocator();
390
391 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
392 defer ctx.deinit(allocator);
393
394 const arith = @import("arith/root.zig");
395 const elem_type = try arith.ArithDialect.getScalarType(&ctx, .f16);
396 const shape = [_]u32{ 16, 8 };
397 const tile_type = try TileDialect.getTileType(&ctx, elem_type, .shared, &shape);
398
399 try testing.expectEqual(TileMemLevel.shared, TileDialect.getTileMemLevel(&ctx, tile_type).?);
400 const parsed_shape = TileDialect.getTileShape(&ctx, tile_type).?;
401 try testing.expectEqual(@as(usize, shape.len), parsed_shape.len);
402 try testing.expectEqual(shape[0], parsed_shape[0]);
403 try testing.expectEqual(shape[1], parsed_shape[1]);
404 try testing.expect(TileDialect.getTileElementType(&ctx, tile_type).?.eql(elem_type));
405 }
406
407 test "TileDialect spec owns tile type params" {
408 const testing = std.testing;
409
410 var arena = alloc_arena.Arena.init(std.testing.allocator);
411 defer arena.deinit();
412 const allocator = arena.allocator();
413
414 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
415 defer ctx.deinit(allocator);
416
417 try ir.dialects.loadDialectSpec(&ctx, TileDialect.spec);
418
419 const type_info = ctx.lookupType(type_names.tile) orelse return error.TestExpectedType;
420 try testing.expect(type_info.hasInterface(interfaces.TypeParamInterface.id));
421 try testing.expect(ctx.lookupType(type_names.barrier) != null);
422 }
423
424 test "TileDialect specs own operation shapes and traits" {
425 const testing = std.testing;
426
427 var arena = alloc_arena.Arena.init(std.testing.allocator);
428 defer arena.deinit();
429 const allocator = arena.allocator();
430
431 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
432 defer ctx.deinit(allocator);
433
434 try ir.dialects.loadDialectSpec(&ctx, TileDialect.spec);
435
436 const Helpers = struct {
437 fn expectLeaf(op_info: *const interfaces.OperationInfo, operands: usize, results: usize) !void {
438 try testing.expect(op_info.shape.operands.allows(operands));
439 try testing.expect(!op_info.shape.operands.allows(operands + 1));
440 if (operands > 0) try testing.expect(!op_info.shape.operands.allows(operands - 1));
441 try testing.expect(op_info.shape.results.allows(results));
442 try testing.expect(!op_info.shape.results.allows(results + 1));
443 if (results > 0) try testing.expect(!op_info.shape.results.allows(results - 1));
444 try testing.expect(op_info.shape.regions.allows(0));
445 try testing.expect(!op_info.shape.regions.allows(1));
446 try testing.expect(op_info.shape.successors.allows(0));
447 try testing.expect(!op_info.shape.successors.allows(1));
448 }
449
450 fn lookup(context: *ir.Context, name: []const u8) !*const interfaces.OperationInfo {
451 return context.lookupOperation(name) orelse error.TestExpectedOperation;
452 }
453 };
454
455 const load_info = try Helpers.lookup(&ctx, TileDialect.LoadOp.operation_name);
456 try Helpers.expectLeaf(load_info, 2, 1);
457 try testing.expect(load_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);
458
459 const store_info = try Helpers.lookup(&ctx, TileDialect.StoreOp.operation_name);
460 try Helpers.expectLeaf(store_info, 3, 0);
461 try testing.expect(store_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);
462
463 const mma_info = try Helpers.lookup(&ctx, TileDialect.MmaOp.operation_name);
464 try Helpers.expectLeaf(mma_info, 3, 1);
465 try testing.expect(mma_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);
466
467 const copy_info = try Helpers.lookup(&ctx, TileDialect.CopyOp.operation_name);
468 try Helpers.expectLeaf(copy_info, 2, 0);
469 try testing.expect(copy_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);
470
471 const barrier_info = try Helpers.lookup(&ctx, TileDialect.BarrierOp.operation_name);
472 try Helpers.expectLeaf(barrier_info, 0, 0);
473 try testing.expect(barrier_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);
474
475 const arrive_info = try Helpers.lookup(&ctx, TileDialect.ArriveOp.operation_name);
476 try Helpers.expectLeaf(arrive_info, 1, 0);
477 try testing.expect(arrive_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);
478
479 const wait_info = try Helpers.lookup(&ctx, TileDialect.WaitOp.operation_name);
480 try Helpers.expectLeaf(wait_info, 2, 0);
481 try testing.expect(wait_info.getInterface(ir.interfaces.EffectOpInterface.id) == null);
482 }
483
484 test "TileDialect verifier rejects malformed operation shapes" {
485 const testing = std.testing;
486
487 var arena = alloc_arena.Arena.init(std.testing.allocator);
488 defer arena.deinit();
489 const allocator = arena.allocator();
490
491 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
492 defer ctx.deinit(allocator);
493
494 const arith = @import("arith/root.zig");
495 const loc = ir.Location.getUnknown();
496 const elem_type = try arith.ArithDialect.getScalarType(&ctx, .f32);
497 const shape = [_]u32{ 16, 8 };
498 const tile_type = try TileDialect.getTileType(&ctx, elem_type, .shared, &shape);
499
500 var state = ir.Operation.State.init(TileDialect.LoadOp.operation_name, loc);
501 state.addTypes(&.{tile_type});
502 const malformed_load = try ctx.createOperation(state);
503
504 try testing.expectError(error.OperandCountMismatch, ir.verifyOperation(malformed_load, .{ .recursive = false }));
505 }
506
507 test "TileDialect ops capture operands and results" {
508 const testing = std.testing;
509
510 var arena = alloc_arena.Arena.init(std.testing.allocator);
511 defer arena.deinit();
512 const allocator = arena.allocator();
513
514 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
515 defer ctx.deinit(allocator);
516
517 const arith = @import("arith/root.zig");
518 const memref = @import("memref.zig");
519
520 const loc = ir.Location.getUnknown();
521 const elem_type = try arith.ArithDialect.getScalarType(&ctx, .f32);
522 const index_type = try arith.ArithDialect.getIndexType(&ctx);
523 const memref_type = try memref.MemrefDialect.getMemrefType1D(&ctx, 128, elem_type, .shared);
524 var alloc = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, memref_type);
525
526 var coord = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 0);
527 const shape = [_]u32{ 16, 8 };
528 const tile_type = try TileDialect.getTileType(&ctx, elem_type, .shared, &shape);
529
530 const load = try TileDialect.LoadOp.create(&ctx, loc, alloc.getResult(), coord.getResult(), tile_type);
531 try testing.expect(load.getResult().type.eql(tile_type));
532 try testing.expect(load.getBuffer() == alloc.getResult());
533 try testing.expect(load.getCoord() == coord.getResult());
534
535 const store = try TileDialect.StoreOp.create(&ctx, loc, load.getResult(), alloc.getResult(), coord.getResult());
536 try testing.expect(store.getTile() == load.getResult());
537 try testing.expect(store.getBuffer() == alloc.getResult());
538 try testing.expect(store.getCoord() == coord.getResult());
539
540 const mma = try TileDialect.MmaOp.create(&ctx, loc, load.getResult(), load.getResult(), load.getResult(), tile_type);
541 try testing.expect(mma.getResult().type.eql(tile_type));
542 try testing.expect(mma.getA() == load.getResult());
543
544 const copy = try TileDialect.CopyOp.create(&ctx, loc, load.getResult(), load.getResult());
545 try testing.expect(copy.getSrc() == load.getResult());
546 try testing.expect(copy.getDst() == load.getResult());
547
548 const barrier_type = try TileDialect.getBarrierType(&ctx);
549 try @import("fixture/root.zig").registerTestDialect(&ctx);
550 _ = try ctx.registerOperation("test.tile.barrier", .{});
551 var builder = ir.OperationBuilder.init(&ctx);
552 var barrier_state = ir.Operation.State.init("test.tile.barrier", loc);
553 barrier_state.addTypes(&.{barrier_type});
554 const barrier_op = try builder.create(barrier_state);
555 const barrier_val = barrier_op.getResult(0).?;
556
557 _ = try TileDialect.BarrierOp.create(&ctx, loc);
558 const arrive = try TileDialect.ArriveOp.create(&ctx, loc, barrier_val);
559 try testing.expect(arrive.getBarrier() == barrier_val);
560
561 var phase = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 1);
562 const wait = try TileDialect.WaitOp.create(&ctx, loc, barrier_val, phase.getResult());
563 try testing.expect(wait.getBarrier() == barrier_val);
564 try testing.expect(wait.getPhase() == phase.getResult());
565 }