Commit 376d0878ec2e366633e2dbc7c91ab7ae2a6ae5b7 | zig

author	Jakub Konka <kubkon@jakubkonka.com> 2022-03-29 18:08:03+02:00
committer	Jakub Konka <kubkon@jakubkonka.com> 2022-03-30 00:37:42+02:00
log	376d0878ec2e366633e2dbc7c91ab7ae2a6ae5b7
tree	d3c55c3acdd58283c85057cf7f351925f5e8969f
parent	d447cd940d7da884f0d699d9da679d8bbabb237a

x64: spill .rdi to stack if expecting return value saved on stack

Since .rdi is not part of the callee saved registers, it needs to be
proactively spilled to the stack so that we don't clobber the
return address where to save the return value.

1 files changed, 45 insertions(+), 68 deletions(-)

src
- arch
  - x86_64
    CodeGen.zig

src/arch/x86_64/CodeGen.zig+45-68

 @@ -58,7 +58,6 @@ arg_index: u32,
 src_loc: Module.SrcLoc,
 stack_align: u32,
-ret_backpatches: std.ArrayListUnmanaged(Mir.Inst.Index) = .{},
 compare_flags_inst: ?Air.Inst.Index = null,
 /// MIR Instructions
 @@ -353,7 +352,6 @@ pub fn generate(
             std.AutoHashMap(Mir.Inst.Index, Air.Inst.Index).init(bin_file.allocator)
         else {},
     };
-    defer function.ret_backpatches.deinit(bin_file.allocator);
     defer function.stack.deinit(bin_file.allocator);
     defer function.blocks.deinit(bin_file.allocator);
     defer function.exitlude_jump_relocs.deinit(bin_file.allocator);
 @@ -482,6 +480,21 @@ fn gen(self: *Self) InnerError!void {
             .data = undefined,
         });
