lib/choir/src/backends/gpu/spirv/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 Scope = gpu.Scope;
10 const ShuffleMode = gpu.ShuffleMode;
11 const WarpOpKind = gpu.WarpOpKind;
12
13 pub const Capability = enum {
14 shader,
15
16 pub fn toString(self: Capability) []const u8 {
17 return @tagName(self);
18 }
19
20 pub fn fromString(s: []const u8) ?Capability {
21 if (std.mem.eql(u8, s, "shader")) return .shader;
22 return null;
23 }
24 };
25
26 pub const AddressingModel = enum {
27 logical,
28
29 pub fn toString(self: AddressingModel) []const u8 {
30 return @tagName(self);
31 }
32
33 pub fn fromString(s: []const u8) ?AddressingModel {
34 if (std.mem.eql(u8, s, "logical")) return .logical;
35 return null;
36 }
37 };
38
39 pub const MemoryModel = enum {
40 glsl450,
41
42 pub fn toString(self: MemoryModel) []const u8 {
43 return @tagName(self);
44 }
45
46 pub fn fromString(s: []const u8) ?MemoryModel {
47 if (std.mem.eql(u8, s, "glsl450")) return .glsl450;
48 return null;
49 }
50 };
51
52 pub const ExecutionModel = enum {
53 vertex,
54 fragment,
55 gl_compute,
56
57 pub fn toString(self: ExecutionModel) []const u8 {
58 return @tagName(self);
59 }
60
61 pub fn fromString(s: []const u8) ?ExecutionModel {
62 if (std.mem.eql(u8, s, "vertex")) return .vertex;
63 if (std.mem.eql(u8, s, "fragment")) return .fragment;
64 if (std.mem.eql(u8, s, "gl_compute")) return .gl_compute;
65 return null;
66 }
67 };
68
69 pub const StorageClass = enum {
70 function,
71 private,
72 workgroup,
73 uniform,
74 storage_buffer,
75 input,
76 output,
77
78 pub fn toString(self: StorageClass) []const u8 {
79 return @tagName(self);
80 }
81
82 pub fn fromString(s: []const u8) ?StorageClass {
83 inline for (
84 @typeInfo(StorageClass).@"enum".field_names,
85 @typeInfo(StorageClass).@"enum".field_values,
86 ) |field_name, field_name_value| {
87 const field = .{ .name = field_name, .value = field_name_value };
88 if (std.mem.eql(u8, s, field.name)) {
89 return @fromBackingInt(@intCast(field.value));
90 }
91 }
92 return null;
93 }
94 };
95
96 pub const SpirvDialect = struct {
97 pub const name = "spirv";
98 const symbol_table_trait = ir.dialects.trait(ir.traits.SymbolTable);
99 const op_specs = ir.dialects.opSpec.dialect(@This());
100 pub const spec = ir.dialects.dialectSpec(@This(), .{
101 .dialect_attributes = &.{
102 "choir.string", "spirv.capability", "spirv.addressing_model",
103 "spirv.memory_model", "spirv.execution_model", "spirv.storage_class",
104 "spirv.ext_inst", "spirv.dim", "spirv.scope",
105 "spirv.warp_op", "spirv.shuffle_mode",
106 },
107 });
108
109 const func_symbol_vtable = ir.interfaces.SymbolOpInterface.VTable{
110 .getSymbolName = getFuncSymbolName,
111 .setSymbolName = setFuncSymbolName,
112 .isDeclaration = isFuncDeclaration,
113 };
114
115 pub const ModuleOp = struct {
116 op: *ir.Operation,
117
118 pub const operation_spec = op_specs.define(.{
119 .mnemonic = "module",
120 .required_attrs = &.{ "addressing_model", "capability", "ext_inst", "memory_model" },
121 .dynamic_traits = &.{symbol_table_trait},
122 });
123 pub const operation_name = operation_spec.name;
124
125 pub fn create(
126 ctx: *ir.Context,
127 loc: ir.Location,
128 addressing_model: AddressingModel,
129 memory_model: MemoryModel,
130 capability: Capability,
131 ext_inst: []const u8,
132 ) !ModuleOp {
133 try loadSpec(ctx);
134 var builder = ir.OperationBuilder.init(ctx);
135 var state = ir.Operation.State.init(operation_name, loc);
136 var body = ir.context.initRegion(ctx);
137 defer body.deinit();
138 var body_builder = ir.OperationBuilder.init(ctx);
139 _ = try body_builder.createBlock(&body, &.{}, &.{});
140 var regions = [_]*ir.Region{&body};
141 state.addRegionBodies(®ions);
142
143 const op = try builder.create(state);
144
145 try setAddressingModelAttr(op, ctx, addressing_model);
146 try setMemoryModelAttr(op, ctx, memory_model);
147 try setCapabilityAttr(op, ctx, capability);
148 try setExtInstAttr(op, ctx, ext_inst);
149
150 return .{ .op = op };
151 }
152
153 pub fn getBody(self: ModuleOp) *ir.Region {
154 return self.op.getRegion(0).?;
155 }
156
157 pub fn getBodyBlock(self: ModuleOp) *ir.Block {
158 return self.getBody().getEntryBlock().?;
159 }
160
161 pub fn getAddressingModel(self: ModuleOp) ?AddressingModel {
162 return getAddressingModelAttr(self.op);
163 }
164
165 pub fn getMemoryModel(self: ModuleOp) ?MemoryModel {
166 return getMemoryModelAttr(self.op);
167 }
168
169 pub fn getCapability(self: ModuleOp) ?Capability {
170 return getCapabilityAttr(self.op);
171 }
172
173 pub fn getExtInst(self: ModuleOp) ?[]const u8 {
174 return getExtInstAttr(self.op);
175 }
176 };
177
178 pub const FuncOp = struct {
179 op: *ir.Operation,
180
181 pub const operation_spec = op_specs.define(.{
182 .mnemonic = "func",
183 .attrs = &.{ "entry_point", "execution_model", "sym_name", ir.SymbolTable.symbol_attr_names.sym_visibility },
184 .required_attrs = &.{"sym_name"},
185 .interfaces = &.{
186 ir.interfaces.SymbolOpInterface.entry(&func_symbol_vtable),
187 },
188 });
189 pub const operation_name = operation_spec.name;
190
191 pub fn create(
192 ctx: *ir.Context,
193 loc: ir.Location,
194 func_name: []const u8,
195 input_types: []const ir.Type,
196 result_types: []const ir.Type,
197 ) !FuncOp {
198 try loadSpec(ctx);
199 var builder = ir.OperationBuilder.init(ctx);
200 var state = ir.Operation.State.init(operation_name, loc);
201 state.addTypes(result_types);
202 var body = ir.context.initRegion(ctx);
203 defer body.deinit();
204 var body_builder = ir.OperationBuilder.init(ctx);
205 _ = try body_builder.createBlockWithLoc(&body, input_types, loc);
206 var regions = [_]*ir.Region{&body};
207 state.addRegionBodies(®ions);
208
209 const op = try builder.create(state);
210
211 const name_attr = try ctx.getDialectAttr("choir.string", func_name);
212 try op.setAttr("sym_name", name_attr);
213
214 return .{ .op = op };
215 }
216
217 pub fn setEntryPoint(self: *FuncOp, ctx: *ir.Context, model: ExecutionModel) !void {
218 const entry_attr = try ctx.getBoolAttr(true);
219 try self.op.setAttr("entry_point", entry_attr);
220 try setExecutionModelAttr(self.op, ctx, model);
221 }
222
223 pub fn isEntryPoint(self: *const FuncOp) bool {
224 const bool_attr = self.op.getAttrAs(ir.Attribute.BoolAttr, "entry_point") orelse return false;
225 return bool_attr.getValue();
226 }
227
228 pub fn getExecutionModel(self: *const FuncOp) ?ExecutionModel {
229 return getExecutionModelAttr(self.op);
230 }
231
232 pub fn getBody(self: FuncOp) *ir.Region {
233 return self.op.getRegion(0).?;
234 }
235
236 pub fn getEntryBlock(self: FuncOp) *ir.Block {
237 return self.getBody().getEntryBlock().?;
238 }
239 };
240
241 pub const ConstantOp = struct {
242 op: *ir.Operation,
243
244 pub const operation_spec = op_specs.define(.{
245 .mnemonic = "constant",
246 .required_attrs = &.{"value"},
247 });
248 pub const operation_name = operation_spec.name;
249
250 pub fn createInt(ctx: *ir.Context, loc: ir.Location, result_type: ir.Type, value: i64) !ConstantOp {
251 try loadSpec(ctx);
252 var builder = ir.OperationBuilder.init(ctx);
253 var state = ir.Operation.State.init(operation_name, loc);
254 state.addTypes(&.{result_type});
255
256 const op = try builder.create(state);
257 const value_attr = try ctx.getI64Attr(value);
258 try op.setAttr("value", value_attr);
259
260 return .{ .op = op };
261 }
262
263 pub fn createFloat(ctx: *ir.Context, loc: ir.Location, result_type: ir.Type, value: f64) !ConstantOp {
264 try loadSpec(ctx);
265 var builder = ir.OperationBuilder.init(ctx);
266 var state = ir.Operation.State.init(operation_name, loc);
267 state.addTypes(&.{result_type});
268
269 const op = try builder.create(state);
270 const value_attr = try ctx.getF64Attr(value);
271 try op.setAttr("value", value_attr);
272
273 return .{ .op = op };
274 }
275
276 pub fn createBool(ctx: *ir.Context, loc: ir.Location, result_type: ir.Type, value: bool) !ConstantOp {
277 try loadSpec(ctx);
278 var builder = ir.OperationBuilder.init(ctx);
279 var state = ir.Operation.State.init(operation_name, loc);
280 state.addTypes(&.{result_type});
281
282 const op = try builder.create(state);
283 const value_attr = try ctx.getBoolAttr(value);
284 try op.setAttr("value", value_attr);
285
286 return .{ .op = op };
287 }
288
289 pub fn getResult(self: *const ConstantOp) *ir.Value {
290 return self.op.getResult(0).?;
291 }
292
293 pub fn getIntValue(self: ConstantOp) ?i64 {
294 const int_attr = self.op.getAttrAs(ir.Attribute.IntegerAttr, "value") orelse return null;
295 return int_attr.getValue();
296 }
297
298 pub fn getFloatValue(self: ConstantOp) ?f64 {
299 const float_attr = self.op.getAttrAs(ir.Attribute.FloatAttr, "value") orelse return null;
300 return float_attr.getValue();
301 }
302
303 pub fn getBoolValue(self: ConstantOp) ?bool {
304 const bool_attr = self.op.getAttrAs(ir.Attribute.BoolAttr, "value") orelse return null;
305 return bool_attr.getValue();
306 }
307 };
308
309 pub const VariableOp = struct {
310 op: *ir.Operation,
311
312 pub const operation_spec = op_specs.define(.{
313 .mnemonic = "variable",
314 .required_attrs = &.{"storage_class"},
315 });
316 pub const operation_name = operation_spec.name;
317
318 pub fn create(
319 ctx: *ir.Context,
320 loc: ir.Location,
321 result_type: ir.Type,
322 storage_class: StorageClass,
323 initializer: ?*ir.Value,
324 ) !VariableOp {
325 try loadSpec(ctx);
326 var builder = ir.OperationBuilder.init(ctx);
327 var state = ir.Operation.State.init(operation_name, loc);
328 if (initializer) |init| {
329 state.addOperands(&.{init});
330 }
331 state.addTypes(&.{result_type});
332
333 const op = try builder.create(state);
334 try setStorageClassAttr(op, ctx, storage_class);
335 return .{ .op = op };
336 }
337
338 pub fn getResult(self: *const VariableOp) *ir.Value {
339 return self.op.getResult(0).?;
340 }
341
342 pub fn getInitializer(self: *const VariableOp) ?*ir.Value {
343 if (self.op.operands.items.len > 0) {
344 return self.op.operands.items[0].value;
345 }
346 return null;
347 }
348
349 pub fn getStorageClass(self: *const VariableOp) ?StorageClass {
350 return getStorageClassAttr(self.op);
351 }
352 };
353
354 pub const LocalInvocationIdOp = struct {
355 op: *ir.Operation,
356
357 pub const operation_spec = op_specs.define(.{
358 .mnemonic = "local_invocation_id",
359 .required_attrs = &.{"dim"},
360 });
361 pub const operation_name = operation_spec.name;
362
363 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !LocalInvocationIdOp {
364 const op = try createIndexOp(ctx, loc, dim, operation_name);
365 return .{ .op = op };
366 }
367
368 pub fn getResult(self: *const LocalInvocationIdOp) *ir.Value {
369 return self.op.getResult(0).?;
370 }
371
372 pub fn getDimension(self: LocalInvocationIdOp) ?Dimension {
373 return getDimensionAttr(self.op);
374 }
375 };
376
377 pub const WorkgroupIdOp = struct {
378 op: *ir.Operation,
379
380 pub const operation_spec = op_specs.define(.{
381 .mnemonic = "workgroup_id",
382 .required_attrs = &.{"dim"},
383 });
384 pub const operation_name = operation_spec.name;
385
386 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !WorkgroupIdOp {
387 const op = try createIndexOp(ctx, loc, dim, operation_name);
388 return .{ .op = op };
389 }
390
391 pub fn getResult(self: *const WorkgroupIdOp) *ir.Value {
392 return self.op.getResult(0).?;
393 }
394
395 pub fn getDimension(self: WorkgroupIdOp) ?Dimension {
396 return getDimensionAttr(self.op);
397 }
398 };
399
400 pub const WorkgroupSizeOp = struct {
401 op: *ir.Operation,
402
403 pub const operation_spec = op_specs.define(.{
404 .mnemonic = "workgroup_size",
405 .required_attrs = &.{"dim"},
406 });
407 pub const operation_name = operation_spec.name;
408
409 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !WorkgroupSizeOp {
410 const op = try createIndexOp(ctx, loc, dim, operation_name);
411 return .{ .op = op };
412 }
413
414 pub fn getResult(self: *const WorkgroupSizeOp) *ir.Value {
415 return self.op.getResult(0).?;
416 }
417
418 pub fn getDimension(self: WorkgroupSizeOp) ?Dimension {
419 return getDimensionAttr(self.op);
420 }
421 };
422
423 pub const NumWorkgroupsOp = struct {
424 op: *ir.Operation,
425
426 pub const operation_spec = op_specs.define(.{
427 .mnemonic = "num_workgroups",
428 .required_attrs = &.{"dim"},
429 });
430 pub const operation_name = operation_spec.name;
431
432 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !NumWorkgroupsOp {
433 const op = try createIndexOp(ctx, loc, dim, operation_name);
434 return .{ .op = op };
435 }
436
437 pub fn getResult(self: *const NumWorkgroupsOp) *ir.Value {
438 return self.op.getResult(0).?;
439 }
440
441 pub fn getDimension(self: NumWorkgroupsOp) ?Dimension {
442 return getDimensionAttr(self.op);
443 }
444 };
445
446 pub const GlobalInvocationIdOp = struct {
447 op: *ir.Operation,
448
449 pub const operation_spec = op_specs.define(.{
450 .mnemonic = "global_invocation_id",
451 .required_attrs = &.{"dim"},
452 });
453 pub const operation_name = operation_spec.name;
454
455 pub fn create(ctx: *ir.Context, loc: ir.Location, dim: Dimension) !GlobalInvocationIdOp {
456 const op = try createIndexOp(ctx, loc, dim, operation_name);
457 return .{ .op = op };
458 }
459
460 pub fn getResult(self: *const GlobalInvocationIdOp) *ir.Value {
461 return self.op.getResult(0).?;
462 }
463
464 pub fn getDimension(self: GlobalInvocationIdOp) ?Dimension {
465 return getDimensionAttr(self.op);
466 }
467 };
468
469 pub const BarrierOp = struct {
470 op: *ir.Operation,
471
472 pub const operation_spec = op_specs.define(.{
473 .mnemonic = "control_barrier",
474 .required_attrs = &.{"scope"},
475 });
476 pub const operation_name = operation_spec.name;
477
478 pub fn create(ctx: *ir.Context, loc: ir.Location, scope: Scope) !BarrierOp {
479 try loadSpec(ctx);
480 var builder = ir.OperationBuilder.init(ctx);
481 const state = ir.Operation.State.init(operation_name, loc);
482
483 const op = try builder.create(state);
484 try setScopeAttr(op, ctx, scope);
485
486 return .{ .op = op };
487 }
488
489 pub fn getScope(self: BarrierOp) ?Scope {
490 return getScopeAttr(self.op);
491 }
492 };
493
494 pub const SyncWarpOp = struct {
495 op: *ir.Operation,
496
497 pub const operation_spec = op_specs.define(.{ .mnemonic = "sync_warp" });
498 pub const operation_name = operation_spec.name;
499
500 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value) !SyncWarpOp {
501 try loadSpec(ctx);
502 var builder = ir.OperationBuilder.init(ctx);
503 var state = ir.Operation.State.init(operation_name, loc);
504 state.addOperands(&.{mask});
505 const op = try builder.create(state);
506 return .{ .op = op };
507 }
508
509 pub fn getMask(self: SyncWarpOp) *ir.Value {
510 return self.op.operands.items[0].value;
511 }
512 };
513
514 pub const ActiveMaskOp = struct {
515 op: *ir.Operation,
516
517 pub const operation_spec = op_specs.define(.{ .mnemonic = "active_mask" });
518 pub const operation_name = operation_spec.name;
519
520 pub fn create(ctx: *ir.Context, loc: ir.Location) !ActiveMaskOp {
521 try loadSpec(ctx);
522 var builder = ir.OperationBuilder.init(ctx);
523 const i32_type = try arith.ArithDialect.getI32Type(ctx);
524 var state = ir.Operation.State.init(operation_name, loc);
525 state.addTypes(&.{i32_type});
526
527 const op = try builder.create(state);
528 return .{ .op = op };
529 }
530
531 pub fn getResult(self: *const ActiveMaskOp) *ir.Value {
532 return self.op.getResult(0).?;
533 }
534 };
535
536 pub const AllSyncOp = struct {
537 op: *ir.Operation,
538
539 pub const operation_spec = op_specs.define(.{ .mnemonic = "all_sync" });
540 pub const operation_name = operation_spec.name;
541
542 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AllSyncOp {
543 try loadSpec(ctx);
544 var builder = ir.OperationBuilder.init(ctx);
545 const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);
546 var state = ir.Operation.State.init(operation_name, loc);
547 state.addOperands(&.{ mask, pred });
548 state.addTypes(&.{bool_type});
549
550 const op = try builder.create(state);
551 return .{ .op = op };
552 }
553
554 pub fn getResult(self: *const AllSyncOp) *ir.Value {
555 return self.op.getResult(0).?;
556 }
557
558 pub fn getMask(self: AllSyncOp) *ir.Value {
559 return self.op.operands.items[0].value;
560 }
561
562 pub fn getPredicate(self: AllSyncOp) *ir.Value {
563 return self.op.operands.items[1].value;
564 }
565 };
566
567 pub const AnySyncOp = struct {
568 op: *ir.Operation,
569
570 pub const operation_spec = op_specs.define(.{ .mnemonic = "any_sync" });
571 pub const operation_name = operation_spec.name;
572
573 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !AnySyncOp {
574 try loadSpec(ctx);
575 var builder = ir.OperationBuilder.init(ctx);
576 const bool_type = try arith.ArithDialect.getScalarType(ctx, .bool);
577 var state = ir.Operation.State.init(operation_name, loc);
578 state.addOperands(&.{ mask, pred });
579 state.addTypes(&.{bool_type});
580
581 const op = try builder.create(state);
582 return .{ .op = op };
583 }
584
585 pub fn getResult(self: *const AnySyncOp) *ir.Value {
586 return self.op.getResult(0).?;
587 }
588
589 pub fn getMask(self: AnySyncOp) *ir.Value {
590 return self.op.operands.items[0].value;
591 }
592
593 pub fn getPredicate(self: AnySyncOp) *ir.Value {
594 return self.op.operands.items[1].value;
595 }
596 };
597
598 pub const BallotSyncOp = struct {
599 op: *ir.Operation,
600
601 pub const operation_spec = op_specs.define(.{ .mnemonic = "ballot_sync" });
602 pub const operation_name = operation_spec.name;
603
604 pub fn create(ctx: *ir.Context, loc: ir.Location, mask: *ir.Value, pred: *ir.Value) !BallotSyncOp {
605 try loadSpec(ctx);
606 var builder = ir.OperationBuilder.init(ctx);
607 const i32_type = try arith.ArithDialect.getI32Type(ctx);
608 var state = ir.Operation.State.init(operation_name, loc);
609 state.addOperands(&.{ mask, pred });
610 state.addTypes(&.{i32_type});
611
612 const op = try builder.create(state);
613 return .{ .op = op };
614 }
615
616 pub fn getResult(self: *const BallotSyncOp) *ir.Value {
617 return self.op.getResult(0).?;
618 }
619
620 pub fn getMask(self: BallotSyncOp) *ir.Value {
621 return self.op.operands.items[0].value;
622 }
623
624 pub fn getPredicate(self: BallotSyncOp) *ir.Value {
625 return self.op.operands.items[1].value;
626 }
627 };
628
629 pub const ShflSyncOp = struct {
630 op: *ir.Operation,
631
632 pub const operation_spec = op_specs.define(.{
633 .mnemonic = "shfl_sync",
634 .required_attrs = &.{"mode"},
635 });
636 pub const operation_name = operation_spec.name;
637
638 pub fn create(
639 ctx: *ir.Context,
640 loc: ir.Location,
641 mode: ShuffleMode,
642 mask: *ir.Value,
643 src: *ir.Value,
644 lane_or_delta: *ir.Value,
645 ) !ShflSyncOp {
646 try loadSpec(ctx);
647 var builder = ir.OperationBuilder.init(ctx);
648 var state = ir.Operation.State.init(operation_name, loc);
649 state.addOperands(&.{ mask, src, lane_or_delta });
650 state.addTypes(&.{src.type});
651
652 const op = try builder.create(state);
653 try setShuffleModeAttr(op, ctx, mode);
654 return .{ .op = op };
655 }
656
657 pub fn getResult(self: *const ShflSyncOp) *ir.Value {
658 return self.op.getResult(0).?;
659 }
660
661 pub fn getMask(self: ShflSyncOp) *ir.Value {
662 return self.op.operands.items[0].value;
663 }
664
665 pub fn getSrc(self: ShflSyncOp) *ir.Value {
666 return self.op.operands.items[1].value;
667 }
668
669 pub fn getLaneOrDelta(self: ShflSyncOp) *ir.Value {
670 return self.op.operands.items[2].value;
671 }
672
673 pub fn getMode(self: ShflSyncOp) ?ShuffleMode {
674 return getShuffleModeAttr(self.op);
675 }
676 };
677
678 pub const WarpReduceOp = struct {
679 op: *ir.Operation,
680
681 pub const operation_spec = op_specs.define(.{
682 .mnemonic = "warp_reduce",
683 .required_attrs = &.{"op"},
684 });
685 pub const operation_name = operation_spec.name;
686
687 pub fn create(
688 ctx: *ir.Context,
689 loc: ir.Location,
690 op_kind: WarpOpKind,
691 mask: *ir.Value,
692 value: *ir.Value,
693 ) !WarpReduceOp {
694 try loadSpec(ctx);
695 var builder = ir.OperationBuilder.init(ctx);
696 var state = ir.Operation.State.init(operation_name, loc);
697 state.addOperands(&.{ mask, value });
698 state.addTypes(&.{value.type});
699
700 const op = try builder.create(state);
701 try setWarpOpAttr(op, ctx, op_kind);
702 return .{ .op = op };
703 }
704
705 pub fn getResult(self: *const WarpReduceOp) *ir.Value {
706 return self.op.getResult(0).?;
707 }
708
709 pub fn getMask(self: WarpReduceOp) *ir.Value {
710 return self.op.operands.items[0].value;
711 }
712
713 pub fn getValue(self: WarpReduceOp) *ir.Value {
714 return self.op.operands.items[1].value;
715 }
716
717 pub fn getOpKind(self: WarpReduceOp) ?WarpOpKind {
718 return getWarpOpAttr(self.op);
719 }
720 };
721
722 pub const WarpScanOp = struct {
723 op: *ir.Operation,
724
725 pub const operation_spec = op_specs.define(.{
726 .mnemonic = "warp_scan",
727 .required_attrs = &.{ "inclusive", "op" },
728 });
729 pub const operation_name = operation_spec.name;
730
731 pub fn create(
732 ctx: *ir.Context,
733 loc: ir.Location,
734 op_kind: WarpOpKind,
735 inclusive: bool,
736 mask: *ir.Value,
737 value: *ir.Value,
738 ) !WarpScanOp {
739 try loadSpec(ctx);
740 var builder = ir.OperationBuilder.init(ctx);
741 var state = ir.Operation.State.init(operation_name, loc);
742 state.addOperands(&.{ mask, value });
743 state.addTypes(&.{value.type});
744
745 const op = try builder.create(state);
746 try setWarpOpAttr(op, ctx, op_kind);
747 try setBoolAttr(op, ctx, "inclusive", inclusive);
748 return .{ .op = op };
749 }
750
751 pub fn getResult(self: *const WarpScanOp) *ir.Value {
752 return self.op.getResult(0).?;
753 }
754
755 pub fn getMask(self: WarpScanOp) *ir.Value {
756 return self.op.operands.items[0].value;
757 }
758
759 pub fn getValue(self: WarpScanOp) *ir.Value {
760 return self.op.operands.items[1].value;
761 }
762
763 pub fn getOpKind(self: WarpScanOp) ?WarpOpKind {
764 return getWarpOpAttr(self.op);
765 }
766
767 pub fn isInclusive(self: WarpScanOp) bool {
768 return getBoolAttrValue(self.op, "inclusive");
769 }
770 };
771
772 pub const IAddOp = struct {
773 op: *ir.Operation,
774
775 pub const operation_spec = op_specs.define(.{ .mnemonic = "iadd" });
776 pub const operation_name = operation_spec.name;
777
778 pub fn create(ctx: *ir.Context, loc: ir.Location, lhs: *ir.Value, rhs: *ir.Value) !IAddOp {
779 const op = try createBinary(ctx, loc, lhs, rhs, operation_name);
780 return .{ .op = op };
781 }
782 };
783
784 pub const FAddOp = struct {
785 op: *ir.Operation,
786
787 pub const operation_spec = op_specs.define(.{ .mnemonic = "fadd" });
788 pub const operation_name = operation_spec.name;
789
790 pub fn create(ctx: *ir.Context, loc: ir.Location, lhs: *ir.Value, rhs: *ir.Value) !FAddOp {
791 const op = try createBinary(ctx, loc, lhs, rhs, operation_name);
792 return .{ .op = op };
793 }
794 };
795
796 pub const ISubOp = struct {
797 op: *ir.Operation,
798
799 pub const operation_spec = op_specs.define(.{ .mnemonic = "isub" });
800 pub const operation_name = operation_spec.name;
801
802 pub fn create(ctx: *ir.Context, loc: ir.Location, lhs: *ir.Value, rhs: *ir.Value) !ISubOp {
803 const op = try createBinary(ctx, loc, lhs, rhs, operation_name);
804 return .{ .op = op };
805 }
806 };
807
808 pub const FSubOp = struct {
809 op: *ir.Operation,
810
811 pub const operation_spec = op_specs.define(.{ .mnemonic = "fsub" });
812 pub const operation_name = operation_spec.name;
813
814 pub fn create(ctx: *ir.Context, loc: ir.Location, lhs: *ir.Value, rhs: *ir.Value) !FSubOp {
815 const op = try createBinary(ctx, loc, lhs, rhs, operation_name);
816 return .{ .op = op };
817 }
818 };
819
820 pub const IMulOp = struct {
821 op: *ir.Operation,
822
823 pub const operation_spec = op_specs.define(.{ .mnemonic = "imul" });
824 pub const operation_name = operation_spec.name;
825
826 pub fn create(ctx: *ir.Context, loc: ir.Location, lhs: *ir.Value, rhs: *ir.Value) !IMulOp {
827 const op = try createBinary(ctx, loc, lhs, rhs, operation_name);
828 return .{ .op = op };
829 }
830 };
831
832 pub const FMulOp = struct {
833 op: *ir.Operation,
834
835 pub const operation_spec = op_specs.define(.{ .mnemonic = "fmul" });
836 pub const operation_name = operation_spec.name;
837
838 pub fn create(ctx: *ir.Context, loc: ir.Location, lhs: *ir.Value, rhs: *ir.Value) !FMulOp {
839 const op = try createBinary(ctx, loc, lhs, rhs, operation_name);
840 return .{ .op = op };
841 }
842 };
843
844 pub const UDivOp = struct {
845 op: *ir.Operation,
846
847 pub const operation_spec = op_specs.define(.{ .mnemonic = "udiv" });
848 pub const operation_name = operation_spec.name;
849
850 pub fn create(ctx: *ir.Context, loc: ir.Location, lhs: *ir.Value, rhs: *ir.Value) !UDivOp {
851 const op = try createBinary(ctx, loc, lhs, rhs, operation_name);
852 return .{ .op = op };
853 }
854 };
855
856 pub const SDivOp = struct {
857 op: *ir.Operation,
858
859 pub const operation_spec = op_specs.define(.{ .mnemonic = "sdiv" });
860 pub const operation_name = operation_spec.name;
861
862 pub fn create(ctx: *ir.Context, loc: ir.Location, lhs: *ir.Value, rhs: *ir.Value) !SDivOp {
863 const op = try createBinary(ctx, loc, lhs, rhs, operation_name);
864 return .{ .op = op };
865 }
866 };
867
868 pub const FDivOp = struct {
869 op: *ir.Operation,
870
871 pub const operation_spec = op_specs.define(.{ .mnemonic = "fdiv" });
872 pub const operation_name = operation_spec.name;
873
874 pub fn create(ctx: *ir.Context, loc: ir.Location, lhs: *ir.Value, rhs: *ir.Value) !FDivOp {
875 const op = try createBinary(ctx, loc, lhs, rhs, operation_name);
876 return .{ .op = op };
877 }
878 };
879
880 fn loadSpec(ctx: *ir.Context) !void {
881 ir.dialects.loadDialectSpec(ctx, spec) catch |err| switch (err) {
882 error.ContextFrozen => {},
883 else => return err,
884 };
885 }
886
887 fn getFuncSymbolName(op_ptr: *const anyopaque) ?[]const u8 {
888 const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));
889 if (op.getAttrAs(ir.Attribute.StringAttr, "sym_name")) |string_attr| {
890 return string_attr.getValue();
891 }
892 const attr = op.getAttr("sym_name") orelse return null;
893 if (std.mem.eql(u8, attr.abstract.name, "choir.string")) {
894 const dialect_attr = attr.cast(ir.Attribute.DialectAttr) orelse return null;
895 return dialect_attr.payload;
896 }
897 return null;
898 }
899
900 fn setFuncSymbolName(op_ptr: *const anyopaque, symbol_name: []const u8) anyerror!void {
901 const op: *ir.Operation = @ptrCast(@alignCast(@constCast(op_ptr)));
902 try op.setAttr("sym_name", try op.getContext().getDialectAttr("choir.string", symbol_name));
903 }
904
905 fn isFuncDeclaration(_: *const anyopaque) bool {
906 return false;
907 }
908
909 fn setCapabilityAttr(op: *ir.Operation, ctx: *ir.Context, capability: Capability) !void {
910 const cap_attr = try ctx.getDialectAttr("spirv.capability", capability.toString());
911 try op.setAttr("capability", cap_attr);
912 }
913
914 fn getCapabilityAttr(op: *const ir.Operation) ?Capability {
915 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "capability") orelse return null;
916 return Capability.fromString(dialect_attr.payload);
917 }
918
919 fn setAddressingModelAttr(op: *ir.Operation, ctx: *ir.Context, model: AddressingModel) !void {
920 const model_attr = try ctx.getDialectAttr("spirv.addressing_model", model.toString());
921 try op.setAttr("addressing_model", model_attr);
922 }
923
924 fn getAddressingModelAttr(op: *const ir.Operation) ?AddressingModel {
925 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "addressing_model") orelse return null;
926 return AddressingModel.fromString(dialect_attr.payload);
927 }
928
929 fn setMemoryModelAttr(op: *ir.Operation, ctx: *ir.Context, model: MemoryModel) !void {
930 const model_attr = try ctx.getDialectAttr("spirv.memory_model", model.toString());
931 try op.setAttr("memory_model", model_attr);
932 }
933
934 fn getMemoryModelAttr(op: *const ir.Operation) ?MemoryModel {
935 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "memory_model") orelse return null;
936 return MemoryModel.fromString(dialect_attr.payload);
937 }
938
939 fn setExecutionModelAttr(op: *ir.Operation, ctx: *ir.Context, model: ExecutionModel) !void {
940 const model_attr = try ctx.getDialectAttr("spirv.execution_model", model.toString());
941 try op.setAttr("execution_model", model_attr);
942 }
943
944 fn getExecutionModelAttr(op: *const ir.Operation) ?ExecutionModel {
945 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "execution_model") orelse return null;
946 return ExecutionModel.fromString(dialect_attr.payload);
947 }
948
949 fn setStorageClassAttr(op: *ir.Operation, ctx: *ir.Context, storage: StorageClass) !void {
950 const storage_attr = try ctx.getDialectAttr("spirv.storage_class", storage.toString());
951 try op.setAttr("storage_class", storage_attr);
952 }
953
954 fn getStorageClassAttr(op: *const ir.Operation) ?StorageClass {
955 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "storage_class") orelse return null;
956 return StorageClass.fromString(dialect_attr.payload);
957 }
958
959 fn setExtInstAttr(op: *ir.Operation, ctx: *ir.Context, ext_inst_name: []const u8) !void {
960 const ext_attr = try ctx.getDialectAttr("spirv.ext_inst", ext_inst_name);
961 try op.setAttr("ext_inst", ext_attr);
962 }
963
964 fn getExtInstAttr(op: *const ir.Operation) ?[]const u8 {
965 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "ext_inst") orelse return null;
966 return dialect_attr.payload;
967 }
968
969 fn setDimensionAttr(op: *ir.Operation, ctx: *ir.Context, dim: Dimension) !void {
970 const dim_attr = try ctx.getDialectAttr("spirv.dim", dim.toString());
971 try op.setAttr("dim", dim_attr);
972 }
973
974 fn getDimensionAttr(op: *const ir.Operation) ?Dimension {
975 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "dim") orelse return null;
976 return Dimension.fromString(dialect_attr.payload);
977 }
978
979 fn setScopeAttr(op: *ir.Operation, ctx: *ir.Context, scope: Scope) !void {
980 const scope_attr = try ctx.getDialectAttr("spirv.scope", scope.toString());
981 try op.setAttr("scope", scope_attr);
982 }
983
984 fn getScopeAttr(op: *const ir.Operation) ?Scope {
985 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "scope") orelse return null;
986 return Scope.fromString(dialect_attr.payload);
987 }
988
989 fn setWarpOpAttr(op: *ir.Operation, ctx: *ir.Context, op_kind: WarpOpKind) !void {
990 const op_attr = try ctx.getDialectAttr("spirv.warp_op", op_kind.toString());
991 try op.setAttr("op", op_attr);
992 }
993
994 fn getWarpOpAttr(op: *const ir.Operation) ?WarpOpKind {
995 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "op") orelse return null;
996 return WarpOpKind.fromString(dialect_attr.payload);
997 }
998
999 fn setShuffleModeAttr(op: *ir.Operation, ctx: *ir.Context, mode: ShuffleMode) !void {
1000 const mode_attr = try ctx.getDialectAttr("spirv.shuffle_mode", mode.toString());
1001 try op.setAttr("mode", mode_attr);
1002 }
1003
1004 fn getShuffleModeAttr(op: *const ir.Operation) ?ShuffleMode {
1005 const dialect_attr = op.getAttrAs(ir.Attribute.DialectAttr, "mode") orelse return null;
1006 return ShuffleMode.fromString(dialect_attr.payload);
1007 }
1008
1009 fn setBoolAttr(op: *ir.Operation, ctx: *ir.Context, attr_name: []const u8, value: bool) !void {
1010 const bool_attr = try ctx.getBoolAttr(value);
1011 try op.setAttr(attr_name, bool_attr);
1012 }
1013
1014 fn getBoolAttrValue(op: *const ir.Operation, attr_name: []const u8) bool {
1015 const bool_attr = op.getAttrAs(ir.Attribute.BoolAttr, attr_name) orelse return false;
1016 return bool_attr.getValue();
1017 }
1018
1019 fn createIndexOp(
1020 ctx: *ir.Context,
1021 loc: ir.Location,
1022 dim: Dimension,
1023 comptime op_name: []const u8,
1024 ) !*ir.Operation {
1025 try loadSpec(ctx);
1026 var builder = ir.OperationBuilder.init(ctx);
1027 const index_type = try arith.ArithDialect.getIndexType(ctx);
1028 var state = ir.Operation.State.init(op_name, loc);
1029 state.addTypes(&.{index_type});
1030 const op = try builder.create(state);
1031 try setDimensionAttr(op, ctx, dim);
1032 return op;
1033 }
1034
1035 fn createBinary(
1036 ctx: *ir.Context,
1037 loc: ir.Location,
1038 lhs: *ir.Value,
1039 rhs: *ir.Value,
1040 comptime op_name: []const u8,
1041 ) !*ir.Operation {
1042 try loadSpec(ctx);
1043 var builder = ir.OperationBuilder.init(ctx);
1044 var state = ir.Operation.State.init(op_name, loc);
1045 state.addOperands(&.{ lhs, rhs });
1046 state.addTypes(&.{lhs.type});
1047 return builder.create(state);
1048 }
1049 };
1050
1051 test "SpirvDialect.ModuleOp creates container with module attributes" {
1052 const testing = std.testing;
1053 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1054 defer arena.deinit();
1055 const allocator = arena.allocator();
1056
1057 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1058 defer ctx.deinit(allocator);
1059
1060 const loc = ir.Location.getUnknown();
1061 const module = try SpirvDialect.ModuleOp.create(
1062 &ctx,
1063 loc,
1064 .logical,
1065 .glsl450,
1066 .shader,
1067 "GLSL.std.450",
1068 );
1069
1070 try testing.expectEqualStrings(SpirvDialect.ModuleOp.operation_name, module.op.name.name);
1071 try testing.expect(module.getBody().getEntryBlock() != null);
1072 try testing.expectEqual(AddressingModel.logical, module.getAddressingModel().?);
1073 try testing.expectEqual(MemoryModel.glsl450, module.getMemoryModel().?);
1074 try testing.expectEqual(Capability.shader, module.getCapability().?);
1075 try testing.expectEqualStrings("GLSL.std.450", module.getExtInst().?);
1076 }
1077
1078 test "SpirvDialect.ModuleOp owns a symbol table" {
1079 const testing = std.testing;
1080 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1081 defer arena.deinit();
1082 const allocator = arena.allocator();
1083
1084 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1085 defer ctx.deinit(allocator);
1086 try ir.dialects.loadDialectSpec(&ctx, SpirvDialect.spec);
1087
1088 const loc = ir.Location.getUnknown();
1089 const module = try SpirvDialect.ModuleOp.create(
1090 &ctx,
1091 loc,
1092 .logical,
1093 .glsl450,
1094 .shader,
1095 "GLSL.std.450",
1096 );
1097 const block = module.getBodyBlock();
1098 const kernel = try SpirvDialect.FuncOp.create(&ctx, loc, "kernel", &.{}, &.{});
1099 try block.addOperation(kernel.op);
1100
1101 var table = ir.SymbolTable.init(allocator);
1102 defer table.deinit();
1103 try table.buildFromOperation(module.op);
1104
1105 try testing.expect(module.op.getTraits().is_symbol_table);
1106 try testing.expect(kernel.op.interface(ir.interfaces.SymbolOpInterface) != null);
1107 try testing.expect(table.lookup("kernel") == kernel.op);
1108 }
1109
1110 test "SpirvDialect.FuncOp marks entry points" {
1111 const testing = std.testing;
1112 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1113 defer arena.deinit();
1114 const allocator = arena.allocator();
1115
1116 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1117 defer ctx.deinit(allocator);
1118
1119 const loc = ir.Location.getUnknown();
1120 const arith_dialect = choir.dialects.arith.ArithDialect;
1121 const i32_type = try arith_dialect.getScalarType(&ctx, .i32);
1122
1123 var func = try SpirvDialect.FuncOp.create(&ctx, loc, "kernel", &.{i32_type}, &.{});
1124 try testing.expect(!func.isEntryPoint());
1125 try func.setEntryPoint(&ctx, .gl_compute);
1126 try testing.expect(func.isEntryPoint());
1127 try testing.expectEqual(ExecutionModel.gl_compute, func.getExecutionModel().?);
1128 }
1129
1130 test "spirv execution models round trip through their names" {
1131 const models = .{ ExecutionModel.vertex, ExecutionModel.fragment, ExecutionModel.gl_compute };
1132 inline for (models) |model| {
1133 try std.testing.expectEqual(model, ExecutionModel.fromString(model.toString()).?);
1134 }
1135 try std.testing.expectEqual(@as(?ExecutionModel, null), ExecutionModel.fromString("kernel"));
1136 }
1137
1138 test "SpirvDialect.ConstantOp stores typed value" {
1139 const testing = std.testing;
1140 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1141 defer arena.deinit();
1142 const allocator = arena.allocator();
1143
1144 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1145 defer ctx.deinit(allocator);
1146
1147 const loc = ir.Location.getUnknown();
1148 const arith_dialect = choir.dialects.arith.ArithDialect;
1149 const i32_type = try arith_dialect.getScalarType(&ctx, .i32);
1150
1151 const const_op = try SpirvDialect.ConstantOp.createInt(&ctx, loc, i32_type, 42);
1152 try testing.expectEqual(@as(i64, 42), const_op.getIntValue().?);
1153 }
1154
1155 test "SpirvDialect.VariableOp sets storage class" {
1156 const testing = std.testing;
1157 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1158 defer arena.deinit();
1159 const allocator = arena.allocator();
1160
1161 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1162 defer ctx.deinit(allocator);
1163
1164 const loc = ir.Location.getUnknown();
1165 const arith_dialect = choir.dialects.arith.ArithDialect;
1166 const i32_type = try arith_dialect.getScalarType(&ctx, .i32);
1167
1168 const var_op = try SpirvDialect.VariableOp.create(&ctx, loc, i32_type, .workgroup, null);
1169 try testing.expectEqual(StorageClass.workgroup, var_op.getStorageClass().?);
1170 }
1171
1172 test "SpirvDialect.IAddOp creates binary op" {
1173 const testing = std.testing;
1174 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1175 defer arena.deinit();
1176 const allocator = arena.allocator();
1177
1178 var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);
1179 defer ctx.deinit(allocator);
1180
1181 const loc = ir.Location.getUnknown();
1182 const arith_dialect = choir.dialects.arith.ArithDialect;
1183 const i32_type = try arith_dialect.getScalarType(&ctx, .i32);
1184
1185 var c1 = try SpirvDialect.ConstantOp.createInt(&ctx, loc, i32_type, 10);
1186 var c2 = try SpirvDialect.ConstantOp.createInt(&ctx, loc, i32_type, 20);
1187 const add = try SpirvDialect.IAddOp.create(&ctx, loc, c1.getResult(), c2.getResult());
1188 try testing.expectEqualStrings(SpirvDialect.IAddOp.operation_name, add.op.name.name);
1189 }