lib/choir/src/backends/gpu/nvptx/dialect.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("../../../root.zig");
3
4 const ir = choir.ir;
5 const arith = choir.dialects.arith;
6 const gpu = @import("../../../dialects/gpu/root.zig");
7
8 const Dimension = gpu.Dimension;
9 const MmaShape = gpu.MmaShape;
10 const ShuffleMode = gpu.ShuffleMode;
11 const WarpOpKind = gpu.WarpOpKind;
12
13 pub const NvptxDialect = struct {
14 pub const name = "nvptx";
15 const op_specs = ir.dialects.opSpec.dialect(@This());
16 pub const spec = ir.dialects.dialectSpec(@This(), .{
17 .dialect_attributes = &.{
18 "nvptx.dim", "nvptx.shuffle_mode", "nvptx.warp_op",
19 "nvptx.mma_shape", "nvptx.atomic_kind",
20 },
21 });
22
23 pub const ThreadIdxOp = struct {
24 op: *ir.Operation,
25 pub const operation_spec = op_specs.leaf(.{
26 .mnemonic = "tid",
27 .operands = 0,
28 .results = 1,
29 .required_attrs = &.{"dim"},
30 });
31 pub const operation_name = operation_spec.name;
32
33 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !ThreadIdxOp {
34 try loadSpec(ctx);
35 var builder = ir.OperationBuilder.init(ctx);
36 const index_type = try arith.ArithDialect.getIndexType(ctx);
37 var state = op_specs.state(@This(), loc);
38 state.addTypes(&.{index_type});
39
40 const op = try builder.create(state);
41 try setDimensionAttr(op, ctx, dim);
42 return .{ .op = op };
43 }
44
45 pub fn getResult(self: *const ThreadIdxOp) *ir.Value {
46 return self.op.getResult(0).?;
47 }
48
49 pub fn getDimension(self: ThreadIdxOp) ?Dimension {
50 return getDimensionAttr(self.op);
51 }
52 };
53
54 pub const BlockIdxOp = struct {
55 op: *ir.Operation,
56 pub const operation_spec = op_specs.leaf(.{
57 .mnemonic = "ctaid",
58 .operands = 0,
59 .results = 1,
60 .required_attrs = &.{"dim"},
61 });
62 pub const operation_name = operation_spec.name;
63
64 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !BlockIdxOp {
65 try loadSpec(ctx);
66 var builder = ir.OperationBuilder.init(ctx);
67 const index_type = try arith.ArithDialect.getIndexType(ctx);
68 var state = op_specs.state(@This(), loc);
69 state.addTypes(&.{index_type});
70
71 const op = try builder.create(state);
72 try setDimensionAttr(op, ctx, dim);
73 return .{ .op = op };
74 }
75
76 pub fn getResult(self: *const BlockIdxOp) *ir.Value {
77 return self.op.getResult(0).?;
78 }
79
80 pub fn getDimension(self: BlockIdxOp) ?Dimension {
81 return getDimensionAttr(self.op);
82 }
83 };
84
85 pub const BlockDimOp = struct {
86 op: *ir.Operation,
87 pub const operation_spec = op_specs.leaf(.{
88 .mnemonic = "ntid",
89 .operands = 0,
90 .results = 1,
91 .required_attrs = &.{"dim"},
92 });
93 pub const operation_name = operation_spec.name;
94
95 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !BlockDimOp {
96 try loadSpec(ctx);
97 var builder = ir.OperationBuilder.init(ctx);
98 const index_type = try arith.ArithDialect.getIndexType(ctx);
99 var state = op_specs.state(@This(), loc);
100 state.addTypes(&.{index_type});
101
102 const op = try builder.create(state);
103 try setDimensionAttr(op, ctx, dim);
104 return .{ .op = op };
105 }
106
107 pub fn getResult(self: *const BlockDimOp) *ir.Value {
108 return self.op.getResult(0).?;
109 }
110
111 pub fn getDimension(self: BlockDimOp) ?Dimension {
112 return getDimensionAttr(self.op);
113 }
114 };
115
116 pub const GridDimOp = struct {
117 op: *ir.Operation,
118 pub const operation_spec = op_specs.leaf(.{
119 .mnemonic = "nctaid",
120 .operands = 0,
121 .results = 1,
122 .required_attrs = &.{"dim"},
123 });
124 pub const operation_name = operation_spec.name;
125
126 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !GridDimOp {
127 try loadSpec(ctx);
128 var builder = ir.OperationBuilder.init(ctx);
129 const index_type = try arith.ArithDialect.getIndexType(ctx);
130 var state = op_specs.state(@This(), loc);
131 state.addTypes(&.{index_type});
132
133 const op = try builder.create(state);
134 try setDimensionAttr(op, ctx, dim);
135 return .{ .op = op };
136 }
137
138 pub fn getResult(self: *const GridDimOp) *ir.Value {
139 return self.op.getResult(0).?;
140 }
141
142 pub fn getDimension(self: GridDimOp) ?Dimension {
143 return getDimensionAttr(self.op);
144 }
145 };
146
147 pub const Barrier0Op = struct {
148 op: *ir.Operation,
149 pub const operation_spec = op_specs.leaf(.{
150 .mnemonic = "barrier0",
151 .operands = 0,
152 .results = 0,
153 });
154 pub const operation_name = operation_spec.name;
155
156 pub fn create(ctx: *ir.Context, loc: ir.Location) !Barrier0Op {
157 try loadSpec(ctx);
158 var builder = ir.OperationBuilder.init(ctx);
159 const state = op_specs.state(@This(), loc);
160 const op = try builder.create(state);
161 return .{ .op = op };
162 }
163 };
164
165 pub const WarpBarrierAllOp = struct {
166 op: *ir.Operation,
167 pub const operation_spec = op_specs.leaf(.{
168 .mnemonic = "bar.warp.sync.all",
169 .operands = 0,
170 .results = 0,
171 });
172 pub const operation_name = operation_spec.name;
173
174 pub fn create(ctx: *ir.Context, loc: ir.Location) !WarpBarrierAllOp {
175 try loadSpec(ctx);
176 var builder = ir.OperationBuilder.init(ctx);
177 const state = op_specs.state(@This(), loc);
178 const op = try builder.create(state);
179 return .{ .op = op };
180 }
181 };
182
183 pub const LaneIdOp = struct {
184 op: *ir.Operation,
185 pub const operation_spec = op_specs.leaf(.{
186 .mnemonic = "laneid",
187 .operands = 0,
188 .results = 1,
189 });
190 pub const operation_name = operation_spec.name;
191
192 pub fn create(ctx: *ir.Context, loc: ir.Location) !LaneIdOp {
193 try loadSpec(ctx);
194 var builder = ir.OperationBuilder.init(ctx);
195 const index_type = try arith.ArithDialect.getIndexType(ctx);
196 var state = op_specs.state(@This(), loc);
197 state.addTypes(&.{index_type});
198 const op = try builder.create(state);
199 return .{ .op = op };
200 }
201
202 pub fn getResult(self: *const LaneIdOp) *ir.Value {
203 return self.op.getResult(0).?;
204 }
205 };
206
207 pub const WarpIdOp = struct {
208 op: *ir.Operation,
209 pub const operation_spec = op_specs.leaf(.{
210 .mnemonic = "warpid",
211 .operands = 0,
212 .results = 1,
213 });
214 pub const operation_name = operation_spec.name;
215
216 pub fn create(ctx: *ir.Context, loc: ir.Location) !WarpIdOp {
217 try loadSpec(ctx);
218 var builder = ir.OperationBuilder.init(ctx);
219 const index_type = try arith.ArithDialect.getIndexType(ctx);
220 var state = op_specs.state(@This(), loc);
221 state.addTypes(&.{index_type});
222 const op = try builder.create(state);
223 return .{ .op = op };
224 }
225
226 pub fn getResult(self: *const WarpIdOp) *ir.Value {
227 return self.op.getResult(0).?;
228 }
229 };
230
231 pub const BarrierSyncOp = struct {
232 op: *ir.Operation,
233 pub const operation_spec = op_specs.leaf(.{
234 .mnemonic = "bar.sync",
235 .operands = 1,
236 .results = 0,
237 });
238 pub const operation_name = operation_spec.name;
239
240 pub fn create(ctx: *ir.Context, loc: ir.Location, barrier_id: *ir.Value) !BarrierSyncOp {
241 try loadSpec(ctx);
242 var builder = ir.OperationBuilder.init(ctx);
243 var state = op_specs.state(@This(), loc);
244 state.addOperands(&.{barrier_id});
245 const op = try builder.create(state);
246 return .{ .op = op };
247 }
248
249 pub fn getBarrierId(self: BarrierSyncOp) *ir.Value {
250 return self.op.operands.items[0].value;
251 }
252 };
253
254 pub const SyncWarpOp = struct {
255 op: *ir.Operation,
256 pub const operation_spec = op_specs.leaf(.{
257 .mnemonic = "bar.warp.sync",
258 .operands = 1,
259 .results = 0,
260 });
261 pub const operation_name = operation_spec.name;
262
263 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value) !SyncWarpOp {
264 try loadSpec(ctx);
265 var builder = ir.OperationBuilder.init(ctx);
266 var state = op_specs.state(@This(), loc);
267 state.addOperands(&.{mask});
268 const op = try builder.create(state);
269 return .{ .op = op };
270 }
271
272 pub fn getMask(self: SyncWarpOp) *ir.Value {
273 return self.op.operands.items[0].value;
274 }
275 };
276
277 pub const ActiveMaskOp = struct {
278 op: *ir.Operation,
279 pub const operation_spec = op_specs.leaf(.{
280 .mnemonic = "activemask",
281 .operands = 0,
282 .results = 1,
283 });
284 pub const operation_name = operation_spec.name;
285
286 pub fn create(ctx: *ir.Context, loc: ir.Location) !ActiveMaskOp {
287 try loadSpec(ctx);
288 var builder = ir.OperationBuilder.init(ctx);
289 const i32_type = try arith.ArithDialect.getI32Type(ctx);
290 var state = op_specs.state(@This(), loc);
291 state.addTypes(&.{i32_type});
292 const op = try builder.create(state);
293 return .{ .op = op };
294 }
295
296 pub fn getResult(self: *const ActiveMaskOp) *ir.Value {
297 return self.op.getResult(0).?;
298 }
299 };
300
301 pub const AllSyncOp = struct {
302 op: *ir.Operation,
303 pub const operation_spec = op_specs.leaf(.{
304 .mnemonic = "vote.sync.all",
305 .operands = .{ "mask", "predicate" },
306 .results = .{"result"},
307 });
308 pub const operation_name = operation_spec.name;
309
310 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AllSyncOp {
311 try loadSpec(ctx);
312 var builder = ir.OperationBuilder.init(ctx);
313 const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);
314 var state = op_specs.state(@This(), loc);
315 state.addOperands(&.{ mask, pred });
316 state.addTypes(&.{bool_type});
317 const op = try builder.create(state);
318 return .{ .op = op };
319 }
320
321 pub fn getResult(self: *const AllSyncOp) *ir.Value {
322 return self.op.getResult(0).?;
323 }
324
325 pub fn getMask(self: AllSyncOp) *ir.Value {
326 return self.op.operands.items[0].value;
327 }
328
329 pub fn getPredicate(self: AllSyncOp) *ir.Value {
330 return self.op.operands.items[1].value;
331 }
332 };
333
334 pub const AnySyncOp = struct {
335 op: *ir.Operation,
336 pub const operation_spec = op_specs.leaf(.{
337 .mnemonic = "vote.sync.any",
338 .operands = .{ "mask", "predicate" },
339 .results = .{"result"},
340 });
341 pub const operation_name = operation_spec.name;
342
343 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AnySyncOp {
344 try loadSpec(ctx);
345 var builder = ir.OperationBuilder.init(ctx);
346 const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);
347 var state = op_specs.state(@This(), loc);
348 state.addOperands(&.{ mask, pred });
349 state.addTypes(&.{bool_type});
350 const op = try builder.create(state);
351 return .{ .op = op };
352 }
353
354 pub fn getResult(self: *const AnySyncOp) *ir.Value {
355 return self.op.getResult(0).?;
356 }
357
358 pub fn getMask(self: AnySyncOp) *ir.Value {
359 return self.op.operands.items[0].value;
360 }
361
362 pub fn getPredicate(self: AnySyncOp) *ir.Value {
363 return self.op.operands.items[1].value;
364 }
365 };
366
367 pub const BallotSyncOp = struct {
368 op: *ir.Operation,
369 pub const operation_spec = op_specs.leaf(.{
370 .mnemonic = "vote.sync.ballot",
371 .operands = .{ "mask", "predicate" },
372 .results = .{"result"},
373 });
374 pub const operation_name = operation_spec.name;
375
376 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !BallotSyncOp {
377 try loadSpec(ctx);
378 var builder = ir.OperationBuilder.init(ctx);
379 const i32_type = try arith.ArithDialect.getI32Type(ctx);
380 var state = op_specs.state(@This(), loc);
381 state.addOperands(&.{ mask, pred });
382 state.addTypes(&.{i32_type});
383 const op = try builder.create(state);
384 return .{ .op = op };
385 }
386
387 pub fn getResult(self: *const BallotSyncOp) *ir.Value {
388 return self.op.getResult(0).?;
389 }
390
391 pub fn getMask(self: BallotSyncOp) *ir.Value {
392 return self.op.operands.items[0].value;
393 }
394
395 pub fn getPredicate(self: BallotSyncOp) *ir.Value {
396 return self.op.operands.items[1].value;
397 }
398 };
399
400 pub const ShflSyncOp = struct {
401 op: *ir.Operation,
402 pub const operation_spec = op_specs.leaf(.{
403 .mnemonic = "shfl.sync",
404 .operands = .{ "mask", "src", "lane_or_delta" },
405 .results = .{"result"},
406 .required_attrs = &.{"mode"},
407 });
408 pub const operation_name = operation_spec.name;
409
410 pub fn create(
411 ctx: *ir.Context,
412 loc: ir.Location,
413 mode: ShuffleMode,
414 mask: *ir.Value,
415 src: *ir.Value,
416 lane_or_delta: *ir.Value,
417 ) !ShflSyncOp {
418 try loadSpec(ctx);
419 var builder = ir.OperationBuilder.init(ctx);
420 var state = op_specs.state(@This(), loc);
421 state.addOperands(&.{ mask, src, lane_or_delta });
422 state.addTypes(&.{src.type});
423 const op = try builder.create(state);
424 try setShuffleModeAttr(op, ctx, mode);
425 return .{ .op = op };
426 }
427
428 pub fn getResult(self: *const ShflSyncOp) *ir.Value {
429 return self.op.getResult(0).?;
430 }
431
432 pub fn getMask(self: ShflSyncOp) *ir.Value {
433 return self.op.operands.items[0].value;
434 }
435
436 pub fn getSrc(self: ShflSyncOp) *ir.Value {
437 return self.op.operands.items[1].value;
438 }
439
440 pub fn getLaneOrDelta(self: ShflSyncOp) *ir.Value {
441 return self.op.operands.items[2].value;
442 }
443
444 pub fn getMode(self: ShflSyncOp) ?ShuffleMode {
445 return getShuffleModeAttr(self.op);
446 }
447 };
448
449 pub const WarpReduceOp = struct {
450 op: *ir.Operation,
451 pub const operation_spec = op_specs.leaf(.{
452 .mnemonic = "warp.reduce",
453 .operands = .{ "mask", "value" },
454 .results = .{"result"},
455 .required_attrs = &.{"op"},
456 });
457 pub const operation_name = operation_spec.name;
458
459 pub fn create(
460 ctx: *ir.Context,
461 loc: ir.Location,
462 op_kind: WarpOpKind,
463 mask: *ir.Value,
464 value: *ir.Value,
465 ) !WarpReduceOp {
466 try loadSpec(ctx);
467 var builder = ir.OperationBuilder.init(ctx);
468 var state = op_specs.state(@This(), loc);
469 state.addOperands(&.{ mask, value });
470 state.addTypes(&.{value.type});
471 const op = try builder.create(state);
472 try setWarpOpAttr(op, ctx, op_kind);
473 return .{ .op = op };
474 }
475
476 pub fn getResult(self: *const WarpReduceOp) *ir.Value {
477 return self.op.getResult(0).?;
478 }
479
480 pub fn getMask(self: WarpReduceOp) *ir.Value {
481 return self.op.operands.items[0].value;
482 }
483
484 pub fn getValue(self: WarpReduceOp) *ir.Value {
485 return self.op.operands.items[1].value;
486 }
487
488 pub fn getOpKind(self: WarpReduceOp) ?WarpOpKind {
489 return getWarpOpAttr(self.op);
490 }
491 };
492
493 pub const MmaSyncOp = struct {
494 op: *ir.Operation,
495 pub const operation_spec = op_specs.leaf(.{
496 .mnemonic = "mma.sync",
497 .operands = .{ "a0", "a1", "a2", "a3", "b0", "b1", "c0", "c1", "c2", "c3" },
498 .results = .{ "d0", "d1", "d2", "d3" },
499 .required_attrs = &.{"shape"},
500 });
501 pub const operation_name = operation_spec.name;
502
503 pub fn create(
504 ctx: *ir.Context,
505 loc: ir.Location,
506 operands: [10]*ir.Value,
507 shape: MmaShape,
508 ) !MmaSyncOp {
509 try loadSpec(ctx);
510 var builder = ir.OperationBuilder.init(ctx);
511 var state = op_specs.state(@This(), loc);
512 state.addOperands(&operands);
513 state.addTypes(&.{ operands[6].type, operands[7].type, operands[8].type, operands[9].type });
514 const op = try builder.create(state);
515 try setMmaShapeAttr(op, ctx, shape);
516 return .{ .op = op };
517 }
518
519 pub fn getOperandValue(self: MmaSyncOp, index: usize) *ir.Value {
520 return self.op.operands.items[index].value;
521 }
522
523 pub fn getD(self: *const MmaSyncOp, index: usize) *ir.Value {
524 return self.op.getResult(index).?;
525 }
526
527 pub fn getShape(self: MmaSyncOp) ?MmaShape {
528 return getMmaShapeAttr(self.op);
529 }
530 };
531
532 pub const CpAsyncSharedOp = struct {
533 op: *ir.Operation,
534 pub const operation_spec = op_specs.leaf(.{
535 .mnemonic = "cp.async.shared",
536 .operands = .{ "dst", "dst_index", "src", "src_index" },
537 .results = 0,
538 .required_attrs = &.{"bytes"},
539 });
540 pub const operation_name = operation_spec.name;
541
542 pub fn create(
543 ctx: *ir.Context,
544 loc: ir.Location,
545 dst: *ir.Value,
546 dst_index: *ir.Value,
547 src: *ir.Value,
548 src_index: *ir.Value,
549 bytes: u32,
550 ) !CpAsyncSharedOp {
551 try loadSpec(ctx);
552 var builder = ir.OperationBuilder.init(ctx);
553 var state = op_specs.state(@This(), loc);
554 state.addOperands(&.{ dst, dst_index, src, src_index });
555 const op = try builder.create(state);
556 const bytes_attr = try ctx.getI64Attr(@intCast(bytes));
557 try op.setAttr("bytes", bytes_attr);
558 return .{ .op = op };
559 }
560
561 pub fn getDst(self: CpAsyncSharedOp) *ir.Value {
562 return ir.dialects.operand(operation_spec, self.op, "dst");
563 }
564
565 pub fn getDstIndex(self: CpAsyncSharedOp) *ir.Value {
566 return ir.dialects.operand(operation_spec, self.op, "dst_index");
567 }
568
569 pub fn getSrc(self: CpAsyncSharedOp) *ir.Value {
570 return ir.dialects.operand(operation_spec, self.op, "src");
571 }
572
573 pub fn getSrcIndex(self: CpAsyncSharedOp) *ir.Value {
574 return ir.dialects.operand(operation_spec, self.op, "src_index");
575 }
576
577 pub fn getBytes(self: CpAsyncSharedOp) ?u32 {
578 const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "bytes") orelse return null;
579 const raw = int_attr.getUnsignedValue();
580 if (raw > std.math.maxInt(u32)) return null;
581 return @intCast(raw);
582 }
583 };
584
585 pub const FenceDeviceOp = struct {
586 op: *ir.Operation,
587 pub const operation_spec = op_specs.leaf(.{
588 .mnemonic = "membar.gl",
589 .operands = 0,
590 .results = 0,
591 });
592 pub const operation_name = operation_spec.name;
593
594 pub fn create(ctx: *ir.Context, loc: ir.Location) !FenceDeviceOp {
595 try loadSpec(ctx);
596 var builder = ir.OperationBuilder.init(ctx);
597 const state = op_specs.state(@This(), loc);
598 const op = try builder.create(state);
599 return .{ .op = op };
600 }
601 };
602
603 pub const CpAsyncCommitOp = struct {
604 op: *ir.Operation,
605 pub const operation_spec = op_specs.leaf(.{
606 .mnemonic = "cp.async.commit",
607 .operands = 0,
608 .results = 0,
609 });
610 pub const operation_name = operation_spec.name;
611
612 pub fn create(ctx: *ir.Context, loc: ir.Location) !CpAsyncCommitOp {
613 try loadSpec(ctx);
614 var builder = ir.OperationBuilder.init(ctx);
615 const state = op_specs.state(@This(), loc);
616 const op = try builder.create(state);
617 return .{ .op = op };
618 }
619 };
620
621 pub const CpAsyncWaitOp = struct {
622 op: *ir.Operation,
623 pub const operation_spec = op_specs.leaf(.{
624 .mnemonic = "cp.async.wait",
625 .operands = 0,
626 .results = 0,
627 .required_attrs = &.{"groups"},
628 });
629 pub const operation_name = operation_spec.name;
630
631 pub fn create(ctx: *ir.Context, loc: ir.Location, groups: u32) !CpAsyncWaitOp {
632 try loadSpec(ctx);
633 var builder = ir.OperationBuilder.init(ctx);
634 const state = op_specs.state(@This(), loc);
635 const op = try builder.create(state);
636 const groups_attr = try ctx.getI64Attr(@intCast(groups));
637 try op.setAttr("groups", groups_attr);
638 return .{ .op = op };
639 }
640
641 pub fn getGroups(self: CpAsyncWaitOp) ?u32 {
642 const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "groups") orelse return null;
643 const raw = int_attr.getUnsignedValue();
644 if (raw > std.math.maxInt(u32)) return null;
645 return @intCast(raw);
646 }
647 };
648
649 pub const WarpScanOp = struct {
650 op: *ir.Operation,
651 pub const operation_spec = op_specs.leaf(.{
652 .mnemonic = "warp.scan",
653 .operands = .{ "mask", "value" },
654 .results = .{"result"},
655 .required_attrs = &.{ "inclusive", "op" },
656 });
657 pub const operation_name = operation_spec.name;
658
659 pub fn create(
660 ctx: *ir.Context,
661 loc: ir.Location,
662 op_kind: WarpOpKind,
663 inclusive: bool,
664 mask: *ir.Value,
665 value: *ir.Value,
666 ) !WarpScanOp {
667 try loadSpec(ctx);
668 var builder = ir.OperationBuilder.init(ctx);
669 var state = op_specs.state(@This(), loc);
670 state.addOperands(&.{ mask, value });
671 state.addTypes(&.{value.type});
672 const op = try builder.create(state);
673 try setWarpOpAttr(op, ctx, op_kind);
674 try setBoolAttr(op, ctx, "inclusive", inclusive);
675 return .{ .op = op };
676 }
677
678 pub fn getResult(self: *const WarpScanOp) *ir.Value {
679 return self.op.getResult(0).?;
680 }
681
682 pub fn getMask(self: WarpScanOp) *ir.Value {
683 return self.op.operands.items[0].value;
684 }
685
686 pub fn getValue(self: WarpScanOp) *ir.Value {
687 return self.op.operands.items[1].value;
688 }
689
690 pub fn getOpKind(self: WarpScanOp) ?WarpOpKind {
691 return getWarpOpAttr(self.op);
692 }
693
694 pub fn isInclusive(self: WarpScanOp) bool {
695 return getBoolAttrValue(self.op, "inclusive");
696 }
697 };
698
699 pub const LoadLocalOp: type = loadOp("local");
700 pub const StoreLocalOp: type = storeOp("local");
701
702 pub const LoadGlobalOp: type = loadOp("global");
703
704 pub const StoreGlobalOp: type = storeOp("global");
705
706 pub const AtomicGlobalOp = struct {
707 op: *ir.Operation,
708 pub const operation_spec = op_specs.leaf(.{
709 .mnemonic = "atom.global",
710 .operands = 3,
711 .results = 1,
712 .attrs = &.{"kind"},
713 });
714 pub const operation_name = operation_spec.name;
715
716 pub fn create(
717 ctx: *ir.Context,
718 loc: ir.Location,
719 kind: choir.dialects.AtomicRmwKind,
720 value: *ir.Value,
721 memref: *ir.Value,
722 index: *ir.Value,
723 result_type: ir.Type,
724 ) !AtomicGlobalOp {
725 try loadSpec(ctx);
726 var builder = ir.OperationBuilder.init(ctx);
727 var state = op_specs.state(@This(), loc);
728 state.addOperands(&.{ value, memref, index });
729 state.addTypes(&.{result_type});
730 const op = try builder.create(state);
731 const kind_attr = try ctx.getDialectAttr("nvptx.atomic_kind", kind.toString());
732 try op.setAttr("kind", kind_attr);
733 return .{ .op = op };
734 }
735
736 pub fn getKind(self: AtomicGlobalOp) ?choir.dialects.AtomicRmwKind {
737 const dialect_attr = self.op.getAttrAs(ir.Attribute.DialectAttr, "kind") orelse return null;
738 return choir.dialects.AtomicRmwKind.fromString(dialect_attr.payload);
739 }
740
741 pub fn getValue(self: AtomicGlobalOp) *ir.Value {
742 return self.op.operands.items[0].value;
743 }
744
745 pub fn getMemref(self: AtomicGlobalOp) *ir.Value {
746 return self.op.operands.items[1].value;
747 }
748
749 pub fn getIndex(self: AtomicGlobalOp) *ir.Value {
750 return self.op.operands.items[2].value;
751 }
752
753 pub fn getResult(self: *const AtomicGlobalOp) *ir.Value {
754 return self.op.getResult(0).?;
755 }
756 };
757
758 pub const AtomicSharedOp = struct {
759 op: *ir.Operation,
760 pub const operation_spec = op_specs.leaf(.{
761 .mnemonic = "atom.shared",
762 .operands = 3,
763 .results = 1,
764 .attrs = &.{"kind"},
765 });
766 pub const operation_name = operation_spec.name;
767
768 pub fn create(
769 ctx: *ir.Context,
770 loc: ir.Location,
771 kind: choir.dialects.AtomicRmwKind,
772 value: *ir.Value,
773 memref: *ir.Value,
774 index: *ir.Value,
775 result_type: ir.Type,
776 ) !AtomicSharedOp {
777 try loadSpec(ctx);
778 var builder = ir.OperationBuilder.init(ctx);
779 var state = op_specs.state(@This(), loc);
780 state.addOperands(&.{ value, memref, index });
781 state.addTypes(&.{result_type});
782 const op = try builder.create(state);
783 const kind_attr = try ctx.getDialectAttr("nvptx.atomic_kind", kind.toString());
784 try op.setAttr("kind", kind_attr);
785 return .{ .op = op };
786 }
787
788 pub fn getKind(self: AtomicSharedOp) ?choir.dialects.AtomicRmwKind {
789 const dialect_attr = self.op.getAttrAs(ir.Attribute.DialectAttr, "kind") orelse return null;
790 return choir.dialects.AtomicRmwKind.fromString(dialect_attr.payload);
791 }
792
793 pub fn getValue(self: AtomicSharedOp) *ir.Value {
794 return self.op.operands.items[0].value;
795 }
796
797 pub fn getMemref(self: AtomicSharedOp) *ir.Value {
798 return self.op.operands.items[1].value;
799 }
800
801 pub fn getIndex(self: AtomicSharedOp) *ir.Value {
802 return self.op.operands.items[2].value;
803 }
804
805 pub fn getResult(self: *const AtomicSharedOp) *ir.Value {
806 return self.op.getResult(0).?;
807 }
808 };
809
810 pub const AtomicCasGlobalOp = struct {
811 op: *ir.Operation,
812 pub const operation_spec = op_specs.leaf(.{
813 .mnemonic = "atom.global.cas",
814 .operands = 4,
815 .results = 1,
816 });
817 pub const operation_name = operation_spec.name;
818
819 pub fn create(
820 ctx: *ir.Context,
821 loc: ir.Location,
822 expected: *ir.Value,
823 desired: *ir.Value,
824 memref: *ir.Value,
825 index: *ir.Value,
826 result_type: ir.Type,
827 ) !AtomicCasGlobalOp {
828 try loadSpec(ctx);
829 var builder = ir.OperationBuilder.init(ctx);
830 var state = op_specs.state(@This(), loc);
831 state.addOperands(&.{ expected, desired, memref, index });
832 state.addTypes(&.{result_type});
833 const op = try builder.create(state);
834 return .{ .op = op };
835 }
836
837 pub fn getExpected(self: AtomicCasGlobalOp) *ir.Value {
838 return self.op.operands.items[0].value;
839 }
840
841 pub fn getDesired(self: AtomicCasGlobalOp) *ir.Value {
842 return self.op.operands.items[1].value;
843 }
844
845 pub fn getMemref(self: AtomicCasGlobalOp) *ir.Value {
846 return self.op.operands.items[2].value;
847 }
848
849 pub fn getIndex(self: AtomicCasGlobalOp) *ir.Value {
850 return self.op.operands.items[3].value;
851 }
852
853 pub fn getResult(self: *const AtomicCasGlobalOp) *ir.Value {
854 return self.op.getResult(0).?;
855 }
856 };
857
858 pub const AtomicCasSharedOp = struct {
859 op: *ir.Operation,
860 pub const operation_spec = op_specs.leaf(.{
861 .mnemonic = "atom.shared.cas",
862 .operands = 4,
863 .results = 1,
864 });
865 pub const operation_name = operation_spec.name;
866
867 pub fn create(
868 ctx: *ir.Context,
869 loc: ir.Location,
870 expected: *ir.Value,
871 desired: *ir.Value,
872 memref: *ir.Value,
873 index: *ir.Value,
874 result_type: ir.Type,
875 ) !AtomicCasSharedOp {
876 try loadSpec(ctx);
877 var builder = ir.OperationBuilder.init(ctx);
878 var state = op_specs.state(@This(), loc);
879 state.addOperands(&.{ expected, desired, memref, index });
880 state.addTypes(&.{result_type});
881 const op = try builder.create(state);
882 return .{ .op = op };
883 }
884
885 pub fn getExpected(self: AtomicCasSharedOp) *ir.Value {
886 return self.op.operands.items[0].value;
887 }
888
889 pub fn getDesired(self: AtomicCasSharedOp) *ir.Value {
890 return self.op.operands.items[1].value;
891 }
892
893 pub fn getMemref(self: AtomicCasSharedOp) *ir.Value {
894 return self.op.operands.items[2].value;
895 }
896
897 pub fn getIndex(self: AtomicCasSharedOp) *ir.Value {
898 return self.op.operands.items[3].value;
899 }
900
901 pub fn getResult(self: *const AtomicCasSharedOp) *ir.Value {
902 return self.op.getResult(0).?;
903 }
904 };
905
906 pub const LoadSharedOp: type = loadOp("shared");
907
908 pub const StoreSharedOp: type = storeOp("shared");
909
910 pub const SinApproxF32Op = struct {
911 op: *ir.Operation,
912 pub const operation_spec = op_specs.leaf(.{
913 .mnemonic = "sin.approx.f32",
914 .operands = 1,
915 .results = 1,
916 });
917 pub const operation_name = operation_spec.name;
918
919 pub fn create(ctx: *ir.Context, loc: ir.Location, value: *ir.Value) !SinApproxF32Op {
920 try loadSpec(ctx);
921 var builder = ir.OperationBuilder.init(ctx);
922 var state = op_specs.state(@This(), loc);
923 state.addOperands(&.{value});
924 state.addTypes(&.{value.type});
925 const op = try builder.create(state);
926 return .{ .op = op };
927 }
928
929 pub fn getValue(self: SinApproxF32Op) *ir.Value {
930 return self.op.operands.items[0].value;
931 }
932
933 pub fn getResult(self: *const SinApproxF32Op) *ir.Value {
934 return self.op.getResult(0).?;
935 }
936 };
937
938 pub const CosApproxF32Op = struct {
939 op: *ir.Operation,
940 pub const operation_spec = op_specs.leaf(.{
941 .mnemonic = "cos.approx.f32",
942 .operands = 1,
943 .results = 1,
944 });
945 pub const operation_name = operation_spec.name;
946
947 pub fn create(ctx: *ir.Context, loc: ir.Location, value: *ir.Value) !CosApproxF32Op {
948 try loadSpec(ctx);
949 var builder = ir.OperationBuilder.init(ctx);
950 var state = op_specs.state(@This(), loc);
951 state.addOperands(&.{value});
952 state.addTypes(&.{value.type});
953 const op = try builder.create(state);
954 return .{ .op = op };
955 }
956
957 pub fn getValue(self: CosApproxF32Op) *ir.Value {
958 return self.op.operands.items[0].value;
959 }
960
961 pub fn getResult(self: *const CosApproxF32Op) *ir.Value {
962 return self.op.getResult(0).?;
963 }
964 };
965
966 pub const FmaRnF32Op = struct {
967 op: *ir.Operation,
968 pub const operation_spec = op_specs.leaf(.{
969 .mnemonic = "fma.rn.f32",
970 .operands = 3,
971 .results = 1,
972 });
973 pub const operation_name = operation_spec.name;
974
975 pub fn create(ctx: *ir.Context, loc: ir.Location, a: *ir.Value, b: *ir.Value, c: *ir.Value) !FmaRnF32Op {
976 try loadSpec(ctx);
977 var builder = ir.OperationBuilder.init(ctx);
978 var state = op_specs.state(@This(), loc);
979 state.addOperands(&.{ a, b, c });
980 state.addTypes(&.{a.type});
981 const op = try builder.create(state);
982 return .{ .op = op };
983 }
984
985 pub fn getA(self: FmaRnF32Op) *ir.Value {
986 return self.op.operands.items[0].value;
987 }
988
989 pub fn getB(self: FmaRnF32Op) *ir.Value {
990 return self.op.operands.items[1].value;
991 }
992
993 pub fn getC(self: FmaRnF32Op) *ir.Value {
994 return self.op.operands.items[2].value;
995 }
996
997 pub fn getResult(self: *const FmaRnF32Op) *ir.Value {
998 return self.op.getResult(0).?;
999 }
1000 };
1001
1002 fn loadOp(comptime space: []const u8) type {
1003 return struct {
1004 op: *ir.Operation,
1005 pub const operation_spec = op_specs.leaf(.{
1006 .mnemonic = "ld." ++ space,
1007 .operands = 2,
1008 .results = 1,
1009 });
1010 pub const operation_name = operation_spec.name;
1011
1012 pub fn create(
1013 ctx: *ir.Context,
1014 loc: ir.Location,
1015 memref: *ir.Value,
1016 index: *ir.Value,
1017 result_type: ir.Type,
1018 ) !@This() {
1019 try loadSpec(ctx);
1020 var builder = ir.OperationBuilder.init(ctx);
1021 var state = op_specs.state(@This(), loc);
1022 state.addOperands(&.{ memref, index });
1023 state.addTypes(&.{result_type});
1024 const op = try builder.create(state);
1025 return .{ .op = op };
1026 }
1027
1028 pub fn getResult(self: *const @This()) *ir.Value {
1029 return self.op.getResult(0).?;
1030 }
1031
1032 pub fn getMemref(self: @This()) *ir.Value {
1033 return self.op.operands.items[0].value;
1034 }
1035
1036 pub fn getIndex(self: @This()) *ir.Value {
1037 return self.op.operands.items[1].value;
1038 }
1039 };
1040 }
1041
1042 fn storeOp(comptime space: []const u8) type {
1043 return struct {
1044 op: *ir.Operation,
1045 pub const operation_spec = op_specs.leaf(.{
1046 .mnemonic = "st." ++ space,
1047 .operands = 3,
1048 .results = 0,
1049 });
1050 pub const operation_name = operation_spec.name;
1051
1052 pub fn create(
1053 ctx: *ir.Context,
1054 loc: ir.Location,
1055 value: *ir.Value,
1056 memref: *ir.Value,
1057 index: *ir.Value,
1058 ) !@This() {
1059 try loadSpec(ctx);
1060 var builder = ir.OperationBuilder.init(ctx);
1061 var state = op_specs.state(@This(), loc);
1062 state.addOperands(&.{ value, memref, index });
1063 const op = try builder.create(state);
1064 return .{ .op = op };
1065 }
1066
1067 pub fn getValue(self: @This()) *ir.Value {
1068 return self.op.operands.items[0].value;
1069 }
1070
1071 pub fn getMemref(self: @This()) *ir.Value {
1072 return self.op.operands.items[1].value;
1073 }
1074
1075 pub fn getIndex(self: @This()) *ir.Value {
1076 return self.op.operands.items[2].value;
1077 }
1078 };
1079 }
1080
1081 fn loadSpec(ctx: *ir.Context) !void {
1082 try ir.dialects.loadDialectSpec(ctx, spec);
1083 }
1084
1085 fn setDimensionAttr(op: *ir.Operation, ctx: *ir.Context, dim: Dimension) !void {
1086 const attr = try ctx.getDialectAttr("nvptx.dim", dim.toString());
1087 try op.setAttr("dim", attr);
1088 }
1089
1090 fn getDimensionAttr(op: *const ir.Operation) ?Dimension {
1091 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "dim") orelse return null;
1092 return Dimension.fromString(dialect_attr.payload);
1093 }
1094
1095 fn setShuffleModeAttr(op: *ir.Operation, ctx: *ir.Context, mode: ShuffleMode) !void {
1096 const attr = try ctx.getDialectAttr("nvptx.shuffle_mode", mode.toString());
1097 try op.setAttr("mode", attr);
1098 }
1099
1100 fn getShuffleModeAttr(op: *const ir.Operation) ?ShuffleMode {
1101 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "mode") orelse return null;
1102 return ShuffleMode.fromString(dialect_attr.payload);
1103 }
1104
1105 fn setWarpOpAttr(op: *ir.Operation, ctx: *ir.Context, op_kind: WarpOpKind) !void {
1106 const attr = try ctx.getDialectAttr("nvptx.warp_op", op_kind.toString());
1107 try op.setAttr("op", attr);
1108 }
1109
1110 fn getWarpOpAttr(op: *const ir.Operation) ?WarpOpKind {
1111 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "op") orelse return null;
1112 return WarpOpKind.fromString(dialect_attr.payload);
1113 }
1114
1115 fn setMmaShapeAttr(op: *ir.Operation, ctx: *ir.Context, shape: MmaShape) !void {
1116 var buf: [32]u8 = undefined;
1117 const shape_str = try shape.toString(buf[0..]);
1118 const shape_attr = try ctx.getDialectAttr("nvptx.mma_shape", shape_str);
1119 try op.setAttr("shape", shape_attr);
1120 }
1121
1122 fn getMmaShapeAttr(op: *const ir.Operation) ?MmaShape {
1123 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "shape") orelse return null;
1124 return MmaShape.parse(dialect_attr.payload);
1125 }
1126
1127 fn setBoolAttr(op: *ir.Operation, ctx: *ir.Context, attr_name: []const u8, value: bool) !void {
1128 const bool_attr = try ctx.getBoolAttr(value);
1129 try op.setAttr(attr_name, bool_attr);
1130 }
1131
1132 fn getBoolAttrValue(op: *const ir.Operation, attr_name: []const u8) bool {
1133 const bool_attr = op.getAttrAs(ir.Attribute.BoolAttr, attr_name) orelse return false;
1134 return bool_attr.getValue();
1135 }
1136 };
1137
1138 test {
1139 std.testing.refAllDecls(@This());
1140 }
1141
1142 test "NvptxDialect operation specs register shapes and attributes" {
1143 const testing = std.testing;
1144
1145 var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);
1146 defer ctx.deinit(testing.allocator);
1147
1148 try ir.dialects.loadDialectSpec(&ctx, NvptxDialect.spec);
1149
1150 const tid_info = ctx.lookupOperation(NvptxDialect.ThreadIdxOp.operation_name) orelse return error.TestExpectedOperation;
1151 try testing.expect(tid_info.shape.operands.allows(0));
1152 try testing.expect(!tid_info.shape.operands.allows(1));
1153 try testing.expect(tid_info.shape.results.allows(1));
1154 try testing.expect(!tid_info.shape.results.allows(0));
1155 try testing.expect(tid_info.shape.regions.allows(0));
1156 try testing.expect(!tid_info.shape.regions.allows(1));
1157 try testing.expect(tid_info.shape.successors.allows(0));
1158 try testing.expect(!tid_info.shape.successors.allows(1));
1159 try testing.expect(tid_info.hasInherentAttributeName("dim"));
1160
1161 const load_info = ctx.lookupOperation(NvptxDialect.LoadGlobalOp.operation_name) orelse return error.TestExpectedOperation;
1162 try testing.expect(load_info.shape.operands.allows(2));
1163 try testing.expect(!load_info.shape.operands.allows(1));
1164 try testing.expect(load_info.shape.results.allows(1));
1165 try testing.expect(!load_info.shape.results.allows(0));
1166
1167 const store_info = ctx.lookupOperation(NvptxDialect.StoreGlobalOp.operation_name) orelse return error.TestExpectedOperation;
1168 try testing.expect(store_info.shape.operands.allows(3));
1169 try testing.expect(!store_info.shape.operands.allows(2));
1170 try testing.expect(store_info.shape.results.allows(0));
1171 try testing.expect(!store_info.shape.results.allows(1));
1172
1173 const cas_info = ctx.lookupOperation(NvptxDialect.AtomicCasGlobalOp.operation_name) orelse return error.TestExpectedOperation;
1174 try testing.expect(cas_info.shape.operands.allows(4));
1175 try testing.expect(!cas_info.shape.operands.allows(3));
1176 try testing.expect(cas_info.shape.results.allows(1));
1177 try testing.expect(!cas_info.shape.results.allows(0));
1178 }
1179
1180 test "NvptxDialect op-spec state supports strict contexts and verifier shape checks" {
1181 const testing = std.testing;
1182
1183 var ctx = try ir.Context.init(testing.allocator, ir.Context.Limits.testing);
1184 defer ctx.deinit(testing.allocator);
1185
1186 try ctx.requireRegistered();
1187 try ir.dialects.loadDialectSpec(&ctx, choir.dialects.arith.spec);
1188 try ir.dialects.loadDialectSpec(&ctx, NvptxDialect.spec);
1189
1190 const loc = ir.Location.getUnknown();
1191 const tid = try NvptxDialect.ThreadIdxOp.create(&ctx, loc, .x);
1192 try testing.expectEqual(Dimension.x, tid.getDimension().?);
1193 try ir.verifyOperation(tid.op, .{ .recursive = false });
1194
1195 const index_type = try choir.dialects.arith.ArithDialect.getIndexType(&ctx);
1196 var bad_load_state = ir.Operation.State.init(NvptxDialect.LoadGlobalOp.operation_name, loc);
1197 bad_load_state.addTypes(&.{index_type});
1198 const bad_load = try ctx.createOperation(bad_load_state);
1199 try testing.expectError(error.OperandCountMismatch, ir.verifyOperation(bad_load, .{ .recursive = false }));
1200 }