+        if (self.ret_mcv == .stack_offset) {
+            // The address where to store the return value for the caller is in `.rdi`
+            // register which the callee is free to clobber. Therefore, we purposely
+            // spill it to stack immediately.
+            const ptr_ty = Type.usize;
+            const abi_size = @intCast(u32, ptr_ty.abiSize(self.target.*));
+            const abi_align = ptr_ty.abiAlignment(self.target.*);
+            const stack_offset = mem.alignForwardGeneric(u32, self.next_stack_offset + abi_size, abi_align);
+            self.next_stack_offset = stack_offset;
+            self.max_end_stack = @maximum(self.max_end_stack, self.next_stack_offset);
+            try self.genSetStack(ptr_ty, @intCast(i32, stack_offset), MCValue{ .register = .rdi }, .{});
+            self.ret_mcv = MCValue{ .stack_offset = @intCast(i32, stack_offset) };
+            log.debug("gen: spilling .rdi to stack at offset {}", .{stack_offset});
+        }
         _ = try self.addInst(.{
             .tag = .dbg_prologue_end,
             .ops = undefined,
 @@ -519,23 +532,6 @@ fn gen(self: *Self) InnerError!void {
             var disp = data.disp + 8;
             inline for (callee_preserved_regs) |reg, i| {
                 if (self.register_manager.isRegAllocated(reg)) {
-                    if (reg.to64() == .rdi) {
-                        for (self.ret_backpatches.items) |inst| {
-                            log.debug(".rdi was spilled, backpatching with mov from stack at offset {}", .{
-                                -@intCast(i32, disp),
-                            });
-                            const ops = Mir.Ops.decode(self.mir_instructions.items(.ops)[inst]);
-                            self.mir_instructions.set(inst, Mir.Inst{
-                                .tag = .mov,
-                                .ops = (Mir.Ops{
-                                    .reg1 = ops.reg1,
-                                    .reg2 = .rbp,
-                                    .flags = 0b01,
-                                }).encode(),
-                                .data = .{ .imm = @bitCast(u32, -@intCast(i32, disp)) },
-                            });
-                        }
-                    }
                     data.regs |= 1 << @intCast(u5, i);
                     self.max_end_stack += 8;
                     disp += 8;
 @@ -912,8 +908,7 @@ fn allocMem(self: *Self, inst: Air.Inst.Index, abi_size: u32, abi_align: u32) !u
     // TODO find a free slot instead of always appending
     const offset = mem.alignForwardGeneric(u32, self.next_stack_offset + abi_size, abi_align);
     self.next_stack_offset = offset;
-    if (self.next_stack_offset > self.max_end_stack)
+    self.max_end_stack = @maximum(self.max_end_stack, self.next_stack_offset);
-        self.max_end_stack = self.next_stack_offset;
     try self.stack.putNoClobber(self.gpa, offset, .{
         .inst = inst,
         .size = abi_size,
 @@ -3316,7 +3311,6 @@ fn genBinMathOpMir(self: *Self, mir_tag: Mir.Inst.Tag, dst_ty: Type, dst_mcv: MC
 }
 /// Performs multi-operand integer multiplication between dst_mcv and src_mcv, storing the result in dst_mcv.
-/// Does not use/spill .rax/.rdx.
 /// Does not support byte-size operands.
 fn genIntMulComplexOpMir(self: *Self, dst_ty: Type, dst_mcv: MCValue, src_mcv: MCValue) !void {
     const abi_size = @intCast(u32, dst_ty.abiSize(self.target.*));
 @@ -3530,10 +3524,11 @@ fn airCall(self: *Self, inst: Air.Inst.Index, modifier: std.builtin.CallOptions.
         const ret_abi_size = @intCast(u32, ret_ty.abiSize(self.target.*));
         const ret_abi_align = @intCast(u32, ret_ty.abiAlignment(self.target.*));
         const stack_offset = @intCast(i32, try self.allocMem(inst, ret_abi_size, ret_abi_align));
+        log.debug("airCall: return value on stack at offset {}", .{stack_offset});
         try self.register_manager.getReg(.rdi, null);
-        self.register_manager.freezeRegs(&.{.rdi});
         try self.genSetReg(Type.usize, .rdi, .{ .ptr_stack_offset = stack_offset });
+        self.register_manager.freezeRegs(&.{.rdi});
         info.return_value.stack_offset = stack_offset;
     }
 @@ -3720,15 +3715,13 @@ fn airCall(self: *Self, inst: Air.Inst.Index, modifier: std.builtin.CallOptions.
     const result: MCValue = result: {
         switch (info.return_value) {
-            .register => |reg| {
+            .register => {
-                if (RegisterManager.indexOfRegIntoTracked(reg) == null) {
+                // Save function return value in a new register
-                    // Save function return value in a callee saved register
+                break :result try self.copyToRegisterWithInstTracking(
-                    break :result try self.copyToRegisterWithInstTracking(
+                    inst,
-                        inst,
+                    self.air.typeOfIndex(inst),
-                        self.air.typeOfIndex(inst),
+                    info.return_value,
-                        info.return_value,
+                );
-                    );
-                }
             },
             else => {},
         }
 @@ -3755,19 +3748,11 @@ fn airRet(self: *Self, inst: Air.Inst.Index) !void {
     const ret_ty = self.fn_type.fnReturnType();
     switch (self.ret_mcv) {
         .stack_offset => {
-            // TODO audit register allocation!
+            self.register_manager.freezeRegs(&.{ .rax, .rcx });
-            self.register_manager.freezeRegs(&.{ .rax, .rcx, .rdi });
+            defer self.register_manager.unfreezeRegs(&.{ .rax, .rcx });
-            defer self.register_manager.unfreezeRegs(&.{ .rax, .rcx, .rdi });
+            const reg = try self.copyToTmpRegister(Type.usize, self.ret_mcv);
-            const reg = try self.register_manager.allocReg(null);
+            self.register_manager.freezeRegs(&.{reg});
-            const backpatch = try self.addInst(.{
+            defer self.register_manager.unfreezeRegs(&.{reg});
-                .tag = .mov,
-                .ops = (Mir.Ops{
-                    .reg1 = reg,
-                    .reg2 = .rdi,
-                }).encode(),
-                .data = undefined,
-            });
-            try self.ret_backpatches.append(self.gpa, backpatch);
             try self.genSetStack(ret_ty, 0, operand, .{
                 .source_stack_base = .rbp,
                 .dest_stack_base = reg,
 @@ -3798,19 +3783,11 @@ fn airRetLoad(self: *Self, inst: Air.Inst.Index) !void {
     const elem_ty = ptr_ty.elemType();
     switch (self.ret_mcv) {
         .stack_offset => {
-            // TODO audit register allocation!
+            self.register_manager.freezeRegs(&.{ .rax, .rcx });
-            self.register_manager.freezeRegs(&.{ .rax, .rcx, .rdi });
+            defer self.register_manager.unfreezeRegs(&.{ .rax, .rcx });
-            defer self.register_manager.unfreezeRegs(&.{ .rax, .rcx, .rdi });
+            const reg = try self.copyToTmpRegister(Type.usize, self.ret_mcv);
-            const reg = try self.register_manager.allocReg(null);
+            self.register_manager.freezeRegs(&.{reg});
-            const backpatch = try self.addInst(.{
+            defer self.register_manager.unfreezeRegs(&.{reg});
-                .tag = .mov,
-                .ops = (Mir.Ops{
-                    .reg1 = reg,
-                    .reg2 = .rdi,
-                }).encode(),
-                .data = undefined,
-            });
-            try self.ret_backpatches.append(self.gpa, backpatch);
             try self.genInlineMemcpy(.{ .stack_offset = 0 }, ptr, .{ .immediate = elem_ty.abiSize(self.target.*) }, .{
                 .source_stack_base = .rbp,
                 .dest_stack_base = reg,
 @@ -5092,6 +5069,7 @@ const InlineMemcpyOpts = struct {
     dest_stack_base: ?Register = null,
 };
+/// Spills .rax and .rcx.
 fn genInlineMemcpy(
     self: *Self,
     dst_ptr: MCValue,
 @@ -5099,13 +5077,7 @@ fn genInlineMemcpy(
     len: MCValue,
     opts: InlineMemcpyOpts,
 ) InnerError!void {
-    // TODO this is wrong. We should check first if any of the operands is in `.rax` or `.rcx` before
-    // spilling. Consolidate with other TODOs regarding register allocation mechanics.
-    try self.register_manager.getReg(.rax, null);
-    try self.register_manager.getReg(.rcx, null);
     self.register_manager.freezeRegs(&.{ .rax, .rcx });
-    defer self.register_manager.unfreezeRegs(&.{ .rax, .rcx });
     if (opts.source_stack_base) |reg| self.register_manager.freezeRegs(&.{reg});
     defer if (opts.source_stack_base) |reg| self.register_manager.unfreezeRegs(&.{reg});
 @@ -5145,7 +5117,6 @@ fn genInlineMemcpy(
             return self.fail("TODO implement memcpy for setting stack when dest is {}", .{dst_ptr});
         },
     }
     self.register_manager.freezeRegs(&.{dst_addr_reg});
     defer self.register_manager.unfreezeRegs(&.{dst_addr_reg});
 @@ -5181,7 +5152,6 @@ fn genInlineMemcpy(
             return self.fail("TODO implement memcpy for setting stack when src is {}", .{src_ptr});
         },
     }
     self.register_manager.freezeRegs(&.{src_addr_reg});
     defer self.register_manager.unfreezeRegs(&.{src_addr_reg});
 @@ -5189,6 +5159,11 @@ fn genInlineMemcpy(
     const count_reg = regs[0].to64();
     const tmp_reg = regs[1].to8();
+    self.register_manager.unfreezeRegs(&.{ .rax, .rcx });
+    try self.register_manager.getReg(.rax, null);
+    try self.register_manager.getReg(.rcx, null);
     try self.genSetReg(Type.usize, count_reg, len);
     // mov rcx, 0
 @@ -5284,6 +5259,7 @@ fn genInlineMemcpy(
     try self.performReloc(loop_reloc);
 }
+/// Spills .rax register.
 fn genInlineMemset(
     self: *Self,
     dst_ptr: MCValue,
 @@ -5291,9 +5267,7 @@ fn genInlineMemset(
     len: MCValue,
     opts: InlineMemcpyOpts,
 ) InnerError!void {
-    try self.register_manager.getReg(.rax, null);
     self.register_manager.freezeRegs(&.{.rax});
-    defer self.register_manager.unfreezeRegs(&.{.rax});
     const addr_reg = try self.register_manager.allocReg(null);
     switch (dst_ptr) {
 @@ -5330,6 +5304,9 @@ fn genInlineMemset(
     self.register_manager.freezeRegs(&.{addr_reg});
     defer self.register_manager.unfreezeRegs(&.{addr_reg});
+    self.register_manager.unfreezeRegs(&.{.rax});
+    try self.register_manager.getReg(.rax, null);
     try self.genSetReg(Type.usize, .rax, len);
     try self.genBinMathOpMir(.sub, Type.usize, .{ .register = .rax }, .{ .immediate = 1 });