stage2 register manager: clean up API and add more unit tests

This commit is contained in:
joachimschmidt557 2021-05-14 17:09:11 +08:00 committed by Andrew Kelley
parent 5185b5619a
commit 8b4e91e18c
2 changed files with 228 additions and 165 deletions

View File

@ -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);
}

View File

@ -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);
}