diff --git a/libs/jit/include/pslang/jit/arch/linux_x86_64/instruction_builder.hpp b/libs/jit/include/pslang/jit/arch/linux_x86_64/instruction_builder.hpp index af6588f..1d311ed 100644 --- a/libs/jit/include/pslang/jit/arch/linux_x86_64/instruction_builder.hpp +++ b/libs/jit/include/pslang/jit/arch/linux_x86_64/instruction_builder.hpp @@ -72,6 +72,9 @@ namespace pslang::jit::linux_x86_64 // Store the zero-extended 16-bit GPR @reg_src in 64-bit register @reg_dst void movzx_16(reg reg_src, reg reg_dst); + // Store the zero-extended 32-bit GPR @reg_src in 64-bit register @reg_dst + void movzx_32(reg reg_src, reg reg_dst); + // Store the sign-extended 8-bit GPR @reg_src in 64-bit register @reg_dst void movsx_8(reg reg_src, reg reg_dst); @@ -198,6 +201,9 @@ namespace pslang::jit::linux_x86_64 // Bitwise and the 64-bit GPR @reg_src with the 64-bit GPR @reg_dst and store the result in @reg_dst void and_(reg reg_src, reg reg_dst); + // Bitwise and the 64-bit GPR @reg_src with the 8-bit @value (sign-extended to 64 bits) and store the result in @reg_dst + void and_imm8(reg reg_src, std::int8_t value); + // Bitwise or the 64-bit GPR @reg_src with the 64-bit GPR @reg_dst and store the result in @reg_dst void or_(reg reg_src, reg reg_dst); @@ -210,6 +216,9 @@ namespace pslang::jit::linux_x86_64 // Right-shift with zero-filling the 64-bit GPR @reg_dst by the value of the low 6 bits of RCX and store the result in @reg_dst void shr(reg reg_dst); + // Right-shift with zero-filling the 64-bit GPR @reg_dst by 1 bit and store the result in @reg_dst + void shr1(reg reg_dst); + // Right-shift with sign-extending the 64-bit GPR @reg_dst by the value of the low 6 bits of RCX and store the result in @reg_dst void sar(reg reg_dst); @@ -309,6 +318,11 @@ namespace pslang::jit::linux_x86_64 // NB: RIP holds the address of the _next_ instruction void jump_if_nonzero(std::int32_t offset); + // Conditional jump to 32-bit signed @offset relative to RIP + // if SF flag is set + // NB: RIP holds the address of the _next_ instruction + void jump_if_sign(std::int32_t offset); + // Assuming that @opcode refers to the location of a JUMP instruction, // replace its 32-bit jump offset with @offset void jump_inject(std::uint8_t * opcode, std::int32_t offset); @@ -336,6 +350,22 @@ namespace pslang::jit::linux_x86_64 // of the LEA instruction itself void lea_rip_inject_prev(std::uint8_t * opcode, std::int32_t offset); + // Call to 32-bit signed @offset relative to RIP + // NB: RIP holds the address of the _next_ instruction + void call_imm(std::int32_t offset); + + // Assuming that @opcode refers to the location of a CALL instruction, + // replace its 32-bit call offset with @offset + void call_imm_inject(std::uint8_t * opcode, std::int32_t offset); + + // Assuming that @opcode refers to the location of a CALL instruction, + // replace its 32-bit call offset with @offset, compensating for the size + // of the CALL instruction itself + void call_imm_inject_prev(std::uint8_t * opcode, std::int32_t offset); + + // Call to address specified in @reg_addr + void call_reg(reg reg_addr); + private: template void do_push(Args ... values); diff --git a/libs/jit/source/arch/linux_x86_64/compiler.cpp b/libs/jit/source/arch/linux_x86_64/compiler.cpp index 58e5950..603c2bc 100644 --- a/libs/jit/source/arch/linux_x86_64/compiler.cpp +++ b/libs/jit/source/arch/linux_x86_64/compiler.cpp @@ -77,6 +77,7 @@ namespace pslang::jit::linux_x86_64 std::vector jump_resolve; std::vector cjump_resolve; std::vector node_resolve; + std::vector call_resolve; }; std::optional classify_small_struct(local_context & lcontext, types::type_ptr type) @@ -148,7 +149,7 @@ namespace pslang::jit::linux_x86_64 void apply(ir::extern_symbol const & node, types::type_ptr const &) { - std::int32_t offset = pcontext.storage.size(); + std::int32_t offset = pcontext.storage.size(); lcontext.extern_symbols[node.name] = offset; pcontext.foreign_resolve.push_back({node.name, offset}); push_bytes(pcontext.storage.storage, (void *)nullptr); @@ -212,6 +213,7 @@ namespace pslang::jit::linux_x86_64 std::vector argument_position; std::unordered_map stack_position; std::int32_t stack_size = 0; + bool return_value_is_large_struct = false; void apply(ir::node_ref, ir::label const &, types::type_ptr const &) {} @@ -335,6 +337,7 @@ namespace pslang::jit::linux_x86_64 void apply(ir::node_ref it, ir::binary_operation const & node, types::type_ptr const & type) { auto const arg1_type = node.arg1->inferred_type; + auto const arg2_type = node.arg2->inferred_type; bool const is_fp = types::is_floating_point_type(*arg1_type); bool const result_is_fp = types::is_floating_point_type(*type); auto const arg_size = ast::type_size(*arg1_type); @@ -444,7 +447,7 @@ namespace pslang::jit::linux_x86_64 break; case ast::binary_operation_type::right_shift: reg_extend(reg::rax, reg::rax, *arg1_type); - reg_extend(reg::rbx, reg::rcx, *arg1_type); + reg_extend(reg::rbx, reg::rcx, *arg2_type); if (types::is_unsigned_integer_type(*arg1_type)) builder.shr(reg::rax); else @@ -625,24 +628,66 @@ namespace pslang::jit::linux_x86_64 } else if (types::is_floating_point_type(*dst_type)) { - if (types::is_unsigned_integer_type(*src_type) && ast::type_size(*src_type) == 8) - throw std::runtime_error("64-bit unsigned to floating-point conversion is not implemented"); - reg_extend(reg::rax, reg::rax, *src_type); auto dst_size = ast::type_size(*dst_type); - if (dst_size <= 4) - builder.gpr_to_xmm_32(reg::rax, reg::xmm0); + + if (types::is_unsigned_integer_type(*src_type) && ast::type_size(*src_type) == 8) + { + // No u64 -> f32/f64 conversion on x86_64 + // Emulate it by splitting x = 2 * (x >> 1) + (x & 1) + + reg_extend(reg::rax, reg::rax, *src_type); + + builder.test(reg::rax, reg::rax); + auto first_jump = builder.code.size(); + builder.jump_if_sign(0); + + if (dst_size <= 4) + builder.gpr_to_xmm_32(reg::rax, reg::xmm0); + else + builder.gpr_to_xmm(reg::rax, reg::xmm0); + + auto second_jump = builder.code.size(); + builder.jump(0); + + auto first_jump_target = builder.code.size(); + + builder.mov(reg::rax, reg::rcx); + builder.shr1(reg::rcx); + builder.and_imm8(reg::rax, 1); + builder.or_(reg::rax, reg::rcx); + + if (dst_size <= 4) + { + builder.gpr_to_xmm_32(reg::rcx, reg::xmm0); + builder.add_xmm_32(reg::xmm0, reg::xmm0); + } + else + { + builder.gpr_to_xmm(reg::rcx, reg::xmm0); + builder.add_xmm(reg::xmm0, reg::xmm0); + } + + auto second_jump_target = builder.code.size(); + + builder.cjump_inject_prev(builder.code.data() + first_jump, first_jump_target - first_jump); + builder.jump_inject_prev(builder.code.data() + second_jump, second_jump_target - second_jump); + } else - builder.gpr_to_xmm(reg::rax, reg::xmm0); + { + if (dst_size <= 4) + builder.gpr_to_xmm_32(reg::rax, reg::xmm0); + else + builder.gpr_to_xmm(reg::rax, reg::xmm0); + } } } else if (types::is_floating_point_type(*src_type)) { if (types::is_integer_type(*dst_type)) { - if (types::is_unsigned_integer_type(*dst_type) && ast::type_size(*dst_type) == 8) - throw std::runtime_error("Floating-point to 64-bit unsigned conversion is not implemented"); + // TODO: correct floating-point to 64-bit unsigned conversion auto src_size = ast::type_size(*src_type); if (src_size <= 4) @@ -691,7 +736,7 @@ namespace pslang::jit::linux_x86_64 void apply(ir::node_ref it, ir::extern_symbol const & node, types::type_ptr const & type) { - builder.lea_rip(lcontext.extern_symbols[node.name] - (std::int32_t)pcontext.storage.size(), reg::rax); + builder.mov_read_rip_prev(lcontext.extern_symbols[node.name] - (std::int32_t)pcontext.storage.size(), reg::rax); store(it, reg::rax); } @@ -730,6 +775,7 @@ namespace pslang::jit::linux_x86_64 void apply(ir::node_ref, ir::jump_if_zero const & node, types::type_ptr const & type) { load(node.condition, reg::rax); + reg_extend(reg::rax, reg::rax, *node.condition->inferred_type); builder.test(reg::rax, reg::rax); lcontext.cjump_resolve.push_back({(std::int32_t)pcontext.storage.size(), node.target}); builder.jump_if_zero(0); @@ -738,19 +784,165 @@ namespace pslang::jit::linux_x86_64 void apply(ir::node_ref, ir::jump_if_nonzero const & node, types::type_ptr const & type) { load(node.condition, reg::rax); + reg_extend(reg::rax, reg::rax, *node.condition->inferred_type); builder.test(reg::rax, reg::rax); lcontext.cjump_resolve.push_back({(std::int32_t)pcontext.storage.size(), node.target}); builder.jump_if_nonzero(0); } - void apply(ir::node_ref, ir::call const & node, types::type_ptr const & type) + template + void apply_call(ir::node_ref it, Node const & node, types::type_ptr const & type, DoCall && do_call) { - throw std::runtime_error("Not implemented"); + bool return_value_is_large_struct = false; + if (std::holds_alternative(*type) || std::holds_alternative(*type)) + return_value_is_large_struct = !classify_small_struct(lcontext, type); + + if (this->return_value_is_large_struct) + { + builder.push(reg::rdi); + stack_size += 8; + } + + // Compute space required for stack arguments + std::int32_t arguments_stack_size = 0; + for (auto const & argument : node.arguments) + { + auto size = ast::type_size(*argument->inferred_type); + auto alignment = ast::type_size(*argument->inferred_type); + auto struct_type = std::get_if(argument->inferred_type.get()); + auto array_type = std::get_if(argument->inferred_type.get()); + if (struct_type || array_type) + { + if (!classify_small_struct(lcontext, argument->inferred_type)) + { + arguments_stack_size += (alignment - (arguments_stack_size % alignment)) % alignment; + arguments_stack_size += size; + } + } + } + + // Ensure stack is 16-byte aligned after CALL pushes return address on stack + // I.e. enforce (stack_size + arguments_stack_size) % 16 == 8 + arguments_stack_size += (24 - ((arguments_stack_size + stack_size) % 16)) % 16; + + // TODO: handle the case when there weren't enough registers + std::uint8_t reg_index = return_value_is_large_struct ? 1 : 0; + std::uint8_t fp_reg = 0; + std::int32_t stack_offset = 0; + for (auto const & argument : node.arguments) + { + auto size = ast::type_size(*argument->inferred_type); + auto alignment = ast::type_alignment(*argument->inferred_type); + auto struct_type = std::get_if(argument->inferred_type.get()); + auto array_type = std::get_if(argument->inferred_type.get()); + if (struct_type || array_type) + { + if (auto small_struct = classify_small_struct(lcontext, argument->inferred_type)) + { + auto address = node_address(argument); + + for (int o : {0, 1}) + { + auto & octet = small_struct->octets[o]; + if (!octet) break; + + if (octet->has_integer) + load(argument, integer_arg_reg[reg_index++], 8 * o); + else if (octet->has_floating_point) + load_xmm(argument, (reg)(fp_reg++), 8, 8 * o); + } + } + else + { + stack_offset += (alignment - (stack_offset % alignment)) % alignment; + copy_memory(node_address(argument), {.base = reg::rsp, .offset = stack_offset - arguments_stack_size}, size); + stack_offset += size; + } + } + else if (types::is_integer_like_type(*argument->inferred_type)) + load(argument, integer_arg_reg[reg_index++]); + else if (types::is_floating_point_type(*argument->inferred_type)) + load_xmm(argument, (reg)(fp_reg++), size); + else + throw std::runtime_error("Unsupported function argument type"); + } + + if (return_value_is_large_struct) + { + auto address = node_address(it); + // Function call node cannot have RIP-relative address + builder.lea(address.base.value(), address.offset, reg::rdi); + } + + if (arguments_stack_size > 0) + { + builder.sub_imm(reg::rsp, arguments_stack_size); + stack_size += arguments_stack_size; + } + + do_call(); + + if (arguments_stack_size > 0) + { + builder.add_imm(reg::rsp, arguments_stack_size); + stack_size -= arguments_stack_size; + } + + std::size_t const size = ast::type_size(*type); + auto struct_type = std::get_if(type.get()); + auto array_type = std::get_if(type.get()); + if (size == 0) + {} + else if (struct_type || array_type) + { + if (auto small_struct = classify_small_struct(lcontext, type)) + { + std::uint8_t reg_index = 0; + std::uint8_t fp_reg = 0; + for (int o : {0, 1}) + { + auto & octet = small_struct->octets[o]; + if (!octet) break; + + if (octet->has_integer) + store(it, integer_ret_reg[reg_index++], 8 * o); + else if (octet->has_floating_point) + store_xmm(it, (reg)(fp_reg++), 8, 8 * o); + } + } + else + { + // Nothing to be done - return value should already be in place + } + } + else if (types::is_integer_like_type(*type)) + store(it, reg::rax); + else if (types::is_floating_point_type(*type)) + store_xmm(it, reg::xmm0, size); + else + throw std::runtime_error("Unsupported return value type"); + + if (this->return_value_is_large_struct) + { + builder.pop(reg::rdi); + stack_size -= 8; + } } - void apply(ir::node_ref, ir::call_pointer const & node, types::type_ptr const & type) + void apply(ir::node_ref it, ir::call const & node, types::type_ptr const & type) { - throw std::runtime_error("Not implemented"); + apply_call(it, node, type, [&]{ + lcontext.call_resolve.emplace_back(pcontext.storage.size(), node.target); + builder.call_imm(0); + }); + } + + void apply(ir::node_ref it, ir::call_pointer const & node, types::type_ptr const & type) + { + apply_call(it, node, type, [&]{ + load(node.pointer, reg::rax); + builder.call_reg(reg::rax); + }); } void apply(ir::node_ref, ir::return_value const & node, types::type_ptr const & type) @@ -795,11 +987,11 @@ namespace pslang::jit::linux_x86_64 throw std::runtime_error("Unsupported return value type"); } if (stack_size > 0) - { builder.add_imm(reg::rsp, stack_size); - if (lcontext.use_frame_pointer) - builder.pop(reg::rbp); - } + + builder.pop(reg::rbx); + if (lcontext.use_frame_pointer) + builder.pop(reg::rbp); builder.ret(); } @@ -807,11 +999,9 @@ namespace pslang::jit::linux_x86_64 { auto result_type = function_definition->inferred_result_type; - if (lcontext.use_frame_pointer) - { - builder.push(reg::rbp); - builder.mov(reg::rsp, reg::rbp); - } + if (result_type) + if (std::holds_alternative(*result_type) || std::holds_alternative(*result_type)) + return_value_is_large_struct = !classify_small_struct(lcontext, result_type); stack_size = 0; @@ -842,25 +1032,33 @@ namespace pslang::jit::linux_x86_64 } } - // Ensure (RSP % 16) == 8 for future function calls - stack_size = ((stack_size + 7) / 16) * 16 + 8; - if (!std::holds_alternative(begin->instruction)) throw std::runtime_error("First IR node of a function must be a label"); auto it = begin; lcontext.nodes[it] = pcontext.storage.size(); + + if (lcontext.use_frame_pointer) + { + builder.push(reg::rbp); + builder.mov(reg::rsp, reg::rbp); + } + + builder.push(reg::rbx); + if (stack_size > 0) builder.sub_imm(reg::rsp, stack_size); // TODO: handle the case when there weren't enough registers - std::uint8_t reg_index = 0; + // If return value is large struct, RDI stores the pointer to return value + std::uint8_t reg_index = return_value_is_large_struct ? 1 : 0; std::uint8_t fp_reg = 0; // 8 for the return address // Another 8 for RBP if using frame pointers - std::int32_t stack_offset = stack_size + (lcontext.use_frame_pointer ? 16 : 8); + // Another 8 for saved RBX + std::int32_t stack_offset = stack_size + (lcontext.use_frame_pointer ? 16 : 8) + 8; for (std::size_t i = 0; i < function_definition->arguments.size(); ++i) { auto const & argument = function_definition->arguments[i]; @@ -916,7 +1114,7 @@ namespace pslang::jit::linux_x86_64 continue; // Uncomment to debug per-node instruction generation: - // builder.nop(); + builder.nop(); lcontext.nodes[it] = pcontext.storage.size(); std::visit([&](auto const & instruction){ apply(it, instruction, it->inferred_type); }, it->instruction); @@ -934,6 +1132,10 @@ namespace pslang::jit::linux_x86_64 builder.movzx_8(reg_src, reg_dst); else if (size == 2) builder.movzx_16(reg_src, reg_dst); + else if (size == 4) + builder.movzx_32(reg_src, reg_dst); + else if (reg_src != reg_dst) + builder.mov(reg_src, reg_dst); } else if (types::is_signed_integer_type(type)) { @@ -943,6 +1145,8 @@ namespace pslang::jit::linux_x86_64 builder.movsx_16(reg_src, reg_dst); else if (size == 4) builder.movsx_32(reg_src, reg_dst); + else if (reg_src != reg_dst) + builder.mov(reg_src, reg_dst); } else if (types::is_bool_type(type)) { @@ -1128,6 +1332,9 @@ namespace pslang::jit::linux_x86_64 for (auto const & resolve : lcontext.node_resolve) builder.lea_rip_inject_prev(pcontext.storage.storage.data() + resolve.offset, lcontext.nodes.at(resolve.target) - resolve.offset); + for (auto const & resolve : lcontext.call_resolve) + builder.call_imm_inject_prev(pcontext.storage.storage.data() + resolve.offset, lcontext.nodes.at(resolve.target) - resolve.offset); + auto code_end = pcontext.storage.align(); pcontext.storage.code.push_back({code_begin, code_end}); } diff --git a/libs/jit/source/arch/linux_x86_64/instruction_builder.cpp b/libs/jit/source/arch/linux_x86_64/instruction_builder.cpp index 8546959..dac2c8c 100644 --- a/libs/jit/source/arch/linux_x86_64/instruction_builder.cpp +++ b/libs/jit/source/arch/linux_x86_64/instruction_builder.cpp @@ -163,6 +163,11 @@ namespace pslang::jit::linux_x86_64 do_push(rex(1, reg_dst, {}, reg_src), 0x0f_ub, 0xb7_ub, modrm(MOD_REG, reg_dst, reg_src)); } + void instruction_builder::movzx_32(reg reg_src, reg reg_dst) + { + do_push(rex(0, reg_src, {}, reg_dst), 0x89_ub, modrm(MOD_REG, reg_src, reg_dst)); + } + void instruction_builder::movsx_8(reg reg_src, reg reg_dst) { do_push(rex(1, reg_dst, {}, reg_src), 0x0f_ub, 0xbe_ub, modrm(MOD_REG, reg_dst, reg_src)); @@ -392,6 +397,11 @@ namespace pslang::jit::linux_x86_64 do_push(rex(1, reg_dst, {}, reg_src), 0x23_ub, modrm(MOD_REG, reg_dst, reg_src)); } + void instruction_builder::and_imm8(reg reg_src, std::int8_t value) + { + do_push(rex(1, {}, {}, reg_src), 0x83_ub, modrm(MOD_REG, 0b100, reg_src), value); + } + void instruction_builder::or_(reg reg_src, reg reg_dst) { do_push(rex(1, reg_dst, {}, reg_src), 0x0b_ub, modrm(MOD_REG, reg_dst, reg_src)); @@ -412,6 +422,11 @@ namespace pslang::jit::linux_x86_64 do_push(rex(1, {}, {}, reg_dst), 0xd3_ub, modrm(MOD_REG, 0b101, reg_dst)); } + void instruction_builder::shr1(reg reg_dst) + { + do_push(rex(1, {}, {}, reg_dst), 0xd1_ub, modrm(MOD_REG, 0b101, reg_dst)); + } + void instruction_builder::sar(reg reg_dst) { do_push(rex(1, {}, {}, reg_dst), 0xd3_ub, modrm(MOD_REG, 0b111, reg_dst)); @@ -567,6 +582,11 @@ namespace pslang::jit::linux_x86_64 do_push(0x0f_ub, 0x85_ub, offset); } + void instruction_builder::jump_if_sign(std::int32_t offset) + { + do_push(0x0f_ub, 0x88_ub, offset); + } + void instruction_builder::jump_inject(std::uint8_t * opcode, std::int32_t offset) { auto src = (std::uint8_t const *)(&offset); @@ -600,6 +620,30 @@ namespace pslang::jit::linux_x86_64 lea_rip_inject(opcode, offset - 7); } + void instruction_builder::call_imm(std::int32_t offset) + { + do_push(0xe8_ub, offset); + } + + void instruction_builder::call_imm_inject(std::uint8_t * opcode, std::int32_t offset) + { + auto src = (std::uint8_t const *)(&offset); + std::copy(src, src + 4, opcode + 1); + } + + void instruction_builder::call_imm_inject_prev(std::uint8_t * opcode, std::int32_t offset) + { + call_imm_inject(opcode, offset - 5); + } + + void instruction_builder::call_reg(reg reg_addr) + { + if ((std::uint8_t(reg_addr) & 0x8) == 0) + do_push(0xff_ub, modrm(MOD_REG, 0b010, reg_addr)); + else + do_push(rex(0, {}, {}, reg_addr), 0xff_ub, modrm(MOD_REG, 0b010, reg_addr)); + } + template void instruction_builder::do_push(Args ... values) {