| ... | ... | @@ -4,6 +4,9 @@ const assert = std.debug.assert; |
| 4 | 4 | const Allocator = std.mem.Allocator; |
| 5 | 5 | const ir = @import("ir.zig"); |
| 6 | 6 | const Type = @import("type.zig").Type; |
| 7 | const Module = @import("Module.zig"); |
| 8 | const LazySrcLoc = Module.LazySrcLoc; |
| 9 | |
| 7 | 10 | const log = std.log.scoped(.register_manager); |
| 8 | 11 | |
| 9 | 12 | pub fn RegisterManager( |
| ... | ... | @@ -22,6 +25,7 @@ pub fn RegisterManager( |
| 22 | 25 | |
| 23 | 26 | /// An integer whose bits represent all the registers and whether they are free. |
| 24 | 27 | const FreeRegInt = std.meta.Int(.unsigned, callee_preserved_regs.len); |
| 28 | const ShiftInt = math.Log2Int(FreeRegInt); |
| 25 | 29 | |
| 26 | 30 | fn getFunction(self: *Self) *Function { |
| 27 | 31 | return @fieldParentPtr(Function, "register_manager", self); |
| ... | ... | @@ -34,7 +38,6 @@ pub fn RegisterManager( |
| 34 | 38 | fn markRegUsed(self: *Self, reg: Register) void { |
| 35 | 39 | if (FreeRegInt == u0) return; |
| 36 | 40 | const index = reg.allocIndex() orelse return; |
| 37 | | const ShiftInt = math.Log2Int(FreeRegInt); |
| 38 | 41 | const shift = @intCast(ShiftInt, index); |
| 39 | 42 | const mask = @as(FreeRegInt, 1) << shift; |
| 40 | 43 | self.free_registers &= ~mask; |
| ... | ... | @@ -44,7 +47,6 @@ pub fn RegisterManager( |
| 44 | 47 | fn markRegFree(self: *Self, reg: Register) void { |
| 45 | 48 | if (FreeRegInt == u0) return; |
| 46 | 49 | const index = reg.allocIndex() orelse return; |
| 47 | | const ShiftInt = math.Log2Int(FreeRegInt); |
| 48 | 50 | const shift = @intCast(ShiftInt, index); |
| 49 | 51 | self.free_registers |= @as(FreeRegInt, 1) << shift; |
| 50 | 52 | } |
| ... | ... | @@ -54,9 +56,8 @@ pub fn RegisterManager( |
| 54 | 56 | pub fn isRegAllocated(self: Self, reg: Register) bool { |
| 55 | 57 | if (FreeRegInt == u0) return false; |
| 56 | 58 | const index = reg.allocIndex() orelse return false; |
| 57 | | const ShiftInt = math.Log2Int(FreeRegInt); |
| 58 | 59 | const shift = @intCast(ShiftInt, index); |
| 59 | | return self.free_registers & @as(FreeRegInt, 1) << shift != 0; |
| 60 | return self.allocated_registers & @as(FreeRegInt, 1) << shift != 0; |
| 60 | 61 | } |
| 61 | 62 | |
| 62 | 63 | /// Before calling, must ensureCapacity + 1 on self.registers. |
| ... | ... | @@ -66,9 +67,16 @@ pub fn RegisterManager( |
| 66 | 67 | if (free_index >= callee_preserved_regs.len) { |
| 67 | 68 | return null; |
| 68 | 69 | } |
| 69 | | const mask = @as(FreeRegInt, 1) << free_index; |
| 70 | |
| 71 | // This is necessary because the return type of @ctz is 1 |
| 72 | // bit longer than ShiftInt if callee_preserved_regs.len |
| 73 | // is a power of two. This int cast is always safe because |
| 74 | // free_index < callee_preserved_regs.len |
| 75 | const shift = @intCast(ShiftInt, free_index); |
| 76 | const mask = @as(FreeRegInt, 1) << shift; |
| 70 | 77 | self.free_registers &= ~mask; |
| 71 | 78 | self.allocated_registers |= mask; |
| 79 | |
| 72 | 80 | const reg = callee_preserved_regs[free_index]; |
| 73 | 81 | self.registers.putAssumeCapacityNoClobber(reg, inst); |
| 74 | 82 | log.debug("alloc {} => {*}", .{ reg, inst }); |
| ... | ... | @@ -125,3 +133,96 @@ pub fn RegisterManager( |
| 125 | 133 | } |
| 126 | 134 | }; |
| 127 | 135 | } |
| 136 | |
| 137 | const MockRegister = enum(u2) { |
| 138 | r0, r1, r2, r3, |
| 139 | |
| 140 | pub fn allocIndex(self: MockRegister) ?u2 { |
| 141 | inline for (mock_callee_preserved_regs) |cpreg, i| { |
| 142 | if (self == cpreg) return i; |
| 143 | } |
| 144 | return null; |
| 145 | } |
| 146 | }; |
| 147 | |
| 148 | const mock_callee_preserved_regs = [_]MockRegister{ .r2, .r3 }; |
| 149 | |
| 150 | const MockFunction = struct { |
| 151 | allocator: *Allocator, |
| 152 | register_manager: RegisterManager(Self, MockRegister, &mock_callee_preserved_regs) = .{}, |
| 153 | spilled: std.ArrayListUnmanaged(MockRegister) = .{}, |
| 154 | |
| 155 | const Self = @This(); |
| 156 | |
| 157 | pub fn deinit(self: *Self) void { |
| 158 | self.register_manager.deinit(self.allocator); |
| 159 | self.spilled.deinit(self.allocator); |
| 160 | } |
| 161 | |
| 162 | pub fn spillInstruction(self: *Self, src: LazySrcLoc, reg: MockRegister, inst: *ir.Inst) !void { |
| 163 | try self.spilled.append(self.allocator, reg); |
| 164 | } |
| 165 | }; |
| 166 | |
| 167 | test "tryAllocReg: no spilling" { |
| 168 | const allocator = std.testing.allocator; |
| 169 | |
| 170 | var function = MockFunction{ |
| 171 | .allocator = allocator, |
| 172 | }; |
| 173 | defer function.deinit(); |
| 174 | |
| 175 | var mock_instruction = ir.Inst{ |
| 176 | .tag = .breakpoint, |
| 177 | .ty = Type.initTag(.void), |
| 178 | .src = .unneeded, |
| 179 | }; |
| 180 | |
| 181 | std.testing.expect(!function.register_manager.isRegAllocated(.r2)); |
| 182 | std.testing.expect(!function.register_manager.isRegAllocated(.r3)); |
| 183 | |
| 184 | try function.register_manager.registers.ensureCapacity(allocator, function.register_manager.registers.count() + 2); |
| 185 | std.testing.expectEqual(@as(?MockRegister, .r2), function.register_manager.tryAllocReg(&mock_instruction)); |
| 186 | std.testing.expectEqual(@as(?MockRegister, .r3), function.register_manager.tryAllocReg(&mock_instruction)); |
| 187 | std.testing.expectEqual(@as(?MockRegister, null), function.register_manager.tryAllocReg(&mock_instruction)); |
| 188 | |
| 189 | std.testing.expect(function.register_manager.isRegAllocated(.r2)); |
| 190 | std.testing.expect(function.register_manager.isRegAllocated(.r3)); |
| 191 | |
| 192 | function.register_manager.freeReg(.r2); |
| 193 | function.register_manager.freeReg(.r3); |
| 194 | |
| 195 | std.testing.expect(function.register_manager.isRegAllocated(.r2)); |
| 196 | std.testing.expect(function.register_manager.isRegAllocated(.r3)); |
| 197 | } |
| 198 | |
| 199 | test "allocReg: spilling" { |
| 200 | const allocator = std.testing.allocator; |
| 201 | |
| 202 | var function = MockFunction{ |
| 203 | .allocator = allocator, |
| 204 | }; |
| 205 | defer function.deinit(); |
| 206 | |
| 207 | var mock_instruction = ir.Inst{ |
| 208 | .tag = .breakpoint, |
| 209 | .ty = Type.initTag(.void), |
| 210 | .src = .unneeded, |
| 211 | }; |
| 212 | |
| 213 | std.testing.expect(!function.register_manager.isRegAllocated(.r2)); |
| 214 | std.testing.expect(!function.register_manager.isRegAllocated(.r3)); |
| 215 | |
| 216 | try function.register_manager.registers.ensureCapacity(allocator, function.register_manager.registers.count() + 2); |
| 217 | std.testing.expectEqual(@as(?MockRegister, .r2), try function.register_manager.allocReg(&mock_instruction)); |
| 218 | std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction)); |
| 219 | |
| 220 | // Spill a register |
| 221 | std.testing.expectEqual(@as(?MockRegister, .r2), try function.register_manager.allocReg(&mock_instruction)); |
| 222 | std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items); |
| 223 | |
| 224 | // No spilling necessary |
| 225 | function.register_manager.freeReg(.r3); |
| 226 | std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction)); |
| 227 | std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items); |
| 228 | } |