From 35503b3d3fe1bfce19f1ea3e78a75ce87b0ed646 Mon Sep 17 00:00:00 2001 From: joachimschmidt557 Date: Wed, 26 Jan 2022 11:47:39 +0100 Subject: [PATCH] stage2 regalloc: Add freezeRegs/unfreezeRegs API The freeze/unfreeze API replaces the exceptions API for hopefully preventing bugs in codegen code using the RegisterManager. The exceptions API is still available for backwards compatibility and will be removed once all backends transition to the new freeze/unfreeze API. --- src/register_manager.zig | 132 +++++++++++++++++++++++++++++++++------ 1 file changed, 114 insertions(+), 18 deletions(-) diff --git a/src/register_manager.zig b/src/register_manager.zig index add0335374..886a717342 100644 --- a/src/register_manager.zig +++ b/src/register_manager.zig @@ -6,7 +6,6 @@ const Allocator = std.mem.Allocator; const Air = @import("Air.zig"); const Type = @import("type.zig").Type; const Module = @import("Module.zig"); -const LazySrcLoc = Module.LazySrcLoc; const expect = std.testing.expect; const expectEqual = std.testing.expectEqual; const expectEqualSlices = std.testing.expectEqualSlices; @@ -19,15 +18,25 @@ pub fn RegisterManager( comptime callee_preserved_regs: []const Register, ) type { return struct { + /// Tracks the AIR instruction allocated to every register or + /// `null` if no instruction is allocated to a register + /// /// The key must be canonical register. registers: [callee_preserved_regs.len]?Air.Inst.Index = [_]?Air.Inst.Index{null} ** callee_preserved_regs.len, + /// Tracks which registers are free (in which case the + /// corresponding bit is set to 1) free_registers: FreeRegInt = math.maxInt(FreeRegInt), - /// Tracks all registers allocated in the course of this function + /// Tracks all registers allocated in the course of this + /// function allocated_registers: FreeRegInt = 0, + /// Tracks registers which are temporarily blocked from being + /// allocated + frozen_registers: FreeRegInt = 0, const Self = @This(); - /// An integer whose bits represent all the registers and whether they are free. + /// An integer whose bits represent all the registers and + /// whether they are free. const FreeRegInt = std.meta.Int(.unsigned, callee_preserved_regs.len); const ShiftInt = math.Log2Int(FreeRegInt); @@ -35,43 +44,76 @@ pub fn RegisterManager( return @fieldParentPtr(Function, "register_manager", self); } - fn markRegUsed(self: *Self, reg: Register) void { - if (FreeRegInt == u0) return; - const index = reg.allocIndex() orelse return; + fn getRegisterMask(reg: Register) ?FreeRegInt { + if (FreeRegInt == u0) return null; + const index = reg.allocIndex() orelse return null; const shift = @intCast(ShiftInt, index); const mask = @as(FreeRegInt, 1) << shift; + return mask; + } + + fn markRegUsed(self: *Self, reg: Register) void { + const mask = getRegisterMask(reg) orelse return; self.free_registers &= ~mask; self.allocated_registers |= mask; } fn markRegFree(self: *Self, reg: Register) void { - if (FreeRegInt == u0) return; - const index = reg.allocIndex() orelse return; - const shift = @intCast(ShiftInt, index); - self.free_registers |= @as(FreeRegInt, 1) << shift; + const mask = getRegisterMask(reg) orelse return; + self.free_registers |= mask; } /// Returns true when this register is not tracked pub fn isRegFree(self: Self, reg: Register) bool { - if (FreeRegInt == u0) return true; - const index = reg.allocIndex() orelse return true; - const shift = @intCast(ShiftInt, index); - return self.free_registers & @as(FreeRegInt, 1) << shift != 0; + const mask = getRegisterMask(reg) orelse return true; + return self.free_registers & mask != 0; } /// Returns whether this register was allocated in the course /// of this function. + /// /// Returns false when this register is not tracked pub fn isRegAllocated(self: Self, reg: Register) bool { - if (FreeRegInt == u0) return false; - const index = reg.allocIndex() orelse return false; - const shift = @intCast(ShiftInt, index); - return self.allocated_registers & @as(FreeRegInt, 1) << shift != 0; + const mask = getRegisterMask(reg) orelse return false; + return self.allocated_registers & mask != 0; + } + + /// Returns whether this register is frozen + /// + /// Returns false when this register is not tracked + pub fn isRegFrozen(self: Self, reg: Register) bool { + const mask = getRegisterMask(reg) orelse return false; + return self.frozen_registers & mask != 0; + } + + /// Prevents the registers from being allocated until they are + /// unfrozen again + pub fn freezeRegs(self: *Self, regs: []const Register) void { + for (regs) |reg| { + const mask = getRegisterMask(reg) orelse continue; + self.frozen_registers |= mask; + } + } + + /// Enables the allocation of the registers + pub fn unfreezeRegs(self: *Self, regs: []const Register) void { + for (regs) |reg| { + const mask = getRegisterMask(reg) orelse continue; + self.frozen_registers &= ~mask; + } + } + + /// Returns true when at least one register is frozen + pub fn frozenRegsExist(self: Self) bool { + return self.frozen_registers != 0; } /// Allocates a specified number of registers, optionally /// tracking them. Returns `null` if not enough registers are /// free. + /// + /// Exceptions are deprecated, use freezeRegs and unfreezeRegs + /// instead. pub fn tryAllocRegs( self: *Self, comptime count: comptime_int, @@ -90,6 +132,7 @@ pub fn RegisterManager( for (callee_preserved_regs) |reg| { if (i >= count) break; if (mem.indexOfScalar(Register, exceptions, reg) != null) continue; + if (self.isRegFrozen(reg)) continue; if (self.isRegFree(reg)) { regs[i] = reg; i += 1; @@ -113,6 +156,9 @@ pub fn RegisterManager( /// Allocates a register and optionally tracks it with a /// corresponding instruction. Returns `null` if all registers /// are allocated. + /// + /// Exceptions are deprecated, use freezeRegs and unfreezeRegs + /// instead. pub fn tryAllocReg(self: *Self, inst: ?Air.Inst.Index, exceptions: []const Register) ?Register { return if (tryAllocRegs(self, 1, .{inst}, exceptions)) |regs| regs[0] else null; } @@ -120,6 +166,9 @@ pub fn RegisterManager( /// Allocates a specified number of registers, optionally /// tracking them. Asserts that count + exceptions.len is not /// larger than the total number of registers available. + /// + /// Exceptions are deprecated, use freezeRegs and unfreezeRegs + /// instead. pub fn allocRegs( self: *Self, comptime count: comptime_int, @@ -138,6 +187,7 @@ pub fn RegisterManager( for (callee_preserved_regs) |reg| { if (i >= count) break; if (mem.indexOfScalar(Register, exceptions, reg) != null) continue; + if (self.isRegFrozen(reg)) continue; regs[i] = reg; const index = reg.allocIndex().?; // allocIndex() on a callee-preserved reg should never return null @@ -171,6 +221,9 @@ pub fn RegisterManager( /// Allocates a register and optionally tracks it with a /// corresponding instruction. + /// + /// Exceptions are deprecated, use freezeRegs and unfreezeRegs + /// instead. pub fn allocReg(self: *Self, inst: ?Air.Inst.Index, exceptions: []const Register) !Register { return (try self.allocRegs(1, .{inst}, exceptions))[0]; } @@ -343,9 +396,22 @@ test "allocReg: spilling" { try expectEqualSlices(MockRegister1, &[_]MockRegister1{.r2}, function.spilled.items); // Exceptions + // + // TODO deprecated, remove test once no backend uses exceptions + // anymore function.register_manager.freeReg(.r2); function.register_manager.freeReg(.r3); try expectEqual(@as(?MockRegister1, .r3), try function.register_manager.allocReg(mock_instruction, &.{.r2})); + + // Frozen registers + function.register_manager.freeReg(.r3); + { + function.register_manager.freezeRegs(&.{.r2}); + defer function.register_manager.unfreezeRegs(&.{.r2}); + + try expectEqual(@as(?MockRegister1, .r3), try function.register_manager.allocReg(mock_instruction, &.{})); + } + try expect(!function.register_manager.frozenRegsExist()); } test "tryAllocRegs" { @@ -359,10 +425,25 @@ test "tryAllocRegs" { try expectEqual([_]MockRegister2{ .r0, .r1, .r2 }, function.register_manager.tryAllocRegs(3, .{ null, null, null }, &.{}).?); // Exceptions + // + // TODO deprecated, remove test once no backend uses exceptions + // anymore function.register_manager.freeReg(.r0); function.register_manager.freeReg(.r1); function.register_manager.freeReg(.r2); try expectEqual([_]MockRegister2{ .r0, .r2, .r3 }, function.register_manager.tryAllocRegs(3, .{ null, null, null }, &.{.r1}).?); + + // Frozen registers + function.register_manager.freeReg(.r0); + function.register_manager.freeReg(.r2); + function.register_manager.freeReg(.r3); + { + function.register_manager.freezeRegs(&.{.r1}); + defer function.register_manager.unfreezeRegs(&.{.r1}); + + try expectEqual([_]MockRegister2{ .r0, .r2, .r3 }, function.register_manager.tryAllocRegs(3, .{ null, null, null }, &.{}).?); + } + try expect(!function.register_manager.frozenRegsExist()); } test "allocRegs" { @@ -382,8 +463,23 @@ test "allocRegs" { }, &.{})); // Exceptions + // + // TODO deprecated, remove test once no backend uses exceptions + // anymore try expectEqual([_]MockRegister2{ .r0, .r2, .r3 }, try function.register_manager.allocRegs(3, .{ null, null, null }, &.{.r1})); try expectEqualSlices(MockRegister2, &[_]MockRegister2{ .r0, .r2 }, function.spilled.items); + + // Frozen registers + function.register_manager.freeReg(.r0); + function.register_manager.freeReg(.r2); + function.register_manager.freeReg(.r3); + { + function.register_manager.freezeRegs(&.{.r1}); + defer function.register_manager.unfreezeRegs(&.{.r1}); + + try expectEqual([_]MockRegister2{ .r0, .r2, .r3 }, try function.register_manager.allocRegs(3, .{ null, null, null }, &.{})); + } + try expect(!function.register_manager.frozenRegsExist()); } test "getReg" {