| ... | ... | @@ -0,0 +1,400 @@ |
| 1 | const std = @import("std"); |
| 2 | const Allocator = std.mem.Allocator; |
| 3 | const log = std.log.scoped(.spirv_link); |
| 4 | const assert = std.debug.assert; |
| 5 | |
| 6 | const BinaryModule = @import("BinaryModule.zig"); |
| 7 | const Section = @import("../../codegen/spirv/Section.zig"); |
| 8 | const spec = @import("../../codegen/spirv/spec.zig"); |
| 9 | const Opcode = spec.Opcode; |
| 10 | const ResultId = spec.IdResult; |
| 11 | const Word = spec.Word; |
| 12 | |
| 13 | fn canDeduplicate(opcode: Opcode) bool { |
| 14 | return switch (opcode) { |
| 15 | .OpTypeForwardPointer => false, // Don't need to handle these |
| 16 | .OpGroupDecorate, .OpGroupMemberDecorate => { |
| 17 | // These are deprecated, so don't bother supporting them for now. |
| 18 | return false; |
| 19 | }, |
| 20 | .OpName, .OpMemberName => true, // Debug decoration-style instructions |
| 21 | else => switch (opcode.class()) { |
| 22 | .TypeDeclaration, |
| 23 | .ConstantCreation, |
| 24 | .Annotation, |
| 25 | => true, |
| 26 | else => false, |
| 27 | }, |
| 28 | }; |
| 29 | } |
| 30 | |
| 31 | const ModuleInfo = struct { |
| 32 | /// This models a type, decoration or constant instruction |
| 33 | /// and its dependencies. |
| 34 | const Entity = struct { |
| 35 | /// The type that this entity represents. This is just |
| 36 | /// the instruction opcode. |
| 37 | kind: Opcode, |
| 38 | /// Offset of first child result-id, stored in entity_children. |
| 39 | /// These are the shallow entities appearing directly in the |
| 40 | /// type's instruction. |
| 41 | first_child: u32, |
| 42 | /// Offset to the first word of extra-data: Data in the instruction |
| 43 | /// that must be considered for uniqueness, but doesn't include |
| 44 | /// any IDs. |
| 45 | first_extra_data: u32, |
| 46 | }; |
| 47 | |
| 48 | /// Maps result-id to Entity's |
| 49 | entities: std.AutoArrayHashMapUnmanaged(ResultId, Entity), |
| 50 | /// The list of children per instruction. |
| 51 | entity_children: []const ResultId, |
| 52 | /// The list of extra data per instruction. |
| 53 | /// TODO: This is a bit awkward, maybe we need to store it some |
| 54 | /// other way? |
| 55 | extra_data: []const u32, |
| 56 | |
| 57 | pub fn parse( |
| 58 | arena: Allocator, |
| 59 | parser: *BinaryModule.Parser, |
| 60 | binary: BinaryModule, |
| 61 | ) !ModuleInfo { |
| 62 | var entities = std.AutoArrayHashMap(ResultId, Entity).init(arena); |
| 63 | var entity_children = std.ArrayList(ResultId).init(arena); |
| 64 | var extra_data = std.ArrayList(u32).init(arena); |
| 65 | var id_offsets = std.ArrayList(u16).init(arena); |
| 66 | |
| 67 | var it = binary.iterateInstructions(); |
| 68 | while (it.next()) |inst| { |
| 69 | if (inst.opcode == .OpFunction) break; // No more declarations are possible |
| 70 | if (!canDeduplicate(inst.opcode)) continue; |
| 71 | |
| 72 | id_offsets.items.len = 0; |
| 73 | try parser.parseInstructionResultIds(binary, inst, &id_offsets); |
| 74 | |
| 75 | const result_id_index: u32 = switch (inst.opcode.class()) { |
| 76 | .TypeDeclaration, .Annotation, .Debug => 0, |
| 77 | .ConstantCreation => 1, |
| 78 | else => unreachable, |
| 79 | }; |
| 80 | |
| 81 | const result_id: ResultId = @enumFromInt(inst.operands[id_offsets.items[result_id_index]]); |
| 82 | |
| 83 | const first_child: u32 = @intCast(entity_children.items.len); |
| 84 | const first_extra_data: u32 = @intCast(extra_data.items.len); |
| 85 | |
| 86 | try entity_children.ensureUnusedCapacity(id_offsets.items.len - 1); |
| 87 | try extra_data.ensureUnusedCapacity(inst.operands.len - id_offsets.items.len); |
| 88 | |
| 89 | var id_i: usize = 0; |
| 90 | for (inst.operands, 0..) |operand, i| { |
| 91 | assert(id_i == id_offsets.items.len or id_offsets.items[id_i] >= i); |
| 92 | if (id_i != id_offsets.items.len and id_offsets.items[id_i] == i) { |
| 93 | // Skip .IdResult / .IdResultType. |
| 94 | if (id_i != result_id_index) { |
| 95 | entity_children.appendAssumeCapacity(@enumFromInt(operand)); |
| 96 | } |
| 97 | id_i += 1; |
| 98 | } else { |
| 99 | // Non-id operand, add it to extra data. |
| 100 | extra_data.appendAssumeCapacity(operand); |
| 101 | } |
| 102 | } |
| 103 | |
| 104 | switch (inst.opcode.class()) { |
| 105 | .Annotation, .Debug => { |
| 106 | // TODO |
| 107 | }, |
| 108 | .TypeDeclaration, .ConstantCreation => { |
| 109 | const entry = try entities.getOrPut(result_id); |
| 110 | if (entry.found_existing) { |
| 111 | log.err("type or constant {} has duplicate definition", .{result_id}); |
| 112 | return error.DuplicateId; |
| 113 | } |
| 114 | entry.value_ptr.* = .{ |
| 115 | .kind = inst.opcode, |
| 116 | .first_child = first_child, |
| 117 | .first_extra_data = first_extra_data, |
| 118 | }; |
| 119 | }, |
| 120 | else => unreachable, |
| 121 | } |
| 122 | } |
| 123 | |
| 124 | return ModuleInfo{ |
| 125 | .entities = entities.unmanaged, |
| 126 | .entity_children = entity_children.items, |
| 127 | .extra_data = extra_data.items, |
| 128 | }; |
| 129 | } |
| 130 | |
| 131 | /// Fetch a slice of children for the index corresponding to an entity. |
| 132 | fn childrenByIndex(self: ModuleInfo, index: usize) []const ResultId { |
| 133 | const values = self.entities.values(); |
| 134 | const first_child = values[index].first_child; |
| 135 | if (index == values.len - 1) { |
| 136 | return self.entity_children[first_child..]; |
| 137 | } else { |
| 138 | const next_first_child = values[index + 1].first_child; |
| 139 | return self.entity_children[first_child..next_first_child]; |
| 140 | } |
| 141 | } |
| 142 | |
| 143 | /// Fetch the slice of extra-data for the index corresponding to an entity. |
| 144 | fn extraDataByIndex(self: ModuleInfo, index: usize) []const u32 { |
| 145 | const values = self.entities.values(); |
| 146 | const first_extra_data = values[index].first_extra_data; |
| 147 | if (index == values.len - 1) { |
| 148 | return self.extra_data[first_extra_data..]; |
| 149 | } else { |
| 150 | const next_extra_data = values[index + 1].first_extra_data; |
| 151 | return self.extra_data[first_extra_data..next_extra_data]; |
| 152 | } |
| 153 | } |
| 154 | }; |
| 155 | |
| 156 | const EntityContext = struct { |
| 157 | a: Allocator, |
| 158 | ptr_map_a: std.AutoArrayHashMapUnmanaged(ResultId, void) = .{}, |
| 159 | ptr_map_b: std.AutoArrayHashMapUnmanaged(ResultId, void) = .{}, |
| 160 | info: *const ModuleInfo, |
| 161 | |
| 162 | fn init(a: Allocator, info: *const ModuleInfo) EntityContext { |
| 163 | return .{ |
| 164 | .a = a, |
| 165 | .info = info, |
| 166 | }; |
| 167 | } |
| 168 | |
| 169 | fn deinit(self: *EntityContext) void { |
| 170 | self.ptr_map_a.deinit(self.a); |
| 171 | self.ptr_map_b.deinit(self.a); |
| 172 | |
| 173 | self.* = undefined; |
| 174 | } |
| 175 | |
| 176 | fn equalizeMapCapacity(self: *EntityContext) !void { |
| 177 | const cap = @max(self.ptr_map_a.capacity(), self.ptr_map_b.capacity()); |
| 178 | try self.ptr_map_a.ensureTotalCapacity(self.a, cap); |
| 179 | try self.ptr_map_b.ensureTotalCapacity(self.a, cap); |
| 180 | } |
| 181 | |
| 182 | fn hash(self: *EntityContext, id: ResultId) !u64 { |
| 183 | var hasher = std.hash.Wyhash.init(0); |
| 184 | self.ptr_map_a.clearRetainingCapacity(); |
| 185 | try self.hashInner(&hasher, id); |
| 186 | return hasher.final(); |
| 187 | } |
| 188 | |
| 189 | fn hashInner(self: *EntityContext, hasher: *std.hash.Wyhash, id: ResultId) !void { |
| 190 | const index = self.info.entities.getIndex(id).?; |
| 191 | const entity = self.info.entities.values()[index]; |
| 192 | |
| 193 | std.hash.autoHash(hasher, entity.kind); |
| 194 | if (entity.kind == .OpTypePointer) { |
| 195 | // This may be either a pointer that is forward-referenced in the future, |
| 196 | // or a forward reference to a pointer. |
| 197 | const entry = try self.ptr_map_a.getOrPut(self.a, id); |
| 198 | if (entry.found_existing) { |
| 199 | // Pointer already seen. Hash the index instead of recursing into its children. |
| 200 | // TODO: Discriminate this path somehow? |
| 201 | std.hash.autoHash(hasher, entry.index); |
| 202 | return; |
| 203 | } |
| 204 | } |
| 205 | |
| 206 | // Hash extra data |
| 207 | for (self.info.extraDataByIndex(index)) |data| { |
| 208 | std.hash.autoHash(hasher, data); |
| 209 | } |
| 210 | |
| 211 | // Hash children |
| 212 | for (self.info.childrenByIndex(index)) |child| { |
| 213 | try self.hashInner(hasher, child); |
| 214 | } |
| 215 | } |
| 216 | |
| 217 | fn eql(self: *EntityContext, a: ResultId, b: ResultId) !bool { |
| 218 | self.ptr_map_a.clearRetainingCapacity(); |
| 219 | self.ptr_map_b.clearRetainingCapacity(); |
| 220 | |
| 221 | return try self.eqlInner(a, b); |
| 222 | } |
| 223 | |
| 224 | fn eqlInner(self: *EntityContext, id_a: ResultId, id_b: ResultId) !bool { |
| 225 | const index_a = self.info.entities.getIndex(id_a).?; |
| 226 | const index_b = self.info.entities.getIndex(id_b).?; |
| 227 | |
| 228 | const entity_a = self.info.entities.values()[index_a]; |
| 229 | const entity_b = self.info.entities.values()[index_b]; |
| 230 | |
| 231 | if (entity_a.kind != entity_b.kind) return false; |
| 232 | |
| 233 | if (entity_a.kind == .OpTypePointer) { |
| 234 | // May be a forward reference, or should be saved as a potential |
| 235 | // forward reference in the future. Whatever the case, it should |
| 236 | // be the same for both a and b. |
| 237 | const entry_a = try self.ptr_map_a.getOrPut(self.a, id_a); |
| 238 | const entry_b = try self.ptr_map_b.getOrPut(self.a, id_b); |
| 239 | |
| 240 | if (entry_a.found_existing != entry_b.found_existing) return false; |
| 241 | if (entry_a.index != entry_b.index) return false; |
| 242 | |
| 243 | if (entry_a.found_existing) { |
| 244 | // No need to recurse. |
| 245 | return true; |
| 246 | } |
| 247 | } |
| 248 | |
| 249 | // Check if extra data is the same. |
| 250 | if (!std.mem.eql(u32, self.info.extraDataByIndex(index_a), self.info.extraDataByIndex(index_b))) { |
| 251 | return false; |
| 252 | } |
| 253 | |
| 254 | // Recursively check if children are the same |
| 255 | const children_a = self.info.childrenByIndex(index_a); |
| 256 | const children_b = self.info.childrenByIndex(index_b); |
| 257 | if (children_a.len != children_b.len) return false; |
| 258 | |
| 259 | for (children_a, children_b) |child_a, child_b| { |
| 260 | if (!try self.eqlInner(child_a, child_b)) { |
| 261 | return false; |
| 262 | } |
| 263 | } |
| 264 | |
| 265 | return true; |
| 266 | } |
| 267 | }; |
| 268 | |
| 269 | /// This struct is a wrapper around EntityContext that adapts it for |
| 270 | /// use in a hash map. Because EntityContext allocates, it cannot be |
| 271 | /// used. This wrapper simply assumes that the maps have been allocated |
| 272 | /// the max amount of memory they are going to use. |
| 273 | /// This is done by pre-hashing all keys. |
| 274 | const EntityHashContext = struct { |
| 275 | entity_context: *EntityContext, |
| 276 | |
| 277 | pub fn hash(self: EntityHashContext, key: ResultId) u64 { |
| 278 | return self.entity_context.hash(key) catch unreachable; |
| 279 | } |
| 280 | |
| 281 | pub fn eql(self: EntityHashContext, a: ResultId, b: ResultId) bool { |
| 282 | return self.entity_context.eql(a, b) catch unreachable; |
| 283 | } |
| 284 | }; |
| 285 | |
| 286 | pub fn run(parser: *BinaryModule.Parser, binary: *BinaryModule) !void { |
| 287 | var arena = std.heap.ArenaAllocator.init(parser.a); |
| 288 | defer arena.deinit(); |
| 289 | const a = arena.allocator(); |
| 290 | |
| 291 | const info = try ModuleInfo.parse(a, parser, binary.*); |
| 292 | log.info("added {} entities", .{info.entities.count()}); |
| 293 | log.info("children size: {}", .{info.entity_children.len}); |
| 294 | log.info("extra data size: {}", .{info.extra_data.len}); |
| 295 | |
| 296 | // Hash all keys once so that the maps can be allocated the right size. |
| 297 | var ctx = EntityContext.init(a, &info); |
| 298 | for (info.entities.keys()) |id| { |
| 299 | _ = try ctx.hash(id); |
| 300 | } |
| 301 | |
| 302 | // hash only uses ptr_map_a, so allocate ptr_map_b too |
| 303 | try ctx.equalizeMapCapacity(); |
| 304 | |
| 305 | // Figure out which entities can be deduplicated. |
| 306 | var map = std.HashMap(ResultId, void, EntityHashContext, 80).initContext(a, .{ |
| 307 | .entity_context = &ctx, |
| 308 | }); |
| 309 | var replace = std.AutoArrayHashMap(ResultId, ResultId).init(a); |
| 310 | for (info.entities.keys(), info.entities.values()) |id, entity| { |
| 311 | const entry = try map.getOrPut(id); |
| 312 | if (entry.found_existing) { |
| 313 | log.info("deduplicating {} - {s} (prior definition: {})", .{ id, @tagName(entity.kind), entry.key_ptr.* }); |
| 314 | try replace.putNoClobber(id, entry.key_ptr.*); |
| 315 | } |
| 316 | } |
| 317 | |
| 318 | // Now process the module, and replace instructions where needed. |
| 319 | var section = Section{}; |
| 320 | var it = binary.iterateInstructions(); |
| 321 | var id_offsets = std.ArrayList(u16).init(a); |
| 322 | var new_functions_section: ?usize = null; |
| 323 | var new_operands = std.ArrayList(u32).init(a); |
| 324 | var emitted_ptrs = std.AutoHashMap(ResultId, void).init(a); |
| 325 | while (it.next()) |inst| { |
| 326 | // Result-id can only be the first or second operand |
| 327 | const inst_spec = parser.getInstSpec(inst.opcode).?; |
| 328 | const maybe_result_id: ?ResultId = for (0..2) |i| { |
| 329 | if (inst_spec.operands.len > i and inst_spec.operands[i].kind == .IdResult) { |
| 330 | break @enumFromInt(inst.operands[i]); |
| 331 | } |
| 332 | } else null; |
| 333 | |
| 334 | if (maybe_result_id) |result_id| { |
| 335 | if (replace.contains(result_id)) continue; |
| 336 | } |
| 337 | |
| 338 | switch (inst.opcode) { |
| 339 | .OpFunction => if (new_functions_section == null) { |
| 340 | new_functions_section = section.instructions.items.len; |
| 341 | }, |
| 342 | .OpTypeForwardPointer => continue, // We re-emit these where needed |
| 343 | // TODO: These aren't supported yet, strip them out for testing purposes. |
| 344 | .OpName, .OpMemberName => continue, |
| 345 | else => {}, |
| 346 | } |
| 347 | |
| 348 | // Re-emit the instruction, but replace all the IDs. |
| 349 | |
| 350 | id_offsets.items.len = 0; |
| 351 | try parser.parseInstructionResultIds(binary.*, inst, &id_offsets); |
| 352 | |
| 353 | new_operands.items.len = 0; |
| 354 | try new_operands.appendSlice(inst.operands); |
| 355 | for (id_offsets.items) |offset| { |
| 356 | { |
| 357 | const id: ResultId = @enumFromInt(inst.operands[offset]); |
| 358 | if (replace.get(id)) |new_id| { |
| 359 | new_operands.items[offset] = @intFromEnum(new_id); |
| 360 | } |
| 361 | } |
| 362 | |
| 363 | // TODO: Does this logic work? Maybe it will emit an OpTypeForwardPointer to |
| 364 | // something thats not a struct... |
| 365 | // It seems to work correctly on behavior.zig at least |
| 366 | const id: ResultId = @enumFromInt(new_operands.items[offset]); |
| 367 | if (maybe_result_id == null or maybe_result_id.? != id) { |
| 368 | const index = info.entities.getIndex(id) orelse continue; |
| 369 | const entity = info.entities.values()[index]; |
| 370 | if (entity.kind == .OpTypePointer) { |
| 371 | if (!emitted_ptrs.contains(id)) { |
| 372 | // The storage class is in the extra data |
| 373 | // TODO: This is kind of hacky... |
| 374 | const extra_data = info.extraDataByIndex(index); |
| 375 | const storage_class: spec.StorageClass = @enumFromInt(extra_data[0]); |
| 376 | try section.emit(a, .OpTypeForwardPointer, .{ |
| 377 | .pointer_type = id, |
| 378 | .storage_class = storage_class, |
| 379 | }); |
| 380 | try emitted_ptrs.put(id, {}); |
| 381 | } |
| 382 | } |
| 383 | } |
| 384 | } |
| 385 | |
| 386 | if (inst.opcode == .OpTypePointer) { |
| 387 | try emitted_ptrs.put(maybe_result_id.?, {}); |
| 388 | } |
| 389 | |
| 390 | try section.emitRawInstruction(a, inst.opcode, new_operands.items); |
| 391 | } |
| 392 | |
| 393 | for (replace.keys()) |key| { |
| 394 | _ = binary.ext_inst_map.remove(key); |
| 395 | _ = binary.arith_type_width.remove(key); |
| 396 | } |
| 397 | |
| 398 | binary.instructions = try parser.a.dupe(Word, section.toWords()); |
| 399 | binary.sections.functions = new_functions_section orelse binary.instructions.len; |
| 400 | } |
Robin Voetter <robin@voetter.nl> 2024-03-29 13:54:28+01:00
Signed by SSH key SHA256:ZS52FNyUv2WUXvO4njmVaFVO46RHojFuOrxRc4LuKzg
src