lib/choir/src/backends/gpu/spirv/emitter/module.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const SpirvOp = @import("ops.zig").SpirvOp;
3 const ExecutionMode = @import("spec.zig").ExecutionMode;
4 const Decoration = @import("spec.zig").Decoration;
5
6 pub const Version = struct {
7 pub const v10: u32 = 0x00010000;
8 pub const v13: u32 = 0x00010300;
9 };
10
11 pub const Header = struct {
12 pub const magic: u32 = 0x07230203;
13 pub const version: u32 = Version.v10;
14 pub const generator: u32 = 0;
15 pub const schema: u32 = 0;
16 };
17
18 pub const Section = std.ArrayListUnmanaged(u32);
19
20 pub const Builder = struct {
21 allocator: std.mem.Allocator,
22 version: u32,
23 next_id: u32,
24 capabilities: Section,
25 extensions: Section,
26 ext_inst_imports: Section,
27 memory_model: Section,
28 entry_points: Section,
29 execution_modes: Section,
30 annotations: Section,
31 types: Section,
32 globals: Section,
33 functions: Section,
34 glsl_std_450_id: u32,
35
36 pub fn init(allocator: std.mem.Allocator) Builder {
37 return .{
38 .allocator = allocator,
39 .version = Header.version,
40 .next_id = 1,
41 .capabilities = .empty,
42 .extensions = .empty,
43 .ext_inst_imports = .empty,
44 .memory_model = .empty,
45 .entry_points = .empty,
46 .execution_modes = .empty,
47 .annotations = .empty,
48 .types = .empty,
49 .globals = .empty,
50 .functions = .empty,
51 .glsl_std_450_id = 0,
52 };
53 }
54
55 pub fn deinit(self: *Builder) void {
56 self.capabilities.deinit(self.allocator);
57 self.extensions.deinit(self.allocator);
58 self.ext_inst_imports.deinit(self.allocator);
59 self.memory_model.deinit(self.allocator);
60 self.entry_points.deinit(self.allocator);
61 self.execution_modes.deinit(self.allocator);
62 self.annotations.deinit(self.allocator);
63 self.types.deinit(self.allocator);
64 self.globals.deinit(self.allocator);
65 self.functions.deinit(self.allocator);
66 }
67
68 pub fn importGlslStd450(self: *Builder) !u32 {
69 if (self.glsl_std_450_id != 0) return self.glsl_std_450_id;
70 const id = self.newId();
71 const set_name = "GLSL.std.450";
72 const total_len = set_name.len + 1;
73 const string_words = (total_len + 3) / 4;
74 const word_count: u32 = @intCast(2 + string_words);
75 const first_word: u32 = (word_count << 16) | SpirvOp.ExtInstImport;
76 try self.ext_inst_imports.append(self.allocator, first_word);
77 try self.ext_inst_imports.append(self.allocator, id);
78 try self.emitString(&self.ext_inst_imports, set_name);
79 self.glsl_std_450_id = id;
80 return id;
81 }
82
83 pub fn newId(self: *Builder) u32 {
84 const id = self.next_id;
85 self.next_id += 1;
86 return id;
87 }
88
89 pub fn emit(self: *Builder, section: *Section, opcode: u16, operands: []const u32) !void {
90 const word_count: u32 = @intCast(operands.len + 1);
91 const first_word: u32 = (word_count << 16) | opcode;
92 try section.append(self.allocator, first_word);
93 if (operands.len > 0) {
94 try section.appendSlice(self.allocator, operands);
95 }
96 if (section == &self.functions and isFloatArithmetic(opcode)) {
97 std.debug.assert(operands.len >= 2);
98 try self.emitNoContraction(operands[1]);
99 }
100 }
101
102 pub fn emitNoContraction(self: *Builder, result_id: u32) !void {
103 try self.annotations.appendSlice(self.allocator, &.{
104 (@as(u32, 3) << 16) | SpirvOp.Decorate,
105 result_id,
106 Decoration.NoContraction,
107 });
108 }
109
110 fn isFloatArithmetic(opcode: u16) bool {
111 return switch (opcode) {
112 SpirvOp.FNegate,
113 SpirvOp.FAdd,
114 SpirvOp.FSub,
115 SpirvOp.FMul,
116 SpirvOp.FDiv,
117 SpirvOp.GroupNonUniformFAdd,
118 SpirvOp.GroupNonUniformFMin,
119 SpirvOp.GroupNonUniformFMax,
120 SpirvOp.DPdxFine,
121 SpirvOp.DPdyFine,
122 SpirvOp.FwidthFine,
123 => true,
124 else => false,
125 };
126 }
127
128 fn emitString(self: *Builder, section: *Section, str: []const u8) !void {
129 const total_len = str.len + 1;
130 const word_len = (total_len + 3) / 4;
131 try section.ensureTotalCapacity(self.allocator, section.items.len + word_len);
132
133 var i: usize = 0;
134 while (i < word_len) : (i += 1) {
135 var word: u32 = 0;
136 var j: usize = 0;
137 while (j < 4) : (j += 1) {
138 const idx = i * 4 + j;
139 var byte: u8 = 0;
140 if (idx < str.len) {
141 byte = str[idx];
142 } else if (idx == str.len) {
143 byte = 0;
144 }
145 const shift: u5 = @intCast(8 * j);
146 word |= (@as(u32, byte) << shift);
147 }
148 section.appendAssumeCapacity(word);
149 }
150 }
151
152 pub fn emitCapability(self: *Builder, cap: u32) !void {
153 try self.emit(&self.capabilities, SpirvOp.Capability, &.{cap});
154 }
155
156 pub fn emitExtension(self: *Builder, name: []const u8) !void {
157 const string_words = (name.len + 1 + 3) / 4;
158 const word_count: u32 = @intCast(1 + string_words);
159 try self.extensions.append(self.allocator, (word_count << 16) | SpirvOp.Extension);
160 try self.emitString(&self.extensions, name);
161 }
162
163 pub fn requireVersion(self: *Builder, version: u32) void {
164 if (version > self.version) {
165 self.version = version;
166 }
167 }
168
169 pub fn emitMemoryModel(self: *Builder, addressing_model: u32, memory_model: u32) !void {
170 try self.emit(&self.memory_model, SpirvOp.MemoryModel, &.{ addressing_model, memory_model });
171 }
172
173 pub fn emitEntryPoint(
174 self: *Builder,
175 exec_model: u32,
176 func_id: u32,
177 name: []const u8,
178 interface_ids: []const u32,
179 ) !void {
180 var operands = Section.empty;
181 defer operands.deinit(self.allocator);
182
183 try operands.append(self.allocator, exec_model);
184 try operands.append(self.allocator, func_id);
185
186 try self.emitString(&operands, name);
187 if (interface_ids.len > 0) {
188 try operands.appendSlice(self.allocator, interface_ids);
189 }
190
191 try self.emit(&self.entry_points, SpirvOp.EntryPoint, operands.items);
192 }
193
194 pub fn emitExecutionModeLocalSize(self: *Builder, func_id: u32, x: u32, y: u32, z: u32) !void {
195 try self.emit(&self.execution_modes, SpirvOp.ExecutionMode, &.{
196 func_id,
197 ExecutionMode.LocalSize,
198 x,
199 y,
200 z,
201 });
202 }
203
204 pub fn emitExecutionMode(self: *Builder, func_id: u32, mode: u32) !void {
205 try self.emit(&self.execution_modes, SpirvOp.ExecutionMode, &.{ func_id, mode });
206 }
207
208 pub fn emitFloatExecutionMode(self: *Builder, func_id: u32, mode: u32, width: u32) !void {
209 switch (width) {
210 16, 32, 64 => {},
211 else => unreachable,
212 }
213 try self.emit(&self.execution_modes, SpirvOp.ExecutionMode, &.{ func_id, mode, width });
214 }
215
216 pub fn toWords(self: *Builder, allocator: std.mem.Allocator) ![]u32 {
217 var out: Section = .empty;
218 errdefer out.deinit(allocator);
219
220 try out.appendSlice(allocator, &.{
221 Header.magic,
222 self.version,
223 Header.generator,
224 self.next_id,
225 Header.schema,
226 });
227
228 const sections = [_]*Section{
229 &self.capabilities,
230 &self.extensions,
231 &self.ext_inst_imports,
232 &self.memory_model,
233 &self.entry_points,
234 &self.execution_modes,
235 &self.annotations,
236 &self.types,
237 &self.globals,
238 &self.functions,
239 };
240
241 for (sections) |section| {
242 if (section.items.len > 0) {
243 try out.appendSlice(allocator, section.items);
244 }
245 }
246
247 return out.toOwnedSlice(allocator);
248 }
249 };
250
251 fn containsOpcode(words: []const u32, opcode: u16) bool {
252 var i: usize = 5;
253 while (i < words.len) {
254 const word = words[i];
255 const word_count = word >> 16;
256 const op: u16 = @intCast(word & 0xffff);
257 if (op == opcode) return true;
258 if (word_count == 0) return false;
259 i += word_count;
260 }
261 return false;
262 }
263
264 test "module builder serializes header and ordered sections" {
265 const testing = std.testing;
266 const allocator = testing.allocator;
267
268 var builder = Builder.init(allocator);
269 defer builder.deinit();
270
271 try builder.emitCapability(1);
272 try builder.emitExtension("SPV_KHR_8bit_storage");
273 try builder.emitMemoryModel(0, 1);
274
275 const words = try builder.toWords(allocator);
276 defer allocator.free(words);
277
278 try testing.expectEqual(Header.magic, words[0]);
279 try testing.expectEqual(Header.version, words[1]);
280 try testing.expect(containsOpcode(words, SpirvOp.Capability));
281 try testing.expect(containsOpcode(words, SpirvOp.Extension));
282 try testing.expect(containsOpcode(words, SpirvOp.MemoryModel));
283 }
284
285 test "module builder caches GLSL import id" {
286 const testing = std.testing;
287 const allocator = testing.allocator;
288
289 var builder = Builder.init(allocator);
290 defer builder.deinit();
291
292 const first = try builder.importGlslStd450();
293 const second = try builder.importGlslStd450();
294
295 try testing.expectEqual(first, second);
296 try testing.expectEqual(@as(u32, 1), builder.ext_inst_imports.items[1]);
297 }