lisyarus/pslang
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

x86_64 backend: tons of fixes + raytracer example works!

Browse source
This commit is contained in:
Nikita Lisitsa 2026-08-02 15:22:35 +03:00
parent 4c583c8b3b
commit b16894ebd1
3 changed files with 311 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
libs/jit/include/pslang/jit/arch/linux_x86_64/instruction_builder.hpp
View file

@ -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 <typename ... Args>
void do_push(Args ... values);

255
libs/jit/source/arch/linux_x86_64/compiler.cpp
View file

@ -77,6 +77,7 @@ namespace pslang::jit::linux_x86_64
std::vector<resolve_data> jump_resolve;
std::vector<resolve_data> cjump_resolve;
std::vector<resolve_data> node_resolve;
std::vector<resolve_data> call_resolve;
};
std::optional<small_struct_data> classify_small_struct(local_context & lcontext, types::type_ptr type)
@ -212,6 +213,7 @@ namespace pslang::jit::linux_x86_64
std::vector<std::int32_t> argument_position;
std::unordered_map<ir::node_ref, std::int32_t> 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 (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
{
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 <typename Node, typename DoCall>
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<types::struct_type>(*type) || std::holds_alternative<types::array_type>(*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;
}
void apply(ir::node_ref, ir::call_pointer const & node, types::type_ptr const & type)
// Compute space required for stack arguments
std::int32_t arguments_stack_size = 0;
for (auto const & argument : node.arguments)
{
throw std::runtime_error("Not implemented");
auto size = ast::type_size(*argument->inferred_type);
auto alignment = ast::type_size(*argument->inferred_type);
auto struct_type = std::get_if<types::struct_type>(argument->inferred_type.get());
auto array_type = std::get_if<types::array_type>(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<types::struct_type>(argument->inferred_type.get());
auto array_type = std::get_if<types::array_type>(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<types::struct_type>(type.get());
auto array_type = std::get_if<types::array_type>(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 it, ir::call const & node, types::type_ptr const & type)
{
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);
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<types::struct_type>(*result_type) || std::holds_alternative<types::array_type>(*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<ir::label>(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});
}

44
libs/jit/source/arch/linux_x86_64/instruction_builder.cpp
View file

@ -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 <typename ... Args>
void instruction_builder::do_push(Args ... values)
{