| ... | ... | @@ -6,7 +6,6 @@ const Allocator = std.mem.Allocator; |
| 6 | 6 | const Air = @import("Air.zig"); |
| 7 | 7 | const Type = @import("type.zig").Type; |
| 8 | 8 | const Module = @import("Module.zig"); |
| 9 | | const LazySrcLoc = Module.LazySrcLoc; |
| 10 | 9 | const expect = std.testing.expect; |
| 11 | 10 | const expectEqual = std.testing.expectEqual; |
| 12 | 11 | const expectEqualSlices = std.testing.expectEqualSlices; |
| ... | ... | @@ -19,15 +18,25 @@ pub fn RegisterManager( |
| 19 | 18 | comptime callee_preserved_regs: []const Register, |
| 20 | 19 | ) type { |
| 21 | 20 | return struct { |
| 21 | /// Tracks the AIR instruction allocated to every register or |
| 22 | /// `null` if no instruction is allocated to a register |
| 23 | /// |
| 22 | 24 | /// The key must be canonical register. |
| 23 | 25 | registers: [callee_preserved_regs.len]?Air.Inst.Index = [_]?Air.Inst.Index{null} ** callee_preserved_regs.len, |
| 26 | /// Tracks which registers are free (in which case the |
| 27 | /// corresponding bit is set to 1) |
| 24 | 28 | free_registers: FreeRegInt = math.maxInt(FreeRegInt), |
| 25 | | /// Tracks all registers allocated in the course of this function |
| 29 | /// Tracks all registers allocated in the course of this |
| 30 | /// function |
| 26 | 31 | allocated_registers: FreeRegInt = 0, |
| 32 | /// Tracks registers which are temporarily blocked from being |
| 33 | /// allocated |
| 34 | frozen_registers: FreeRegInt = 0, |
| 27 | 35 | |
| 28 | 36 | const Self = @This(); |
| 29 | 37 | |
| 30 | | /// An integer whose bits represent all the registers and whether they are free. |
| 38 | /// An integer whose bits represent all the registers and |
| 39 | /// whether they are free. |
| 31 | 40 | const FreeRegInt = std.meta.Int(.unsigned, callee_preserved_regs.len); |
| 32 | 41 | const ShiftInt = math.Log2Int(FreeRegInt); |
| 33 | 42 | |
| ... | ... | @@ -35,43 +44,76 @@ pub fn RegisterManager( |
| 35 | 44 | return @fieldParentPtr(Function, "register_manager", self); |
| 36 | 45 | } |
| 37 | 46 | |
| 38 | | fn markRegUsed(self: *Self, reg: Register) void { |
| 39 | | if (FreeRegInt == u0) return; |
| 40 | | const index = reg.allocIndex() orelse return; |
| 47 | fn getRegisterMask(reg: Register) ?FreeRegInt { |
| 48 | if (FreeRegInt == u0) return null; |
| 49 | const index = reg.allocIndex() orelse return null; |
| 41 | 50 | const shift = @intCast(ShiftInt, index); |
| 42 | 51 | const mask = @as(FreeRegInt, 1) << shift; |
| 52 | return mask; |
| 53 | } |
| 54 | |
| 55 | fn markRegUsed(self: *Self, reg: Register) void { |
| 56 | const mask = getRegisterMask(reg) orelse return; |
| 43 | 57 | self.free_registers &= ~mask; |
| 44 | 58 | self.allocated_registers |= mask; |
| 45 | 59 | } |
| 46 | 60 | |
| 47 | 61 | fn markRegFree(self: *Self, reg: Register) void { |
| 48 | | if (FreeRegInt == u0) return; |
| 49 | | const index = reg.allocIndex() orelse return; |
| 50 | | const shift = @intCast(ShiftInt, index); |
| 51 | | self.free_registers |= @as(FreeRegInt, 1) << shift; |
| 62 | const mask = getRegisterMask(reg) orelse return; |
| 63 | self.free_registers |= mask; |
| 52 | 64 | } |
| 53 | 65 | |
| 54 | 66 | /// Returns true when this register is not tracked |
| 55 | 67 | pub fn isRegFree(self: Self, reg: Register) bool { |
| 56 | | if (FreeRegInt == u0) return true; |
| 57 | | const index = reg.allocIndex() orelse return true; |
| 58 | | const shift = @intCast(ShiftInt, index); |
| 59 | | return self.free_registers & @as(FreeRegInt, 1) << shift != 0; |
| 68 | const mask = getRegisterMask(reg) orelse return true; |
| 69 | return self.free_registers & mask != 0; |
| 60 | 70 | } |
| 61 | 71 | |
| 62 | 72 | /// Returns whether this register was allocated in the course |
| 63 | 73 | /// of this function. |
| 74 | /// |
| 64 | 75 | /// Returns false when this register is not tracked |
| 65 | 76 | pub fn isRegAllocated(self: Self, reg: Register) bool { |
| 66 | | if (FreeRegInt == u0) return false; |
| 67 | | const index = reg.allocIndex() orelse return false; |
| 68 | | const shift = @intCast(ShiftInt, index); |
| 69 | | return self.allocated_registers & @as(FreeRegInt, 1) << shift != 0; |
| 77 | const mask = getRegisterMask(reg) orelse return false; |
| 78 | return self.allocated_registers & mask != 0; |
| 79 | } |
| 80 | |
| 81 | /// Returns whether this register is frozen |
| 82 | /// |
| 83 | /// Returns false when this register is not tracked |
| 84 | pub fn isRegFrozen(self: Self, reg: Register) bool { |
| 85 | const mask = getRegisterMask(reg) orelse return false; |
| 86 | return self.frozen_registers & mask != 0; |
| 87 | } |
| 88 | |
| 89 | /// Prevents the registers from being allocated until they are |
| 90 | /// unfrozen again |
| 91 | pub fn freezeRegs(self: *Self, regs: []const Register) void { |
| 92 | for (regs) |reg| { |
| 93 | const mask = getRegisterMask(reg) orelse continue; |
| 94 | self.frozen_registers |= mask; |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | /// Enables the allocation of the registers |
| 99 | pub fn unfreezeRegs(self: *Self, regs: []const Register) void { |
| 100 | for (regs) |reg| { |
| 101 | const mask = getRegisterMask(reg) orelse continue; |
| 102 | self.frozen_registers &= ~mask; |
| 103 | } |
| 104 | } |
| 105 | |
| 106 | /// Returns true when at least one register is frozen |
| 107 | pub fn frozenRegsExist(self: Self) bool { |
| 108 | return self.frozen_registers != 0; |
| 70 | 109 | } |
| 71 | 110 | |
| 72 | 111 | /// Allocates a specified number of registers, optionally |
| 73 | 112 | /// tracking them. Returns `null` if not enough registers are |
| 74 | 113 | /// free. |
| 114 | /// |
| 115 | /// Exceptions are deprecated, use freezeRegs and unfreezeRegs |
| 116 | /// instead. |
| 75 | 117 | pub fn tryAllocRegs( |
| 76 | 118 | self: *Self, |
| 77 | 119 | comptime count: comptime_int, |
| ... | ... | @@ -90,6 +132,7 @@ pub fn RegisterManager( |
| 90 | 132 | for (callee_preserved_regs) |reg| { |
| 91 | 133 | if (i >= count) break; |
| 92 | 134 | if (mem.indexOfScalar(Register, exceptions, reg) != null) continue; |
| 135 | if (self.isRegFrozen(reg)) continue; |
| 93 | 136 | if (self.isRegFree(reg)) { |
| 94 | 137 | regs[i] = reg; |
| 95 | 138 | i += 1; |
| ... | ... | @@ -113,6 +156,9 @@ pub fn RegisterManager( |
| 113 | 156 | /// Allocates a register and optionally tracks it with a |
| 114 | 157 | /// corresponding instruction. Returns `null` if all registers |
| 115 | 158 | /// are allocated. |
| 159 | /// |
| 160 | /// Exceptions are deprecated, use freezeRegs and unfreezeRegs |
| 161 | /// instead. |
| 116 | 162 | pub fn tryAllocReg(self: *Self, inst: ?Air.Inst.Index, exceptions: []const Register) ?Register { |
| 117 | 163 | return if (tryAllocRegs(self, 1, .{inst}, exceptions)) |regs| regs[0] else null; |
| 118 | 164 | } |
| ... | ... | @@ -120,6 +166,9 @@ pub fn RegisterManager( |
| 120 | 166 | /// Allocates a specified number of registers, optionally |
| 121 | 167 | /// tracking them. Asserts that count + exceptions.len is not |
| 122 | 168 | /// larger than the total number of registers available. |
| 169 | /// |
| 170 | /// Exceptions are deprecated, use freezeRegs and unfreezeRegs |
| 171 | /// instead. |
| 123 | 172 | pub fn allocRegs( |
| 124 | 173 | self: *Self, |
| 125 | 174 | comptime count: comptime_int, |
| ... | ... | @@ -138,6 +187,7 @@ pub fn RegisterManager( |
| 138 | 187 | for (callee_preserved_regs) |reg| { |
| 139 | 188 | if (i >= count) break; |
| 140 | 189 | if (mem.indexOfScalar(Register, exceptions, reg) != null) continue; |
| 190 | if (self.isRegFrozen(reg)) continue; |
| 141 | 191 | regs[i] = reg; |
| 142 | 192 | |
| 143 | 193 | const index = reg.allocIndex().?; // allocIndex() on a callee-preserved reg should never return null |
| ... | ... | @@ -171,6 +221,9 @@ pub fn RegisterManager( |
| 171 | 221 | |
| 172 | 222 | /// Allocates a register and optionally tracks it with a |
| 173 | 223 | /// corresponding instruction. |
| 224 | /// |
| 225 | /// Exceptions are deprecated, use freezeRegs and unfreezeRegs |
| 226 | /// instead. |
| 174 | 227 | pub fn allocReg(self: *Self, inst: ?Air.Inst.Index, exceptions: []const Register) !Register { |
| 175 | 228 | return (try self.allocRegs(1, .{inst}, exceptions))[0]; |
| 176 | 229 | } |
| ... | ... | @@ -343,9 +396,22 @@ test "allocReg: spilling" { |
| 343 | 396 | try expectEqualSlices(MockRegister1, &[_]MockRegister1{.r2}, function.spilled.items); |
| 344 | 397 | |
| 345 | 398 | // Exceptions |
| 399 | // |
| 400 | // TODO deprecated, remove test once no backend uses exceptions |
| 401 | // anymore |
| 346 | 402 | function.register_manager.freeReg(.r2); |
| 347 | 403 | function.register_manager.freeReg(.r3); |
| 348 | 404 | try expectEqual(@as(?MockRegister1, .r3), try function.register_manager.allocReg(mock_instruction, &.{.r2})); |
| 405 | |
| 406 | // Frozen registers |
| 407 | function.register_manager.freeReg(.r3); |
| 408 | { |
| 409 | function.register_manager.freezeRegs(&.{.r2}); |
| 410 | defer function.register_manager.unfreezeRegs(&.{.r2}); |
| 411 | |
| 412 | try expectEqual(@as(?MockRegister1, .r3), try function.register_manager.allocReg(mock_instruction, &.{})); |
| 413 | } |
| 414 | try expect(!function.register_manager.frozenRegsExist()); |
| 349 | 415 | } |
| 350 | 416 | |
| 351 | 417 | test "tryAllocRegs" { |
| ... | ... | @@ -359,10 +425,25 @@ test "tryAllocRegs" { |
| 359 | 425 | try expectEqual([_]MockRegister2{ .r0, .r1, .r2 }, function.register_manager.tryAllocRegs(3, .{ null, null, null }, &.{}).?); |
| 360 | 426 | |
| 361 | 427 | // Exceptions |
| 428 | // |
| 429 | // TODO deprecated, remove test once no backend uses exceptions |
| 430 | // anymore |
| 362 | 431 | function.register_manager.freeReg(.r0); |
| 363 | 432 | function.register_manager.freeReg(.r1); |
| 364 | 433 | function.register_manager.freeReg(.r2); |
| 365 | 434 | try expectEqual([_]MockRegister2{ .r0, .r2, .r3 }, function.register_manager.tryAllocRegs(3, .{ null, null, null }, &.{.r1}).?); |
| 435 | |
| 436 | // Frozen registers |
| 437 | function.register_manager.freeReg(.r0); |
| 438 | function.register_manager.freeReg(.r2); |
| 439 | function.register_manager.freeReg(.r3); |
| 440 | { |
| 441 | function.register_manager.freezeRegs(&.{.r1}); |
| 442 | defer function.register_manager.unfreezeRegs(&.{.r1}); |
| 443 | |
| 444 | try expectEqual([_]MockRegister2{ .r0, .r2, .r3 }, function.register_manager.tryAllocRegs(3, .{ null, null, null }, &.{}).?); |
| 445 | } |
| 446 | try expect(!function.register_manager.frozenRegsExist()); |
| 366 | 447 | } |
| 367 | 448 | |
| 368 | 449 | test "allocRegs" { |
| ... | ... | @@ -382,8 +463,23 @@ test "allocRegs" { |
| 382 | 463 | }, &.{})); |
| 383 | 464 | |
| 384 | 465 | // Exceptions |
| 466 | // |
| 467 | // TODO deprecated, remove test once no backend uses exceptions |
| 468 | // anymore |
| 385 | 469 | try expectEqual([_]MockRegister2{ .r0, .r2, .r3 }, try function.register_manager.allocRegs(3, .{ null, null, null }, &.{.r1})); |
| 386 | 470 | try expectEqualSlices(MockRegister2, &[_]MockRegister2{ .r0, .r2 }, function.spilled.items); |
| 471 | |
| 472 | // Frozen registers |
| 473 | function.register_manager.freeReg(.r0); |
| 474 | function.register_manager.freeReg(.r2); |
| 475 | function.register_manager.freeReg(.r3); |
| 476 | { |
| 477 | function.register_manager.freezeRegs(&.{.r1}); |
| 478 | defer function.register_manager.unfreezeRegs(&.{.r1}); |
| 479 | |
| 480 | try expectEqual([_]MockRegister2{ .r0, .r2, .r3 }, try function.register_manager.allocRegs(3, .{ null, null, null }, &.{})); |
| 481 | } |
| 482 | try expect(!function.register_manager.frozenRegsExist()); |
| 387 | 483 | } |
| 388 | 484 | |
| 389 | 485 | test "getReg" { |