1const std = @import("std");
2const Allocator = std.mem.Allocator;
3const BinaryModule = @import("BinaryModule.zig");
4
5const spec = @import("../../codegen/spirv/spec.zig");
6const Word = spec.Word;
7const Id = spec.Id;
8const Opcode = spec.Opcode;
9const Instruction = BinaryModule.Instruction;
10
11/// Deduplicate types and constants in a SPIR-V binary module.
12///
13/// The SPIR-V spec requires that non-aggregate types be unique.
14/// When merging fragments from parallel codegen, duplicate type definitions
15/// may exist. This pass identifies structurally identical types/constants,
16/// keeps one canonical instance, and remaps all references to duplicates.
17pub fn run(parser: *BinaryModule.Parser, binary: *BinaryModule) !void {
18 const gpa = parser.gpa;
19
20 const Decoration = struct { offset: usize, len: usize };
21 var decorations_by_id: std.array_hash_map.Auto(Id, std.ArrayList(Decoration)) = .empty;
22 defer {
23 for (decorations_by_id.values()) |*list| list.deinit(gpa);
24 decorations_by_id.deinit(gpa);
25 }
26
27 var it = binary.iterateInstructions();
28 while (it.next()) |inst| {
29 if (inst.offset >= binary.functions_start) break;
30 switch (inst.opcode) {
31 .OpName, .OpMemberName => continue,
32 else => switch (inst.opcode.class()) {
33 .annotation => {},
34 else => continue,
35 },
36 }
37 if (inst.operands.len == 0) continue;
38 const target_id: Id = @fromBackingInt(@intCast(inst.operands[0]));
39
40 const gop = try decorations_by_id.getOrPut(gpa, target_id);
41 if (!gop.found_existing) gop.value_ptr.* = .empty;
42 try gop.value_ptr.append(gpa, .{
43 .offset = inst.offset,
44 .len = 1 + inst.operands.len,
45 });
46 }
47
48 var canonical_map: std.array_hash_map.Custom(TypeKey, Id, TypeKey.HashContext, true) = .empty;
49 defer {
50 for (canonical_map.keys()) |key| gpa.free(key.words);
51 canonical_map.deinit(gpa);
52 }
53
54 var id_remap: std.AutoHashMapUnmanaged(Id, Id) = .empty;
55 defer id_remap.deinit(gpa);
56
57 var id_offsets: std.ArrayList(u16) = .empty;
58 defer id_offsets.deinit(gpa);
59
60 var key_words: std.ArrayList(Word) = .empty;
61 defer key_words.deinit(gpa);
62
63 var dec_hashes: std.ArrayList(u64) = .empty;
64 defer dec_hashes.deinit(gpa);
65
66 // first pass: build canonical map, identify duplicates
67 it = binary.iterateInstructions();
68 while (it.next()) |inst| {
69 if (inst.offset >= binary.functions_start) break;
70 if (!canDeduplicate(inst.opcode)) continue;
71
72 const result_id_index: usize = switch (inst.opcode.class()) {
73 .type_declaration, .extension => 0,
74 .constant_creation => 1,
75 else => continue,
76 };
77 if (result_id_index >= inst.operands.len) continue;
78 const result_id: Id = @fromBackingInt(@intCast(inst.operands[result_id_index]));
79
80 key_words.items.len = 0;
81 try key_words.append(gpa, @backingInt(inst.opcode));
82
83 id_offsets.items.len = 0;
84 parser.parseInstructionResultIds(binary.*, inst, &id_offsets) catch continue;
85
86 for (inst.operands, 0..) |word, i| {
87 if (i == result_id_index) continue;
88 if (std.mem.findScalar(u16, id_offsets.items, @intCast(i)) != null) {
89 const canonical = id_remap.get(@fromBackingInt(@intCast(word))) orelse @as(Id, @fromBackingInt(@intCast(word)));
90 try key_words.append(gpa, @backingInt(canonical));
91 } else {
92 try key_words.append(gpa, word);
93 }
94 }
95
96 if (decorations_by_id.getPtr(result_id)) |dec_list| {
97 dec_hashes.items.len = 0;
98 for (dec_list.items) |dec| {
99 const dec_words = binary.instructions[dec.offset..][0..dec.len];
100 var hasher = std.hash.Wyhash.init(0);
101 hasher.update(std.mem.asBytes(&dec_words[0]));
102 for (dec_words[2..]) |w| {
103 const w_val = if (id_remap.get(@fromBackingInt(@intCast(w)))) |c| @backingInt(c) else w;
104 hasher.update(std.mem.asBytes(&w_val));
105 }
106 try dec_hashes.append(gpa, hasher.final());
107 }
108 std.mem.sort(u64, dec_hashes.items, {}, std.sort.asc(u64));
109 var prev: u64 = 0;
110 for (dec_hashes.items) |h| {
111 if (h == prev) continue;
112 prev = h;
113 try key_words.append(gpa, @truncate(h));
114 try key_words.append(gpa, @truncate(h >> 32));
115 }
116 }
117
118 const key = TypeKey{ .words = try gpa.dupe(Word, key_words.items) };
119 const gop = try canonical_map.getOrPut(gpa, key);
120 if (gop.found_existing) {
121 try id_remap.put(gpa, result_id, gop.value_ptr.*);
122 gpa.free(key.words);
123 } else {
124 gop.value_ptr.* = result_id;
125 }
126 }
127
128 if (id_remap.count() == 0) return;
129
130 // second pass: rewrite id references, remove duplicates and redundant annotations
131 var new_words: std.ArrayList(Word) = .empty;
132 defer new_words.deinit(gpa);
133 try new_words.ensureTotalCapacity(gpa, binary.instructions.len);
134
135 var emitted_annotations: std.AutoHashMapUnmanaged(u64, void) = .empty;
136 defer emitted_annotations.deinit(gpa);
137
138 var new_functions_offset: ?usize = null;
139 var max_id: Word = 0;
140
141 it = binary.iterateInstructions();
142 while (it.next()) |inst| {
143 if (new_functions_offset == null and inst.offset >= binary.functions_start) {
144 new_functions_offset = new_words.items.len;
145 }
146
147 if (canDeduplicate(inst.opcode)) {
148 const result_id_index: usize = switch (inst.opcode.class()) {
149 .type_declaration, .extension => 0,
150 .constant_creation => 1,
151 else => unreachable,
152 };
153 if (result_id_index < inst.operands.len) {
154 const result_id: Id = @fromBackingInt(@intCast(inst.operands[result_id_index]));
155 if (id_remap.contains(result_id)) continue;
156 }
157 }
158
159 switch (inst.opcode.class()) {
160 .annotation, .debug => {
161 if (inst.operands.len > 0) {
162 const target: Id = @fromBackingInt(@intCast(inst.operands[0]));
163 if (id_remap.contains(target)) continue;
164 }
165 },
166 else => {},
167 }
168
169 const inst_start = new_words.items.len;
170 new_words.appendAssumeCapacity(binary.instructions[inst.offset]);
171 new_words.appendSliceAssumeCapacity(inst.operands);
172 const inst_slice = new_words.items[inst_start + 1 ..];
173
174 id_offsets.items.len = 0;
175 parser.parseInstructionResultIds(binary.*, inst, &id_offsets) catch continue;
176
177 const inst_spec = parser.getInstSpec(inst.opcode);
178 const maybe_result_id_index: ?usize = if (inst_spec) |ispec| blk: {
179 break :blk for (0..@min(2, ispec.operands.len)) |i| {
180 if (ispec.operands[i].kind == .id_result) break @intCast(i);
181 } else null;
182 } else null;
183
184 for (inst_slice, 0..) |*word, i| {
185 if (std.mem.findScalar(u16, id_offsets.items, @intCast(i)) == null) continue;
186 max_id = @max(max_id, word.*);
187 if (maybe_result_id_index != null and i == maybe_result_id_index.?) continue;
188
189 if (id_remap.get(@fromBackingInt(@intCast(word.*)))) |canonical| {
190 word.* = @backingInt(canonical);
191 max_id = @max(max_id, word.*);
192 }
193 }
194
195 switch (inst.opcode.class()) {
196 .annotation, .debug => {
197 const ann_words = new_words.items[inst_start..];
198 const ann_hash = std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(ann_words));
199 const gop = try emitted_annotations.getOrPut(gpa, ann_hash);
200 if (gop.found_existing) {
201 new_words.items.len = inst_start;
202 continue;
203 }
204 },
205 else => {},
206 }
207 }
208
209 var remap_it = id_remap.iterator();
210 while (remap_it.next()) |entry| {
211 _ = binary.ext_inst_map.remove(entry.key_ptr.*);
212 _ = binary.arith_type_width.remove(entry.key_ptr.*);
213 }
214
215 binary.instructions = try gpa.dupe(Word, new_words.items);
216 binary.functions_start = new_functions_offset orelse new_words.items.len;
217 binary.id_bound = max_id + 1;
218}
219
220fn canDeduplicate(opcode: Opcode) bool {
221 return switch (opcode) {
222 .OpTypeForwardPointer => false,
223 .OpGroupDecorate, .OpGroupMemberDecorate => false,
224 else => switch (opcode.class()) {
225 .type_declaration, .constant_creation => true,
226 .extension => opcode == .OpExtInstImport,
227 else => false,
228 },
229 };
230}
231
232const TypeKey = struct {
233 words: []const Word,
234
235 const HashContext = struct {
236 pub fn hash(_: @This(), key: TypeKey) u32 {
237 return @truncate(std.hash.Wyhash.hash(0, std.mem.sliceAsBytes(key.words)));
238 }
239
240 pub fn eql(_: @This(), a: TypeKey, b: TypeKey, _: usize) bool {
241 return std.mem.eql(Word, a.words, b.words);
242 }
243 };
244};