lib/choir/src/backends/gpu/spirv/emitter/dialect.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("../../../../root.zig");
3
4 const ir = choir.ir;
5 const dialects = choir.dialects;
6 const spirv_target = @import("../root.zig");
7 const gpu = @import("gpu.zig");
8 const binary = @import("module.zig");
9 const scalar = @import("scalar.zig");
10 const spec = @import("spec.zig");
11 const spirv_ops = @import("ops.zig");
12
13 const FuncDialect = dialects.func.FuncDialect;
14 const SpirvDialect = spirv_target.SpirvDialect;
15 const SpirvExecutionModel = spirv_target.ExecutionModel;
16 const SpirvStorageClass = spirv_target.StorageClass;
17 const Section = binary.Section;
18 const SpirvOp = spirv_ops.SpirvOp;
19
20 pub fn functionName(op: *ir.Operation) ?[]const u8 {
21 return ir.SymbolTable.getSymbolName(op);
22 }
23
24 fn storageClassToSpv(storage: SpirvStorageClass) u32 {
25 return switch (storage) {
26 .function => spec.StorageClass.Function,
27 .private => spec.StorageClass.Private,
28 .workgroup => spec.StorageClass.Workgroup,
29 .uniform => spec.StorageClass.Uniform,
30 .storage_buffer => spec.StorageClass.StorageBuffer,
31 .input => spec.StorageClass.Input,
32 .output => spec.StorageClass.Output,
33 };
34 }
35
36 pub fn emitModuleWords(self: anytype, module: *ir.Operation) ![]u32 {
37 if (!std.mem.eql(u8, module.name.name, SpirvDialect.ModuleOp.operation_name)) {
38 return error.InvalidModule;
39 }
40
41 const spirv_module = SpirvDialect.ModuleOp{ .op = module };
42 const addressing_model = spirv_module.getAddressingModel() orelse return error.MissingAttribute;
43 const memory_model = spirv_module.getMemoryModel() orelse return error.MissingAttribute;
44 const capability = spirv_module.getCapability() orelse return error.MissingAttribute;
45
46 try self.requireCapability(switch (capability) {
47 .shader => spec.Capability.Shader,
48 });
49 try self.builder.emitMemoryModel(
50 switch (addressing_model) {
51 .logical => spec.AddressingModel.Logical,
52 },
53 switch (memory_model) {
54 .glsl450 => spec.MemoryModel.GLSL450,
55 },
56 );
57
58 const region = module.getRegion(0) orelse return error.InvalidModule;
59 const block = region.getEntryBlock() orelse return error.InvalidModule;
60
61 var funcs: std.ArrayListUnmanaged(*ir.Operation) = .empty;
62 var vars: std.ArrayListUnmanaged(*ir.Operation) = .empty;
63 var consts: std.ArrayListUnmanaged(*ir.Operation) = .empty;
64 defer funcs.deinit(self.allocator);
65 defer vars.deinit(self.allocator);
66 defer consts.deinit(self.allocator);
67
68 var op_iter = block.operations.head;
69 while (op_iter) |op_ptr| {
70 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
71 const name = op.name.name;
72 if (std.mem.eql(u8, name, SpirvDialect.FuncOp.operation_name)) {
73 try funcs.append(self.allocator, op);
74 } else if (std.mem.eql(u8, name, SpirvDialect.VariableOp.operation_name)) {
75 try vars.append(self.allocator, op);
76 } else if (std.mem.eql(u8, name, SpirvDialect.ConstantOp.operation_name)) {
77 try consts.append(self.allocator, op);
78 } else {
79 return error.UnsupportedOperation;
80 }
81 op_iter = op.next_op;
82 }
83
84 for (consts.items) |const_op| {
85 try emitConstant(self, const_op);
86 }
87
88 for (vars.items) |var_op| {
89 try emitVariable(self, var_op, true);
90 }
91
92 for (funcs.items) |func_op| {
93 try emitFunction(self, func_op);
94 }
95
96 return self.builder.toWords(self.allocator);
97 }
98
99 fn emitTopLevel(self: anytype, op: *ir.Operation) !void {
100 const name = op.name.name;
101 if (std.mem.eql(u8, name, SpirvDialect.FuncOp.operation_name)) {
102 return emitFunction(self, op);
103 }
104 if (std.mem.eql(u8, name, SpirvDialect.VariableOp.operation_name)) {
105 return emitVariable(self, op, true);
106 }
107 if (std.mem.eql(u8, name, SpirvDialect.ConstantOp.operation_name)) {
108 return emitConstant(self, op);
109 }
110 return error.UnsupportedOperation;
111 }
112
113 fn emitFunction(self: anytype, func_op: *ir.Operation) !void {
114 const name = functionName(func_op) orelse return error.MissingFunctionName;
115
116 const region = func_op.getRegion(0) orelse return error.InvalidModule;
117 if (!region.hasOneBlock()) return error.UnsupportedControlFlow;
118 const entry = region.getEntryBlock() orelse return error.InvalidModule;
119
120 const result_types = func_op.getResultTypes();
121 if (result_types.len > 1) return error.UnsupportedFunctionSignature;
122
123 const returns_void = result_types.len == 0;
124 const return_kind: scalar.Kind = if (returns_void)
125 .void
126 else
127 (scalar.kindFromType(result_types[0]) orelse return error.UnsupportedType);
128 const return_type_id = try self.getScalarType(return_kind);
129
130 var param_type_ids = std.ArrayListUnmanaged(u32).empty;
131 defer param_type_ids.deinit(self.allocator);
132 for (entry.arguments.items) |arg| {
133 const kind = scalar.kindFromType(arg.type) orelse return error.UnsupportedType;
134 const type_id = try self.getScalarType(kind);
135 try param_type_ids.append(self.allocator, type_id);
136 }
137
138 const func_type_id = try emitFunctionType(self, return_type_id, param_type_ids.items);
139 const func_id = self.builder.newId();
140 try self.builder.emit(&self.builder.functions, SpirvOp.Function, &.{
141 return_type_id,
142 func_id,
143 0,
144 func_type_id,
145 });
146
147 for (entry.arguments.items, 0..) |arg, idx| {
148 const param_id = self.builder.newId();
149 try self.builder.emit(&self.builder.functions, SpirvOp.FunctionParameter, &.{
150 param_type_ids.items[idx],
151 param_id,
152 });
153 try self.bindValue(arg, param_id);
154 }
155
156 const label_id = self.builder.newId();
157 try self.emitLabel(label_id);
158
159 var op_iter = entry.operations.head;
160 while (op_iter) |op_ptr| {
161 const op: *ir.Operation = @ptrCast(@alignCast(op_ptr));
162 try emitOp(self, op, returns_void);
163 op_iter = op.next_op;
164 }
165
166 try self.builder.emit(&self.builder.functions, SpirvOp.FunctionEnd, &.{});
167 self.current_block = null;
168
169 const spirv_func = SpirvDialect.FuncOp{ .op = func_op };
170 if (spirv_func.isEntryPoint()) {
171 const exec_model = spirv_func.getExecutionModel() orelse return error.MissingAttribute;
172 try self.builder.emitEntryPoint(
173 switch (exec_model) {
174 SpirvExecutionModel.vertex => spec.ExecutionModel.Vertex,
175 SpirvExecutionModel.fragment => spec.ExecutionModel.Fragment,
176 SpirvExecutionModel.gl_compute => spec.ExecutionModel.GLCompute,
177 },
178 func_id,
179 name,
180 self.interface_vars.items,
181 );
182 switch (exec_model) {
183 SpirvExecutionModel.vertex => {},
184 SpirvExecutionModel.fragment => try self.builder.emitExecutionMode(
185 func_id,
186 spec.ExecutionMode.OriginUpperLeft,
187 ),
188 SpirvExecutionModel.gl_compute => {
189 try self.builder.emitExecutionModeLocalSize(func_id, 1, 1, 1);
190 },
191 }
192 try self.emitFloatExecutionModes(func_id);
193 }
194 }
195
196 fn emitOp(self: anytype, op: *ir.Operation, returns_void: bool) !void {
197 const name = op.name.name;
198 if (std.mem.eql(u8, name, SpirvDialect.ConstantOp.operation_name)) {
199 return emitConstant(self, op);
200 }
201 if (std.mem.eql(u8, name, SpirvDialect.VariableOp.operation_name)) {
202 return emitVariable(self, op, false);
203 }
204 if (std.mem.eql(u8, name, SpirvDialect.IAddOp.operation_name)) {
205 return emitBinary(self, op, SpirvOp.IAdd, .int_any);
206 }
207 if (std.mem.eql(u8, name, SpirvDialect.FAddOp.operation_name)) {
208 return emitBinary(self, op, SpirvOp.FAdd, .float);
209 }
210 if (std.mem.eql(u8, name, SpirvDialect.ISubOp.operation_name)) {
211 return emitBinary(self, op, SpirvOp.ISub, .int_any);
212 }
213 if (std.mem.eql(u8, name, SpirvDialect.FSubOp.operation_name)) {
214 return emitBinary(self, op, SpirvOp.FSub, .float);
215 }
216 if (std.mem.eql(u8, name, SpirvDialect.IMulOp.operation_name)) {
217 return emitBinary(self, op, SpirvOp.IMul, .int_any);
218 }
219 if (std.mem.eql(u8, name, SpirvDialect.FMulOp.operation_name)) {
220 return emitBinary(self, op, SpirvOp.FMul, .float);
221 }
222 if (std.mem.eql(u8, name, SpirvDialect.UDivOp.operation_name)) {
223 return emitBinary(self, op, SpirvOp.UDiv, .int_unsigned);
224 }
225 if (std.mem.eql(u8, name, SpirvDialect.SDivOp.operation_name)) {
226 return emitBinary(self, op, SpirvOp.SDiv, .int_signed);
227 }
228 if (std.mem.eql(u8, name, SpirvDialect.FDivOp.operation_name)) {
229 return emitBinary(self, op, SpirvOp.FDiv, .float);
230 }
231 if (std.mem.eql(u8, name, SpirvDialect.LocalInvocationIdOp.operation_name)) {
232 return gpu.emitIndex(self, op, .local_invocation_id);
233 }
234 if (std.mem.eql(u8, name, SpirvDialect.WorkgroupIdOp.operation_name)) {
235 return gpu.emitIndex(self, op, .workgroup_id);
236 }
237 if (std.mem.eql(u8, name, SpirvDialect.WorkgroupSizeOp.operation_name)) {
238 return gpu.emitIndex(self, op, .workgroup_size);
239 }
240 if (std.mem.eql(u8, name, SpirvDialect.NumWorkgroupsOp.operation_name)) {
241 return gpu.emitIndex(self, op, .num_workgroups);
242 }
243 if (std.mem.eql(u8, name, SpirvDialect.GlobalInvocationIdOp.operation_name)) {
244 return gpu.emitIndex(self, op, .global_invocation_id);
245 }
246 if (std.mem.eql(u8, name, SpirvDialect.BarrierOp.operation_name)) {
247 return gpu.emitBarrier(self, op);
248 }
249 if (std.mem.eql(u8, name, SpirvDialect.SyncWarpOp.operation_name)) {
250 return gpu.emitSyncWarp(self, op);
251 }
252 if (std.mem.eql(u8, name, SpirvDialect.ActiveMaskOp.operation_name)) {
253 return gpu.emitActiveMask(self, op);
254 }
255 if (std.mem.eql(u8, name, SpirvDialect.AllSyncOp.operation_name)) {
256 return gpu.emitAllAny(self, op, .all);
257 }
258 if (std.mem.eql(u8, name, SpirvDialect.AnySyncOp.operation_name)) {
259 return gpu.emitAllAny(self, op, .any);
260 }
261 if (std.mem.eql(u8, name, SpirvDialect.BallotSyncOp.operation_name)) {
262 return gpu.emitBallot(self, op);
263 }
264 if (std.mem.eql(u8, name, SpirvDialect.ShflSyncOp.operation_name)) {
265 return gpu.emitShuffle(self, op);
266 }
267 if (std.mem.eql(u8, name, SpirvDialect.WarpReduceOp.operation_name)) {
268 return gpu.emitWarpReduce(self, op);
269 }
270 if (std.mem.eql(u8, name, SpirvDialect.WarpScanOp.operation_name)) {
271 return gpu.emitWarpScan(self, op);
272 }
273 if (std.mem.eql(u8, name, FuncDialect.ReturnOp.operation_name)) {
274 return emitReturn(self, op, returns_void);
275 }
276 return error.UnsupportedOperation;
277 }
278
279 pub fn emitConstant(self: anytype, op: *ir.Operation) !void {
280 const result = op.getResult(0) orelse return error.UnsupportedOperation;
281 const type_id = try self.getTypeForValue(result);
282
283 const constant = SpirvDialect.ConstantOp{ .op = op };
284 if (constant.getIntValue()) |int_value| {
285 const kind = scalar.kindFromType(result.type) orelse return error.UnsupportedType;
286 const const_id = try self.getIntConstant(type_id, kind, int_value);
287 try self.bindValue(result, const_id);
288 return;
289 }
290
291 if (constant.getFloatValue()) |float_value| {
292 const kind = scalar.kindFromType(result.type) orelse return error.UnsupportedType;
293 const const_id = try self.getFloatConstant(type_id, kind, float_value);
294 try self.bindValue(result, const_id);
295 return;
296 }
297
298 if (op.getAttrAs(ir.Attribute.BoolAttr, "value")) |bool_attr| {
299 const kind = scalar.kindFromType(result.type) orelse return error.UnsupportedType;
300 if (kind != .bool) return error.UnsupportedType;
301 const const_id = try self.getBoolConstant(bool_attr.getValue());
302 try self.bindValue(result, const_id);
303 return;
304 }
305
306 return error.UnsupportedOperation;
307 }
308
309 pub const SpirvBinaryExpectation = enum {
310 int_any,
311 int_signed,
312 int_unsigned,
313 float,
314 };
315
316 pub fn emitBinary(
317 self: anytype,
318 op: *ir.Operation,
319 opcode: u16,
320 expectation: SpirvBinaryExpectation,
321 ) !void {
322 if (op.operands.items.len != 2) return error.UnsupportedOperation;
323 const lhs = op.operands.items[0].value;
324 const rhs = op.operands.items[1].value;
325 const result = op.getResult(0) orelse return error.UnsupportedOperation;
326
327 const kind = scalar.kindFromType(result.type) orelse return error.UnsupportedType;
328 switch (expectation) {
329 .float => if (!scalar.isFloat(kind)) return error.UnsupportedType,
330 .int_any => if (!scalar.isInt(kind)) return error.UnsupportedType,
331 .int_signed => if (!scalar.isSignedInt(kind)) return error.UnsupportedType,
332 .int_unsigned => if (!scalar.isUnsignedInt(kind)) return error.UnsupportedType,
333 }
334
335 const lhs_id = try self.getValue(lhs);
336 const rhs_id = try self.getValue(rhs);
337 const result_type_id = try self.getScalarType(kind);
338
339 const result_id = self.builder.newId();
340 try self.builder.emit(&self.builder.functions, opcode, &.{
341 result_type_id,
342 result_id,
343 lhs_id,
344 rhs_id,
345 });
346 try self.bindValue(result, result_id);
347 }
348
349 fn emitVariable(
350 self: anytype,
351 op: *ir.Operation,
352 module_scope: bool,
353 ) !void {
354 const var_op = SpirvDialect.VariableOp{ .op = op };
355 const result = op.getResult(0) orelse return error.UnsupportedOperation;
356 const storage = var_op.getStorageClass() orelse return error.MissingAttribute;
357 const storage_class = storageClassToSpv(storage);
358
359 const elem_kind = scalar.kindFromType(result.type) orelse return error.UnsupportedType;
360 const elem_type = try self.getScalarType(elem_kind);
361 const ptr_type = try self.getPointerType(storage_class, elem_type);
362
363 const var_id = self.builder.newId();
364 const initializer = var_op.getInitializer();
365
366 const section = if (module_scope) &self.builder.globals else &self.builder.functions;
367 if (initializer) |init_val| {
368 const init_id = try self.getValue(init_val);
369 try self.builder.emit(section, SpirvOp.Variable, &.{
370 ptr_type,
371 var_id,
372 storage_class,
373 init_id,
374 });
375 } else {
376 try self.builder.emit(section, SpirvOp.Variable, &.{
377 ptr_type,
378 var_id,
379 storage_class,
380 });
381 }
382
383 try self.bindValue(result, var_id);
384 if (module_scope) {
385 try self.addInterfaceVar(var_id);
386 }
387 }
388
389 fn emitReturn(self: anytype, op: *ir.Operation, returns_void: bool) !void {
390 if (returns_void) {
391 if (op.operands.items.len != 0) return error.UnsupportedFunctionSignature;
392 try self.builder.emit(&self.builder.functions, SpirvOp.Return, &.{});
393 return;
394 }
395
396 if (op.operands.items.len != 1) return error.UnsupportedFunctionSignature;
397 const value_id = try self.getValue(op.operands.items[0].value);
398 try self.builder.emit(&self.builder.functions, SpirvOp.ReturnValue, &.{value_id});
399 }
400
401 /// The function type of a signature, declared once per module as SPIR-V
402 /// requires of a non-aggregate type.
403 pub fn emitFunctionType(
404 self: anytype,
405 return_type: u32,
406 param_types: []const u32,
407 ) !u32 {
408 const types = self.builder.types.items;
409 if (findFunctionType(types, return_type, param_types)) |existing| return existing;
410 const id = self.builder.newId();
411 var operands = Section.empty;
412 defer operands.deinit(self.allocator);
413
414 try operands.append(self.allocator, id);
415 try operands.append(self.allocator, return_type);
416 if (param_types.len > 0) {
417 try operands.appendSlice(self.allocator, param_types);
418 }
419
420 try self.builder.emit(&self.builder.types, SpirvOp.TypeFunction, operands.items);
421 return id;
422 }
423
424 fn findFunctionType(types: []const u32, return_type: u32, param_types: []const u32) ?u32 {
425 var index: usize = 0;
426 while (index < types.len) {
427 const count = types[index] >> 16;
428 std.debug.assert(count > 0);
429 const opcode: u16 = @truncate(types[index]);
430 if (opcode == SpirvOp.TypeFunction and count == 3 + param_types.len and
431 types[index + 2] == return_type and
432 std.mem.eql(u32, types[index + 3 ..][0..param_types.len], param_types))
433 {
434 return types[index + 1];
435 }
436 index += count;
437 }
438 return null;
439 }
440
441 test "spirv dialect writer maps storage classes" {
442 try std.testing.expectEqual(@as(u32, spec.StorageClass.Workgroup), storageClassToSpv(.workgroup));
443 try std.testing.expectEqual(@as(u32, spec.StorageClass.StorageBuffer), storageClassToSpv(.storage_buffer));
444 }