Commit 4efbcad26d088a0b40cae308ca2990a71ea1f3f9 | zig

author	Joachim Schmidt <joachim.schmidt557@outlook.com> 2021-03-19 10:11:45+01:00
committer	Andrew Kelley <andrew@ziglang.org> 2021-04-02 12:01:52-07:00
log	4efbcad26d088a0b40cae308ca2990a71ea1f3f9
tree	299e234e4354a6a7da0e27d90e7476aa8eb27e34
parent	fc9c1b4e4ae86c36a7ff29859cc589a9816ab7d5

stage2 codegen: Extract register management code into separate file

2 files changed, 155 insertions(+), 104 deletions(-)

src
- codegen.zig
- register_manager.zig

src/codegen.zig+28-104

 @@ -18,6 +18,7 @@ const leb128 = std.leb;
 const log = std.log.scoped(.codegen);
 const build_options = @import("build_options");
 const LazySrcLoc = Module.LazySrcLoc;
+const RegisterManager = @import("register_manager.zig").RegisterManager;
 /// The codegen-related data that is stored in `ir.Inst.Block` instructions.
 pub const BlockData = struct {
 @@ -286,11 +287,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
         /// across each runtime branch upon joining.
         branch_stack: *std.ArrayList(Branch),
-        /// The key must be canonical register.
+        register_manager: RegisterManager(Self, Register, &callee_preserved_regs) = .{},
-        registers: std.AutoHashMapUnmanaged(Register, *ir.Inst) = .{},
-        free_registers: FreeRegInt = math.maxInt(FreeRegInt),
-        /// Tracks all registers allocated in the course of this function
-        allocated_registers: FreeRegInt = 0,
         /// Maps offset to what is stored there.
         stack: std.AutoHashMapUnmanaged(u32, StackAllocation) = .{},
 @@ -382,49 +379,6 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             }
         };
-        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;
-            self.allocated_registers |= mask;
-        }
-        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;
-        }
-        /// Before calling, must ensureCapacity + 1 on self.registers.
-        /// Returns `null` if all registers are allocated.
-        fn allocReg(self: *Self, inst: *ir.Inst) ?Register {
-            const free_index = @ctz(FreeRegInt, self.free_registers);
-            if (free_index >= callee_preserved_regs.len) {
-                return null;
-            }
-            const mask = @as(FreeRegInt, 1) << free_index;
-            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 });
-            return reg;
-        }
-        /// Does not track the register.
-        fn findUnusedReg(self: *Self) ?Register {
-            const free_index = @ctz(FreeRegInt, self.free_registers);
-            if (free_index >= callee_preserved_regs.len) {
-                return null;
-            }
-            return callee_preserved_regs[free_index];
-        }
         const StackAllocation = struct {
             inst: *ir.Inst,
             /// TODO do we need size? should be determined by inst.ty.abiSize()
 @@ -495,7 +449,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                 .rbrace_src = src_data.rbrace_src,
                 .source = src_data.source,
             };
