From 228a1ce3e8d112a7710fa47c6b9486cf320b5d6f Mon Sep 17 00:00:00 2001 From: joachimschmidt557 Date: Sat, 27 Mar 2021 17:24:51 +0100 Subject: [PATCH] stage2 register_manager: Add unit tests for tryAllocReg and allocReg --- src/codegen.zig | 2 +- src/register_manager.zig | 111 +++++++++++++++++++++++++++++++++++++-- 2 files changed, 107 insertions(+), 6 deletions(-) diff --git a/src/codegen.zig b/src/codegen.zig index d5282e2797..beb3540d37 100644 --- a/src/codegen.zig +++ b/src/codegen.zig @@ -929,7 +929,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type { return MCValue{ .stack_offset = stack_offset }; } - pub fn spillInstruction(self: *Self, src: usize, reg: Register, inst: *ir.Inst) !void { + pub fn spillInstruction(self: *Self, src: LazySrcLoc, reg: Register, inst: *ir.Inst) !void { const stack_mcv = try self.allocRegOrMem(inst, false); const reg_mcv = self.getResolvedInstValue(inst); assert(reg == toCanonicalReg(reg_mcv.register)); diff --git a/src/register_manager.zig b/src/register_manager.zig index 165b33870b..abb1632165 100644 --- a/src/register_manager.zig +++ b/src/register_manager.zig @@ -4,6 +4,9 @@ const assert = std.debug.assert; const Allocator = std.mem.Allocator; const ir = @import("ir.zig"); const Type = @import("type.zig").Type; +const Module = @import("Module.zig"); +const LazySrcLoc = Module.LazySrcLoc; + const log = std.log.scoped(.register_manager); pub fn RegisterManager( @@ -22,6 +25,7 @@ pub fn RegisterManager( /// 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); fn getFunction(self: *Self) *Function { return @fieldParentPtr(Function, "register_manager", self); @@ -34,7 +38,6 @@ pub fn RegisterManager( fn markRegUsed(self: *Self, reg: Register) void { if (FreeRegInt == u0) return; const index = reg.allocIndex() orelse return; - const ShiftInt = math.Log2Int(FreeRegInt); const shift = @intCast(ShiftInt, index); const mask = @as(FreeRegInt, 1) << shift; self.free_registers &= ~mask; @@ -44,7 +47,6 @@ pub fn RegisterManager( fn markRegFree(self: *Self, reg: Register) void { if (FreeRegInt == u0) return; const index = reg.allocIndex() orelse return; - const ShiftInt = math.Log2Int(FreeRegInt); const shift = @intCast(ShiftInt, index); self.free_registers |= @as(FreeRegInt, 1) << shift; } @@ -54,9 +56,8 @@ pub fn RegisterManager( pub fn isRegAllocated(self: Self, reg: Register) bool { if (FreeRegInt == u0) return false; const index = reg.allocIndex() orelse return false; - const ShiftInt = math.Log2Int(FreeRegInt); const shift = @intCast(ShiftInt, index); - return self.free_registers & @as(FreeRegInt, 1) << shift != 0; + return self.allocated_registers & @as(FreeRegInt, 1) << shift != 0; } /// Before calling, must ensureCapacity + 1 on self.registers. @@ -66,9 +67,16 @@ pub fn RegisterManager( if (free_index >= callee_preserved_regs.len) { return null; } - const mask = @as(FreeRegInt, 1) << free_index; + + // This is necessary because the return type of @ctz is 1 + // bit longer than ShiftInt if callee_preserved_regs.len + // is a power of two. This int cast is always safe because + // free_index < callee_preserved_regs.len + const shift = @intCast(ShiftInt, free_index); + const mask = @as(FreeRegInt, 1) << shift; self.free_registers &= ~mask; self.allocated_registers |= mask; + const reg = callee_preserved_regs[free_index]; self.registers.putAssumeCapacityNoClobber(reg, inst); log.debug("alloc {} => {*}", .{ reg, inst }); @@ -125,3 +133,96 @@ pub fn RegisterManager( } }; } + +const MockRegister = enum(u2) { + r0, r1, r2, r3, + + pub fn allocIndex(self: MockRegister) ?u2 { + inline for (mock_callee_preserved_regs) |cpreg, i| { + if (self == cpreg) return i; + } + return null; + } +}; + +const mock_callee_preserved_regs = [_]MockRegister{ .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.register_manager.deinit(self.allocator); + self.spilled.deinit(self.allocator); + } + + pub fn spillInstruction(self: *Self, src: LazySrcLoc, reg: MockRegister, inst: *ir.Inst) !void { + try self.spilled.append(self.allocator, reg); + } +}; + +test "tryAllocReg: no spilling" { + const allocator = std.testing.allocator; + + var function = MockFunction{ + .allocator = allocator, + }; + defer function.deinit(); + + var mock_instruction = ir.Inst{ + .tag = .breakpoint, + .ty = Type.initTag(.void), + .src = .unneeded, + }; + + std.testing.expect(!function.register_manager.isRegAllocated(.r2)); + std.testing.expect(!function.register_manager.isRegAllocated(.r3)); + + try function.register_manager.registers.ensureCapacity(allocator, function.register_manager.registers.count() + 2); + std.testing.expectEqual(@as(?MockRegister, .r2), function.register_manager.tryAllocReg(&mock_instruction)); + std.testing.expectEqual(@as(?MockRegister, .r3), function.register_manager.tryAllocReg(&mock_instruction)); + std.testing.expectEqual(@as(?MockRegister, null), function.register_manager.tryAllocReg(&mock_instruction)); + + std.testing.expect(function.register_manager.isRegAllocated(.r2)); + std.testing.expect(function.register_manager.isRegAllocated(.r3)); + + function.register_manager.freeReg(.r2); + function.register_manager.freeReg(.r3); + + std.testing.expect(function.register_manager.isRegAllocated(.r2)); + std.testing.expect(function.register_manager.isRegAllocated(.r3)); +} + +test "allocReg: spilling" { + const allocator = std.testing.allocator; + + var function = MockFunction{ + .allocator = allocator, + }; + defer function.deinit(); + + var mock_instruction = ir.Inst{ + .tag = .breakpoint, + .ty = Type.initTag(.void), + .src = .unneeded, + }; + + std.testing.expect(!function.register_manager.isRegAllocated(.r2)); + std.testing.expect(!function.register_manager.isRegAllocated(.r3)); + + try function.register_manager.registers.ensureCapacity(allocator, function.register_manager.registers.count() + 2); + std.testing.expectEqual(@as(?MockRegister, .r2), try function.register_manager.allocReg(&mock_instruction)); + std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction)); + + // Spill a register + std.testing.expectEqual(@as(?MockRegister, .r2), try function.register_manager.allocReg(&mock_instruction)); + std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items); + + // No spilling necessary + function.register_manager.freeReg(.r3); + std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction)); + std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items); +}