1const std = @import("std");
2const Allocator = std.mem.Allocator;
3const assert = std.debug.assert;
4const log = std.log.scoped(.spirv_link);
5
6const BinaryModule = @import("BinaryModule.zig");
7const Section = @import("../../codegen/spirv/Section.zig");
8const spec = @import("../../codegen/spirv/spec.zig");
9const ResultId = spec.Id;
10const Word = spec.Word;
11
12/// This structure contains all the stuff that we need to parse from the module in
13/// order to run this pass, as well as some functions to ease its use.
14const ModuleInfo = struct {
15 /// Information about a particular function.
16 const Fn = struct {
17 /// The index of the first callee in `callee_store`.
18 first_callee: usize,
19 /// The return type id of this function
20 return_type: ResultId,
21 /// The parameter types of this function
22 param_types: []const ResultId,
23 /// The set of (result-id's of) invocation globals that are accessed
24 /// in this function, or after resolution, that are accessed in this
25 /// function or any of it's callees.
26 invocation_globals: std.array_hash_map.Auto(ResultId, void),
27 };
28
29 /// Information about a particular invocation global
30 const InvocationGlobal = struct {
31 /// The list of invocation globals that this invocation global
32 /// depends on.
33 dependencies: std.array_hash_map.Auto(ResultId, void),
34 /// The invocation global's type
35 ty: ResultId,
36 /// Initializer function. May be `none`.
37 /// Note that if the initializer is `none`, then `dependencies` is empty.
38 initializer: ResultId,
39 };
40
41 /// Maps function result-id -> Fn information structure.
42 functions: std.array_hash_map.Auto(ResultId, Fn),
43 /// Set of OpFunction result-ids in this module.
44 entry_points: std.array_hash_map.Auto(ResultId, void),
45 /// For each function, a list of function result-ids that it calls.
46 callee_store: []const ResultId,
47 /// Maps each invocation global result-id to a type-id.
48 invocation_globals: std.array_hash_map.Auto(ResultId, InvocationGlobal),
49 /// Subset of `invocation_globals` reachable from any entry point.
50 live_invocation_globals: std.array_hash_map.Auto(ResultId, void),
51 /// Initializer functions of unreachable invocation globals. Their
52 /// OpFunction...OpFunctionEnd ranges are skipped during rewriteFunctions.
53 dead_initializers: std.array_hash_map.Auto(ResultId, void),
54
55 /// Fetch the list of callees per function. Guaranteed to contain only unique IDs.
56 fn callees(self: ModuleInfo, fn_id: ResultId) []const ResultId {
57 const fn_index = self.functions.getIndex(fn_id).?;
58 const values = self.functions.values();
59 const first_callee = values[fn_index].first_callee;
60 if (fn_index == values.len - 1) {
61 return self.callee_store[first_callee..];
62 } else {
63 const next_first_callee = values[fn_index + 1].first_callee;
64 return self.callee_store[first_callee..next_first_callee];
65 }
66 }
67
68 /// Extract most of the required information from the binary. The remaining info is
69 /// constructed by `resolve()`.
70 fn parse(
71 arena: Allocator,
72 parser: *BinaryModule.Parser,
73 binary: BinaryModule,
74 ) !ModuleInfo {
75 var entry_points: std.array_hash_map.Auto(ResultId, void) = .empty;
76 var functions: std.array_hash_map.Auto(ResultId, Fn) = .empty;
77 var fn_types = std.AutoHashMap(ResultId, struct {
78 return_type: ResultId,
79 param_types: []const ResultId,
80 }).init(arena);
81 var calls: std.array_hash_map.Auto(ResultId, void) = .empty;
82 var callee_store: std.ArrayList(ResultId) = .empty;
83 var function_invocation_globals: std.array_hash_map.Auto(ResultId, void) = .empty;
84 var result_id_offsets: std.ArrayList(u16) = .empty;
85 var invocation_globals: std.array_hash_map.Auto(ResultId, InvocationGlobal) = .empty;
86
87 var maybe_current_function: ?ResultId = null;
88 var fn_ty_id: ResultId = undefined;
89
90 var it = binary.iterateInstructions();
91 while (it.next()) |inst| {
92 result_id_offsets.items.len = 0;
93 try parser.parseInstructionResultIds(binary, inst, &result_id_offsets);
94
95 switch (inst.opcode) {
96 .OpEntryPoint => {
97 const entry_point: ResultId = @fromBackingInt(@intCast(inst.operands[1]));
98 const entry = try entry_points.getOrPut(arena, entry_point);
99 if (entry.found_existing) {
100 log.err("Entry point type {f} has duplicate definition", .{entry_point});
101 return error.DuplicateId;
102 }
103 },
104 .OpTypeFunction => {
105 const fn_type: ResultId = @fromBackingInt(@intCast(inst.operands[0]));
106 const return_type: ResultId = @fromBackingInt(@intCast(inst.operands[1]));
107 const param_types: []const ResultId = @ptrCast(inst.operands[2..]);
108
109 const entry = try fn_types.getOrPut(fn_type);
110 if (entry.found_existing) {
111 log.err("Function type {f} has duplicate definition", .{fn_type});
112 return error.DuplicateId;
113 }
114
115 entry.value_ptr.* = .{
116 .return_type = return_type,
117 .param_types = param_types,
118 };
119 },
120 .OpExtInst => {
121 // Note: format and set are already verified by parseInstructionResultIds().
122 const global_type: ResultId = @fromBackingInt(@intCast(inst.operands[0]));
123 const result_id: ResultId = @fromBackingInt(@intCast(inst.operands[1]));
124 const set_id: ResultId = @fromBackingInt(@intCast(inst.operands[2]));
125 const set_inst = inst.operands[3];
126
127 const set = binary.ext_inst_map.get(set_id).?;
128 if (set == .zig and set_inst == 0) {
129 const initializer: ResultId = if (inst.operands.len >= 5)
130 @fromBackingInt(@intCast(inst.operands[4]))
131 else
132 .none;
133
134 try invocation_globals.put(arena, result_id, .{
135 .dependencies = .{},
136 .ty = global_type,
137 .initializer = initializer,
138 });
139 }
140 },
141 .OpFunction => {
142 if (maybe_current_function) |current_function| {
143 log.err("OpFunction {f} does not have an OpFunctionEnd", .{current_function});
144 return error.InvalidPhysicalFormat;
145 }
146
147 maybe_current_function = @fromBackingInt(@intCast(inst.operands[1]));
148 fn_ty_id = @fromBackingInt(@intCast(inst.operands[3]));
149 function_invocation_globals.clearRetainingCapacity();
150 },
151 .OpFunctionCall => {
152 const callee: ResultId = @fromBackingInt(@intCast(inst.operands[2]));
153 try calls.put(arena, callee, {});
154 },
155 .OpFunctionEnd => {
156 const current_function = maybe_current_function orelse {
157 log.err("encountered OpFunctionEnd without corresponding OpFunction", .{});
158 return error.InvalidPhysicalFormat;
159 };
160 const entry = try functions.getOrPut(arena, current_function);
161 if (entry.found_existing) {
162 log.err("Function {f} has duplicate definition", .{current_function});
163 return error.DuplicateId;
164 }
165
166 const first_callee = callee_store.items.len;
167 try callee_store.appendSlice(arena, calls.keys());
168
169 const fn_type = fn_types.get(fn_ty_id) orelse {
170 log.err("Function {f} has invalid OpFunction type", .{current_function});
171 return error.InvalidId;
172 };
173
174 entry.value_ptr.* = .{
175 .first_callee = first_callee,
176 .return_type = fn_type.return_type,
177 .param_types = fn_type.param_types,
178 .invocation_globals = try function_invocation_globals.clone(arena),
179 };
180 maybe_current_function = null;
181 calls.clearRetainingCapacity();
182 },
183 else => {},
184 }
185
186 for (result_id_offsets.items) |off| {
187 const result_id: ResultId = @fromBackingInt(@intCast(inst.operands[off]));
188 if (invocation_globals.contains(result_id)) {
189 try function_invocation_globals.put(arena, result_id, {});
190 }
191 }
192 }
193
194 if (maybe_current_function) |current_function| {
195 log.err("OpFunction {f} does not have an OpFunctionEnd", .{current_function});
196 return error.InvalidPhysicalFormat;
197 }
198
199 return .{
200 .functions = functions,
201 .entry_points = entry_points,
202 .callee_store = callee_store.items,
203 .invocation_globals = invocation_globals,
204 .live_invocation_globals = .empty,
205 .dead_initializers = .empty,
206 };
207 }
208
209 /// Derive the remaining info from the structures filled in by parsing.
210 fn resolve(self: *ModuleInfo, arena: Allocator) !void {
211 try self.resolveInvocationGlobalUsage(arena);
212 try self.resolveInvocationGlobalDependencies(arena);
213 try self.resolveLiveSet(arena);
214 }
215
216 fn resolveLiveSet(self: *ModuleInfo, arena: Allocator) !void {
217 for (self.entry_points.keys()) |ep_id| {
218 const ep_info = self.functions.get(ep_id) orelse continue;
219 for (ep_info.invocation_globals.keys()) |g| {
220 try self.live_invocation_globals.put(arena, g, {});
221 const g_info = self.invocation_globals.get(g).?;
222 for (g_info.dependencies.keys()) |dep| {
223 try self.live_invocation_globals.put(arena, dep, {});
224 }
225 }
226 }
227 for (self.invocation_globals.keys(), self.invocation_globals.values()) |g, info| {
228 if (info.initializer == .none) continue;
229 if (self.live_invocation_globals.contains(g)) continue;
230 try self.dead_initializers.put(arena, info.initializer, {});
231 }
232 }
233
234 /// For each function, extend the list of `invocation_globals` with the
235 /// invocation globals that ALL of its dependencies use.
236 fn resolveInvocationGlobalUsage(self: *ModuleInfo, arena: Allocator) !void {
237 var seen: std.bit_set.Dynamic = try .initEmpty(arena, self.functions.count());
238
239 for (self.functions.keys()) |id| {
240 try self.resolveInvocationGlobalUsageStep(arena, id, &seen);
241 }
242 }
243
244 fn resolveInvocationGlobalUsageStep(
245 self: *ModuleInfo,
246 arena: Allocator,
247 id: ResultId,
248 seen: *std.bit_set.Dynamic,
249 ) !void {
250 const index = self.functions.getIndex(id) orelse {
251 log.err("function calls invalid function {f}", .{id});
252 return error.InvalidId;
253 };
254
255 if (seen.isSet(index)) {
256 return;
257 }
258 seen.set(index);
259
260 const info = &self.functions.values()[index];
261 for (self.callees(id)) |callee| {
262 try self.resolveInvocationGlobalUsageStep(arena, callee, seen);
263 const callee_info = self.functions.get(callee).?;
264 for (callee_info.invocation_globals.keys()) |global| {
265 try info.invocation_globals.put(arena, global, {});
266 }
267 }
268 }
269
270 /// For each invocation global, populate and fully resolve the `dependencies` set.
271 /// This requires `resolveInvocationGlobalUsage()` to be already done.
272 fn resolveInvocationGlobalDependencies(
273 self: *ModuleInfo,
274 arena: Allocator,
275 ) !void {
276 var seen: std.bit_set.Dynamic = try .initEmpty(arena, self.invocation_globals.count());
277
278 for (self.invocation_globals.keys()) |id| {
279 try self.resolveInvocationGlobalDependenciesStep(arena, id, &seen);
280 }
281 }
282
283 fn resolveInvocationGlobalDependenciesStep(
284 self: *ModuleInfo,
285 arena: Allocator,
286 id: ResultId,
287 seen: *std.bit_set.Dynamic,
288 ) !void {
289 const index = self.invocation_globals.getIndex(id) orelse {
290 log.err("invalid invocation global {f}", .{id});
291 return error.InvalidId;
292 };
293
294 if (seen.isSet(index)) {
295 return;
296 }
297 seen.set(index);
298
299 const info = &self.invocation_globals.values()[index];
300 if (info.initializer == .none) {
301 return;
302 }
303
304 const initializer = self.functions.get(info.initializer) orelse {
305 log.err("invocation global {f} has invalid initializer {f}", .{ id, info.initializer });
306 return error.InvalidId;
307 };
308
309 for (initializer.invocation_globals.keys()) |dependency| {
310 if (dependency == id) {
311 // The set of invocation global dependencies includes the dependency itself,
312 // so we need to skip that case.
313 continue;
314 }
315
316 try info.dependencies.put(arena, dependency, {});
317 try self.resolveInvocationGlobalDependenciesStep(arena, dependency, seen);
318
319 const dep_info = self.invocation_globals.getPtr(dependency).?;
320
321 for (dep_info.dependencies.keys()) |global| {
322 try info.dependencies.put(arena, global, {});
323 }
324 }
325 }
326};
327
328const ModuleBuilder = struct {
329 const FunctionType = struct {
330 return_type: ResultId,
331 param_types: []const ResultId,
332
333 const Context = struct {
334 pub fn hash(_: @This(), ty: FunctionType) u32 {
335 var hasher = std.hash.Wyhash.init(0);
336 hasher.update(std.mem.asBytes(&ty.return_type));
337 hasher.update(std.mem.sliceAsBytes(ty.param_types));
338 return @truncate(hasher.final());
339 }
340
341 pub fn eql(_: @This(), a: FunctionType, b: FunctionType, _: usize) bool {
342 if (a.return_type != b.return_type) return false;
343 return std.mem.eql(ResultId, a.param_types, b.param_types);
344 }
345 };
346 };
347
348 const FunctionNewInfo = struct {
349 /// This is here just so that we don't need to allocate the new
350 /// param_types multiple times.
351 new_function_type: ResultId,
352 /// The first ID of the parameters for the invocation globals.
353 /// Each global is allocate here according to the index in
354 /// `ModuleInfo.Fn.invocation_globals`.
355 global_id_base: u32,
356
357 fn invocationGlobalId(self: FunctionNewInfo, index: usize) ResultId {
358 return @fromBackingInt(@intCast(self.global_id_base + @as(u32, @intCast(index))));
359 }
360 };
361
362 arena: Allocator,
363 section: Section,
364 /// The ID bound of the new module.
365 id_bound: u32,
366 /// The first ID of the new entry points. Entry points are allocated from
367 /// here according to their index in `info.entry_points`.
368 entry_point_new_id_base: u32,
369 /// OpName operands saved for invocation globals to re-emit.
370 global_names: std.array_hash_map.Auto(ResultId, []const Word) = .empty,
371 /// A set of all function types in the new program. SPIR-V mandates that these are unique,
372 /// and until a general type deduplication pass is programmed, we just handle it here via this.
373 function_types: std.array_hash_map.Custom(FunctionType, ResultId, FunctionType.Context, true) = .empty,
374 /// Maps functions to new information required for creating the module
375 function_new_info: std.array_hash_map.Auto(ResultId, FunctionNewInfo) = .empty,
376 /// Offset of the functions section in the new binary.
377 new_functions_section: ?usize,
378
379 fn init(arena: Allocator, binary: BinaryModule, info: ModuleInfo) !ModuleBuilder {
380 var self = ModuleBuilder{
381 .arena = arena,
382 .section = .{},
383 .id_bound = binary.id_bound,
384 .entry_point_new_id_base = undefined,
385 .new_functions_section = null,
386 };
387 self.entry_point_new_id_base = @backingInt(self.allocIds(@intCast(info.entry_points.count())));
388 return self;
389 }
390
391 fn allocId(self: *ModuleBuilder) ResultId {
392 return self.allocIds(1);
393 }
394
395 fn allocIds(self: *ModuleBuilder, n: u32) ResultId {
396 defer self.id_bound += n;
397 return @fromBackingInt(@intCast(self.id_bound));
398 }
399
400 fn finalize(self: *ModuleBuilder, arena: Allocator, binary: *BinaryModule) !void {
401 binary.id_bound = self.id_bound;
402 binary.instructions = try arena.dupe(Word, self.section.instructions.items);
403 // Nothing is removed in this pass so we don't need to change any of the maps,
404 // just make sure the section is updated.
405 binary.functions_start = self.new_functions_section orelse binary.instructions.len;
406 }
407
408 fn emitGlobalNames(self: *ModuleBuilder, info: ModuleInfo) !void {
409 for (info.functions.keys(), info.functions.values()) |func, fn_info| {
410 if (info.dead_initializers.contains(func)) continue;
411 const new_info = self.function_new_info.get(func) orelse continue;
412 for (fn_info.invocation_globals.keys(), 0..) |global, i| {
413 if (!info.live_invocation_globals.contains(global)) continue;
414 const name_words = self.global_names.get(global) orelse continue;
415 const id = new_info.invocationGlobalId(i);
416 try self.section.emitRaw(self.arena, .OpName, 1 + name_words.len);
417 self.section.writeOperand(ResultId, id);
418 self.section.writeWords(name_words);
419 }
420 }
421 }
422
423 /// Process everything from `binary` up to the first function and emit it into the builder.
424 fn processPreamble(self: *ModuleBuilder, binary: BinaryModule, info: ModuleInfo) !void {
425 var emitted_global_names = false;
426 var it = binary.iterateInstructions();
427 while (it.next()) |inst| {
428 if (!emitted_global_names) switch (inst.opcode.class()) {
429 .annotation, .type_declaration, .constant_creation => {
430 try self.emitGlobalNames(info);
431 emitted_global_names = true;
432 },
433 else => {},
434 };
435
436 switch (inst.opcode) {
437 .OpName => {
438 const id: ResultId = @fromBackingInt(@intCast(inst.operands[0]));
439 if (info.invocation_globals.contains(id)) {
440 try self.global_names.put(self.arena, id, inst.operands[1..]);
441 continue;
442 }
443 if (info.dead_initializers.contains(id)) continue;
444 },
445 .OpExtInstImport => {
446 const set_id: ResultId = @fromBackingInt(@intCast(inst.operands[0]));
447 const set = binary.ext_inst_map.get(set_id).?;
448 if (set == .zig) continue;
449 },
450 .OpExtInst => {
451 const set_id: ResultId = @fromBackingInt(@intCast(inst.operands[2]));
452 const set_inst = inst.operands[3];
453 const set = binary.ext_inst_map.get(set_id).?;
454 if (set == .zig and set_inst == 0) {
455 continue;
456 }
457 },
458 .OpEntryPoint => {
459 const original_id: ResultId = @fromBackingInt(@intCast(inst.operands[1]));
460 const fn_info = info.functions.get(original_id).?;
461 if (fn_info.invocation_globals.count() > 0) {
462 const new_id_index = info.entry_points.getIndex(original_id).?;
463 const new_id: ResultId = @fromBackingInt(@intCast(self.entry_point_new_id_base + new_id_index));
464 try self.section.emitRaw(self.arena, .OpEntryPoint, inst.operands.len);
465 self.section.writeWord(inst.operands[0]);
466 self.section.writeOperand(ResultId, new_id);
467 self.section.writeWords(inst.operands[2..]);
468 } else {
469 try self.section.emitRawInstruction(self.arena, inst.opcode, inst.operands);
470 }
471 continue;
472 },
473 .OpExecutionMode, .OpExecutionModeId => {
474 const original_id: ResultId = @fromBackingInt(@intCast(inst.operands[0]));
475 const fn_info = info.functions.get(original_id).?;
476 if (fn_info.invocation_globals.count() > 0) {
477 const new_id_index = info.entry_points.getIndex(original_id).?;
478 const new_id: ResultId = @fromBackingInt(@intCast(self.entry_point_new_id_base + new_id_index));
479 try self.section.emitRaw(self.arena, inst.opcode, inst.operands.len);
480 self.section.writeOperand(ResultId, new_id);
481 self.section.writeWords(inst.operands[1..]);
482 } else {
483 try self.section.emitRawInstruction(self.arena, inst.opcode, inst.operands);
484 }
485 continue;
486 },
487 .OpTypeFunction => {
488 continue;
489 },
490 .OpFunction => break,
491 else => {},
492 }
493
494 try self.section.emitRawInstruction(self.arena, inst.opcode, inst.operands);
495 }
496
497 if (!emitted_global_names) {
498 try self.emitGlobalNames(info);
499 }
500 }
501
502 /// Derive new information required for further emitting this module,
503 fn deriveNewFnInfo(self: *ModuleBuilder, info: ModuleInfo) !void {
504 for (info.functions.keys(), info.functions.values()) |func, fn_info| {
505 const invocation_global_count = fn_info.invocation_globals.count();
506 const new_param_types = try self.arena.alloc(ResultId, fn_info.param_types.len + invocation_global_count);
507 for (fn_info.invocation_globals.keys(), 0..) |global, i| {
508 new_param_types[i] = info.invocation_globals.get(global).?.ty;
509 }
510 @memcpy(new_param_types[invocation_global_count..], fn_info.param_types);
511
512 const new_type = try self.internFunctionType(fn_info.return_type, new_param_types);
513 try self.function_new_info.put(self.arena, func, .{
514 .new_function_type = new_type,
515 .global_id_base = @backingInt(self.allocIds(@intCast(invocation_global_count))),
516 });
517 }
518 }
519
520 /// Emit the new function types, which include the parameters for the invocation globals.
521 /// Currently, this function re-emits ALL function types to ensure that there are
522 /// no duplicates in the final program.
523 /// TODO: The above should be resolved by a generalized deduplication pass, and then
524 /// we only need to emit the new function pointers type here.
525 fn emitFunctionTypes(self: *ModuleBuilder, info: ModuleInfo) !void {
526 // TODO: Handle decorators. Function types usually don't have those
527 // though, but stuff like OpName could be a possibility.
528
529 // Entry points retain their old function type, so make sure to emit
530 // those in the `function_types` set.
531 for (info.entry_points.keys()) |func| {
532 const fn_info = info.functions.get(func).?;
533 _ = try self.internFunctionType(fn_info.return_type, fn_info.param_types);
534 }
535
536 for (self.function_types.keys(), self.function_types.values()) |fn_type, result_id| {
537 try self.section.emit(self.arena, .OpTypeFunction, .{
538 .id_result = result_id,
539 .return_type = fn_type.return_type,
540 .id_ref_2 = fn_type.param_types,
541 });
542 }
543 }
544
545 fn internFunctionType(self: *ModuleBuilder, return_type: ResultId, param_types: []const ResultId) !ResultId {
546 const entry = try self.function_types.getOrPut(self.arena, .{
547 .return_type = return_type,
548 .param_types = param_types,
549 });
550
551 if (!entry.found_existing) {
552 const new_id = self.allocId();
553 entry.value_ptr.* = new_id;
554 }
555
556 return entry.value_ptr.*;
557 }
558
559 /// Rewrite the modules functions and emit them with the new parameter types.
560 fn rewriteFunctions(
561 self: *ModuleBuilder,
562 parser: *BinaryModule.Parser,
563 binary: BinaryModule,
564 info: ModuleInfo,
565 ) !void {
566 var result_id_offsets: std.ArrayList(u16) = .empty;
567 var operands: std.ArrayList(u32) = .empty;
568
569 var maybe_current_function: ?ResultId = null;
570 var skip_until_end: bool = false;
571 var it = binary.iterateInstructionsFrom(binary.functions_start);
572 self.new_functions_section = self.section.instructions.items.len;
573 while (it.next()) |inst| {
574 if (skip_until_end) {
575 if (inst.opcode == .OpFunctionEnd) skip_until_end = false;
576 continue;
577 }
578 result_id_offsets.items.len = 0;
579 try parser.parseInstructionResultIds(binary, inst, &result_id_offsets);
580
581 operands.items.len = 0;
582 try operands.appendSlice(self.arena, inst.operands);
583
584 // Replace the result-ids with the global's new result-id if required.
585 for (result_id_offsets.items) |off| {
586 const result_id: ResultId = @fromBackingInt(@intCast(operands.items[off]));
587 if (info.invocation_globals.contains(result_id)) {
588 const func = maybe_current_function.?;
589 const new_info = self.function_new_info.get(func).?;
590 const fn_info = info.functions.get(func).?;
591 const index = fn_info.invocation_globals.getIndex(result_id).?;
592 operands.items[off] = @backingInt(new_info.invocationGlobalId(index));
593 }
594 }
595
596 switch (inst.opcode) {
597 .OpFunction => {
598 // Re-declare the function with the new parameters.
599 const func: ResultId = @fromBackingInt(@intCast(operands.items[1]));
600 if (info.dead_initializers.contains(func)) {
601 skip_until_end = true;
602 continue;
603 }
604 const fn_info = info.functions.get(func).?;
605 const new_info = self.function_new_info.get(func).?;
606
607 try self.section.emitRaw(self.arena, .OpFunction, 4);
608 self.section.writeOperand(ResultId, fn_info.return_type);
609 self.section.writeOperand(ResultId, func);
610 self.section.writeWord(operands.items[2]);
611 self.section.writeOperand(ResultId, new_info.new_function_type);
612
613 // Emit the OpFunctionParameters for the invocation globals. The functions
614 // actual parameters are emitted unchanged from their original form, so
615 // we don't need to handle those here.
616
617 for (fn_info.invocation_globals.keys(), 0..) |global, index| {
618 const ty = info.invocation_globals.get(global).?.ty;
619 const id = new_info.invocationGlobalId(index);
620 try self.section.emit(self.arena, .OpFunctionParameter, .{
621 .id_result_type = ty,
622 .id_result = id,
623 });
624 }
625
626 maybe_current_function = func;
627 },
628 .OpFunctionCall => {
629 // Add the required invocation globals to the function's new parameter list.
630 const caller = maybe_current_function.?;
631 const callee: ResultId = @fromBackingInt(@intCast(operands.items[2]));
632 const caller_info = info.functions.get(caller).?;
633 const callee_info = info.functions.get(callee).?;
634 const caller_new_info = self.function_new_info.get(caller).?;
635 const total_params = callee_info.invocation_globals.count() + callee_info.param_types.len;
636
637 try self.section.emitRaw(self.arena, .OpFunctionCall, 3 + total_params);
638 self.section.writeWord(operands.items[0]); // Copy result type-id
639 self.section.writeWord(operands.items[1]); // Copy result-id
640 self.section.writeOperand(ResultId, callee);
641
642 // Add the new arguments
643 for (callee_info.invocation_globals.keys()) |global| {
644 const caller_global_index = caller_info.invocation_globals.getIndex(global).?;
645 const id = caller_new_info.invocationGlobalId(caller_global_index);
646 self.section.writeOperand(ResultId, id);
647 }
648
649 // Add the original arguments
650 self.section.writeWords(operands.items[3..]);
651 },
652 else => {
653 try self.section.emitRawInstruction(self.arena, inst.opcode, operands.items);
654 },
655 }
656 }
657 }
658
659 fn emitNewEntryPoints(self: *ModuleBuilder, info: ModuleInfo) !void {
660 const arena = self.arena;
661 var all_function_invocation_globals: std.array_hash_map.Auto(ResultId, void) = .empty;
662
663 for (info.entry_points.keys(), 0..) |func, entry_point_index| {
664 const fn_info = info.functions.get(func).?;
665 if (fn_info.invocation_globals.count() == 0) continue;
666 const ep_id: ResultId = @fromBackingInt(@intCast(self.entry_point_new_id_base + @as(u32, @intCast(entry_point_index))));
667 const fn_type = self.function_types.get(.{
668 .return_type = fn_info.return_type,
669 .param_types = fn_info.param_types,
670 }).?;
671
672 try self.section.emit(arena, .OpFunction, .{
673 .id_result_type = fn_info.return_type,
674 .id_result = ep_id,
675 .function_control = .{}, // TODO: Copy the attributes from the original function maybe?
676 .function_type = fn_type,
677 });
678
679 // Emit OpFunctionParameter instructions for the original kernel's parameters.
680 const params_id_base: u32 = @backingInt(self.allocIds(@intCast(fn_info.param_types.len)));
681 for (fn_info.param_types, 0..) |param_type, i| {
682 const id: ResultId = @fromBackingInt(@intCast(params_id_base + @as(u32, @intCast(i))));
683 try self.section.emit(arena, .OpFunctionParameter, .{
684 .id_result_type = param_type,
685 .id_result = id,
686 });
687 }
688
689 try self.section.emit(arena, .OpLabel, .{
690 .id_result = self.allocId(),
691 });
692
693 // Besides the IDs of the main kernel, we also need the
694 // dependencies of the globals.
695 // Just quickly construct that set here.
696 all_function_invocation_globals.clearRetainingCapacity();
697 for (fn_info.invocation_globals.keys()) |global| {
698 try all_function_invocation_globals.put(arena, global, {});
699 const global_info = info.invocation_globals.get(global).?;
700 for (global_info.dependencies.keys()) |dependency| {
701 try all_function_invocation_globals.put(arena, dependency, {});
702 }
703 }
704
705 // Declare the IDs of the invocation globals.
706 const global_id_base: u32 = @backingInt(self.allocIds(@intCast(all_function_invocation_globals.count())));
707 for (all_function_invocation_globals.keys(), 0..) |global, i| {
708 const global_info = info.invocation_globals.get(global).?;
709
710 const id: ResultId = @fromBackingInt(@intCast(global_id_base + @as(u32, @intCast(i))));
711 try self.section.emit(arena, .OpVariable, .{
712 .id_result_type = global_info.ty,
713 .id_result = id,
714 .storage_class = .function,
715 .initializer = null,
716 });
717 }
718
719 // Call initializers for invocation globals that need it
720 for (all_function_invocation_globals.keys()) |global| {
721 const global_info = info.invocation_globals.get(global).?;
722 if (global_info.initializer == .none) continue;
723
724 const initializer_info = info.functions.get(global_info.initializer).?;
725 assert(initializer_info.param_types.len == 0);
726
727 try self.callWithGlobalsAndLinearParams(
728 &all_function_invocation_globals,
729 global_info.initializer,
730 initializer_info,
731 global_id_base,
732 undefined,
733 );
734 }
735
736 // Call the main kernel entry
737 try self.callWithGlobalsAndLinearParams(
738 &all_function_invocation_globals,
739 func,
740 fn_info,
741 global_id_base,
742 params_id_base,
743 );
744
745 try self.section.emit(arena, .OpReturn, {});
746 try self.section.emit(arena, .OpFunctionEnd, {});
747 }
748 }
749
750 fn callWithGlobalsAndLinearParams(
751 self: *ModuleBuilder,
752 all_globals: *const std.array_hash_map.Auto(ResultId, void),
753 func: ResultId,
754 callee_info: ModuleInfo.Fn,
755 global_id_base: u32,
756 params_id_base: u32,
757 ) !void {
758 const total_arguments = callee_info.invocation_globals.count() + callee_info.param_types.len;
759 try self.section.emitRaw(self.arena, .OpFunctionCall, 3 + total_arguments);
760 self.section.writeOperand(ResultId, callee_info.return_type);
761 self.section.writeOperand(ResultId, self.allocId());
762 self.section.writeOperand(ResultId, func);
763
764 // Add the invocation globals
765 for (callee_info.invocation_globals.keys()) |global| {
766 const index = all_globals.getIndex(global).?;
767 const id: ResultId = @fromBackingInt(@intCast(global_id_base + @as(u32, @intCast(index))));
768 self.section.writeOperand(ResultId, id);
769 }
770
771 // Add the arguments
772 for (0..callee_info.param_types.len) |index| {
773 const id: ResultId = @fromBackingInt(@intCast(params_id_base + @as(u32, @intCast(index))));
774 self.section.writeOperand(ResultId, id);
775 }
776 }
777};
778
779pub fn run(parser: *BinaryModule.Parser, binary: *BinaryModule, progress: std.Progress.Node) !void {
780 const sub_node = progress.start("Lower invocation globals", 6);
781 defer sub_node.end();
782
783 var arena_state = std.heap.ArenaAllocator.init(parser.gpa);
784 defer arena_state.deinit();
785 const arena = arena_state.allocator();
786
787 var info = try ModuleInfo.parse(arena, parser, binary.*);
788 try info.resolve(arena);
789
790 var builder = try ModuleBuilder.init(arena, binary.*, info);
791 sub_node.completeOne();
792 try builder.deriveNewFnInfo(info);
793 sub_node.completeOne();
794 try builder.processPreamble(binary.*, info);
795 sub_node.completeOne();
796 try builder.emitFunctionTypes(info);
797 sub_node.completeOne();
798 try builder.rewriteFunctions(parser, binary.*, info);
799 sub_node.completeOne();
800 try builder.emitNewEntryPoints(info);
801 sub_node.completeOne();
802 try builder.finalize(parser.gpa, binary);
803}