| ... | @@ -0,0 +1,482 @@ |
| 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 | // Debug decoration-style instructions |
| 21 | .OpName, .OpMemberName => true, |
| 22 | else => switch (opcode.class()) { |
| 23 | .TypeDeclaration, |
| 24 | .ConstantCreation, |
| 25 | .Annotation, |
| 26 | => true, |
| 27 | else => false, |
| 28 | }, |
| 29 | }; |
| 30 | } |
| 31 | |
| 32 | const ModuleInfo = struct { |
| 33 | /// This models a type, decoration or constant instruction |
| 34 | /// and its dependencies. |
| 35 | const Entity = struct { |
| 36 | /// The type that this entity represents. This is just |
| 37 | /// the instruction opcode. |
| 38 | kind: Opcode, |
| 39 | /// The offset of this entity's operands, in |
| 40 | /// `binary.instructions`. |
| 41 | first_operand: u32, |
| 42 | /// The number of operands in this entity |
| 43 | num_operands: u16, |
| 44 | /// The (first_operand-relative) offset of the result-id, |
| 45 | /// or the entity that is affected by this entity if this entity |
| 46 | /// is a decoration. |
| 47 | result_id_index: u16, |
| 48 | /// The first decoration in `self.decorations`. |
| 49 | first_decoration: u32, |
| 50 | }; |
| 51 | |
| 52 | /// Maps result-id to Entity's |
| 53 | entities: std.AutoArrayHashMapUnmanaged(ResultId, Entity), |
| 54 | /// A bit set that keeps track of which operands are result-ids. |
| 55 | /// Note: This also includes any result-id! |
| 56 | /// Because we need these values when recoding the module anyway, |
| 57 | /// it contains the status of ALL operands in the module. |
| 58 | operand_is_id: std.DynamicBitSetUnmanaged, |
| 59 | /// Store of decorations for each entity. |
| 60 | decorations: []const Entity, |
| 61 | |
| 62 | pub fn parse( |
| 63 | arena: Allocator, |
| 64 | parser: *BinaryModule.Parser, |
| 65 | binary: BinaryModule, |
| 66 | ) !ModuleInfo { |
| 67 | var entities = std.AutoArrayHashMap(ResultId, Entity).init(arena); |
| 68 | var id_offsets = std.ArrayList(u16).init(arena); |
| 69 | var operand_is_id = try std.DynamicBitSetUnmanaged.initEmpty(arena, binary.instructions.len); |
| 70 | var decorations = std.MultiArrayList(struct { target_id: ResultId, entity: Entity }){}; |
| 71 | |
| 72 | var it = binary.iterateInstructions(); |
| 73 | while (it.next()) |inst| { |
| 74 | id_offsets.items.len = 0; |
| 75 | try parser.parseInstructionResultIds(binary, inst, &id_offsets); |
| 76 | |
| 77 | const first_operand_offset: u32 = @intCast(inst.offset + 1); |
| 78 | for (id_offsets.items) |offset| { |
| 79 | operand_is_id.set(first_operand_offset + offset); |
| 80 | } |
| 81 | |
| 82 | if (!canDeduplicate(inst.opcode)) continue; |
| 83 | |
| 84 | const result_id_index: u16 = switch (inst.opcode.class()) { |
| 85 | .TypeDeclaration, .Annotation, .Debug => 0, |
| 86 | .ConstantCreation => 1, |
| 87 | else => unreachable, |
| 88 | }; |
| 89 | |
| 90 | const result_id: ResultId = @enumFromInt(inst.operands[id_offsets.items[result_id_index]]); |
| 91 | const entity = Entity{ |
| 92 | .kind = inst.opcode, |
| 93 | .first_operand = first_operand_offset, |
| 94 | .num_operands = @intCast(inst.operands.len), |
| 95 | .result_id_index = result_id_index, |
| 96 | .first_decoration = undefined, // Filled in later |
| 97 | }; |
| 98 | |
| 99 | switch (inst.opcode.class()) { |
| 100 | .Annotation, .Debug => { |
| 101 | try decorations.append(arena, .{ |
| 102 | .target_id = result_id, |
| 103 | .entity = entity, |
| 104 | }); |
| 105 | }, |
| 106 | .TypeDeclaration, .ConstantCreation => { |
| 107 | const entry = try entities.getOrPut(result_id); |
| 108 | if (entry.found_existing) { |
| 109 | log.err("type or constant {} has duplicate definition", .{result_id}); |
| 110 | return error.DuplicateId; |
| 111 | } |
| 112 | entry.value_ptr.* = entity; |
| 113 | }, |
| 114 | else => unreachable, |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | // Sort decorations by the index of the result-id in `entities. |
| 119 | // This ensures not only that the decorations of a particular reuslt-id |
| 120 | // are continuous, but the subsequences also appear in the same order as in `entities`. |
| 121 | |
| 122 | const SortContext = struct { |
| 123 | entities: std.AutoArrayHashMapUnmanaged(ResultId, Entity), |
| 124 | ids: []const ResultId, |
| 125 | |
| 126 | pub fn lessThan(ctx: @This(), a_index: usize, b_index: usize) bool { |
| 127 | // If any index is not in the entities set, its because its not a |
| 128 | // deduplicatable result-id. Those should be considered largest and |
| 129 | // float to the end. |
| 130 | const entity_index_a = ctx.entities.getIndex(ctx.ids[a_index]) orelse return false; |
| 131 | const entity_index_b = ctx.entities.getIndex(ctx.ids[b_index]) orelse return true; |
| 132 | |
| 133 | return entity_index_a < entity_index_b; |
| 134 | } |
| 135 | }; |
| 136 | |
| 137 | decorations.sort(SortContext{ |
| 138 | .entities = entities.unmanaged, |
| 139 | .ids = decorations.items(.target_id), |
| 140 | }); |
| 141 | |
| 142 | // Now go through the decorations and add the offsets to the entities list. |
| 143 | var decoration_i: u32 = 0; |
| 144 | const target_ids = decorations.items(.target_id); |
| 145 | for (entities.keys(), entities.values()) |id, *entity| { |
| 146 | entity.first_decoration = decoration_i; |
| 147 | |
| 148 | // Scan ahead to the next decoration |
| 149 | while (decoration_i < target_ids.len and target_ids[decoration_i] == id) { |
| 150 | decoration_i += 1; |
| 151 | } |
| 152 | } |
| 153 | |
| 154 | return ModuleInfo{ |
| 155 | .entities = entities.unmanaged, |
| 156 | .operand_is_id = operand_is_id, |
| 157 | // There may be unrelated decorations at the end, so make sure to |
| 158 | // slice those off. |
| 159 | .decorations = decorations.items(.entity)[0..decoration_i], |
| 160 | }; |
| 161 | } |
| 162 | |
| 163 | fn entityDecorationsByIndex(self: ModuleInfo, index: usize) []const Entity { |
| 164 | const values = self.entities.values(); |
| 165 | const first_decoration = values[index].first_decoration; |
| 166 | if (index == values.len - 1) { |
| 167 | return self.decorations[first_decoration..]; |
| 168 | } else { |
| 169 | const next_first_decoration = values[index + 1].first_decoration; |
| 170 | return self.decorations[first_decoration..next_first_decoration]; |
| 171 | } |
| 172 | } |
| 173 | }; |
| 174 | |
| 175 | const EntityContext = struct { |
| 176 | a: Allocator, |
| 177 | ptr_map_a: std.AutoArrayHashMapUnmanaged(ResultId, void) = .{}, |
| 178 | ptr_map_b: std.AutoArrayHashMapUnmanaged(ResultId, void) = .{}, |
| 179 | info: *const ModuleInfo, |
| 180 | binary: *const BinaryModule, |
| 181 | |
| 182 | fn deinit(self: *EntityContext) void { |
| 183 | self.ptr_map_a.deinit(self.a); |
| 184 | self.ptr_map_b.deinit(self.a); |
| 185 | |
| 186 | self.* = undefined; |
| 187 | } |
| 188 | |
| 189 | fn equalizeMapCapacity(self: *EntityContext) !void { |
| 190 | const cap = @max(self.ptr_map_a.capacity(), self.ptr_map_b.capacity()); |
| 191 | try self.ptr_map_a.ensureTotalCapacity(self.a, cap); |
| 192 | try self.ptr_map_b.ensureTotalCapacity(self.a, cap); |
| 193 | } |
| 194 | |
| 195 | fn hash(self: *EntityContext, id: ResultId) !u64 { |
| 196 | var hasher = std.hash.Wyhash.init(0); |
| 197 | self.ptr_map_a.clearRetainingCapacity(); |
| 198 | try self.hashInner(&hasher, id); |
| 199 | return hasher.final(); |
| 200 | } |
| 201 | |
| 202 | fn hashInner(self: *EntityContext, hasher: *std.hash.Wyhash, id: ResultId) error{OutOfMemory}!void { |
| 203 | const index = self.info.entities.getIndex(id) orelse { |
| 204 | // Index unknown, the type or constant may depend on another result-id |
| 205 | // that couldn't be deduplicated and so it wasn't added to info.entities. |
| 206 | // In this case, just has the ID itself. |
| 207 | std.hash.autoHash(hasher, id); |
| 208 | return; |
| 209 | }; |
| 210 | |
| 211 | const entity = self.info.entities.values()[index]; |
| 212 | |
| 213 | if (entity.kind == .OpTypePointer) { |
| 214 | // This may be either a pointer that is forward-referenced in the future, |
| 215 | // or a forward reference to a pointer. |
| 216 | const entry = try self.ptr_map_a.getOrPut(self.a, id); |
| 217 | if (entry.found_existing) { |
| 218 | // Pointer already seen. Hash the index instead of recursing into its children. |
| 219 | std.hash.autoHash(hasher, entry.index); |
| 220 | return; |
| 221 | } |
| 222 | } |
| 223 | |
| 224 | try self.hashEntity(hasher, entity); |
| 225 | |
| 226 | // Process decorations. |
| 227 | const decorations = self.info.entityDecorationsByIndex(index); |
| 228 | for (decorations) |decoration| { |
| 229 | try self.hashEntity(hasher, decoration); |
| 230 | } |
| 231 | } |
| 232 | |
| 233 | fn hashEntity(self: *EntityContext, hasher: *std.hash.Wyhash, entity: ModuleInfo.Entity) !void { |
| 234 | std.hash.autoHash(hasher, entity.kind); |
| 235 | // Process operands |
| 236 | const operands = self.binary.instructions[entity.first_operand..][0..entity.num_operands]; |
| 237 | for (operands, 0..) |operand, i| { |
| 238 | if (i == entity.result_id_index) { |
| 239 | // Not relevant, skip... |
| 240 | continue; |
| 241 | } else if (self.info.operand_is_id.isSet(entity.first_operand + i)) { |
| 242 | // Operand is ID |
| 243 | try self.hashInner(hasher, @enumFromInt(operand)); |
| 244 | } else { |
| 245 | // Operand is merely data |
| 246 | std.hash.autoHash(hasher, operand); |
| 247 | } |
| 248 | } |
| 249 | } |
| 250 | |
| 251 | fn eql(self: *EntityContext, a: ResultId, b: ResultId) !bool { |
| 252 | self.ptr_map_a.clearRetainingCapacity(); |
| 253 | self.ptr_map_b.clearRetainingCapacity(); |
| 254 | |
| 255 | return try self.eqlInner(a, b); |
| 256 | } |
| 257 | |
| 258 | fn eqlInner(self: *EntityContext, id_a: ResultId, id_b: ResultId) error{OutOfMemory}!bool { |
| 259 | const maybe_index_a = self.info.entities.getIndex(id_a); |
| 260 | const maybe_index_b = self.info.entities.getIndex(id_b); |
| 261 | |
| 262 | if (maybe_index_a == null and maybe_index_b == null) { |
| 263 | // Both indices unknown. In this case the type or constant |
| 264 | // may depend on another result-id that couldn't be deduplicated |
| 265 | // (so it wasn't added to info.entities). In this case, that particular |
| 266 | // result-id should be the same one. |
| 267 | return id_a == id_b; |
| 268 | } |
| 269 | |
| 270 | const index_a = maybe_index_a orelse return false; |
| 271 | const index_b = maybe_index_b orelse return false; |
| 272 | |
| 273 | const entity_a = self.info.entities.values()[index_a]; |
| 274 | const entity_b = self.info.entities.values()[index_b]; |
| 275 | |
| 276 | if (entity_a.kind == .OpTypePointer) { |
| 277 | // May be a forward reference, or should be saved as a potential |
| 278 | // forward reference in the future. Whatever the case, it should |
| 279 | // be the same for both a and b. |
| 280 | const entry_a = try self.ptr_map_a.getOrPut(self.a, id_a); |
| 281 | const entry_b = try self.ptr_map_b.getOrPut(self.a, id_b); |
| 282 | |
| 283 | if (entry_a.found_existing != entry_b.found_existing) return false; |
| 284 | if (entry_a.index != entry_b.index) return false; |
| 285 | |
| 286 | if (entry_a.found_existing) { |
| 287 | // No need to recurse. |
| 288 | return true; |
| 289 | } |
| 290 | } |
| 291 | |
| 292 | if (!try self.eqlEntities(entity_a, entity_b)) { |
| 293 | return false; |
| 294 | } |
| 295 | |
| 296 | // Compare decorations. |
| 297 | const decorations_a = self.info.entityDecorationsByIndex(index_a); |
| 298 | const decorations_b = self.info.entityDecorationsByIndex(index_b); |
| 299 | if (decorations_a.len != decorations_b.len) { |
| 300 | return false; |
| 301 | } |
| 302 | |
| 303 | for (decorations_a, decorations_b) |decoration_a, decoration_b| { |
| 304 | if (!try self.eqlEntities(decoration_a, decoration_b)) { |
| 305 | return false; |
| 306 | } |
| 307 | } |
| 308 | |
| 309 | return true; |
| 310 | } |
| 311 | |
| 312 | fn eqlEntities(self: *EntityContext, entity_a: ModuleInfo.Entity, entity_b: ModuleInfo.Entity) !bool { |
| 313 | if (entity_a.kind != entity_b.kind) { |
| 314 | return false; |
| 315 | } else if (entity_a.result_id_index != entity_a.result_id_index) { |
| 316 | return false; |
| 317 | } |
| 318 | |
| 319 | const operands_a = self.binary.instructions[entity_a.first_operand..][0..entity_a.num_operands]; |
| 320 | const operands_b = self.binary.instructions[entity_b.first_operand..][0..entity_b.num_operands]; |
| 321 | |
| 322 | // Note: returns false for operands that have explicit defaults in optional operands... oh well |
| 323 | if (operands_a.len != operands_b.len) { |
| 324 | return false; |
| 325 | } |
| 326 | |
| 327 | for (operands_a, operands_b, 0..) |operand_a, operand_b, i| { |
| 328 | const a_is_id = self.info.operand_is_id.isSet(entity_a.first_operand + i); |
| 329 | const b_is_id = self.info.operand_is_id.isSet(entity_b.first_operand + i); |
| 330 | if (a_is_id != b_is_id) { |
| 331 | return false; |
| 332 | } else if (i == entity_a.result_id_index) { |
| 333 | // result-id for both... |
| 334 | continue; |
| 335 | } else if (a_is_id) { |
| 336 | // Both are IDs, so recurse. |
| 337 | if (!try self.eqlInner(@enumFromInt(operand_a), @enumFromInt(operand_b))) { |
| 338 | return false; |
| 339 | } |
| 340 | } else if (operand_a != operand_b) { |
| 341 | return false; |
| 342 | } |
| 343 | } |
| 344 | |
| 345 | return true; |
| 346 | } |
| 347 | }; |
| 348 | |
| 349 | /// This struct is a wrapper around EntityContext that adapts it for |
| 350 | /// use in a hash map. Because EntityContext allocates, it cannot be |
| 351 | /// used. This wrapper simply assumes that the maps have been allocated |
| 352 | /// the max amount of memory they are going to use. |
| 353 | /// This is done by pre-hashing all keys. |
| 354 | const EntityHashContext = struct { |
| 355 | entity_context: *EntityContext, |
| 356 | |
| 357 | pub fn hash(self: EntityHashContext, key: ResultId) u64 { |
| 358 | return self.entity_context.hash(key) catch unreachable; |
| 359 | } |
| 360 | |
| 361 | pub fn eql(self: EntityHashContext, a: ResultId, b: ResultId) bool { |
| 362 | return self.entity_context.eql(a, b) catch unreachable; |
| 363 | } |
| 364 | }; |
| 365 | |
| 366 | pub fn run(parser: *BinaryModule.Parser, binary: *BinaryModule) !void { |
| 367 | var arena = std.heap.ArenaAllocator.init(parser.a); |
| 368 | defer arena.deinit(); |
| 369 | const a = arena.allocator(); |
| 370 | |
| 371 | const info = try ModuleInfo.parse(a, parser, binary.*); |
| 372 | |
| 373 | // Hash all keys once so that the maps can be allocated the right size. |
| 374 | var ctx = EntityContext{ |
| 375 | .a = a, |
| 376 | .info = &info, |
| 377 | .binary = binary, |
| 378 | }; |
| 379 | for (info.entities.keys()) |id| { |
| 380 | _ = try ctx.hash(id); |
| 381 | } |
| 382 | |
| 383 | // hash only uses ptr_map_a, so allocate ptr_map_b too |
| 384 | try ctx.equalizeMapCapacity(); |
| 385 | |
| 386 | // Figure out which entities can be deduplicated. |
| 387 | var map = std.HashMap(ResultId, void, EntityHashContext, 80).initContext(a, .{ |
| 388 | .entity_context = &ctx, |
| 389 | }); |
| 390 | var replace = std.AutoArrayHashMap(ResultId, ResultId).init(a); |
| 391 | for (info.entities.keys()) |id| { |
| 392 | const entry = try map.getOrPut(id); |
| 393 | if (entry.found_existing) { |
| 394 | try replace.putNoClobber(id, entry.key_ptr.*); |
| 395 | } |
| 396 | } |
| 397 | |
| 398 | // Now process the module, and replace instructions where needed. |
| 399 | var section = Section{}; |
| 400 | var it = binary.iterateInstructions(); |
| 401 | var new_functions_section: ?usize = null; |
| 402 | var new_operands = std.ArrayList(u32).init(a); |
| 403 | var emitted_ptrs = std.AutoHashMap(ResultId, void).init(a); |
| 404 | while (it.next()) |inst| { |
| 405 | // Result-id can only be the first or second operand |
| 406 | const inst_spec = parser.getInstSpec(inst.opcode).?; |
| 407 | |
| 408 | const maybe_result_id_offset: ?u16 = for (0..2) |i| { |
| 409 | if (inst_spec.operands.len > i and inst_spec.operands[i].kind == .IdResult) { |
| 410 | break @intCast(i); |
| 411 | } |
| 412 | } else null; |
| 413 | |
| 414 | if (maybe_result_id_offset) |offset| { |
| 415 | const result_id: ResultId = @enumFromInt(inst.operands[offset]); |
| 416 | if (replace.contains(result_id)) continue; |
| 417 | } |
| 418 | |
| 419 | switch (inst.opcode) { |
| 420 | .OpFunction => if (new_functions_section == null) { |
| 421 | new_functions_section = section.instructions.items.len; |
| 422 | }, |
| 423 | .OpTypeForwardPointer => continue, // We re-emit these where needed |
| 424 | else => {}, |
| 425 | } |
| 426 | |
| 427 | switch (inst.opcode.class()) { |
| 428 | .Annotation, .Debug => { |
| 429 | // For decoration-style instructions, only emit them |
| 430 | // if the target is not removed. |
| 431 | const target: ResultId = @enumFromInt(inst.operands[0]); |
| 432 | if (replace.contains(target)) continue; |
| 433 | }, |
| 434 | else => {}, |
| 435 | } |
| 436 | |
| 437 | // Re-emit the instruction, but replace all the IDs. |
| 438 | |
| 439 | new_operands.items.len = 0; |
| 440 | try new_operands.appendSlice(inst.operands); |
| 441 | |
| 442 | for (new_operands.items, 0..) |*operand, i| { |
| 443 | const is_id = info.operand_is_id.isSet(inst.offset + 1 + i); |
| 444 | if (!is_id) continue; |
| 445 | |
| 446 | if (replace.get(@enumFromInt(operand.*))) |new_id| { |
| 447 | operand.* = @intFromEnum(new_id); |
| 448 | } |
| 449 | |
| 450 | if (maybe_result_id_offset == null or maybe_result_id_offset.? != i) { |
| 451 | const id: ResultId = @enumFromInt(operand.*); |
| 452 | const index = info.entities.getIndex(id) orelse continue; |
| 453 | const entity = info.entities.values()[index]; |
| 454 | if (entity.kind == .OpTypePointer and !emitted_ptrs.contains(id)) { |
| 455 | // Grab the pointer's storage class from its operands in the original |
| 456 | // module. |
| 457 | const storage_class: spec.StorageClass = @enumFromInt(binary.instructions[entity.first_operand + 1]); |
| 458 | try section.emit(a, .OpTypeForwardPointer, .{ |
| 459 | .pointer_type = id, |
| 460 | .storage_class = storage_class, |
| 461 | }); |
| 462 | try emitted_ptrs.put(id, {}); |
| 463 | } |
| 464 | } |
| 465 | } |
| 466 | |
| 467 | if (inst.opcode == .OpTypePointer) { |
| 468 | const result_id: ResultId = @enumFromInt(new_operands.items[maybe_result_id_offset.?]); |
| 469 | try emitted_ptrs.put(result_id, {}); |
| 470 | } |
| 471 | |
| 472 | try section.emitRawInstruction(a, inst.opcode, new_operands.items); |
| 473 | } |
| 474 | |
| 475 | for (replace.keys()) |key| { |
| 476 | _ = binary.ext_inst_map.remove(key); |
| 477 | _ = binary.arith_type_width.remove(key); |
| 478 | } |
| 479 | |
| 480 | binary.instructions = try parser.a.dupe(Word, section.toWords()); |
| 481 | binary.sections.functions = new_functions_section orelse binary.instructions.len; |
| 482 | } |
Robin Voetter <robin@voetter.nl> 2024-04-01 09:51:04+02:00
GitHub <noreply@github.com> 2024-04-01 09:51:04+02:00
Signed by PGP key B5690EEEBB952194
src