lib/choir/src/dialects/gpu/dialect.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("../../root.zig");
3
4 const ir = choir.ir;
5 const interfaces = choir.ir.interfaces;
6 const effects = interfaces.effects;
7 const func = choir.dialects.func;
8 const tags = @import("tags.zig");
9 const stage = @import("stage.zig");
10
11 pub const type_names = tags.type_names;
12 pub const Dimension = tags.Dimension;
13 pub const Scope = tags.Scope;
14 pub const MemoryOrder = tags.MemoryOrder;
15 pub const ShuffleMode = tags.ShuffleMode;
16 pub const WarpOpKind = tags.WarpOpKind;
17 pub const MmaShape = tags.MmaShape;
18 pub const Stage = tags.Stage;
19
20 pub const GpuDialect = struct {
21 pub const name = "gpu";
22 const op_specs = ir.dialects.opSpec.dialect(@This());
23 pub const spec = ir.dialects.dialectSpec(@This(), .{
24 .types = ir.dialects.typeNames(type_specs),
25 });
26 const symbol_table_trait = ir.dialects.trait(ir.traits.SymbolTable);
27
28 const func_symbol_vtable = interfaces.SymbolOpInterface.VTable{
29 .getSymbolName = getFuncSymbolName,
30 .setSymbolName = setFuncSymbolName,
31 .isDeclaration = isFuncDeclaration,
32 };
33
34 const type_specs = struct {
35 pub const tma_desc = type_names.tma_desc;
36 pub const mbarrier = type_names.mbarrier;
37 pub const sampled_texture = type_names.sampled_texture;
38 };
39
40 pub const StageInputOp = stage.StageInputOp;
41 pub const StageOutputOp = stage.StageOutputOp;
42 pub const PositionOp = stage.PositionOp;
43 pub const FragCoordOp = stage.FragCoordOp;
44 pub const VertexIndexOp = stage.VertexIndexOp;
45 pub const InstanceIndexOp = stage.InstanceIndexOp;
46 pub const FrontFacingOp = stage.FrontFacingOp;
47 pub const SampledTextureOp = stage.SampledTextureOp;
48 pub const SampleOp = stage.SampleOp;
49 pub const SampleLodOp = stage.SampleLodOp;
50 pub const DpdxOp = stage.DpdxOp;
51 pub const DpdyOp = stage.DpdyOp;
52 pub const FwidthOp = stage.FwidthOp;
53 pub const PushConstantOp = stage.PushConstantOp;
54 pub const UniformOp = stage.UniformOp;
55
56 pub const ModuleOp = struct {
57 op: *ir.Operation,
58
59 pub const operation_spec = op_specs.define(.{
60 .mnemonic = "module",
61 .operands = 0,
62 .results = 0,
63 .regions = .{"body"},
64 .successors = 0,
65 .dynamic_traits = .{symbol_table_trait},
66 });
67 pub const operation_name = operation_spec.name;
68
69 pub fn create(ctx: *ir.Context, loc: ir.Location) !ModuleOp {
70 try loadSpec(ctx);
71 var builder = ir.OperationBuilder.init(ctx);
72 var state = ir.Operation.State.init(operation_name, loc);
73 var body = ir.context.initRegion(ctx);
74 defer body.deinit();
75 var body_builder = ir.OperationBuilder.init(ctx);
76 _ = try body_builder.createBlock(&body, &.{}, &.{});
77 var regions = [_]*ir.Region{&body};
78 state.addRegionBodies(®ions);
79
80 const op = try builder.create(state);
81 errdefer op.erase();
82
83 return .{ .op = op };
84 }
85
86 pub fn getBody(self: ModuleOp) *ir.Region {
87 return self.op.getRegion(0).?;
88 }
89
90 pub fn getBodyBlock(self: ModuleOp) *ir.Block {
91 return self.getBody().getEntryBlock().?;
92 }
93 };
94
95 pub const FuncOp = struct {
96 op: *ir.Operation,
97
98 pub const operation_spec = op_specs.define(.{
99 .mnemonic = "func",
100 .operands = 0,
101 .regions = .{"body"},
102 .successors = 0,
103 .attrs = &.{ "kernel", ir.SymbolTable.symbol_attr_names.sym_visibility },
104 .required_attrs = &.{"sym_name"},
105 .interfaces = &.{
106 interfaces.SymbolOpInterface.entry(&func_symbol_vtable),
107 effects.EffectOpInterface.entryFor(.{ .facts = &.{.{ .region = .{
108 .index = 0,
109 .execution = .latent,
110 .may_diverge = false,
111 .captures = false,
112 } }} }),
113 },
114 });
115 pub const operation_name = operation_spec.name;
116
117 pub fn create(
118 ctx: *ir.Context,
119 loc: ir.Location,
120 func_name: []const u8,
121 input_types: []const ir.Type,
122 result_types: []const ir.Type,
123 ) !FuncOp {
124 try loadSpec(ctx);
125 var builder = ir.OperationBuilder.init(ctx);
126 var state = ir.Operation.State.init(operation_name, loc);
127 state.addTypes(result_types);
128 var body = ir.context.initRegion(ctx);
129 defer body.deinit();
130 var body_builder = ir.OperationBuilder.init(ctx);
131 _ = try body_builder.createBlockWithLoc(&body, input_types, loc);
132 var regions = [_]*ir.Region{&body};
133 state.addRegionBodies(®ions);
134
135 const op = try builder.create(state);
136 errdefer op.erase();
137
138 const name_attr = try func.FuncDialect.getSymNameAttr(ctx, func_name);
139 try op.setAttr("sym_name", name_attr);
140
141 return .{ .op = op };
142 }
143
144 pub fn createKernel(
145 ctx: *ir.Context,
146 loc: ir.Location,
147 kernel_name: []const u8,
148 input_types: []const ir.Type,
149 ) !FuncOp {
150 var func_op = try create(ctx, loc, kernel_name, input_types, &.{});
151 errdefer func_op.op.erase();
152 const kernel_attr = try func.FuncDialect.getKernelAttr(ctx);
153 try func_op.op.setAttr("kernel", kernel_attr);
154 return func_op;
155 }
156
157 pub fn getName(self: FuncOp) ?[]const u8 {
158 return ir.SymbolTable.getSymbolName(self.op);
159 }
160
161 pub fn isKernel(self: FuncOp) bool {
162 return self.op.getAttr("kernel") != null;
163 }
164
165 pub fn getBody(self: FuncOp) *ir.Region {
166 return self.op.getRegion(0).?;
167 }
168
169 pub fn getEntryBlock(self: FuncOp) *ir.Block {
170 return self.getBody().getEntryBlock().?;
171 }
172
173 pub fn getArguments(self: FuncOp) []*ir.Value {
174 return self.getEntryBlock().arguments.items;
175 }
176
177 pub fn getNumArguments(self: FuncOp) usize {
178 return self.getEntryBlock().arguments.items.len;
179 }
180
181 pub fn getArgument(self: FuncOp, index: usize) *ir.Value {
182 return self.getEntryBlock().arguments.items[index];
183 }
184
185 pub fn getResultTypes(self: FuncOp) []const ir.Type {
186 return self.op.getResultTypes();
187 }
188
189 pub fn getNumResults(self: FuncOp) usize {
190 return self.op.results.items.len;
191 }
192 };
193
194 pub const YieldOp = struct {
195 op: *ir.Operation,
196
197 pub const operation_spec = op_specs.terminator(.{ .mnemonic = "yield" });
198 pub const operation_name = operation_spec.name;
199
200 pub fn create(
201 ctx: *ir.Context,
202 loc: ir.Location,
203 operands: []const *ir.Value,
204 ) !YieldOp {
205 try loadSpec(ctx);
206 var builder = ir.OperationBuilder.init(ctx);
207 var state = ir.Operation.State.init(operation_name, loc);
208 state.addOperands(operands);
209 const op = try builder.create(state);
210 errdefer op.erase();
211 return .{ .op = op };
212 }
213
214 pub fn getOperands(self: YieldOp) []const *ir.Value {
215 return self.op.getOperandValues();
216 }
217 };
218
219 pub const LaunchOp = struct {
220 op: *ir.Operation,
221
222 const dim_attr_keys = struct {
223 pub const grid_x = "grid_x";
224 pub const grid_y = "grid_y";
225 pub const grid_z = "grid_z";
226 pub const block_x = "block_x";
227 pub const block_y = "block_y";
228 pub const block_z = "block_z";
229 };
230
231 pub const operation_spec = op_specs.leaf(.{
232 .mnemonic = "launch",
233 .interfaces = &.{gpuEffects(.launch, &.{}, &.{})},
234 .operands = ir.dialects.shape.any(),
235 .results = 0,
236 .required_attrs = &.{
237 dim_attr_keys.block_x,
238 dim_attr_keys.block_y,
239 dim_attr_keys.block_z,
240 dim_attr_keys.grid_x,
241 dim_attr_keys.grid_y,
242 dim_attr_keys.grid_z,
243 "kernel",
244 "num_kernel_args",
245 },
246 });
247 pub const operation_name = operation_spec.name;
248
249 fn setDimAttr(op: *ir.Operation, ctx: *ir.Context, key: []const u8, value: u32) !void {
250 const attr = try ctx.getI64Attr(@intCast(value));
251 try op.setAttr(key, attr);
252 }
253
254 fn getDimAttr(op: *const ir.Operation, key: []const u8) ?u32 {
255 const int_attr = op.getAttrAs(ir.Attribute.IntegerAttr, key) orelse return null;
256 const raw = int_attr.getUnsignedValue();
257 if (raw > std.math.maxInt(u32)) return null;
258 return @intCast(raw);
259 }
260
261 pub fn create(
262 ctx: *ir.Context,
263 loc: ir.Location,
264 kernel_name: []const u8,
265 kernel_args: []const *ir.Value,
266 grid_dim: [3]u32,
267 block_dim: [3]u32,
268 ) !LaunchOp {
269 try loadSpec(ctx);
270 var builder = ir.OperationBuilder.init(ctx);
271 var state = ir.Operation.State.init(operation_name, loc);
272
273 var all_operands: std.ArrayList(*ir.Value) = .empty;
274 const allocator = ir.context.transientAllocator(ctx);
275 defer all_operands.deinit(allocator);
276 for (kernel_args) |arg| {
277 try all_operands.append(allocator, arg);
278 }
279 state.addOperands(all_operands.items);
280
281 const op = try builder.create(state);
282 errdefer op.erase();
283
284 const kernel_attr = try func.FuncDialect.getSymNameAttr(ctx, kernel_name);
285 try op.setAttr("kernel", kernel_attr);
286
287 var buf: [16]u8 = undefined;
288 const num_args_str = try ir.format.intPayload(buf[0..], kernel_args.len);
289 const num_args_attr = try ctx.getDialectAttr("gpu.num_kernel_args", num_args_str);
290 try op.setAttr("num_kernel_args", num_args_attr);
291
292 try setDimAttr(op, ctx, dim_attr_keys.grid_x, grid_dim[0]);
293 try setDimAttr(op, ctx, dim_attr_keys.grid_y, grid_dim[1]);
294 try setDimAttr(op, ctx, dim_attr_keys.grid_z, grid_dim[2]);
295 try setDimAttr(op, ctx, dim_attr_keys.block_x, block_dim[0]);
296 try setDimAttr(op, ctx, dim_attr_keys.block_y, block_dim[1]);
297 try setDimAttr(op, ctx, dim_attr_keys.block_z, block_dim[2]);
298
299 return .{ .op = op };
300 }
301
302 pub fn getKernelName(self: LaunchOp) ?[]const u8 {
303 if (self.op.getAttr("kernel")) |attr| {
304 return func.FuncDialect.getSymNameValue(attr);
305 }
306 return null;
307 }
308
309 pub fn getNumKernelArgs(self: LaunchOp) usize {
310 const dialect_attr = self.op.getAttrAs(ir.Attribute.DialectAttr, "num_kernel_args") orelse return 0;
311 return std.fmt.parseInt(usize, dialect_attr.payload, 10) catch 0;
312 }
313
314 pub fn getKernelArgs(self: LaunchOp) []const *ir.Value {
315 const num_args = self.getNumKernelArgs();
316 return self.op.getOperandValues()[0..num_args];
317 }
318
319 pub fn getGridDim(self: LaunchOp) ?[3]u32 {
320 const gx = getDimAttr(self.op, dim_attr_keys.grid_x) orelse return null;
321 const gy = getDimAttr(self.op, dim_attr_keys.grid_y) orelse return null;
322 const gz = getDimAttr(self.op, dim_attr_keys.grid_z) orelse return null;
323 return .{ gx, gy, gz };
324 }
325
326 pub fn getBlockDim(self: LaunchOp) ?[3]u32 {
327 const bx = getDimAttr(self.op, dim_attr_keys.block_x) orelse return null;
328 const by = getDimAttr(self.op, dim_attr_keys.block_y) orelse return null;
329 const bz = getDimAttr(self.op, dim_attr_keys.block_z) orelse return null;
330 return .{ bx, by, bz };
331 }
332 };
333
334 pub const ThreadIdxOp = struct {
335 op: *ir.Operation,
336
337 pub const operation_spec = dimIndexSpec("thread_idx");
338 pub const operation_name = operation_spec.name;
339
340 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !ThreadIdxOp {
341 const arith = choir.dialects.arith;
342 try loadSpec(ctx);
343 var builder = ir.OperationBuilder.init(ctx);
344 const index_type = try arith.ArithDialect.getIndexType(ctx);
345 var state = ir.Operation.State.init(operation_name, loc);
346 state.addTypes(&.{index_type});
347
348 const op = try builder.create(state);
349 errdefer op.erase();
350 try setDimensionAttr(op, ctx, dim);
351
352 return .{ .op = op };
353 }
354
355 pub fn getResult(self: *const ThreadIdxOp) *ir.Value {
356 return self.op.getResult(0).?;
357 }
358
359 pub fn getDimension(self: ThreadIdxOp) ?Dimension {
360 return getDimensionAttr(self.op);
361 }
362 };
363
364 pub const BlockIdxOp = struct {
365 op: *ir.Operation,
366
367 pub const operation_spec = dimIndexSpec("block_idx");
368 pub const operation_name = operation_spec.name;
369
370 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !BlockIdxOp {
371 const arith = choir.dialects.arith;
372 try loadSpec(ctx);
373 var builder = ir.OperationBuilder.init(ctx);
374 const index_type = try arith.ArithDialect.getIndexType(ctx);
375 var state = ir.Operation.State.init(operation_name, loc);
376 state.addTypes(&.{index_type});
377
378 const op = try builder.create(state);
379 errdefer op.erase();
380 try setDimensionAttr(op, ctx, dim);
381
382 return .{ .op = op };
383 }
384
385 pub fn getResult(self: *const BlockIdxOp) *ir.Value {
386 return self.op.getResult(0).?;
387 }
388
389 pub fn getDimension(self: BlockIdxOp) ?Dimension {
390 return getDimensionAttr(self.op);
391 }
392 };
393
394 pub const BlockDimOp = struct {
395 op: *ir.Operation,
396
397 pub const operation_spec = dimIndexSpec("block_dim");
398 pub const operation_name = operation_spec.name;
399
400 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !BlockDimOp {
401 const arith = choir.dialects.arith;
402 try loadSpec(ctx);
403 var builder = ir.OperationBuilder.init(ctx);
404 const index_type = try arith.ArithDialect.getIndexType(ctx);
405 var state = ir.Operation.State.init(operation_name, loc);
406 state.addTypes(&.{index_type});
407
408 const op = try builder.create(state);
409 errdefer op.erase();
410 try setDimensionAttr(op, ctx, dim);
411
412 return .{ .op = op };
413 }
414
415 pub fn getResult(self: *const BlockDimOp) *ir.Value {
416 return self.op.getResult(0).?;
417 }
418
419 pub fn getDimension(self: BlockDimOp) ?Dimension {
420 return getDimensionAttr(self.op);
421 }
422 };
423
424 pub const GridDimOp = struct {
425 op: *ir.Operation,
426
427 pub const operation_spec = dimIndexSpec("grid_dim");
428 pub const operation_name = operation_spec.name;
429
430 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !GridDimOp {
431 const arith = choir.dialects.arith;
432 try loadSpec(ctx);
433 var builder = ir.OperationBuilder.init(ctx);
434 const index_type = try arith.ArithDialect.getIndexType(ctx);
435 var state = ir.Operation.State.init(operation_name, loc);
436 state.addTypes(&.{index_type});
437
438 const op = try builder.create(state);
439 errdefer op.erase();
440 try setDimensionAttr(op, ctx, dim);
441
442 return .{ .op = op };
443 }
444
445 pub fn getResult(self: *const GridDimOp) *ir.Value {
446 return self.op.getResult(0).?;
447 }
448
449 pub fn getDimension(self: GridDimOp) ?Dimension {
450 return getDimensionAttr(self.op);
451 }
452 };
453
454 pub const GlobalIdxOp = struct {
455 op: *ir.Operation,
456
457 pub const operation_spec = dimIndexSpec("global_idx");
458 pub const operation_name = operation_spec.name;
459
460 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !GlobalIdxOp {
461 const arith = choir.dialects.arith;
462 try loadSpec(ctx);
463 var builder = ir.OperationBuilder.init(ctx);
464 const index_type = try arith.ArithDialect.getIndexType(ctx);
465 var state = ir.Operation.State.init(operation_name, loc);
466 state.addTypes(&.{index_type});
467
468 const op = try builder.create(state);
469 errdefer op.erase();
470 try setDimensionAttr(op, ctx, dim);
471
472 return .{ .op = op };
473 }
474
475 pub fn getResult(self: *const GlobalIdxOp) *ir.Value {
476 return self.op.getResult(0).?;
477 }
478
479 pub fn getDimension(self: GlobalIdxOp) ?Dimension {
480 return getDimensionAttr(self.op);
481 }
482 };
483
484 pub const LaneIdOp = struct {
485 op: *ir.Operation,
486
487 pub const operation_spec = indexSpec("lane_id");
488 pub const operation_name = operation_spec.name;
489
490 pub fn create(ctx: *ir.Context, loc: ir.Location) !LaneIdOp {
491 const arith = choir.dialects.arith;
492 try loadSpec(ctx);
493 var builder = ir.OperationBuilder.init(ctx);
494 const index_type = try arith.ArithDialect.getIndexType(ctx);
495 var state = ir.Operation.State.init(operation_name, loc);
496 state.addTypes(&.{index_type});
497
498 const op = try builder.create(state);
499 errdefer op.erase();
500 return .{ .op = op };
501 }
502
503 pub fn getResult(self: *const LaneIdOp) *ir.Value {
504 return self.op.getResult(0).?;
505 }
506 };
507
508 pub const WarpIdOp = struct {
509 op: *ir.Operation,
510
511 pub const operation_spec = indexSpec("warp_id");
512 pub const operation_name = operation_spec.name;
513
514 pub fn create(ctx: *ir.Context, loc: ir.Location) !WarpIdOp {
515 const arith = choir.dialects.arith;
516 try loadSpec(ctx);
517 var builder = ir.OperationBuilder.init(ctx);
518 const index_type = try arith.ArithDialect.getIndexType(ctx);
519 var state = ir.Operation.State.init(operation_name, loc);
520 state.addTypes(&.{index_type});
521
522 const op = try builder.create(state);
523 errdefer op.erase();
524 return .{ .op = op };
525 }
526
527 pub fn getResult(self: *const WarpIdOp) *ir.Value {
528 return self.op.getResult(0).?;
529 }
530 };
531
532 pub const BarrierOp = struct {
533 op: *ir.Operation,
534
535 pub const operation_spec = op_specs.leaf(.{
536 .mnemonic = "barrier",
537 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
538 .operands = 0,
539 .results = 0,
540 .required_attrs = &.{"scope"},
541 });
542 pub const operation_name = operation_spec.name;
543
544 pub fn create(ctx: *ir.Context, loc: ir.Location, scope: Scope) !BarrierOp {
545 try loadSpec(ctx);
546 var builder = ir.OperationBuilder.init(ctx);
547 const state = ir.Operation.State.init(operation_name, loc);
548
549 const op = try builder.create(state);
550 errdefer op.erase();
551 try setScopeAttr(op, ctx, scope);
552
553 return .{ .op = op };
554 }
555
556 pub fn getScope(self: BarrierOp) ?Scope {
557 return getScopeAttr(self.op);
558 }
559 };
560
561 pub const FenceOp = struct {
562 op: *ir.Operation,
563
564 pub const operation_spec = op_specs.leaf(.{
565 .mnemonic = "fence",
566 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
567 .operands = 0,
568 .results = 0,
569 .required_attrs = &.{ "ordering", "scope" },
570 });
571 pub const operation_name = operation_spec.name;
572
573 pub fn create(ctx: *ir.Context, loc: ir.Location, scope: Scope, ordering: MemoryOrder) !FenceOp {
574 try loadSpec(ctx);
575 var builder = ir.OperationBuilder.init(ctx);
576 const state = ir.Operation.State.init(operation_name, loc);
577
578 const op = try builder.create(state);
579 errdefer op.erase();
580 try setScopeAttr(op, ctx, scope);
581 try setOrderingAttr(op, ctx, ordering);
582
583 return .{ .op = op };
584 }
585
586 pub fn getScope(self: FenceOp) ?Scope {
587 return getScopeAttr(self.op);
588 }
589
590 pub fn getOrdering(self: FenceOp) ?MemoryOrder {
591 return getOrderingAttr(self.op);
592 }
593 };
594
595 pub const MemcpyAsyncOp = struct {
596 op: *ir.Operation,
597
598 pub const operation_spec = op_specs.leaf(.{
599 .mnemonic = "memcpy_async",
600 .interfaces = &.{gpuEffects(.launch, &.{0}, &.{1})},
601 .operands = ir.dialects.shape.between(3, 5),
602 .operand_names = .{ "src", "dst", "num_bytes", "stream", "event" },
603 .results = 0,
604 .operand_segments = ir.dialects.segments.operands(.{
605 1,
606 1,
607 1,
608 ir.dialects.shape.atMost(1),
609 ir.dialects.shape.atMost(1),
610 }),
611 });
612 pub const operation_name = operation_spec.name;
613
614 pub fn create(
615 ctx: *ir.Context,
616 loc: ir.Location,
617 src: *ir.Value,
618 dst: *ir.Value,
619 num_bytes: *ir.Value,
620 stream: ?*ir.Value,
621 event: ?*ir.Value,
622 ) !MemcpyAsyncOp {
623 try loadSpec(ctx);
624 var builder = ir.OperationBuilder.init(ctx);
625 var state = ir.Operation.State.init(operation_name, loc);
626
627 var operands: [5]*ir.Value = undefined;
628 var operand_count: usize = 0;
629 operands[operand_count] = src;
630 operand_count += 1;
631 operands[operand_count] = dst;
632 operand_count += 1;
633 operands[operand_count] = num_bytes;
634 operand_count += 1;
635
636 if (stream) |stream_val| {
637 operands[operand_count] = stream_val;
638 operand_count += 1;
639 }
640
641 if (event) |event_val| {
642 operands[operand_count] = event_val;
643 operand_count += 1;
644 }
645
646 state.addOperands(operands[0..operand_count]);
647 const op = try builder.create(state);
648 errdefer op.erase();
649 const segment_sizes = [_]usize{
650 1,
651 1,
652 1,
653 if (stream != null) 1 else 0,
654 if (event != null) 1 else 0,
655 };
656 try ir.dialects.setOperandSegmentSizes(operation_spec, op, &segment_sizes);
657
658 return .{ .op = op };
659 }
660
661 pub fn getSrc(self: MemcpyAsyncOp) *ir.Value {
662 return ir.dialects.operand(operation_spec, self.op, "src");
663 }
664
665 pub fn getDst(self: MemcpyAsyncOp) *ir.Value {
666 return ir.dialects.operand(operation_spec, self.op, "dst");
667 }
668
669 pub fn getNumBytes(self: MemcpyAsyncOp) *ir.Value {
670 return ir.dialects.operand(operation_spec, self.op, "num_bytes");
671 }
672
673 pub fn getStream(self: MemcpyAsyncOp) ?*ir.Value {
674 return ir.dialects.operandSegmentValue(operation_spec, self.op, "stream");
675 }
676
677 pub fn getEvent(self: MemcpyAsyncOp) ?*ir.Value {
678 return ir.dialects.operandSegmentValue(operation_spec, self.op, "event");
679 }
680 };
681
682 pub const TmaCreateDescriptorOp = struct {
683 op: *ir.Operation,
684
685 pub const operation_spec = op_specs.leaf(.{
686 .mnemonic = "tma.create_descriptor",
687 .operands = .{ "tensor", "box_shape" },
688 .results = .{"descriptor"},
689 });
690 pub const operation_name = operation_spec.name;
691
692 pub fn create(
693 ctx: *ir.Context,
694 loc: ir.Location,
695 tensor: *ir.Value,
696 box_shape: *ir.Value,
697 ) !TmaCreateDescriptorOp {
698 try loadSpec(ctx);
699 var builder = ir.OperationBuilder.init(ctx);
700 var state = ir.Operation.State.init(operation_name, loc);
701 state.addOperands(&.{ tensor, box_shape });
702
703 const desc_type = try getTmaDescriptorType(ctx);
704 state.addTypes(&.{desc_type});
705
706 const op = try builder.create(state);
707 errdefer op.erase();
708 return .{ .op = op };
709 }
710
711 pub fn getResult(self: *const TmaCreateDescriptorOp) *ir.Value {
712 return self.op.getResult(0).?;
713 }
714
715 pub fn getTensor(self: TmaCreateDescriptorOp) *ir.Value {
716 return self.op.operands.items[0].value;
717 }
718
719 pub fn getBoxShape(self: TmaCreateDescriptorOp) *ir.Value {
720 return self.op.operands.items[1].value;
721 }
722 };
723
724 pub const TmaLoadOp = struct {
725 op: *ir.Operation,
726
727 pub const operation_spec = noResultSpec("tma.load", .{ "descriptor", "shared_mem", "barrier", "coords" });
728 pub const operation_name = operation_spec.name;
729
730 pub fn create(
731 ctx: *ir.Context,
732 loc: ir.Location,
733 desc: *ir.Value,
734 shmem: *ir.Value,
735 mbarrier: *ir.Value,
736 coords: *ir.Value,
737 ) !TmaLoadOp {
738 try loadSpec(ctx);
739 var builder = ir.OperationBuilder.init(ctx);
740 var state = ir.Operation.State.init(operation_name, loc);
741 state.addOperands(&.{ desc, shmem, mbarrier, coords });
742
743 const op = try builder.create(state);
744 errdefer op.erase();
745 return .{ .op = op };
746 }
747
748 pub fn getDescriptor(self: TmaLoadOp) *ir.Value {
749 return self.op.operands.items[0].value;
750 }
751
752 pub fn getSharedMem(self: TmaLoadOp) *ir.Value {
753 return self.op.operands.items[1].value;
754 }
755
756 pub fn getBarrier(self: TmaLoadOp) *ir.Value {
757 return self.op.operands.items[2].value;
758 }
759
760 pub fn getCoords(self: TmaLoadOp) *ir.Value {
761 return self.op.operands.items[3].value;
762 }
763 };
764
765 pub const TmaCommitGroupOp = struct {
766 op: *ir.Operation,
767
768 pub const operation_spec = noResultSpec("tma.commit_group", 0);
769 pub const operation_name = operation_spec.name;
770
771 pub fn create(ctx: *ir.Context, loc: ir.Location) !TmaCommitGroupOp {
772 try loadSpec(ctx);
773 var builder = ir.OperationBuilder.init(ctx);
774 const state = ir.Operation.State.init(operation_name, loc);
775
776 const op = try builder.create(state);
777 errdefer op.erase();
778 return .{ .op = op };
779 }
780 };
781
782 pub const TmaWaitGroupOp = struct {
783 op: *ir.Operation,
784
785 pub const operation_spec = op_specs.leaf(.{
786 .mnemonic = "tma.wait_group",
787 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
788 .operands = 0,
789 .results = 0,
790 .required_attrs = &.{"count"},
791 });
792 pub const operation_name = operation_spec.name;
793
794 pub fn create(ctx: *ir.Context, loc: ir.Location, count: i64) !TmaWaitGroupOp {
795 try loadSpec(ctx);
796 var builder = ir.OperationBuilder.init(ctx);
797 const state = ir.Operation.State.init(operation_name, loc);
798
799 const op = try builder.create(state);
800 errdefer op.erase();
801 try setI64Attr(op, ctx, "count", count);
802 return .{ .op = op };
803 }
804
805 pub fn getCount(self: TmaWaitGroupOp) ?i64 {
806 return getI64AttrValue(self.op, "count");
807 }
808 };
809
810 pub const ShflSyncOp = struct {
811 op: *ir.Operation,
812
813 pub const operation_spec = op_specs.leaf(.{
814 .mnemonic = "shfl_sync",
815 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
816 .operands = .{ "mask", "src", "lane_or_delta" },
817 .results = .{"result"},
818 .required_attrs = &.{"mode"},
819 });
820 pub const operation_name = operation_spec.name;
821
822 pub fn create(
823 ctx: *ir.Context,
824 loc: ir.Location,
825 mode: ShuffleMode,
826 mask: *ir.Value,
827 src: *ir.Value,
828 lane_or_delta: *ir.Value,
829 ) !ShflSyncOp {
830 try loadSpec(ctx);
831 var builder = ir.OperationBuilder.init(ctx);
832 var state = ir.Operation.State.init(operation_name, loc);
833 state.addOperands(&.{ mask, src, lane_or_delta });
834 state.addTypes(&.{src.type});
835
836 const op = try builder.create(state);
837 errdefer op.erase();
838
839 const mode_attr = try ctx.getDialectAttr("gpu.shuffle_mode", mode.toString());
840 try op.setAttr("mode", mode_attr);
841
842 return .{ .op = op };
843 }
844
845 pub fn getResult(self: *const ShflSyncOp) *ir.Value {
846 return self.op.getResult(0).?;
847 }
848
849 pub fn getMask(self: ShflSyncOp) *ir.Value {
850 return self.op.operands.items[0].value;
851 }
852
853 pub fn getSrc(self: ShflSyncOp) *ir.Value {
854 return self.op.operands.items[1].value;
855 }
856
857 pub fn getLaneOrDelta(self: ShflSyncOp) *ir.Value {
858 return self.op.operands.items[2].value;
859 }
860
861 pub fn getMode(self: ShflSyncOp) ?ShuffleMode {
862 const dialect_attr = self.op.getAttrAs(ir.Attribute.DialectAttr, "mode") orelse return null;
863 return ShuffleMode.fromString(dialect_attr.payload);
864 }
865 };
866
867 pub const AllSyncOp = struct {
868 op: *ir.Operation,
869
870 pub const operation_spec = op_specs.leaf(.{
871 .mnemonic = "all_sync",
872 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
873 .operands = .{ "mask", "predicate" },
874 .results = .{"result"},
875 });
876 pub const operation_name = operation_spec.name;
877
878 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AllSyncOp {
879 const arith = choir.dialects.arith;
880 try loadSpec(ctx);
881 var builder = ir.OperationBuilder.init(ctx);
882 const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);
883 var state = ir.Operation.State.init(operation_name, loc);
884 state.addOperands(&.{ mask, pred });
885 state.addTypes(&.{bool_type});
886
887 const op = try builder.create(state);
888 errdefer op.erase();
889 return .{ .op = op };
890 }
891
892 pub fn getResult(self: *const AllSyncOp) *ir.Value {
893 return self.op.getResult(0).?;
894 }
895
896 pub fn getMask(self: AllSyncOp) *ir.Value {
897 return self.op.operands.items[0].value;
898 }
899
900 pub fn getPredicate(self: AllSyncOp) *ir.Value {
901 return self.op.operands.items[1].value;
902 }
903 };
904
905 pub const AnySyncOp = struct {
906 op: *ir.Operation,
907
908 pub const operation_spec = op_specs.leaf(.{
909 .mnemonic = "any_sync",
910 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
911 .operands = .{ "mask", "predicate" },
912 .results = .{"result"},
913 });
914 pub const operation_name = operation_spec.name;
915
916 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AnySyncOp {
917 const arith = choir.dialects.arith;
918 try loadSpec(ctx);
919 var builder = ir.OperationBuilder.init(ctx);
920 const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);
921 var state = ir.Operation.State.init(operation_name, loc);
922 state.addOperands(&.{ mask, pred });
923 state.addTypes(&.{bool_type});
924
925 const op = try builder.create(state);
926 errdefer op.erase();
927 return .{ .op = op };
928 }
929
930 pub fn getResult(self: *const AnySyncOp) *ir.Value {
931 return self.op.getResult(0).?;
932 }
933
934 pub fn getMask(self: AnySyncOp) *ir.Value {
935 return self.op.operands.items[0].value;
936 }
937
938 pub fn getPredicate(self: AnySyncOp) *ir.Value {
939 return self.op.operands.items[1].value;
940 }
941 };
942
943 pub const BallotSyncOp = struct {
944 op: *ir.Operation,
945
946 pub const operation_spec = op_specs.leaf(.{
947 .mnemonic = "ballot_sync",
948 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
949 .operands = .{ "mask", "predicate" },
950 .results = .{"result"},
951 });
952 pub const operation_name = operation_spec.name;
953
954 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !BallotSyncOp {
955 const arith = choir.dialects.arith;
956 try loadSpec(ctx);
957 var builder = ir.OperationBuilder.init(ctx);
958 const i32_type = try arith.ArithDialect.getI32Type(ctx);
959 var state = ir.Operation.State.init(operation_name, loc);
960 state.addOperands(&.{ mask, pred });
961 state.addTypes(&.{i32_type});
962
963 const op = try builder.create(state);
964 errdefer op.erase();
965 return .{ .op = op };
966 }
967
968 pub fn getResult(self: *const BallotSyncOp) *ir.Value {
969 return self.op.getResult(0).?;
970 }
971
972 pub fn getMask(self: BallotSyncOp) *ir.Value {
973 return self.op.operands.items[0].value;
974 }
975
976 pub fn getPredicate(self: BallotSyncOp) *ir.Value {
977 return self.op.operands.items[1].value;
978 }
979 };
980
981 pub const WarpReduceOp = struct {
982 op: *ir.Operation,
983
984 pub const operation_spec = op_specs.leaf(.{
985 .mnemonic = "warp_reduce",
986 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
987 .operands = .{ "mask", "value" },
988 .results = .{"result"},
989 .required_attrs = &.{"op"},
990 });
991 pub const operation_name = operation_spec.name;
992
993 pub fn create(
994 ctx: *ir.Context,
995 loc: ir.Location,
996 op_kind: WarpOpKind,
997 mask: *ir.Value,
998 value: *ir.Value,
999 ) !WarpReduceOp {
1000 try loadSpec(ctx);
1001 var builder = ir.OperationBuilder.init(ctx);
1002 var state = ir.Operation.State.init(operation_name, loc);
1003 state.addOperands(&.{ mask, value });
1004 state.addTypes(&.{value.type});
1005
1006 const op = try builder.create(state);
1007 errdefer op.erase();
1008 try setWarpOpAttr(op, ctx, op_kind);
1009
1010 return .{ .op = op };
1011 }
1012
1013 pub fn getResult(self: *const WarpReduceOp) *ir.Value {
1014 return self.op.getResult(0).?;
1015 }
1016
1017 pub fn getMask(self: WarpReduceOp) *ir.Value {
1018 return self.op.operands.items[0].value;
1019 }
1020
1021 pub fn getValue(self: WarpReduceOp) *ir.Value {
1022 return self.op.operands.items[1].value;
1023 }
1024
1025 pub fn getOpKind(self: WarpReduceOp) ?WarpOpKind {
1026 return getWarpOpAttr(self.op);
1027 }
1028 };
1029
1030 pub const WarpScanOp = struct {
1031 op: *ir.Operation,
1032
1033 pub const operation_spec = op_specs.leaf(.{
1034 .mnemonic = "warp_scan",
1035 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
1036 .operands = .{ "mask", "value" },
1037 .results = .{"result"},
1038 .required_attrs = &.{ "inclusive", "op" },
1039 });
1040 pub const operation_name = operation_spec.name;
1041
1042 pub fn create(
1043 ctx: *ir.Context,
1044 loc: ir.Location,
1045 op_kind: WarpOpKind,
1046 inclusive: bool,
1047 mask: *ir.Value,
1048 value: *ir.Value,
1049 ) !WarpScanOp {
1050 try loadSpec(ctx);
1051 var builder = ir.OperationBuilder.init(ctx);
1052 var state = ir.Operation.State.init(operation_name, loc);
1053 state.addOperands(&.{ mask, value });
1054 state.addTypes(&.{value.type});
1055
1056 const op = try builder.create(state);
1057 errdefer op.erase();
1058 try setWarpOpAttr(op, ctx, op_kind);
1059 try setBoolAttr(op, ctx, "inclusive", inclusive);
1060
1061 return .{ .op = op };
1062 }
1063
1064 pub fn getResult(self: *const WarpScanOp) *ir.Value {
1065 return self.op.getResult(0).?;
1066 }
1067
1068 pub fn getMask(self: WarpScanOp) *ir.Value {
1069 return self.op.operands.items[0].value;
1070 }
1071
1072 pub fn getValue(self: WarpScanOp) *ir.Value {
1073 return self.op.operands.items[1].value;
1074 }
1075
1076 pub fn getOpKind(self: WarpScanOp) ?WarpOpKind {
1077 return getWarpOpAttr(self.op);
1078 }
1079
1080 pub fn isInclusive(self: WarpScanOp) bool {
1081 return getBoolAttrValue(self.op, "inclusive");
1082 }
1083 };
1084
1085 pub const MatchAnyOp = struct {
1086 op: *ir.Operation,
1087
1088 pub const operation_spec = op_specs.leaf(.{
1089 .mnemonic = "match_any",
1090 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
1091 .operands = .{"value"},
1092 .results = .{"mask"},
1093 });
1094 pub const operation_name = operation_spec.name;
1095
1096 pub fn create(ctx: *ir.Context, loc: ir.Location, value: *ir.Value) !MatchAnyOp {
1097 const arith = choir.dialects.arith;
1098 try loadSpec(ctx);
1099 var builder = ir.OperationBuilder.init(ctx);
1100 const i32_type = try arith.ArithDialect.getI32Type(ctx);
1101 var state = ir.Operation.State.init(operation_name, loc);
1102 state.addOperands(&.{value});
1103 state.addTypes(&.{i32_type});
1104
1105 const op = try builder.create(state);
1106 errdefer op.erase();
1107 return .{ .op = op };
1108 }
1109
1110 pub fn getResult(self: *const MatchAnyOp) *ir.Value {
1111 return self.op.getResult(0).?;
1112 }
1113
1114 pub fn getValue(self: MatchAnyOp) *ir.Value {
1115 return self.op.operands.items[0].value;
1116 }
1117 };
1118
1119 pub const MatchAllOp = struct {
1120 op: *ir.Operation,
1121
1122 pub const operation_spec = op_specs.leaf(.{
1123 .mnemonic = "match_all",
1124 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
1125 .operands = .{"value"},
1126 .results = .{ "mask", "all_equal" },
1127 });
1128 pub const operation_name = operation_spec.name;
1129
1130 pub fn create(ctx: *ir.Context, loc: ir.Location, value: *ir.Value) !MatchAllOp {
1131 const arith = choir.dialects.arith;
1132 try loadSpec(ctx);
1133 var builder = ir.OperationBuilder.init(ctx);
1134 const i32_type = try arith.ArithDialect.getI32Type(ctx);
1135 const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);
1136 var state = ir.Operation.State.init(operation_name, loc);
1137 state.addOperands(&.{value});
1138 state.addTypes(&.{ i32_type, bool_type });
1139
1140 const op = try builder.create(state);
1141 errdefer op.erase();
1142 return .{ .op = op };
1143 }
1144
1145 pub fn getMaskResult(self: *const MatchAllOp) *ir.Value {
1146 return self.op.getResult(0).?;
1147 }
1148
1149 pub fn getAllEqualResult(self: *const MatchAllOp) *ir.Value {
1150 return self.op.getResult(1).?;
1151 }
1152
1153 pub fn getValue(self: MatchAllOp) *ir.Value {
1154 return self.op.operands.items[0].value;
1155 }
1156 };
1157
1158 pub const ActiveMaskOp = struct {
1159 op: *ir.Operation,
1160
1161 pub const operation_spec = op_specs.leaf(.{
1162 .mnemonic = "active_mask",
1163 .interfaces = &.{gpuEffects(.state_observe, &.{}, &.{})},
1164 .operands = 0,
1165 .results = .{"mask"},
1166 });
1167 pub const operation_name = operation_spec.name;
1168
1169 pub fn create(ctx: *ir.Context, loc: ir.Location) !ActiveMaskOp {
1170 const arith = choir.dialects.arith;
1171 try loadSpec(ctx);
1172 var builder = ir.OperationBuilder.init(ctx);
1173 const i32_type = try arith.ArithDialect.getI32Type(ctx);
1174 var state = ir.Operation.State.init(operation_name, loc);
1175 state.addTypes(&.{i32_type});
1176
1177 const op = try builder.create(state);
1178 errdefer op.erase();
1179 return .{ .op = op };
1180 }
1181
1182 pub fn getResult(self: *const ActiveMaskOp) *ir.Value {
1183 return self.op.getResult(0).?;
1184 }
1185 };
1186
1187 pub const SyncWarpOp = struct {
1188 op: *ir.Operation,
1189
1190 pub const operation_spec = noResultSpec("sync_warp", .{"mask"});
1191 pub const operation_name = operation_spec.name;
1192
1193 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value) !SyncWarpOp {
1194 try loadSpec(ctx);
1195 var builder = ir.OperationBuilder.init(ctx);
1196 var state = ir.Operation.State.init(operation_name, loc);
1197 state.addOperands(&.{mask});
1198
1199 const op = try builder.create(state);
1200 errdefer op.erase();
1201 return .{ .op = op };
1202 }
1203
1204 pub fn getMask(self: SyncWarpOp) *ir.Value {
1205 return self.op.operands.items[0].value;
1206 }
1207 };
1208
1209 pub const mma_sync_a_count = 4;
1210 pub const mma_sync_b_count = 2;
1211 pub const mma_sync_acc_count = 4;
1212
1213 pub const MmaSyncOp = struct {
1214 op: *ir.Operation,
1215
1216 pub const operation_spec = op_specs.leaf(.{
1217 .mnemonic = "mma_sync",
1218 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
1219 .operands = .{ "a0", "a1", "a2", "a3", "b0", "b1", "c0", "c1", "c2", "c3" },
1220 .results = .{ "d0", "d1", "d2", "d3" },
1221 .required_attrs = &.{"shape"},
1222 });
1223 pub const operation_name = operation_spec.name;
1224
1225 pub fn create(
1226 ctx: *ir.Context,
1227 loc: ir.Location,
1228 a: [mma_sync_a_count]*ir.Value,
1229 b: [mma_sync_b_count]*ir.Value,
1230 c: [mma_sync_acc_count]*ir.Value,
1231 shape: MmaShape,
1232 ) !MmaSyncOp {
1233 try loadSpec(ctx);
1234 var builder = ir.OperationBuilder.init(ctx);
1235 var state = ir.Operation.State.init(operation_name, loc);
1236 state.addOperands(&.{ a[0], a[1], a[2], a[3], b[0], b[1], c[0], c[1], c[2], c[3] });
1237 state.addTypes(&.{ c[0].type, c[1].type, c[2].type, c[3].type });
1238
1239 const op = try builder.create(state);
1240 errdefer op.erase();
1241 try setMmaShapeAttr(op, ctx, shape);
1242 return .{ .op = op };
1243 }
1244
1245 pub fn getA(self: MmaSyncOp, index: usize) *ir.Value {
1246 return self.op.operands.items[index].value;
1247 }
1248
1249 pub fn getB(self: MmaSyncOp, index: usize) *ir.Value {
1250 return self.op.operands.items[mma_sync_a_count + index].value;
1251 }
1252
1253 pub fn getC(self: MmaSyncOp, index: usize) *ir.Value {
1254 return self.op.operands.items[mma_sync_a_count + mma_sync_b_count + index].value;
1255 }
1256
1257 pub fn getD(self: *const MmaSyncOp, index: usize) *ir.Value {
1258 return self.op.getResult(index).?;
1259 }
1260
1261 pub fn getShape(self: MmaSyncOp) ?MmaShape {
1262 return getMmaShapeAttr(self.op);
1263 }
1264 };
1265
1266 pub const CpAsyncSharedOp = struct {
1267 op: *ir.Operation,
1268
1269 pub const operation_spec = op_specs.leaf(.{
1270 .mnemonic = "cp_async_shared",
1271 .interfaces = &.{gpuEffects(.launch, &.{2}, &.{0})},
1272 .operands = .{ "dst", "dst_index", "src", "src_index" },
1273 .results = 0,
1274 .required_attrs = &.{"bytes"},
1275 });
1276 pub const operation_name = operation_spec.name;
1277
1278 pub fn create(
1279 ctx: *ir.Context,
1280 loc: ir.Location,
1281 dst: *ir.Value,
1282 dst_index: *ir.Value,
1283 src: *ir.Value,
1284 src_index: *ir.Value,
1285 bytes: u32,
1286 ) !CpAsyncSharedOp {
1287 try loadSpec(ctx);
1288 var builder = ir.OperationBuilder.init(ctx);
1289 var state = ir.Operation.State.init(operation_name, loc);
1290 state.addOperands(&.{ dst, dst_index, src, src_index });
1291 const op = try builder.create(state);
1292 errdefer op.erase();
1293 const bytes_attr = try ctx.getI64Attr(@intCast(bytes));
1294 try op.setAttr("bytes", bytes_attr);
1295 return .{ .op = op };
1296 }
1297
1298 pub fn getDst(self: CpAsyncSharedOp) *ir.Value {
1299 return ir.dialects.operand(operation_spec, self.op, "dst");
1300 }
1301
1302 pub fn getDstIndex(self: CpAsyncSharedOp) *ir.Value {
1303 return ir.dialects.operand(operation_spec, self.op, "dst_index");
1304 }
1305
1306 pub fn getSrc(self: CpAsyncSharedOp) *ir.Value {
1307 return ir.dialects.operand(operation_spec, self.op, "src");
1308 }
1309
1310 pub fn getSrcIndex(self: CpAsyncSharedOp) *ir.Value {
1311 return ir.dialects.operand(operation_spec, self.op, "src_index");
1312 }
1313
1314 pub fn getBytes(self: CpAsyncSharedOp) ?u32 {
1315 const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "bytes") orelse return null;
1316 const raw = int_attr.getUnsignedValue();
1317 if (raw > std.math.maxInt(u32)) return null;
1318 return @intCast(raw);
1319 }
1320 };
1321
1322 pub const CpAsyncCommitOp = struct {
1323 op: *ir.Operation,
1324
1325 pub const operation_spec = op_specs.leaf(.{
1326 .mnemonic = "cp_async_commit",
1327 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
1328 .operands = 0,
1329 .results = 0,
1330 });
1331 pub const operation_name = operation_spec.name;
1332
1333 pub fn create(ctx: *ir.Context, loc: ir.Location) !CpAsyncCommitOp {
1334 try loadSpec(ctx);
1335 var builder = ir.OperationBuilder.init(ctx);
1336 const state = ir.Operation.State.init(operation_name, loc);
1337 const op = try builder.create(state);
1338 errdefer op.erase();
1339 return .{ .op = op };
1340 }
1341 };
1342
1343 pub const CpAsyncWaitOp = struct {
1344 op: *ir.Operation,
1345
1346 pub const operation_spec = op_specs.leaf(.{
1347 .mnemonic = "cp_async_wait",
1348 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
1349 .operands = 0,
1350 .results = 0,
1351 .required_attrs = &.{"groups"},
1352 });
1353 pub const operation_name = operation_spec.name;
1354
1355 pub fn create(ctx: *ir.Context, loc: ir.Location, groups: u32) !CpAsyncWaitOp {
1356 try loadSpec(ctx);
1357 var builder = ir.OperationBuilder.init(ctx);
1358 const state = ir.Operation.State.init(operation_name, loc);
1359 const op = try builder.create(state);
1360 errdefer op.erase();
1361 const groups_attr = try ctx.getI64Attr(@intCast(groups));
1362 try op.setAttr("groups", groups_attr);
1363 return .{ .op = op };
1364 }
1365
1366 pub fn getGroups(self: CpAsyncWaitOp) ?u32 {
1367 const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "groups") orelse return null;
1368 const raw = int_attr.getUnsignedValue();
1369 if (raw > std.math.maxInt(u32)) return null;
1370 return @intCast(raw);
1371 }
1372 };
1373
1374 pub const AtomicLoadOp = struct {
1375 op: *ir.Operation,
1376
1377 pub const operation_spec = op_specs.leaf(.{
1378 .mnemonic = "atomic_load",
1379 .interfaces = &.{gpuEffects(.synchronize, &.{0}, &.{})},
1380 .operands = .{ "memref", "index" },
1381 .results = .{"value"},
1382 .required_attrs = &.{"ordering"},
1383 });
1384 pub const operation_name = operation_spec.name;
1385
1386 pub fn create(
1387 ctx: *ir.Context,
1388 loc: ir.Location,
1389 memref: *ir.Value,
1390 index: *ir.Value,
1391 result_type: ir.Type,
1392 ordering: MemoryOrder,
1393 ) !AtomicLoadOp {
1394 try loadSpec(ctx);
1395 var builder = ir.OperationBuilder.init(ctx);
1396 var state = ir.Operation.State.init(operation_name, loc);
1397 state.addOperands(&.{ memref, index });
1398 state.addTypes(&.{result_type});
1399
1400 const op = try builder.create(state);
1401 errdefer op.erase();
1402 try setOrderingAttr(op, ctx, ordering);
1403 return .{ .op = op };
1404 }
1405
1406 pub fn getResult(self: *const AtomicLoadOp) *ir.Value {
1407 return self.op.getResult(0).?;
1408 }
1409
1410 pub fn getMemref(self: AtomicLoadOp) *ir.Value {
1411 return self.op.operands.items[0].value;
1412 }
1413
1414 pub fn getIndex(self: AtomicLoadOp) *ir.Value {
1415 return self.op.operands.items[1].value;
1416 }
1417
1418 pub fn getOrdering(self: AtomicLoadOp) ?MemoryOrder {
1419 return getOrderingAttr(self.op);
1420 }
1421 };
1422
1423 pub const AtomicStoreOp = struct {
1424 op: *ir.Operation,
1425
1426 pub const operation_spec = op_specs.leaf(.{
1427 .mnemonic = "atomic_store",
1428 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{1})},
1429 .operands = .{ "value", "memref", "index" },
1430 .results = 0,
1431 .required_attrs = &.{"ordering"},
1432 });
1433 pub const operation_name = operation_spec.name;
1434
1435 pub fn create(
1436 ctx: *ir.Context,
1437 loc: ir.Location,
1438 value: *ir.Value,
1439 memref: *ir.Value,
1440 index: *ir.Value,
1441 ordering: MemoryOrder,
1442 ) !AtomicStoreOp {
1443 try loadSpec(ctx);
1444 var builder = ir.OperationBuilder.init(ctx);
1445 var state = ir.Operation.State.init(operation_name, loc);
1446 state.addOperands(&.{ value, memref, index });
1447
1448 const op = try builder.create(state);
1449 errdefer op.erase();
1450 try setOrderingAttr(op, ctx, ordering);
1451 return .{ .op = op };
1452 }
1453
1454 pub fn getValue(self: AtomicStoreOp) *ir.Value {
1455 return self.op.operands.items[0].value;
1456 }
1457
1458 pub fn getMemref(self: AtomicStoreOp) *ir.Value {
1459 return self.op.operands.items[1].value;
1460 }
1461
1462 pub fn getIndex(self: AtomicStoreOp) *ir.Value {
1463 return self.op.operands.items[2].value;
1464 }
1465
1466 pub fn getOrdering(self: AtomicStoreOp) ?MemoryOrder {
1467 return getOrderingAttr(self.op);
1468 }
1469 };
1470
1471 pub const AtomicAddOp = struct {
1472 op: *ir.Operation,
1473
1474 pub const operation_spec = op_specs.leaf(.{
1475 .mnemonic = "atomic_add",
1476 .interfaces = &.{gpuEffects(.synchronize, &.{0}, &.{0})},
1477 .operands = .{ "memref", "index", "value" },
1478 .results = .{"old_value"},
1479 .attrs = &.{"scope"},
1480 .required_attrs = &.{"ordering"},
1481 });
1482 pub const operation_name = operation_spec.name;
1483
1484 pub fn create(
1485 ctx: *ir.Context,
1486 loc: ir.Location,
1487 memref: *ir.Value,
1488 index: *ir.Value,
1489 val: *ir.Value,
1490 ordering: MemoryOrder,
1491 scope: ?Scope,
1492 ) !AtomicAddOp {
1493 try loadSpec(ctx);
1494 var builder = ir.OperationBuilder.init(ctx);
1495 var state = ir.Operation.State.init(operation_name, loc);
1496 state.addOperands(&.{ memref, index, val });
1497 state.addTypes(&.{val.type});
1498
1499 const op = try builder.create(state);
1500 errdefer op.erase();
1501 try setOrderingAttr(op, ctx, ordering);
1502 if (scope) |scope_value| {
1503 try setScopeAttr(op, ctx, scope_value);
1504 }
1505 return .{ .op = op };
1506 }
1507
1508 pub fn getResult(self: *const AtomicAddOp) *ir.Value {
1509 return self.op.getResult(0).?;
1510 }
1511
1512 pub fn getMemref(self: AtomicAddOp) *ir.Value {
1513 return self.op.operands.items[0].value;
1514 }
1515
1516 pub fn getIndex(self: AtomicAddOp) *ir.Value {
1517 return self.op.operands.items[1].value;
1518 }
1519
1520 pub fn getVal(self: AtomicAddOp) *ir.Value {
1521 return self.op.operands.items[2].value;
1522 }
1523
1524 pub fn getOrdering(self: AtomicAddOp) ?MemoryOrder {
1525 return getOrderingAttr(self.op);
1526 }
1527
1528 pub fn getScope(self: AtomicAddOp) ?Scope {
1529 return getScopeAttr(self.op);
1530 }
1531 };
1532
1533 pub const AtomicMaxOp = struct {
1534 op: *ir.Operation,
1535
1536 pub const operation_spec = op_specs.leaf(.{
1537 .mnemonic = "atomic_max",
1538 .interfaces = &.{gpuEffects(.synchronize, &.{0}, &.{0})},
1539 .operands = .{ "memref", "index", "value" },
1540 .results = .{"old_value"},
1541 });
1542 pub const operation_name = operation_spec.name;
1543
1544 pub fn create(
1545 ctx: *ir.Context,
1546 loc: ir.Location,
1547 memref: *ir.Value,
1548 index: *ir.Value,
1549 val: *ir.Value,
1550 ) !AtomicMaxOp {
1551 try loadSpec(ctx);
1552 var builder = ir.OperationBuilder.init(ctx);
1553 var state = ir.Operation.State.init(operation_name, loc);
1554 state.addOperands(&.{ memref, index, val });
1555 state.addTypes(&.{val.type});
1556
1557 const op = try builder.create(state);
1558 errdefer op.erase();
1559 return .{ .op = op };
1560 }
1561
1562 pub fn getResult(self: *const AtomicMaxOp) *ir.Value {
1563 return self.op.getResult(0).?;
1564 }
1565
1566 pub fn getMemref(self: AtomicMaxOp) *ir.Value {
1567 return self.op.operands.items[0].value;
1568 }
1569
1570 pub fn getIndex(self: AtomicMaxOp) *ir.Value {
1571 return self.op.operands.items[1].value;
1572 }
1573
1574 pub fn getVal(self: AtomicMaxOp) *ir.Value {
1575 return self.op.operands.items[2].value;
1576 }
1577 };
1578
1579 pub const AtomicCasOp = struct {
1580 op: *ir.Operation,
1581
1582 pub const operation_spec = op_specs.leaf(.{
1583 .mnemonic = "atomic_cas",
1584 .interfaces = &.{gpuEffects(.synchronize, &.{0}, &.{0})},
1585 .operands = .{ "memref", "index", "expected", "desired" },
1586 .results = .{"old_value"},
1587 .attrs = &.{"scope"},
1588 .required_attrs = &.{"ordering"},
1589 });
1590 pub const operation_name = operation_spec.name;
1591
1592 pub fn create(
1593 ctx: *ir.Context,
1594 loc: ir.Location,
1595 memref: *ir.Value,
1596 index: *ir.Value,
1597 expected: *ir.Value,
1598 desired: *ir.Value,
1599 ordering: MemoryOrder,
1600 scope: ?Scope,
1601 ) !AtomicCasOp {
1602 try loadSpec(ctx);
1603 var builder = ir.OperationBuilder.init(ctx);
1604 var state = ir.Operation.State.init(operation_name, loc);
1605 state.addOperands(&.{ memref, index, expected, desired });
1606 state.addTypes(&.{expected.type});
1607
1608 const op = try builder.create(state);
1609 errdefer op.erase();
1610 try setOrderingAttr(op, ctx, ordering);
1611 if (scope) |scope_value| {
1612 try setScopeAttr(op, ctx, scope_value);
1613 }
1614 return .{ .op = op };
1615 }
1616
1617 pub fn getResult(self: *const AtomicCasOp) *ir.Value {
1618 return self.op.getResult(0).?;
1619 }
1620
1621 pub fn getMemref(self: AtomicCasOp) *ir.Value {
1622 return self.op.operands.items[0].value;
1623 }
1624
1625 pub fn getIndex(self: AtomicCasOp) *ir.Value {
1626 return self.op.operands.items[1].value;
1627 }
1628
1629 pub fn getExpected(self: AtomicCasOp) *ir.Value {
1630 return self.op.operands.items[2].value;
1631 }
1632
1633 pub fn getDesired(self: AtomicCasOp) *ir.Value {
1634 return self.op.operands.items[3].value;
1635 }
1636
1637 pub fn getOrdering(self: AtomicCasOp) ?MemoryOrder {
1638 return getOrderingAttr(self.op);
1639 }
1640
1641 pub fn getScope(self: AtomicCasOp) ?Scope {
1642 return getScopeAttr(self.op);
1643 }
1644 };
1645
1646 fn loadSpec(ctx: *ir.Context) !void {
1647 try ir.dialects.loadDialectSpec(ctx, spec);
1648 }
1649
1650 fn dimIndexSpec(comptime mnemonic: []const u8) ir.dialects.OperationSpec {
1651 return op_specs.leaf(.{
1652 .mnemonic = mnemonic,
1653 .interfaces = &.{gpuEffects(.state_observe, &.{}, &.{})},
1654 .operands = 0,
1655 .results = .{"index"},
1656 .required_attrs = &.{"dim"},
1657 });
1658 }
1659
1660 fn indexSpec(comptime mnemonic: []const u8) ir.dialects.OperationSpec {
1661 return op_specs.leaf(.{
1662 .mnemonic = mnemonic,
1663 .interfaces = &.{gpuEffects(.state_observe, &.{}, &.{})},
1664 .operands = 0,
1665 .results = .{"index"},
1666 });
1667 }
1668
1669 fn noResultSpec(comptime mnemonic: []const u8, comptime operands: anytype) ir.dialects.OperationSpec {
1670 return op_specs.leaf(.{
1671 .mnemonic = mnemonic,
1672 .interfaces = &.{gpuEffects(.synchronize, &.{}, &.{})},
1673 .operands = operands,
1674 .results = 0,
1675 });
1676 }
1677
1678 fn getFuncSymbolName(op_ptr: *const anyopaque) ?[]const u8 {
1679 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
1680 if (op.getAttr("sym_name")) |attr| {
1681 return func.FuncDialect.getSymNameValue(attr);
1682 }
1683 return null;
1684 }
1685
1686 fn setFuncSymbolName(op_ptr: *const anyopaque, symbol_name: []const u8) anyerror!void {
1687 const op: *ir.Operation = @ptrCast(@alignCast(@constCast(op_ptr)));
1688 try op.setAttr("sym_name", try func.FuncDialect.getSymNameAttr(op.getContext(), symbol_name));
1689 }
1690
1691 fn isFuncDeclaration(_: *const anyopaque) bool {
1692 return false;
1693 }
1694
1695 pub fn getTmaDescriptorType(ctx: *ir.Context) !ir.Type {
1696 try loadSpec(ctx);
1697 return ctx.getDialectTypeFromName(type_names.tma_desc);
1698 }
1699
1700 pub fn getMBarrierType(ctx: *ir.Context) !ir.Type {
1701 try loadSpec(ctx);
1702 return ctx.getDialectTypeFromName(type_names.mbarrier);
1703 }
1704
1705 fn setDimensionAttr(op: *ir.Operation, ctx: *ir.Context, dim: Dimension) !void {
1706 const dim_attr = try ctx.getDialectAttr("gpu.dim", dim.toString());
1707 try op.setAttr("dim", dim_attr);
1708 }
1709
1710 fn getDimensionAttr(op: *const ir.Operation) ?Dimension {
1711 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "dim") orelse return null;
1712 return Dimension.fromString(dialect_attr.payload);
1713 }
1714
1715 fn setScopeAttr(op: *ir.Operation, ctx: *ir.Context, scope: Scope) !void {
1716 const scope_attr = try ctx.getDialectAttr("gpu.scope", scope.toString());
1717 try op.setAttr("scope", scope_attr);
1718 }
1719
1720 fn getScopeAttr(op: *const ir.Operation) ?Scope {
1721 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "scope") orelse return null;
1722 return Scope.fromString(dialect_attr.payload);
1723 }
1724
1725 fn setOrderingAttr(op: *ir.Operation, ctx: *ir.Context, ordering: MemoryOrder) !void {
1726 const order_attr = try ctx.getDialectAttr("gpu.ordering", ordering.toString());
1727 try op.setAttr("ordering", order_attr);
1728 }
1729
1730 fn getOrderingAttr(op: *const ir.Operation) ?MemoryOrder {
1731 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "ordering") orelse return null;
1732 return MemoryOrder.fromString(dialect_attr.payload);
1733 }
1734
1735 fn setBoolAttr(op: *ir.Operation, ctx: *ir.Context, attr_name: []const u8, value: bool) !void {
1736 const bool_attr = try ctx.getBoolAttr(value);
1737 try op.setAttr(attr_name, bool_attr);
1738 }
1739
1740 fn getBoolAttrValue(op: *const ir.Operation, attr_name: []const u8) bool {
1741 const bool_attr = op.getAttrAs(ir.Attribute.BoolAttr, attr_name) orelse return false;
1742 return bool_attr.getValue();
1743 }
1744
1745 fn setI64Attr(op: *ir.Operation, ctx: *ir.Context, attr_name: []const u8, value: i64) !void {
1746 const int_attr = try ctx.getI64Attr(value);
1747 try op.setAttr(attr_name, int_attr);
1748 }
1749
1750 fn getI64AttrValue(op: *const ir.Operation, attr_name: []const u8) ?i64 {
1751 const int_attr = op.getAttrAs(ir.Attribute.IntegerAttr, attr_name) orelse return null;
1752 return int_attr.getValue();
1753 }
1754
1755 fn setWarpOpAttr(op: *ir.Operation, ctx: *ir.Context, op_kind: WarpOpKind) !void {
1756 const op_attr = try ctx.getDialectAttr("gpu.warp_op", op_kind.toString());
1757 try op.setAttr("op", op_attr);
1758 }
1759
1760 fn getWarpOpAttr(op: *const ir.Operation) ?WarpOpKind {
1761 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "op") orelse return null;
1762 return WarpOpKind.fromString(dialect_attr.payload);
1763 }
1764
1765 fn setMmaShapeAttr(op: *ir.Operation, ctx: *ir.Context, shape: MmaShape) !void {
1766 var buf: [32]u8 = undefined;
1767 const shape_str = try shape.toString(buf[0..]);
1768 const shape_attr = try ctx.getDialectAttr("gpu.mma_shape", shape_str);
1769 try op.setAttr("shape", shape_attr);
1770 }
1771
1772 fn getMmaShapeAttr(op: *const ir.Operation) ?MmaShape {
1773 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "shape") orelse return null;
1774 return MmaShape.parse(dialect_attr.payload);
1775 }
1776 };
1777
1778 const GpuFactoryResourceCounts = struct {
1779 operations: usize,
1780
1781 fn capture(ctx: *const ir.Context) GpuFactoryResourceCounts {
1782 return .{
1783 .operations = ctx.operationCount(),
1784 };
1785 }
1786
1787 fn expectEqual(self: GpuFactoryResourceCounts, ctx: *const ir.Context) !void {
1788 try std.testing.expectEqual(self.operations, ctx.operationCount());
1789 }
1790 };
1791
1792 fn checkGpuFactoryAllocationFailures(allocator: std.mem.Allocator) !void {
1793 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1794 defer ctx.deinit(allocator);
1795 try ctx.allowUnregistered();
1796
1797 const arith = choir.dialects.arith.ArithDialect;
1798 const loc = ir.Location.getUnknown();
1799 const i32_type = try arith.getScalarType(&ctx, .i32);
1800 const f32_type = try arith.getScalarType(&ctx, .f32);
1801 const index_type = try arith.getIndexType(&ctx);
1802 var source_builder = ir.OperationBuilder.init(&ctx);
1803 var source_state = ir.Operation.State.init("test.gpu_factory_source", loc);
1804 source_state.addTypes(&.{
1805 i32_type,
1806 f32_type,
1807 index_type,
1808 f32_type,
1809 f32_type,
1810 f32_type,
1811 f32_type,
1812 f32_type,
1813 f32_type,
1814 f32_type,
1815 f32_type,
1816 f32_type,
1817 f32_type,
1818 });
1819 const source = try source_builder.create(source_state);
1820 defer source.erase();
1821 const baseline = GpuFactoryResourceCounts.capture(&ctx);
1822 defer baseline.expectEqual(&ctx) catch unreachable;
1823
1824 const module = try GpuDialect.ModuleOp.create(&ctx, loc);
1825 defer module.op.erase();
1826 const function = try GpuDialect.FuncOp.createKernel(&ctx, loc, "gpu_factory_kernel", &.{f32_type});
1827 defer function.op.erase();
1828 const launch = try GpuDialect.LaunchOp.create(
1829 &ctx,
1830 loc,
1831 "gpu_factory_kernel",
1832 &.{source.getResult(1).?},
1833 .{ 1, 1, 1 },
1834 .{ 32, 1, 1 },
1835 );
1836 defer launch.op.erase();
1837 const thread = try GpuDialect.ThreadIdxOp.create(&ctx, loc, .x);
1838 defer thread.op.erase();
1839 const block = try GpuDialect.BlockIdxOp.create(&ctx, loc, .y);
1840 defer block.op.erase();
1841 const block_dim = try GpuDialect.BlockDimOp.create(&ctx, loc, .z);
1842 defer block_dim.op.erase();
1843 const grid_dim = try GpuDialect.GridDimOp.create(&ctx, loc, .x);
1844 defer grid_dim.op.erase();
1845 const global = try GpuDialect.GlobalIdxOp.create(&ctx, loc, .y);
1846 defer global.op.erase();
1847 const barrier = try GpuDialect.BarrierOp.create(&ctx, loc, .block);
1848 defer barrier.op.erase();
1849 const fence = try GpuDialect.FenceOp.create(&ctx, loc, .device, .seq_cst);
1850 defer fence.op.erase();
1851 const memcpy = try GpuDialect.MemcpyAsyncOp.create(
1852 &ctx,
1853 loc,
1854 source.getResult(1).?,
1855 source.getResult(3).?,
1856 source.getResult(2).?,
1857 null,
1858 null,
1859 );
1860 defer memcpy.op.erase();
1861 const tma_wait = try GpuDialect.TmaWaitGroupOp.create(&ctx, loc, 2);
1862 defer tma_wait.op.erase();
1863 const shuffle = try GpuDialect.ShflSyncOp.create(
1864 &ctx,
1865 loc,
1866 .down,
1867 source.getResult(0).?,
1868 source.getResult(1).?,
1869 source.getResult(2).?,
1870 );
1871 defer shuffle.op.erase();
1872 const warp_reduce = try GpuDialect.WarpReduceOp.create(
1873 &ctx,
1874 loc,
1875 .add,
1876 source.getResult(0).?,
1877 source.getResult(1).?,
1878 );
1879 defer warp_reduce.op.erase();
1880 const warp_scan = try GpuDialect.WarpScanOp.create(
1881 &ctx,
1882 loc,
1883 .add,
1884 true,
1885 source.getResult(0).?,
1886 source.getResult(1).?,
1887 );
1888 defer warp_scan.op.erase();
1889 const mma = try GpuDialect.MmaSyncOp.create(&ctx, loc, .{
1890 source.getResult(3).?,
1891 source.getResult(4).?,
1892 source.getResult(5).?,
1893 source.getResult(6).?,
1894 }, .{
1895 source.getResult(7).?,
1896 source.getResult(8).?,
1897 }, .{
1898 source.getResult(9).?,
1899 source.getResult(10).?,
1900 source.getResult(11).?,
1901 source.getResult(12).?,
1902 }, .{ .m = 16, .n = 8, .k = 8 });
1903 defer mma.op.erase();
1904 const async_copy = try GpuDialect.CpAsyncSharedOp.create(
1905 &ctx,
1906 loc,
1907 source.getResult(1).?,
1908 source.getResult(2).?,
1909 source.getResult(3).?,
1910 source.getResult(2).?,
1911 16,
1912 );
1913 defer async_copy.op.erase();
1914 const async_wait = try GpuDialect.CpAsyncWaitOp.create(&ctx, loc, 1);
1915 defer async_wait.op.erase();
1916 const atomic_load = try GpuDialect.AtomicLoadOp.create(
1917 &ctx,
1918 loc,
1919 source.getResult(1).?,
1920 source.getResult(2).?,
1921 f32_type,
1922 .acquire,
1923 );
1924 defer atomic_load.op.erase();
1925 const atomic_store = try GpuDialect.AtomicStoreOp.create(
1926 &ctx,
1927 loc,
1928 source.getResult(3).?,
1929 source.getResult(1).?,
1930 source.getResult(2).?,
1931 .release,
1932 );
1933 defer atomic_store.op.erase();
1934 const atomic_add = try GpuDialect.AtomicAddOp.create(
1935 &ctx,
1936 loc,
1937 source.getResult(1).?,
1938 source.getResult(2).?,
1939 source.getResult(3).?,
1940 .relaxed,
1941 .block,
1942 );
1943 defer atomic_add.op.erase();
1944 const atomic_cas = try GpuDialect.AtomicCasOp.create(
1945 &ctx,
1946 loc,
1947 source.getResult(1).?,
1948 source.getResult(2).?,
1949 source.getResult(3).?,
1950 source.getResult(4).?,
1951 .seq_cst,
1952 .device,
1953 );
1954 defer atomic_cas.op.erase();
1955 }
1956
1957 test "GpuDialect factories restore resources on allocation failure" {
1958 try std.testing.checkAllAllocationFailures(
1959 std.testing.allocator,
1960 checkGpuFactoryAllocationFailures,
1961 .{},
1962 );
1963 }
1964
1965 test "GpuDialect.ModuleOp creates module container" {
1966 const testing = std.testing;
1967
1968 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1969 defer arena.deinit();
1970 const allocator = arena.allocator();
1971
1972 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1973 defer ctx.deinit(allocator);
1974
1975 const loc = ir.Location.getUnknown();
1976 const module = try GpuDialect.ModuleOp.create(&ctx, loc);
1977
1978 try testing.expectEqualStrings("gpu.module", module.op.name.name);
1979 try testing.expect(module.getBody().getEntryBlock() != null);
1980 }
1981
1982 test "GpuDialect.ModuleOp owns a symbol table" {
1983 const testing = std.testing;
1984
1985 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1986 defer arena.deinit();
1987 const allocator = arena.allocator();
1988
1989 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1990 defer ctx.deinit(allocator);
1991 try ir.dialects.loadDialectSpec(&ctx, GpuDialect.spec);
1992
1993 const loc = ir.Location.getUnknown();
1994 const module = try GpuDialect.ModuleOp.create(&ctx, loc);
1995 const block = module.getBodyBlock();
1996 const kernel = try GpuDialect.FuncOp.create(&ctx, loc, "kernel", &.{}, &.{});
1997 try block.addOperation(kernel.op);
1998
1999 var table = ir.SymbolTable.init(allocator);
2000 defer table.deinit();
2001 try table.buildFromOperation(module.op);
2002
2003 try testing.expect(module.op.getTraits().is_symbol_table);
2004 try testing.expect(kernel.op.interface(ir.interfaces.SymbolOpInterface) != null);
2005 try testing.expect(table.lookup("kernel") == kernel.op);
2006 }
2007
2008 test "GpuDialect operation specs register shapes and attributes" {
2009 const testing = std.testing;
2010
2011 var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);
2012 defer ctx.deinit(testing.allocator);
2013
2014 try ir.dialects.loadDialectSpec(&ctx, GpuDialect.spec);
2015
2016 const launch_info = ctx.lookupOperation(GpuDialect.LaunchOp.operation_name) orelse return error.TestExpectedOperation;
2017 try testing.expect(launch_info.shape.operands.allows(0));
2018 try testing.expect(launch_info.shape.operands.allows(8));
2019 try testing.expect(launch_info.shape.results.allows(0));
2020 try testing.expect(!launch_info.shape.results.allows(1));
2021 try testing.expect(launch_info.hasInherentAttributeName("kernel"));
2022 try testing.expect(launch_info.hasRequiredAttributeName("kernel"));
2023 try testing.expect(launch_info.hasRequiredAttributeName("num_kernel_args"));
2024
2025 const memcpy_info = ctx.lookupOperation(GpuDialect.MemcpyAsyncOp.operation_name) orelse return error.TestExpectedOperation;
2026 try testing.expect(memcpy_info.shape.operands.allows(3));
2027 try testing.expect(memcpy_info.shape.operands.allows(5));
2028 try testing.expect(!memcpy_info.shape.operands.allows(2));
2029 try testing.expect(!memcpy_info.shape.operands.allows(6));
2030 try testing.expect(memcpy_info.shape.results.allows(0));
2031 try testing.expect(!memcpy_info.shape.results.allows(1));
2032 const memcpy_segments = memcpy_info.getOperandSegments() orelse return error.TestExpectedOperationSegments;
2033 try testing.expectEqualStrings("operand_segment_sizes", memcpy_segments.attribute_name);
2034 try testing.expectEqual(@as(usize, 5), memcpy_segments.segments.len);
2035 try testing.expect(memcpy_segments.segments[3].allows(0));
2036 try testing.expect(memcpy_segments.segments[3].allows(1));
2037 try testing.expect(!memcpy_segments.segments[3].allows(2));
2038 try testing.expect(memcpy_segments.segments[4].allows(0));
2039 try testing.expect(memcpy_segments.segments[4].allows(1));
2040 try testing.expect(!memcpy_segments.segments[4].allows(2));
2041 try testing.expectEqualStrings("src", GpuDialect.MemcpyAsyncOp.operation_spec.operand_names[0]);
2042 try testing.expectEqualStrings("stream", GpuDialect.MemcpyAsyncOp.operation_spec.operand_names[3]);
2043 try testing.expectEqualStrings("event", GpuDialect.MemcpyAsyncOp.operation_spec.operand_names[4]);
2044
2045 const shuffle_info = ctx.lookupOperation(GpuDialect.ShflSyncOp.operation_name) orelse return error.TestExpectedOperation;
2046 try testing.expect(shuffle_info.shape.operands.allows(3));
2047 try testing.expect(!shuffle_info.shape.operands.allows(2));
2048 try testing.expect(shuffle_info.shape.results.allows(1));
2049 try testing.expect(shuffle_info.hasRequiredAttributeName("mode"));
2050 try testing.expectEqualStrings("mask", GpuDialect.ShflSyncOp.operation_spec.operand_names[0]);
2051 try testing.expectEqualStrings("lane_or_delta", GpuDialect.ShflSyncOp.operation_spec.operand_names[2]);
2052 try testing.expectEqualStrings("result", GpuDialect.ShflSyncOp.operation_spec.result_names[0]);
2053
2054 const module_info = ctx.lookupOperation(GpuDialect.ModuleOp.operation_name) orelse return error.TestExpectedOperation;
2055 try testing.expect(module_info.shape.operands.allows(0));
2056 try testing.expect(!module_info.shape.operands.allows(1));
2057 try testing.expect(module_info.shape.regions.allows(1));
2058 try testing.expect(!module_info.shape.regions.allows(0));
2059 try testing.expect(module_info.hasTraitId(ir.traits.SymbolTable.id));
2060
2061 const atomic_add_info = ctx.lookupOperation(GpuDialect.AtomicAddOp.operation_name) orelse return error.TestExpectedOperation;
2062 try testing.expect(atomic_add_info.hasInherentAttributeName("scope"));
2063 try testing.expect(!atomic_add_info.hasRequiredAttributeName("scope"));
2064 try testing.expect(atomic_add_info.hasRequiredAttributeName("ordering"));
2065 }
2066
2067 test "GpuDialect.FuncOp creates kernel" {
2068 const testing = std.testing;
2069 const arith = choir.dialects.arith;
2070 const memref = choir.dialects.memref;
2071
2072 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2073 defer arena.deinit();
2074 const allocator = arena.allocator();
2075
2076 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2077 defer ctx.deinit(allocator);
2078
2079 const loc = ir.Location.getUnknown();
2080 const f32_type = try arith.ArithDialect.getScalarType(&ctx, .f32);
2081 const memref_type = try memref.MemrefDialect.getMemrefType1D(&ctx, 1024, f32_type, .device);
2082
2083 const kernel_op = try GpuDialect.FuncOp.createKernel(
2084 &ctx,
2085 loc,
2086 "vector_add",
2087 &.{ memref_type, memref_type, memref_type },
2088 );
2089
2090 try testing.expectEqualStrings("gpu.func", kernel_op.op.name.name);
2091 try testing.expectEqualStrings("vector_add", kernel_op.getName().?);
2092 try testing.expectEqual(@as(usize, 3), kernel_op.getNumArguments());
2093 try testing.expectEqual(@as(usize, 0), kernel_op.getNumResults());
2094 try testing.expect(kernel_op.isKernel());
2095 }
2096
2097 test "GpuDialect.YieldOp captures operands" {
2098 const testing = std.testing;
2099 const arith = choir.dialects.arith;
2100
2101 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2102 defer arena.deinit();
2103 const allocator = arena.allocator();
2104
2105 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2106 defer ctx.deinit(allocator);
2107
2108 const loc = ir.Location.getUnknown();
2109 const i32_type = try arith.ArithDialect.getI32Type(&ctx);
2110 var val = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 7);
2111
2112 const yield_op = try GpuDialect.YieldOp.create(&ctx, loc, &.{val.getResult()});
2113 try testing.expectEqualStrings("gpu.yield", yield_op.op.name.name);
2114 const operands = yield_op.getOperands();
2115 try testing.expectEqual(@as(usize, 1), operands.len);
2116 try testing.expect(operands[0] == val.getResult());
2117 }
2118
2119 test "GpuDialect.LaunchOp creates kernel launch" {
2120 const testing = std.testing;
2121 const arith = choir.dialects.arith;
2122 const memref = choir.dialects.memref;
2123
2124 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2125 defer arena.deinit();
2126 const allocator = arena.allocator();
2127
2128 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2129 defer ctx.deinit(allocator);
2130
2131 const loc = ir.Location.getUnknown();
2132 const f32_type = try arith.ArithDialect.getScalarType(&ctx, .f32);
2133 const memref_type = try memref.MemrefDialect.getMemrefType1D(&ctx, 1024, f32_type, .device);
2134
2135 var alloc1 = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, memref_type);
2136 var alloc2 = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, memref_type);
2137 var alloc3 = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, memref_type);
2138
2139 const launch = try GpuDialect.LaunchOp.create(
2140 &ctx,
2141 loc,
2142 "vector_add",
2143 &.{ alloc1.getResult(), alloc2.getResult(), alloc3.getResult() },
2144 .{ 4, 1, 1 },
2145 .{ 256, 1, 1 },
2146 );
2147
2148 try testing.expectEqualStrings("gpu.launch", launch.op.name.name);
2149 try testing.expectEqualStrings("vector_add", launch.getKernelName().?);
2150 try testing.expectEqual(@as(usize, 3), launch.getNumKernelArgs());
2151 const kernel_args = launch.getKernelArgs();
2152 try testing.expectEqual(@as(usize, 3), kernel_args.len);
2153 try testing.expect(kernel_args[0] == alloc1.getResult());
2154 try testing.expect(kernel_args[1] == alloc2.getResult());
2155 try testing.expect(kernel_args[2] == alloc3.getResult());
2156 const grid = launch.getGridDim().?;
2157 const block = launch.getBlockDim().?;
2158 try testing.expectEqual(@as(u32, 4), grid[0]);
2159 try testing.expectEqual(@as(u32, 1), grid[1]);
2160 try testing.expectEqual(@as(u32, 1), grid[2]);
2161 try testing.expectEqual(@as(u32, 256), block[0]);
2162 try testing.expectEqual(@as(u32, 1), block[1]);
2163 try testing.expectEqual(@as(u32, 1), block[2]);
2164 }
2165
2166 test "GpuDialect.ThreadIdxOp creates thread index" {
2167 const testing = std.testing;
2168
2169 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2170 defer arena.deinit();
2171 const allocator = arena.allocator();
2172
2173 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2174 defer ctx.deinit(allocator);
2175
2176 const loc = ir.Location.getUnknown();
2177 const tid_x = try GpuDialect.ThreadIdxOp.create(&ctx, loc, .x);
2178
2179 try testing.expectEqualStrings("gpu.thread_idx", tid_x.op.name.name);
2180 try testing.expectEqual(Dimension.x, tid_x.getDimension().?);
2181 }
2182
2183 test "GpuDialect.BlockIdxOp creates block index" {
2184 const testing = std.testing;
2185
2186 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2187 defer arena.deinit();
2188 const allocator = arena.allocator();
2189
2190 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2191 defer ctx.deinit(allocator);
2192
2193 const loc = ir.Location.getUnknown();
2194 const bid_y = try GpuDialect.BlockIdxOp.create(&ctx, loc, .y);
2195
2196 try testing.expectEqualStrings("gpu.block_idx", bid_y.op.name.name);
2197 try testing.expectEqual(Dimension.y, bid_y.getDimension().?);
2198 }
2199
2200 test "GpuDialect.BarrierOp creates barrier" {
2201 const testing = std.testing;
2202
2203 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2204 defer arena.deinit();
2205 const allocator = arena.allocator();
2206
2207 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2208 defer ctx.deinit(allocator);
2209
2210 const loc = ir.Location.getUnknown();
2211 const barrier = try GpuDialect.BarrierOp.create(&ctx, loc, .block);
2212
2213 try testing.expectEqualStrings("gpu.barrier", barrier.op.name.name);
2214 try testing.expectEqual(Scope.block, barrier.getScope().?);
2215 }
2216
2217 test "GpuDialect.FenceOp creates fence" {
2218 const testing = std.testing;
2219
2220 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2221 defer arena.deinit();
2222 const allocator = arena.allocator();
2223
2224 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2225 defer ctx.deinit(allocator);
2226
2227 const loc = ir.Location.getUnknown();
2228 const fence = try GpuDialect.FenceOp.create(&ctx, loc, .device, .seq_cst);
2229
2230 try testing.expectEqualStrings("gpu.fence", fence.op.name.name);
2231 try testing.expectEqual(Scope.device, fence.getScope().?);
2232 try testing.expectEqual(MemoryOrder.seq_cst, fence.getOrdering().?);
2233 }
2234
2235 test "GpuDialect.MemcpyAsyncOp resolves optional operands" {
2236 const testing = std.testing;
2237 const arith = choir.dialects.arith;
2238 const memref = choir.dialects.memref;
2239
2240 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2241 defer arena.deinit();
2242 const allocator = arena.allocator();
2243
2244 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2245 defer ctx.deinit(allocator);
2246
2247 const loc = ir.Location.getUnknown();
2248 const i32_type = try arith.ArithDialect.getI32Type(&ctx);
2249 const index_type = try arith.ArithDialect.getIndexType(&ctx);
2250 const memref_type = try memref.MemrefDialect.getMemrefType1D(&ctx, 16, i32_type, .device);
2251
2252 var src = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, memref_type);
2253 var dst = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, memref_type);
2254 var size = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 16);
2255 var stream = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 1);
2256 var event = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 2);
2257
2258 const event_only = try GpuDialect.MemcpyAsyncOp.create(
2259 &ctx,
2260 loc,
2261 src.getResult(),
2262 dst.getResult(),
2263 size.getResult(),
2264 null,
2265 event.getResult(),
2266 );
2267 try testing.expect(event_only.getStream() == null);
2268 try testing.expect(event_only.getEvent().? == event.getResult());
2269 try ir.verifyOperation(event_only.op, .{ .recursive = false });
2270
2271 const stream_only = try GpuDialect.MemcpyAsyncOp.create(
2272 &ctx,
2273 loc,
2274 src.getResult(),
2275 dst.getResult(),
2276 size.getResult(),
2277 stream.getResult(),
2278 null,
2279 );
2280 try testing.expect(stream_only.getStream().? == stream.getResult());
2281 try testing.expect(stream_only.getEvent() == null);
2282 try ir.verifyOperation(stream_only.op, .{ .recursive = false });
2283
2284 const both = try GpuDialect.MemcpyAsyncOp.create(
2285 &ctx,
2286 loc,
2287 src.getResult(),
2288 dst.getResult(),
2289 size.getResult(),
2290 stream.getResult(),
2291 event.getResult(),
2292 );
2293 try testing.expect(both.getStream().? == stream.getResult());
2294 try testing.expect(both.getEvent().? == event.getResult());
2295 try ir.verifyOperation(both.op, .{ .recursive = false });
2296 }
2297
2298 test "GpuDialect.ShflSyncOp creates shuffle" {
2299 const testing = std.testing;
2300 const arith = choir.dialects.arith;
2301
2302 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2303 defer arena.deinit();
2304 const allocator = arena.allocator();
2305
2306 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2307 defer ctx.deinit(allocator);
2308
2309 const loc = ir.Location.getUnknown();
2310 const i32_type = try arith.ArithDialect.getI32Type(&ctx);
2311
2312 var mask = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 0xFFFFFFFF);
2313 var src = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 42);
2314 var delta = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 1);
2315
2316 const shfl = try GpuDialect.ShflSyncOp.create(&ctx, loc, .down, mask.getResult(), src.getResult(), delta.getResult());
2317
2318 try testing.expectEqualStrings("gpu.shfl_sync", shfl.op.name.name);
2319 try testing.expectEqual(ShuffleMode.down, shfl.getMode().?);
2320 }
2321
2322 test "GpuDialect.WarpReduceOp creates warp reduction" {
2323 const testing = std.testing;
2324 const arith = choir.dialects.arith;
2325
2326 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2327 defer arena.deinit();
2328 const allocator = arena.allocator();
2329
2330 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2331 defer ctx.deinit(allocator);
2332
2333 const loc = ir.Location.getUnknown();
2334 const i32_type = try arith.ArithDialect.getI32Type(&ctx);
2335
2336 var mask = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 0xFFFFFFFF);
2337 var value = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 7);
2338
2339 const reduce = try GpuDialect.WarpReduceOp.create(&ctx, loc, .add, mask.getResult(), value.getResult());
2340
2341 try testing.expectEqualStrings("gpu.warp_reduce", reduce.op.name.name);
2342 try testing.expectEqual(WarpOpKind.add, reduce.getOpKind().?);
2343 try testing.expect(reduce.getResult().type.eql(i32_type));
2344 }
2345
2346 test "GpuDialect.WarpScanOp captures op kind and inclusive flag" {
2347 const testing = std.testing;
2348 const arith = choir.dialects.arith;
2349
2350 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2351 defer arena.deinit();
2352 const allocator = arena.allocator();
2353
2354 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2355 defer ctx.deinit(allocator);
2356
2357 const loc = ir.Location.getUnknown();
2358 const i32_type = try arith.ArithDialect.getI32Type(&ctx);
2359
2360 var mask = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 0xFFFFFFFF);
2361 var value = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 3);
2362
2363 const scan = try GpuDialect.WarpScanOp.create(&ctx, loc, .xor, false, mask.getResult(), value.getResult());
2364
2365 try testing.expectEqualStrings("gpu.warp_scan", scan.op.name.name);
2366 try testing.expectEqual(WarpOpKind.xor, scan.getOpKind().?);
2367 try testing.expect(!scan.isInclusive());
2368 try testing.expect(scan.getResult().type.eql(i32_type));
2369 }
2370
2371 test "GpuDialect.Match ops return expected result shapes" {
2372 const testing = std.testing;
2373 const arith = choir.dialects.arith;
2374
2375 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2376 defer arena.deinit();
2377 const allocator = arena.allocator();
2378
2379 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2380 defer ctx.deinit(allocator);
2381
2382 const loc = ir.Location.getUnknown();
2383 const i32_type = try arith.ArithDialect.getI32Type(&ctx);
2384 const bool_type = try arith.ArithDialect.getScalarType(&ctx, .bool);
2385
2386 var value = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 11);
2387
2388 const any = try GpuDialect.MatchAnyOp.create(&ctx, loc, value.getResult());
2389 try testing.expectEqualStrings("gpu.match_any", any.op.name.name);
2390 try testing.expect(any.getResult().type.eql(i32_type));
2391
2392 const all = try GpuDialect.MatchAllOp.create(&ctx, loc, value.getResult());
2393 try testing.expectEqualStrings("gpu.match_all", all.op.name.name);
2394 try testing.expect(all.getMaskResult().type.eql(i32_type));
2395 try testing.expect(all.getAllEqualResult().type.eql(bool_type));
2396 }
2397
2398 test "GpuDialect.ActiveMaskOp and SyncWarpOp create warp control ops" {
2399 const testing = std.testing;
2400 const arith = choir.dialects.arith;
2401
2402 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2403 defer arena.deinit();
2404 const allocator = arena.allocator();
2405
2406 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2407 defer ctx.deinit(allocator);
2408
2409 const loc = ir.Location.getUnknown();
2410 const i32_type = try arith.ArithDialect.getI32Type(&ctx);
2411
2412 const active = try GpuDialect.ActiveMaskOp.create(&ctx, loc);
2413 try testing.expectEqualStrings("gpu.active_mask", active.op.name.name);
2414 try testing.expect(active.getResult().type.eql(i32_type));
2415
2416 var mask = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 0xFFFFFFFF);
2417 const sync = try GpuDialect.SyncWarpOp.create(&ctx, loc, mask.getResult());
2418 try testing.expectEqualStrings("gpu.sync_warp", sync.op.name.name);
2419 try testing.expect(sync.getMask() == mask.getResult());
2420 }
2421
2422 test "GpuDialect.AtomicAddOp creates atomic add" {
2423 const testing = std.testing;
2424 const arith = choir.dialects.arith;
2425 const memref = choir.dialects.memref;
2426
2427 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2428 defer arena.deinit();
2429 const allocator = arena.allocator();
2430
2431 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2432 defer ctx.deinit(allocator);
2433
2434 const loc = ir.Location.getUnknown();
2435 const i32_type = try arith.ArithDialect.getI32Type(&ctx);
2436 const index_type = try arith.ArithDialect.getIndexType(&ctx);
2437 const memref_type = try memref.MemrefDialect.getMemrefType1D(&ctx, 1024, i32_type, .device);
2438
2439 var alloc = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, memref_type);
2440 var idx = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 0);
2441 var val = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 1);
2442
2443 const atomic = try GpuDialect.AtomicAddOp.create(&ctx, loc, alloc.getResult(), idx.getResult(), val.getResult(), .relaxed, .block);
2444
2445 try testing.expectEqualStrings("gpu.atomic_add", atomic.op.name.name);
2446 try testing.expect(atomic.getOrdering() == .relaxed);
2447 try testing.expect(atomic.getScope() == .block);
2448 }
2449
2450 test "GpuDialect.AtomicLoadOp and AtomicStoreOp create atomic ops" {
2451 const testing = std.testing;
2452 const arith = choir.dialects.arith;
2453 const memref = choir.dialects.memref;
2454
2455 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2456 defer arena.deinit();
2457 const allocator = arena.allocator();
2458
2459 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2460 defer ctx.deinit(allocator);
2461
2462 const loc = ir.Location.getUnknown();
2463 const i32_type = try arith.ArithDialect.getI32Type(&ctx);
2464 const index_type = try arith.ArithDialect.getIndexType(&ctx);
2465 const memref_type = try memref.MemrefDialect.getMemrefType1D(&ctx, 64, i32_type, .shared);
2466
2467 var alloc = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, memref_type);
2468 var idx = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 3);
2469 var val = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 7);
2470
2471 const atomic_load = try GpuDialect.AtomicLoadOp.create(&ctx, loc, alloc.getResult(), idx.getResult(), i32_type, .relaxed);
2472 try testing.expectEqualStrings("gpu.atomic_load", atomic_load.op.name.name);
2473 try testing.expect(atomic_load.getOrdering() == .relaxed);
2474
2475 const atomic_store = try GpuDialect.AtomicStoreOp.create(&ctx, loc, val.getResult(), alloc.getResult(), idx.getResult(), .release);
2476 try testing.expectEqualStrings("gpu.atomic_store", atomic_store.op.name.name);
2477 try testing.expect(atomic_store.getOrdering() == .release);
2478 }
2479
2480 test "GpuDialect.AtomicCasOp creates atomic cas" {
2481 const testing = std.testing;
2482 const arith = choir.dialects.arith;
2483 const memref = choir.dialects.memref;
2484
2485 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2486 defer arena.deinit();
2487 const allocator = arena.allocator();
2488
2489 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2490 defer ctx.deinit(allocator);
2491
2492 const loc = ir.Location.getUnknown();
2493 const i32_type = try arith.ArithDialect.getI32Type(&ctx);
2494 const index_type = try arith.ArithDialect.getIndexType(&ctx);
2495 const memref_type = try memref.MemrefDialect.getMemrefType1D(&ctx, 256, i32_type, .device);
2496
2497 var alloc = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, memref_type);
2498 var idx = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 1);
2499 var expected = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 7);
2500 var desired = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, i32_type, 9);
2501
2502 const atomic = try GpuDialect.AtomicCasOp.create(&ctx, loc, alloc.getResult(), idx.getResult(), expected.getResult(), desired.getResult(), .seq_cst, .device);
2503
2504 try testing.expectEqualStrings("gpu.atomic_cas", atomic.op.name.name);
2505 try testing.expect(atomic.getOrdering() == .seq_cst);
2506 try testing.expect(atomic.getScope() == .device);
2507 }
2508
2509 test "GpuDialect.LaneIdOp creates lane id" {
2510 const testing = std.testing;
2511
2512 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2513 defer arena.deinit();
2514 const allocator = arena.allocator();
2515
2516 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2517 defer ctx.deinit(allocator);
2518
2519 const loc = ir.Location.getUnknown();
2520 const lane = try GpuDialect.LaneIdOp.create(&ctx, loc);
2521
2522 try testing.expectEqualStrings("gpu.lane_id", lane.op.name.name);
2523 }
2524
2525 test "GpuDialect.GlobalIdxOp creates global index" {
2526 const testing = std.testing;
2527
2528 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2529 defer arena.deinit();
2530 const allocator = arena.allocator();
2531
2532 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2533 defer ctx.deinit(allocator);
2534
2535 const loc = ir.Location.getUnknown();
2536 const gid = try GpuDialect.GlobalIdxOp.create(&ctx, loc, .x);
2537
2538 try testing.expectEqualStrings("gpu.global_idx", gid.op.name.name);
2539 try testing.expectEqual(Dimension.x, gid.getDimension().?);
2540 }
2541
2542 test "GpuDialect.MmaSyncOp carries lane fragments and shape" {
2543 const testing = std.testing;
2544 const arith = choir.dialects.arith;
2545
2546 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2547 defer arena.deinit();
2548 const allocator = arena.allocator();
2549
2550 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2551 defer ctx.deinit(allocator);
2552
2553 const loc = ir.Location.getUnknown();
2554 const f32_type = try arith.ArithDialect.getScalarType(&ctx, .f32);
2555 const shape = MmaShape{ .m = 16, .n = 8, .k = 8 };
2556
2557 try ctx.allowUnregistered();
2558 var builder = ir.OperationBuilder.init(&ctx);
2559 var source_state = ir.Operation.State.init("test.frag_source", loc);
2560 source_state.addTypes(&.{ f32_type, f32_type, f32_type, f32_type, f32_type, f32_type, f32_type, f32_type, f32_type, f32_type });
2561 const source = try builder.create(source_state);
2562
2563 const mma = try GpuDialect.MmaSyncOp.create(&ctx, loc, .{
2564 source.getResult(0).?,
2565 source.getResult(1).?,
2566 source.getResult(2).?,
2567 source.getResult(3).?,
2568 }, .{
2569 source.getResult(4).?,
2570 source.getResult(5).?,
2571 }, .{
2572 source.getResult(6).?,
2573 source.getResult(7).?,
2574 source.getResult(8).?,
2575 source.getResult(9).?,
2576 }, shape);
2577
2578 try testing.expectEqualStrings("gpu.mma_sync", mma.op.name.name);
2579 try testing.expectEqual(@as(usize, 10), mma.op.operands.items.len);
2580 try testing.expectEqual(@as(usize, 4), mma.op.getNumResults());
2581 try testing.expect(mma.getA(1) == source.getResult(1).?);
2582 try testing.expect(mma.getB(0) == source.getResult(4).?);
2583 try testing.expect(mma.getC(3) == source.getResult(9).?);
2584 try testing.expect(mma.getD(0).type.eql(f32_type));
2585
2586 const parsed = mma.getShape().?;
2587 try testing.expectEqual(@as(u32, 16), parsed.m);
2588 try testing.expectEqual(@as(u32, 8), parsed.n);
2589 try testing.expectEqual(@as(u32, 8), parsed.k);
2590 }
2591
2592 test "GpuDialect.Tma ops create descriptor/load/commit/wait" {
2593 const testing = std.testing;
2594 const arith = choir.dialects.arith;
2595 const memref = choir.dialects.memref;
2596
2597 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2598 defer arena.deinit();
2599 const allocator = arena.allocator();
2600
2601 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
2602 defer ctx.deinit(allocator);
2603
2604 const loc = ir.Location.getUnknown();
2605 const index_type = try arith.ArithDialect.getIndexType(&ctx);
2606 const f16_type = try arith.ArithDialect.getScalarType(&ctx, .f16);
2607
2608 const src_type = try memref.MemrefDialect.getMemrefType1D(&ctx, 64, f16_type, .device);
2609 var src = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, src_type);
2610 var shape = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 64);
2611
2612 const desc = try GpuDialect.TmaCreateDescriptorOp.create(&ctx, loc, src.getResult(), shape.getResult());
2613 const desc_type = try GpuDialect.getTmaDescriptorType(&ctx);
2614 try testing.expect(desc.getResult().type.eql(desc_type));
2615 try testing.expect(desc.getTensor() == src.getResult());
2616 try testing.expect(desc.getBoxShape() == shape.getResult());
2617
2618 const shmem_type = try memref.MemrefDialect.getMemrefType1D(&ctx, 64, f16_type, .shared);
2619 var shmem = try memref.MemrefDialect.AllocOp.createStatic(&ctx, loc, shmem_type);
2620
2621 const barrier_type = try GpuDialect.getMBarrierType(&ctx);
2622 var builder = ir.OperationBuilder.init(&ctx);
2623 _ = try ctx.registerOperation("gpu.test_mbarrier", .{});
2624 var barrier_state = ir.Operation.State.init("gpu.test_mbarrier", loc);
2625 barrier_state.addTypes(&.{barrier_type});
2626 const barrier_op = try builder.create(barrier_state);
2627 const barrier = barrier_op.getResult(0).?;
2628
2629 var coords = try arith.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 0);
2630 const load = try GpuDialect.TmaLoadOp.create(&ctx, loc, desc.getResult(), shmem.getResult(), barrier, coords.getResult());
2631 try testing.expectEqualStrings("gpu.tma.load", load.op.name.name);
2632 try testing.expect(load.getDescriptor() == desc.getResult());
2633 try testing.expect(load.getSharedMem() == shmem.getResult());
2634 try testing.expect(load.getBarrier() == barrier);
2635 try testing.expect(load.getCoords() == coords.getResult());
2636
2637 const commit = try GpuDialect.TmaCommitGroupOp.create(&ctx, loc);
2638 try testing.expectEqualStrings("gpu.tma.commit_group", commit.op.name.name);
2639
2640 const wait = try GpuDialect.TmaWaitGroupOp.create(&ctx, loc, 0);
2641 try testing.expectEqualStrings("gpu.tma.wait_group", wait.op.name.name);
2642 try testing.expectEqual(@as(i64, 0), wait.getCount().?);
2643 }
2644
2645 fn gpuEffects(
2646 comptime kind: effects.EventKind,
2647 comptime reads: []const usize,
2648 comptime writes: []const usize,
2649 ) interfaces.InterfaceEntry {
2650 const Declaration = struct {
2651 fn enumerate(op: *const ir.Operation, collector: *effects.Collector) void {
2652 var resource = effects.Resource{};
2653 if (op.getAttrAs(ir.Attribute.DialectAttr, "scope")) |scope| {
2654 if (scope.payload.len > 0) resource.ordering_scope = .{ .named = scope.payload };
2655 }
2656 if (kind == .state_observe) resource.state_key = "gpu.participants";
2657 collector.append(.{ .event = .{
2658 .kind = kind,
2659 .resource = resource,
2660 .ordered = true,
2661 } });
2662 collector.append(.{ .requirement = .{
2663 .kind = .execution_context,
2664 .subject = .operation,
2665 } });
2666 if (kind == .synchronize) collector.append(.{ .event = .{ .kind = .diverge } });
2667 for (reads) |index| {
2668 if (index >= op.getNumOperands()) continue;
2669 var access = resource;
2670 access.subject = .{ .operand = index };
2671 collector.append(.{ .event = .{
2672 .kind = .read,
2673 .resource = access,
2674 .ordered = true,
2675 } });
2676 }
2677 for (writes) |index| {
2678 if (index >= op.getNumOperands()) continue;
2679 var access = resource;
2680 access.subject = .{ .operand = index };
2681 collector.append(.{ .event = .{
2682 .kind = .write,
2683 .resource = access,
2684 .ordered = true,
2685 } });
2686 }
2687 for (0..op.getNumResults()) |index| {
2688 collector.append(.{ .result = .{ .index = index } });
2689 }
2690 }
2691 };
2692 return effects.EffectOpInterface.entryFor(.{
2693 .capacity = .{ .entries = 3 + reads.len + writes.len, .per_result = 1 },
2694 .enumerate = Declaration.enumerate,
2695 });
2696 }
2697
2698 test "gpu effect declarations retain participant observations and collective ordering" {
2699 var ctx = try ir.Context.init(std.testing.allocator, ir.Context.Limits.testing);
2700 defer ctx.deinit(std.testing.allocator);
2701 const mask = try GpuDialect.ActiveMaskOp.create(&ctx, .unknown);
2702 const barrier = try GpuDialect.BarrierOp.create(&ctx, .unknown, .block);
2703 var mask_facts = try effects.inspect(std.testing.allocator, mask.op);
2704 defer mask_facts.deinit(std.testing.allocator);
2705 try std.testing.expectEqual(
2706 effects.EventKind.state_observe,
2707 mask_facts.facts.records[0].event.kind,
2708 );
2709 try std.testing.expectEqualStrings(
2710 "gpu.participants",
2711 mask_facts.facts.records[0].event.resource.state_key.?,
2712 );
2713 try std.testing.expect(!effects.duplicate(mask_facts.facts, .{}));
2714 var barrier_facts = try effects.inspect(std.testing.allocator, barrier.op);
2715 defer barrier_facts.deinit(std.testing.allocator);
2716 const event = barrier_facts.facts.records[0].event;
2717 try std.testing.expectEqual(effects.EventKind.synchronize, event.kind);
2718 try std.testing.expect(event.ordered);
2719 try std.testing.expectEqualStrings("block", event.resource.ordering_scope.named);
2720 try std.testing.expect(!effects.discard(barrier_facts.facts));
2721 }