diff --git a/src/codegen.zig b/src/codegen.zig index c41802ca00..9d344bf1d0 100644 --- a/src/codegen.zig +++ b/src/codegen.zig @@ -952,7 +952,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type { /// allocated. A second call to `copyToTmpRegister` may return the same register. /// This can have a side effect of spilling instructions to the stack to free up a register. fn copyToTmpRegister(self: *Self, src: LazySrcLoc, ty: Type, mcv: MCValue) !Register { - const reg = try self.register_manager.allocRegWithoutTracking(&.{}); + const reg = try self.register_manager.allocReg(null, &.{}); try self.genSetReg(src, ty, reg, mcv); return reg; } @@ -2228,7 +2228,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type { switch (mc_arg) { .none => continue, .register => |reg| { - try self.register_manager.getRegWithoutTracking(reg); + try self.register_manager.getReg(reg, null); try self.genSetReg(arg.src, arg.ty, reg, arg_mcv); }, .stack_offset => { @@ -2370,7 +2370,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type { .compare_flags_signed => unreachable, .compare_flags_unsigned => unreachable, .register => |reg| { - try self.register_manager.getRegWithoutTracking(reg); + try self.register_manager.getReg(reg, null); try self.genSetReg(arg.src, arg.ty, reg, arg_mcv); }, .stack_offset => { @@ -2433,7 +2433,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type { .compare_flags_signed => unreachable, .compare_flags_unsigned => unreachable, .register => |reg| { - try self.register_manager.getRegWithoutTracking(reg); + try self.register_manager.getReg(reg, null); try self.genSetReg(arg.src, arg.ty, reg, arg_mcv); }, .stack_offset => { @@ -2486,7 +2486,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type { .register => |reg| { // TODO prevent this macho if block to be generated for all archs switch (arch) { - .x86_64, .aarch64 => try self.register_manager.getRegWithoutTracking(reg), + .x86_64, .aarch64 => try self.register_manager.getReg(reg, null), else => unreachable, } try self.genSetReg(arg.src, arg.ty, reg, arg_mcv); @@ -3190,7 +3190,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type { const arg = inst.args[i]; const arg_mcv = try self.resolveInst(arg); - try self.register_manager.getRegWithoutTracking(reg); + try self.register_manager.getReg(reg, null); try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv); } @@ -3223,7 +3223,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type { const arg = inst.args[i]; const arg_mcv = try self.resolveInst(arg); - try self.register_manager.getRegWithoutTracking(reg); + try self.register_manager.getReg(reg, null); try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv); } @@ -3258,7 +3258,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type { const arg = inst.args[i]; const arg_mcv = try self.resolveInst(arg); - try self.register_manager.getRegWithoutTracking(reg); + try self.register_manager.getReg(reg, null); try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv); } @@ -3291,7 +3291,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type { const arg = inst.args[i]; const arg_mcv = try self.resolveInst(arg); - try self.register_manager.getRegWithoutTracking(reg); + try self.register_manager.getReg(reg, null); try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv); } diff --git a/src/register_manager.zig b/src/register_manager.zig index c9235b88c1..5f7cccd60f 100644 --- a/src/register_manager.zig +++ b/src/register_manager.zig @@ -7,6 +7,9 @@ const ir = @import("ir.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; const log = std.log.scoped(.register_manager); @@ -66,77 +69,14 @@ pub fn RegisterManager( return self.allocated_registers & @as(FreeRegInt, 1) << shift != 0; } - /// Returns `null` if all registers are allocated. + /// Allocates a specified number of registers, optionally + /// tracking them. Returns `null` if not enough registers are + /// free. pub fn tryAllocRegs( self: *Self, comptime count: comptime_int, - insts: [count]*ir.Inst, - exceptions: []Register, - ) ?[count]Register { - if (self.tryAllocRegsWithoutTracking(count, exceptions)) |regs| { - for (regs) |reg, i| { - const index = reg.allocIndex().?; // allocIndex() on a callee-preserved reg should never return null - self.registers[index] = insts[i]; - self.markRegUsed(reg); - } - - return regs; - } else { - return null; - } - } - - /// Returns `null` if all registers are allocated. - pub fn tryAllocReg(self: *Self, inst: *ir.Inst, exceptions: []Register) ?Register { - return if (tryAllocRegs(self, 1, .{inst}, exceptions)) |regs| regs[0] else null; - } - - pub fn allocRegs( - self: *Self, - comptime count: comptime_int, - insts: [count]*ir.Inst, - exceptions: []Register, - ) ![count]Register { - comptime assert(count > 0 and count <= callee_preserved_regs.len); - assert(count + exceptions.len <= callee_preserved_regs.len); - - return self.tryAllocRegs(count, insts, exceptions) orelse blk: { - // We'll take over the first count registers. Spill - // the instructions that were previously there to a - // stack allocations. - var regs: [count]Register = undefined; - var i: usize = 0; - for (callee_preserved_regs) |reg| { - if (i >= count) break; - if (mem.indexOfScalar(Register, exceptions, reg) != null) continue; - regs[i] = reg; - - const index = reg.allocIndex().?; // allocIndex() on a callee-preserved reg should never return null - if (self.isRegFree(reg)) { - self.markRegUsed(reg); - } else { - const spilled_inst = self.registers[index].?; - try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst); - } - self.registers[index] = insts[i]; - - i += 1; - } - - break :blk regs; - }; - } - - pub fn allocReg(self: *Self, inst: *ir.Inst, exceptions: []Register) !Register { - return (try self.allocRegs(1, .{inst}, exceptions))[0]; - } - - /// Does not track the registers. - /// Returns `null` if not enough registers are free. - pub fn tryAllocRegsWithoutTracking( - self: *Self, - comptime count: comptime_int, - exceptions: []Register, + insts: [count]?*ir.Inst, + exceptions: []const Register, ) ?[count]Register { comptime if (callee_preserved_regs.len == 0) return null; comptime assert(count > 0 and count <= callee_preserved_regs.len); @@ -156,18 +96,40 @@ pub fn RegisterManager( } } - return if (i < count) null else regs; + if (i == count) { + for (regs) |reg, j| { + if (insts[j]) |inst| { + // Track the register + const index = reg.allocIndex().?; // allocIndex() on a callee-preserved reg should never return null + self.registers[index] = inst; + self.markRegUsed(reg); + } + } + + return regs; + } else return null; } - /// Does not track the register. - /// Returns `null` if all registers are allocated. - pub fn tryAllocRegWithoutTracking(self: *Self, exceptions: []Register) ?Register { - return if (self.tryAllocRegsWithoutTracking(1, exceptions)) |regs| regs[0] else null; + /// Allocates a register and optionally tracks it with a + /// corresponding instruction. Returns `null` if all registers + /// are allocated. + pub fn tryAllocReg(self: *Self, inst: ?*ir.Inst, exceptions: []const Register) ?Register { + return if (tryAllocRegs(self, 1, .{inst}, exceptions)) |regs| regs[0] else null; } - /// Does not track the registers - pub fn allocRegsWithoutTracking(self: *Self, comptime count: comptime_int, exceptions: []Register) ![count]Register { - return self.tryAllocRegsWithoutTracking(count, exceptions) orelse blk: { + /// Allocates a specified number of registers, optionally + /// tracking them. Asserts that count + exceptions.len is not + /// larger than the total number of registers available. + pub fn allocRegs( + self: *Self, + comptime count: comptime_int, + insts: [count]?*ir.Inst, + exceptions: []const Register, + ) ![count]Register { + comptime assert(count > 0 and count <= callee_preserved_regs.len); + assert(count + exceptions.len <= callee_preserved_regs.len); + + return self.tryAllocRegs(count, insts, exceptions) orelse blk: { // We'll take over the first count registers. Spill // the instructions that were previously there to a // stack allocations. @@ -179,11 +141,22 @@ pub fn RegisterManager( regs[i] = reg; const index = reg.allocIndex().?; // allocIndex() on a callee-preserved reg should never return null - if (!self.isRegFree(reg)) { - const spilled_inst = self.registers[index].?; - try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst); - self.registers[index] = null; - self.markRegFree(reg); + if (insts[i]) |inst| { + // Track the register + if (self.isRegFree(reg)) { + self.markRegUsed(reg); + } else { + const spilled_inst = self.registers[index].?; + try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst); + } + self.registers[index] = inst; + } else { + // Don't track the register + if (!self.isRegFree(reg)) { + const spilled_inst = self.registers[index].?; + try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst); + self.freeReg(reg); + } } i += 1; @@ -193,39 +166,36 @@ pub fn RegisterManager( }; } - /// Does not track the register. - pub fn allocRegWithoutTracking(self: *Self, exceptions: []Register) !Register { - return (try self.allocRegsWithoutTracking(1, exceptions))[0]; + /// Allocates a register and optionally tracks it with a + /// corresponding instruction. + pub fn allocReg(self: *Self, inst: ?*ir.Inst, exceptions: []const Register) !Register { + return (try self.allocRegs(1, .{inst}, exceptions))[0]; } - /// Allocates the specified register with the specified - /// instruction. Spills the register if it is currently - /// allocated. - pub fn getReg(self: *Self, reg: Register, inst: *ir.Inst) !void { + /// Spills the register if it is currently allocated. If a + /// corresponding instruction is passed, will also track this + /// register. + pub fn getReg(self: *Self, reg: Register, inst: ?*ir.Inst) !void { const index = reg.allocIndex() orelse return; - if (!self.isRegFree(reg)) { - // Move the instruction that was previously there to a - // stack allocation. - const spilled_inst = self.registers[index].?; - self.registers[index] = inst; - try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst); - } else { - self.getRegAssumeFree(reg, inst); - } - } - - /// Spills the register if it is currently allocated. - /// Does not track the register. - pub fn getRegWithoutTracking(self: *Self, reg: Register) !void { - const index = reg.allocIndex() orelse return; - - if (!self.isRegFree(reg)) { - // Move the instruction that was previously there to a - // stack allocation. - const spilled_inst = self.registers[index].?; - try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst); - self.markRegFree(reg); + if (inst) |tracked_inst| + if (!self.isRegFree(reg)) { + // Move the instruction that was previously there to a + // stack allocation. + const spilled_inst = self.registers[index].?; + self.registers[index] = tracked_inst; + try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst); + } else { + self.getRegAssumeFree(reg, tracked_inst); + } + else { + if (!self.isRegFree(reg)) { + // Move the instruction that was previously there to a + // stack allocation. + const spilled_inst = self.registers[index].?; + try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst); + self.freeReg(reg); + } } } @@ -250,42 +220,63 @@ pub fn RegisterManager( }; } -const MockRegister = enum(u2) { +const MockRegister1 = enum(u2) { r0, r1, r2, r3, - pub fn allocIndex(self: MockRegister) ?u2 { - inline for (mock_callee_preserved_regs) |cpreg, i| { + pub fn allocIndex(self: MockRegister1) ?u2 { + inline for (callee_preserved_regs) |cpreg, i| { if (self == cpreg) return i; } return null; } + + const callee_preserved_regs = [_]MockRegister1{ .r2, .r3 }; }; -const mock_callee_preserved_regs = [_]MockRegister{ .r2, .r3 }; +const MockRegister2 = enum(u2) { + r0, + r1, + r2, + r3, -const MockFunction = struct { - allocator: *Allocator, - register_manager: RegisterManager(Self, MockRegister, &mock_callee_preserved_regs) = .{}, - spilled: std.ArrayListUnmanaged(MockRegister) = .{}, - - const Self = @This(); - - pub fn deinit(self: *Self) void { - self.spilled.deinit(self.allocator); + pub fn allocIndex(self: MockRegister2) ?u2 { + inline for (callee_preserved_regs) |cpreg, i| { + if (self == cpreg) return i; + } + return null; } - pub fn spillInstruction(self: *Self, src: LazySrcLoc, reg: MockRegister, inst: *ir.Inst) !void { - try self.spilled.append(self.allocator, reg); - } + const callee_preserved_regs = [_]MockRegister2{ .r0, .r1, .r2, .r3 }; }; -test "tryAllocReg: no spilling" { +fn MockFunction(comptime Register: type) type { + return struct { + allocator: *Allocator, + register_manager: RegisterManager(Self, Register, &Register.callee_preserved_regs) = .{}, + spilled: std.ArrayListUnmanaged(Register) = .{}, + + const Self = @This(); + + pub fn deinit(self: *Self) void { + self.spilled.deinit(self.allocator); + } + + pub fn spillInstruction(self: *Self, src: LazySrcLoc, reg: Register, inst: *ir.Inst) !void { + try self.spilled.append(self.allocator, reg); + } + }; +} + +const MockFunction1 = MockFunction(MockRegister1); +const MockFunction2 = MockFunction(MockRegister2); + +test "default state" { const allocator = std.testing.allocator; - var function = MockFunction{ + var function = MockFunction1{ .allocator = allocator, }; defer function.deinit(); @@ -296,27 +287,48 @@ test "tryAllocReg: no spilling" { .src = .unneeded, }; - try std.testing.expect(!function.register_manager.isRegAllocated(.r2)); - try std.testing.expect(!function.register_manager.isRegAllocated(.r3)); + try expect(!function.register_manager.isRegAllocated(.r2)); + try expect(!function.register_manager.isRegAllocated(.r3)); + try expect(function.register_manager.isRegFree(.r2)); + try expect(function.register_manager.isRegFree(.r3)); +} - try std.testing.expectEqual(@as(?MockRegister, .r2), function.register_manager.tryAllocReg(&mock_instruction, &.{})); - try std.testing.expectEqual(@as(?MockRegister, .r3), function.register_manager.tryAllocReg(&mock_instruction, &.{})); - try std.testing.expectEqual(@as(?MockRegister, null), function.register_manager.tryAllocReg(&mock_instruction, &.{})); +test "tryAllocReg: no spilling" { + const allocator = std.testing.allocator; - try std.testing.expect(function.register_manager.isRegAllocated(.r2)); - try std.testing.expect(function.register_manager.isRegAllocated(.r3)); + var function = MockFunction1{ + .allocator = allocator, + }; + defer function.deinit(); + + var mock_instruction = ir.Inst{ + .tag = .breakpoint, + .ty = Type.initTag(.void), + .src = .unneeded, + }; + + try expectEqual(@as(?MockRegister1, .r2), function.register_manager.tryAllocReg(&mock_instruction, &.{})); + try expectEqual(@as(?MockRegister1, .r3), function.register_manager.tryAllocReg(&mock_instruction, &.{})); + try expectEqual(@as(?MockRegister1, null), function.register_manager.tryAllocReg(&mock_instruction, &.{})); + + try expect(function.register_manager.isRegAllocated(.r2)); + try expect(function.register_manager.isRegAllocated(.r3)); + try expect(!function.register_manager.isRegFree(.r2)); + try expect(!function.register_manager.isRegFree(.r3)); function.register_manager.freeReg(.r2); function.register_manager.freeReg(.r3); - try std.testing.expect(function.register_manager.isRegAllocated(.r2)); - try std.testing.expect(function.register_manager.isRegAllocated(.r3)); + try expect(function.register_manager.isRegAllocated(.r2)); + try expect(function.register_manager.isRegAllocated(.r3)); + try expect(function.register_manager.isRegFree(.r2)); + try expect(function.register_manager.isRegFree(.r3)); } test "allocReg: spilling" { const allocator = std.testing.allocator; - var function = MockFunction{ + var function = MockFunction1{ .allocator = allocator, }; defer function.deinit(); @@ -327,26 +339,76 @@ test "allocReg: spilling" { .src = .unneeded, }; - try std.testing.expect(!function.register_manager.isRegAllocated(.r2)); - try std.testing.expect(!function.register_manager.isRegAllocated(.r3)); - - try std.testing.expectEqual(@as(?MockRegister, .r2), try function.register_manager.allocReg(&mock_instruction, &.{})); - try std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction, &.{})); + try expectEqual(@as(?MockRegister1, .r2), try function.register_manager.allocReg(&mock_instruction, &.{})); + try expectEqual(@as(?MockRegister1, .r3), try function.register_manager.allocReg(&mock_instruction, &.{})); // Spill a register - try std.testing.expectEqual(@as(?MockRegister, .r2), try function.register_manager.allocReg(&mock_instruction, &.{})); - try std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items); + try expectEqual(@as(?MockRegister1, .r2), try function.register_manager.allocReg(&mock_instruction, &.{})); + try expectEqualSlices(MockRegister1, &[_]MockRegister1{.r2}, function.spilled.items); // No spilling necessary function.register_manager.freeReg(.r3); - try std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction, &.{})); - try std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items); + try expectEqual(@as(?MockRegister1, .r3), try function.register_manager.allocReg(&mock_instruction, &.{})); + try expectEqualSlices(MockRegister1, &[_]MockRegister1{.r2}, function.spilled.items); + + // Exceptions + function.register_manager.freeReg(.r2); + function.register_manager.freeReg(.r3); + try expectEqual(@as(?MockRegister1, .r3), try function.register_manager.allocReg(&mock_instruction, &.{.r2})); +} + +test "tryAllocRegs" { + const allocator = std.testing.allocator; + + var function = MockFunction2{ + .allocator = allocator, + }; + defer function.deinit(); + + var mock_instruction = ir.Inst{ + .tag = .breakpoint, + .ty = Type.initTag(.void), + .src = .unneeded, + }; + + try expectEqual([_]MockRegister2{ .r0, .r1, .r2 }, function.register_manager.tryAllocRegs(3, .{ null, null, null }, &.{}).?); + + // Exceptions + 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}).?); +} + +test "allocRegs" { + const allocator = std.testing.allocator; + + var function = MockFunction2{ + .allocator = allocator, + }; + defer function.deinit(); + + var mock_instruction = ir.Inst{ + .tag = .breakpoint, + .ty = Type.initTag(.void), + .src = .unneeded, + }; + + try expectEqual([_]MockRegister2{ .r0, .r1, .r2 }, try function.register_manager.allocRegs(3, .{ + &mock_instruction, + &mock_instruction, + &mock_instruction, + }, &.{})); + + // Exceptions + 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); } test "getReg" { const allocator = std.testing.allocator; - var function = MockFunction{ + var function = MockFunction1{ .allocator = allocator, }; defer function.deinit(); @@ -357,18 +419,19 @@ test "getReg" { .src = .unneeded, }; - try std.testing.expect(!function.register_manager.isRegAllocated(.r2)); - try std.testing.expect(!function.register_manager.isRegAllocated(.r3)); - try function.register_manager.getReg(.r3, &mock_instruction); - try std.testing.expect(!function.register_manager.isRegAllocated(.r2)); - try std.testing.expect(function.register_manager.isRegAllocated(.r3)); + try expect(!function.register_manager.isRegAllocated(.r2)); + try expect(function.register_manager.isRegAllocated(.r3)); + try expect(function.register_manager.isRegFree(.r2)); + try expect(!function.register_manager.isRegFree(.r3)); // Spill r3 try function.register_manager.getReg(.r3, &mock_instruction); - try std.testing.expect(!function.register_manager.isRegAllocated(.r2)); - try std.testing.expect(function.register_manager.isRegAllocated(.r3)); - try std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r3}, function.spilled.items); + try expect(!function.register_manager.isRegAllocated(.r2)); + try expect(function.register_manager.isRegAllocated(.r3)); + try expect(function.register_manager.isRegFree(.r2)); + try expect(!function.register_manager.isRegFree(.r3)); + try expectEqualSlices(MockRegister1, &[_]MockRegister1{.r3}, function.spilled.items); }