-            defer function.registers.deinit(bin_file.allocator);
+            defer function.register_manager.deinit(bin_file.allocator);
             defer function.stack.deinit(bin_file.allocator);
             defer function.exitlude_jump_relocs.deinit(bin_file.allocator);
 @@ -607,10 +561,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                             .r14 = true, // lr
                         };
                         inline for (callee_preserved_regs) |reg, i| {
-                            const ShiftInt = math.Log2Int(FreeRegInt);
+                            if (self.register_manager.isRegAllocated(reg)) {
-                            const shift = @intCast(ShiftInt, i);
-                            const mask = @as(FreeRegInt, 1) << shift;
-                            if (self.allocated_registers & mask != 0) {
                                 @field(saved_regs, @tagName(reg)) = true;
                             }
                         }
 @@ -829,8 +780,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             switch (prev_value) {
                 .register => |reg| {
                     const canon_reg = toCanonicalReg(reg);
-                    _ = self.registers.remove(canon_reg);
+                    self.register_manager.freeReg(canon_reg);
-                    self.markRegFree(canon_reg);
                 },
                 else => {}, // TODO process stack allocation death
             }
 @@ -969,8 +919,8 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                 const ptr_bits = arch.ptrBitWidth();
                 const ptr_bytes: u64 = @divExact(ptr_bits, 8);
                 if (abi_size <= ptr_bytes) {
-                    try self.registers.ensureCapacity(self.gpa, self.registers.count() + 1);
+                    try self.register_manager.registers.ensureCapacity(self.gpa, self.register_manager.registers.count() + 1);
-                    if (self.allocReg(inst)) |reg| {
+                    if (self.register_manager.tryAllocReg(inst)) |reg| {
                         return MCValue{ .register = registerAlias(reg, abi_size) };
                     }
                 }
 @@ -979,26 +929,20 @@ 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 {
+            const stack_mcv = try self.allocRegOrMem(inst, false);
+            const reg_mcv = self.getResolvedInstValue(inst);
+            assert(reg == toCanonicalReg(reg_mcv.register));
+            const branch = &self.branch_stack.items[self.branch_stack.items.len - 1];
+            try branch.inst_table.put(self.gpa, inst, stack_mcv);
+            try self.genSetStack(src, inst.ty, stack_mcv.stack_offset, reg_mcv);
+        }
         /// Copies a value to a register without tracking the register. The register is not considered
         /// 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 = self.findUnusedReg() orelse b: {
+            const reg = try self.register_manager.allocRegWithoutTracking();
-                // We'll take over the first register. Move the instruction that was previously
-                // there to a stack allocation.
-                const reg = callee_preserved_regs[0];
-                const regs_entry = self.registers.remove(reg).?;
-                const spilled_inst = regs_entry.value;
-                const stack_mcv = try self.allocRegOrMem(spilled_inst, false);
-                const reg_mcv = self.getResolvedInstValue(spilled_inst);
-                assert(reg == toCanonicalReg(reg_mcv.register));
-                const branch = &self.branch_stack.items[self.branch_stack.items.len - 1];
-                try branch.inst_table.put(self.gpa, spilled_inst, stack_mcv);
-                try self.genSetStack(src, spilled_inst.ty, stack_mcv.stack_offset, reg_mcv);
-                break :b reg;
-            };
             try self.genSetReg(src, ty, reg, mcv);
             return reg;
         }
 @@ -1007,25 +951,9 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
         /// `reg_owner` is the instruction that gets associated with the register in the register table.
         /// This can have a side effect of spilling instructions to the stack to free up a register.
         fn copyToNewRegister(self: *Self, reg_owner: *ir.Inst, mcv: MCValue) !MCValue {
-            try self.registers.ensureCapacity(self.gpa, @intCast(u32, self.registers.count() + 1));
+            try self.register_manager.registers.ensureCapacity(self.gpa, @intCast(u32, self.register_manager.registers.count() + 1));
-            const reg = self.allocReg(reg_owner) orelse b: {
+            const reg = try self.register_manager.allocReg(reg_owner);
-                // We'll take over the first register. Move the instruction that was previously
-                // there to a stack allocation.
-                const reg = callee_preserved_regs[0];
-                const regs_entry = self.registers.getEntry(reg).?;
-                const spilled_inst = regs_entry.value;
-                regs_entry.value = reg_owner;
-                const stack_mcv = try self.allocRegOrMem(spilled_inst, false);
-                const reg_mcv = self.getResolvedInstValue(spilled_inst);
-                assert(reg == toCanonicalReg(reg_mcv.register));
-                const branch = &self.branch_stack.items[self.branch_stack.items.len - 1];
-                try branch.inst_table.put(self.gpa, spilled_inst, stack_mcv);
-                try self.genSetStack(reg_owner.src, spilled_inst.ty, stack_mcv.stack_offset, reg_mcv);
-                break :b reg;
-            };
             try self.genSetReg(reg_owner.src, reg_owner.ty, reg, mcv);
             return MCValue{ .register = reg };
         }
 @@ -1302,7 +1230,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                 .register => |reg| {
                     // If it's in the registers table, need to associate the register with the
                     // new instruction.
-                    if (self.registers.getEntry(toCanonicalReg(reg))) |entry| {
+                    if (self.register_manager.registers.getEntry(toCanonicalReg(reg))) |entry| {
                         entry.value = inst;
                     }
                     log.debug("reusing {} => {*}", .{ reg, inst });
 @@ -1795,7 +1723,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             const arg_index = self.arg_index;
             self.arg_index += 1;
-            if (FreeRegInt == u0) {
+            if (callee_preserved_regs.len == 0) {
                 return self.fail(inst.base.src, "TODO implement Register enum for {}", .{self.target.cpu.arch});
             }
 @@ -1807,8 +1735,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             switch (result) {
                 .register => |reg| {
-                    try self.registers.putNoClobber(self.gpa, toCanonicalReg(reg), &inst.base);
+                    try self.register_manager.getRegAssumeFree(toCanonicalReg(reg), &inst.base);
-                    self.markRegUsed(reg);
                 },
                 else => {},
             }
 @@ -2431,10 +2358,10 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             // Capture the state of register and stack allocation state so that we can revert to it.
             const parent_next_stack_offset = self.next_stack_offset;
-            const parent_free_registers = self.free_registers;
+            const parent_free_registers = self.register_manager.free_registers;
             var parent_stack = try self.stack.clone(self.gpa);
             defer parent_stack.deinit(self.gpa);
-            var parent_registers = try self.registers.clone(self.gpa);
+            var parent_registers = try self.register_manager.registers.clone(self.gpa);
             defer parent_registers.deinit(self.gpa);
             try self.branch_stack.append(.{});
 @@ -2451,8 +2378,8 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             var saved_then_branch = self.branch_stack.pop();
             defer saved_then_branch.deinit(self.gpa);
-            self.registers.deinit(self.gpa);
+            self.register_manager.registers.deinit(self.gpa);
-            self.registers = parent_registers;
+            self.register_manager.registers = parent_registers;
             parent_registers = .{};
             self.stack.deinit(self.gpa);
 @@ -2460,7 +2387,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             parent_stack = .{};
             self.next_stack_offset = parent_next_stack_offset;
-            self.free_registers = parent_free_registers;
+            self.register_manager.free_registers = parent_free_registers;
             try self.performReloc(inst.base.src, reloc);
             const else_branch = self.branch_stack.addOneAssumeCapacity();
 @@ -4049,9 +3976,6 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             },
         };
-        /// An integer whose bits represent all the registers and whether they are free.
-        const FreeRegInt = std.meta.Int(.unsigned, callee_preserved_regs.len);
         fn parseRegName(name: []const u8) ?Register {
             if (@hasDecl(Register, "parseRegName")) {
                 return Register.parseRegName(name);

src/register_manager.zig created+127

...	@@ -0,0 +1,127 @@
1	const std = @import("std");
2	const math = std.math;
3	const assert = std.debug.assert;
4	const Allocator = std.mem.Allocator;
5	const ir = @import("ir.zig");
6	const Type = @import("type.zig").Type;
7	const log = std.log.scoped(.register_manager);
8
9	pub fn RegisterManager(
10	comptime Function: type,
11	comptime Register: type,
12	comptime callee_preserved_regs: []const Register,
13	) type {
14	return struct {
15	/// The key must be canonical register.
16	registers: std.AutoHashMapUnmanaged(Register, *ir.Inst) = .{},
17	free_registers: FreeRegInt = math.maxInt(FreeRegInt),
18	/// Tracks all registers allocated in the course of this function
19	allocated_registers: FreeRegInt = 0,
20
21	const Self = @This();
22
23	/// An integer whose bits represent all the registers and whether they are free.
24	const FreeRegInt = std.meta.Int(.unsigned, callee_preserved_regs.len);
25
26	fn getFunction(self: Self) Function {
27	return @fieldParentPtr(Function, "register_manager", self);
28	}
29
30	pub fn deinit(self: Self, allocator: Allocator) void {
31	self.registers.deinit(allocator);
32	}
33
34	fn markRegUsed(self: *Self, reg: Register) void {
35	if (FreeRegInt == u0) return;
36	const index = reg.allocIndex() orelse return;
37	const ShiftInt = math.Log2Int(FreeRegInt);
38	const shift = @intCast(ShiftInt, index);
39	const mask = @as(FreeRegInt, 1) << shift;
40	self.free_registers &= ~mask;
41	self.allocated_registers \|= mask;
42	}
43
44	fn markRegFree(self: *Self, reg: Register) void {
45	if (FreeRegInt == u0) return;
46	const index = reg.allocIndex() orelse return;
47	const ShiftInt = math.Log2Int(FreeRegInt);
48	const shift = @intCast(ShiftInt, index);
49	self.free_registers \|= @as(FreeRegInt, 1) << shift;
50	}
51
52	/// Returns whether this register was allocated in the course
53	/// of this function
54	pub fn isRegAllocated(self: Self, reg: Register) bool {
55	if (FreeRegInt == u0) return false;
56	const index = reg.allocIndex() orelse return false;
57	const ShiftInt = math.Log2Int(FreeRegInt);
58	const shift = @intCast(ShiftInt, index);
59	return self.free_registers & @as(FreeRegInt, 1) << shift != 0;
60	}
61
62	/// Before calling, must ensureCapacity + 1 on self.registers.
63	/// Returns `null` if all registers are allocated.
64	pub fn tryAllocReg(self: Self, inst: ir.Inst) ?Register {
65	const free_index = @ctz(FreeRegInt, self.free_registers);
66	if (free_index >= callee_preserved_regs.len) {
67	return null;
68	}
69	const mask = @as(FreeRegInt, 1) << free_index;
70	self.free_registers &= ~mask;
71	self.allocated_registers \|= mask;
72	const reg = callee_preserved_regs[free_index];
73	self.registers.putAssumeCapacityNoClobber(reg, inst);
74	log.debug("alloc {} => {*}", .{ reg, inst });
75	return reg;
76	}
77
78	/// Before calling, must ensureCapacity + 1 on self.registers.
79	pub fn allocReg(self: Self, inst: ir.Inst) !Register {
80	return self.tryAllocReg(inst) orelse b: {
81	// We'll take over the first register. Move the instruction that was previously
82	// there to a stack allocation.
83	const reg = callee_preserved_regs[0];
84	const regs_entry = self.registers.getEntry(reg).?;
85	const spilled_inst = regs_entry.value;
86	regs_entry.value = inst;
87	try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
88
89	break :b reg;
90	};
91	}
92
93	/// Does not track the register.
94	/// Returns `null` if all registers are allocated.
95	pub fn findUnusedReg(self: *Self) ?Register {
96	const free_index = @ctz(FreeRegInt, self.free_registers);
97	if (free_index >= callee_preserved_regs.len) {
98	return null;
99	}
100	return callee_preserved_regs[free_index];
101	}
102
103	/// Does not track the register.
104	pub fn allocRegWithoutTracking(self: *Self) !Register {
105	return self.findUnusedReg() orelse b: {
106	// We'll take over the first register. Move the instruction that was previously
107	// there to a stack allocation.
108	const reg = callee_preserved_regs[0];
109	const regs_entry = self.registers.remove(reg).?;
110	const spilled_inst = regs_entry.value;
111	try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
112
113	break :b reg;
114	};
115	}
116
117	pub fn getRegAssumeFree(self: Self, reg: Register, inst: ir.Inst) !void {
118	try self.registers.putNoClobber(self.getFunction().gpa, reg, inst);
119	self.markRegUsed(reg);
120	}
121
122	pub fn freeReg(self: *Self, reg: Register) void {
123	_ = self.registers.remove(reg);
124	self.markRegFree(reg);
125	}
126	};
127	}