From a0dd6f9ca6a6bfc8a31d2ed05d28f15428d6f935 Mon Sep 17 00:00:00 2001 From: lisyarus Date: Wed, 5 Aug 2026 14:30:38 +0300 Subject: [PATCH] Modules wip: huge refactor + support linking several input files into a single executable memory blob --- apps/interpreter/source/main.cpp | 32 ++-- libs/ir/include/pslang/ir/compiler.hpp | 3 +- libs/ir/include/pslang/ir/module.hpp | 14 +- libs/ir/include/pslang/ir/node.hpp | 15 -- libs/ir/include/pslang/ir/node_fwd.hpp | 18 ++ libs/ir/source/compiler.cpp | 41 ++-- libs/jit/include/pslang/jit/abi.hpp | 2 + .../pslang/jit/arch/linux_x86_64/compiler.hpp | 4 +- .../arch/linux_x86_64/instruction_builder.hpp | 28 ++- .../pslang/jit/arch/linux_x86_64/resolver.hpp | 10 + .../jit/arch/macos_aarch64/compiler.hpp | 4 +- .../macos_aarch64/instruction_builder.hpp | 8 +- .../jit/arch/macos_aarch64/resolver.hpp | 10 + libs/jit/include/pslang/jit/compiler.hpp | 12 ++ libs/jit/include/pslang/jit/executable.hpp | 13 -- .../include/pslang/jit/host_executable.hpp | 25 +++ libs/jit/include/pslang/jit/jit.hpp | 12 -- .../jit/include/pslang/jit/linked_program.hpp | 24 +++ libs/jit/include/pslang/jit/linker.hpp | 49 +++++ libs/jit/include/pslang/jit/module.hpp | 77 ++++++++ .../include/pslang/jit/program_context.hpp | 40 ---- libs/jit/include/pslang/jit/resolver.hpp | 22 +++ libs/jit/include/pslang/jit/storage.hpp | 28 --- .../jit/source/arch/linux_x86_64/compiler.cpp | 175 +++++++++--------- .../arch/linux_x86_64/instruction_builder.cpp | 45 +++++ .../jit/source/arch/linux_x86_64/resolver.cpp | 26 +++ .../source/arch/macos_aarch64/compiler.cpp | 92 ++++----- .../macos_aarch64/instruction_builder.cpp | 9 + .../source/arch/macos_aarch64/resolver.cpp | 26 +++ libs/jit/source/{jit.cpp => compiler.cpp} | 18 +- libs/jit/source/executable.cpp | 57 ------ libs/jit/source/host_executable.cpp | 71 +++++++ libs/jit/source/linker.cpp | 110 +++++++++++ libs/jit/source/resolver.cpp | 21 +++ libs/jit/source/storage.cpp | 32 ---- 35 files changed, 775 insertions(+), 398 deletions(-) create mode 100644 libs/jit/include/pslang/jit/arch/linux_x86_64/resolver.hpp create mode 100644 libs/jit/include/pslang/jit/arch/macos_aarch64/resolver.hpp create mode 100644 libs/jit/include/pslang/jit/compiler.hpp delete mode 100644 libs/jit/include/pslang/jit/executable.hpp create mode 100644 libs/jit/include/pslang/jit/host_executable.hpp delete mode 100644 libs/jit/include/pslang/jit/jit.hpp create mode 100644 libs/jit/include/pslang/jit/linked_program.hpp create mode 100644 libs/jit/include/pslang/jit/linker.hpp create mode 100644 libs/jit/include/pslang/jit/module.hpp delete mode 100644 libs/jit/include/pslang/jit/program_context.hpp create mode 100644 libs/jit/include/pslang/jit/resolver.hpp delete mode 100644 libs/jit/include/pslang/jit/storage.hpp create mode 100644 libs/jit/source/arch/linux_x86_64/resolver.cpp create mode 100644 libs/jit/source/arch/macos_aarch64/resolver.cpp rename libs/jit/source/{jit.cpp => compiler.cpp} (65%) delete mode 100644 libs/jit/source/executable.cpp create mode 100644 libs/jit/source/host_executable.cpp create mode 100644 libs/jit/source/linker.cpp create mode 100644 libs/jit/source/resolver.cpp delete mode 100644 libs/jit/source/storage.cpp diff --git a/apps/interpreter/source/main.cpp b/apps/interpreter/source/main.cpp index 8c3e91c..fe27e1e 100644 --- a/apps/interpreter/source/main.cpp +++ b/apps/interpreter/source/main.cpp @@ -8,8 +8,8 @@ #include #include #include -#include -#include +#include +#include #include #include @@ -198,8 +198,7 @@ int main(int argc, char ** argv) semantic::validate(module); modules.push_back(std::move(module)); - ir_compiled.emplace_back(); - ir::compile(ir_compiled.back(), modules.back()); + ir_compiled.push_back(ir::compile(modules.back())); if (dump_ir) { @@ -248,25 +247,18 @@ int main(int argc, char ** argv) if (jit) { - // TODO: treat all input files as modules combined into a single program + std::vector modules; for (std::size_t i = 0; i < filenames.size(); ++i) + modules.push_back(jit::compile(ir_compiled[i], jit::host_abi())); + + auto executable = jit::make_host_executable(modules); + + for (std::size_t i = 0; i < modules.size(); ++i) { - jit::program_context pcontext - { - .abi = jit::host_abi(), - }; + if (!executable.entry_points[i]) continue; - jit::compile(pcontext, ir_compiled[i]); - - for (auto const & resolve : pcontext.foreign_resolve) - { - auto fptr = jit::load_foreign(resolve.name); - std::copy_n((std::uint8_t const *)(&fptr), 8, pcontext.storage.storage.data() + resolve.offset); - } - - auto executable = jit::make_host_executable(pcontext.storage); - - auto entry_point = (void(*)())(executable.get() + pcontext.entry_point); + using entry_point_t = void(*)(); + auto entry_point = entry_point_t(executable.mapping.get() + *executable.entry_points[i]); entry_point(); } } diff --git a/libs/ir/include/pslang/ir/compiler.hpp b/libs/ir/include/pslang/ir/compiler.hpp index 7bbf0b3..53f0aaa 100644 --- a/libs/ir/include/pslang/ir/compiler.hpp +++ b/libs/ir/include/pslang/ir/compiler.hpp @@ -1,10 +1,11 @@ #pragma once #include +#include namespace pslang::ir { - void compile(compiled_module & result, semantic::module const & module); + compiled_module compile(semantic::module const & module); } diff --git a/libs/ir/include/pslang/ir/module.hpp b/libs/ir/include/pslang/ir/module.hpp index 591340c..4004bcb 100644 --- a/libs/ir/include/pslang/ir/module.hpp +++ b/libs/ir/include/pslang/ir/module.hpp @@ -1,7 +1,7 @@ #pragma once -#include #include +#include #include @@ -9,6 +9,7 @@ namespace pslang::ast { struct function_definition; + struct variable_declaration; } @@ -17,20 +18,25 @@ namespace pslang::ir struct compiled_module { + // Fake root AST node for module entry point + ast::statement_ptr root; + node_list_ptr nodes; // Needed for debug IR dump std::unordered_map labels; - struct function_info + struct node_range { node_ref begin; node_ref end; }; - std::unordered_map functions; + std::unordered_map functions; + std::unordered_map globals; - node_ref entry_point; + // Can be null if the entry point is empty + ast::function_definition const * entry_point; }; } diff --git a/libs/ir/include/pslang/ir/node.hpp b/libs/ir/include/pslang/ir/node.hpp index 95ed72a..1473036 100644 --- a/libs/ir/include/pslang/ir/node.hpp +++ b/libs/ir/include/pslang/ir/node.hpp @@ -6,7 +6,6 @@ #include #include -#include namespace pslang::ir { @@ -161,17 +160,3 @@ namespace pslang::ir bool is_value_instruction(instruction const & instruction); } - -namespace std -{ - - template <> - struct hash<::pslang::ir::node_ref> - { - std::size_t operator()(pslang::ir::node_ref const & ref) const - { - return std::hash()(ref.operator->()); - } - }; - -} diff --git a/libs/ir/include/pslang/ir/node_fwd.hpp b/libs/ir/include/pslang/ir/node_fwd.hpp index d2e919f..35a22df 100644 --- a/libs/ir/include/pslang/ir/node_fwd.hpp +++ b/libs/ir/include/pslang/ir/node_fwd.hpp @@ -2,6 +2,7 @@ #include #include +#include namespace pslang::ir { @@ -12,3 +13,20 @@ namespace pslang::ir using node_list_ptr = std::shared_ptr; } + +namespace std +{ + + template <> + struct hash<::pslang::ir::node_ref> + { + std::size_t operator()(pslang::ir::node_ref const & ref) const + { + // NON STANDARD!!! + // Use the knowledge that list iterator is just a node pointer in + // pretty much all implementations + return std::hash{}(*(void const **)(&ref)); + } + }; + +} diff --git a/libs/ir/source/compiler.cpp b/libs/ir/source/compiler.cpp index 9ddf8f8..3b9a2fe 100644 --- a/libs/ir/source/compiler.cpp +++ b/libs/ir/source/compiler.cpp @@ -21,6 +21,7 @@ namespace pslang::ir struct scope { std::string label_prefix; + bool single_execution_scope = true; }; std::vector scopes; @@ -685,6 +686,19 @@ namespace pslang::ir { if (node.global) { + if (lcontext.scopes.back().single_execution_scope) + { + // Optimization: top-level globals will be initialized unconditionally during module entry point execution + // Skip the initialization guard and just initialize it immediately + module.nodes->emplace_back(global{.initializer = {}}, node.inferred_type); + auto result = last(); + auto value = apply(*node.initializer); + module.nodes->emplace_back(assignment{result, value}); + lcontext.globals[&node] = result; + return result; + } + + // Non top-level global: add boolean initialization guard to ensure once-only initialization module.nodes->emplace_back(global{.initializer = {0}}, std::make_unique(types::primitive_type{types::bool_type{}})); auto guard = last(); module.nodes->emplace_back(global{.initializer = {}}, node.inferred_type); @@ -757,7 +771,7 @@ namespace pslang::ir auto jump_to_end = last(); lcontext.loop_scopes.emplace_back(); - lcontext.scopes.emplace_back(); + lcontext.scopes.emplace_back().single_execution_scope = false; apply(*node.statements); lcontext.scopes.pop_back(); @@ -824,7 +838,8 @@ namespace pslang::ir module.nodes->emplace_back(label{}); auto begin = last(); - lcontext.scopes.emplace_back(); + bool const is_root = lcontext.scopes.empty(); + lcontext.scopes.emplace_back().single_execution_scope = is_root; for (std::size_t i = 0; i < node.arguments.size(); ++i) { module.nodes->emplace_back(argument{i}, ast::get_type(*node.arguments[i].type)); @@ -896,10 +911,6 @@ namespace pslang::ir label_prefix = lcontext.scopes.back().label_prefix + node.name + "."; lcontext.scopes.emplace_back(std::move(label_prefix)); apply(*node.statements); - - // Don't pop_back entry point scope - if (lcontext.scopes.size() > 1) - lcontext.scopes.pop_back(); } void apply(ast::foreign_function_declaration const &) {} @@ -911,14 +922,15 @@ namespace pslang::ir } - void compile(compiled_module & module, semantic::module const & module_in) + compiled_module compile(semantic::module const & module_in) { - if (!module.nodes) - module.nodes = std::make_shared(); + compiled_module module; + + module.nodes = std::make_shared(); // Add a fake root AST function node // for module entry point - auto root = std::make_shared( + module.root = std::make_shared( ast::function_definition { { "[entry point]", @@ -937,7 +949,7 @@ namespace pslang::ir module.nodes->emplace_back(label{}); auto extra_label = std::prev(module.nodes->end()); - compile_visitor{{}, module, lcontext}.apply(*root); + compile_visitor{{}, module, lcontext}.apply(*module.root); module.nodes->erase(extra_label); for (auto & function : lcontext.functions) @@ -952,7 +964,12 @@ namespace pslang::ir for (auto const & function : lcontext.functions) module.functions[function.first] = {.begin = function.second.first, .end = function.second.second}; - module.entry_point = lcontext.functions.at(std::get_if(root.get())).first; + for (auto const & global : lcontext.globals) + module.globals[global.first] = global.second; + + module.entry_point = std::get_if(module.root.get()); + + return module; } } diff --git a/libs/jit/include/pslang/jit/abi.hpp b/libs/jit/include/pslang/jit/abi.hpp index ada6f4b..721702f 100644 --- a/libs/jit/include/pslang/jit/abi.hpp +++ b/libs/jit/include/pslang/jit/abi.hpp @@ -24,6 +24,8 @@ namespace pslang::jit { jit::isa isa; jit::platform platform; + + friend bool operator == (abi, abi) = default; }; abi host_abi(); diff --git a/libs/jit/include/pslang/jit/arch/linux_x86_64/compiler.hpp b/libs/jit/include/pslang/jit/arch/linux_x86_64/compiler.hpp index 77a3c11..279a304 100644 --- a/libs/jit/include/pslang/jit/arch/linux_x86_64/compiler.hpp +++ b/libs/jit/include/pslang/jit/arch/linux_x86_64/compiler.hpp @@ -1,11 +1,11 @@ #pragma once -#include +#include #include namespace pslang::jit::linux_x86_64 { - void compile(program_context & pcontext, ir::compiled_module const & module_in); + compiled_module compile(ir::compiled_module const & module_in); } 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 1d311ed..95cb337 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 @@ -325,30 +325,39 @@ namespace pslang::jit::linux_x86_64 // 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); + static void jump_inject(std::uint8_t * opcode, std::int32_t offset); // Assuming that @opcode refers to the location of a JUMP instruction, // replace its 32-bit jump offset with @offset, compensating for the size // of the JUMP instruction itself - void jump_inject_prev(std::uint8_t * opcode, std::int32_t offset); + static void jump_inject_prev(std::uint8_t * opcode, std::int32_t offset); // Assuming that @opcode refers to the location of a conditional JUMP instruction // (JZ or JNZ), replace its 32-bit jump offset with @offset - void cjump_inject(std::uint8_t * opcode, std::int32_t offset); + static void cjump_inject(std::uint8_t * opcode, std::int32_t offset); // Assuming that @opcode refers to the location of a conditional JUMP instruction // (JZ or JNZ), replace its 32-bit jump offset with @offset, compensating for the size // of the JUMP instruction itself - void cjump_inject_prev(std::uint8_t * opcode, std::int32_t offset); + static void cjump_inject_prev(std::uint8_t * opcode, std::int32_t offset); + + // Assuming that @opcode refers to the location of a RIP-relative MOV instruction, + // replace its 32-bit jump offset with @offset + static void mov_rip_inject(std::uint8_t * opcode, std::int32_t offset); + + // Assuming that @opcode refers to the location of a RIP-relative MOV instruction, + // replace its 32-bit jump offset with @offset, compensating for the size + // of the MOV instruction itself + static void mov_rip_inject_prev(std::uint8_t * opcode, std::int32_t offset); // Assuming that @opcode refers to the location of a RIP-relative LEA instruction, // replace its 32-bit jump offset with @offset - void lea_rip_inject(std::uint8_t * opcode, std::int32_t offset); + static void lea_rip_inject(std::uint8_t * opcode, std::int32_t offset); // Assuming that @opcode refers to the location of a RIP-relative LEA instruction, // replace its 32-bit jump offset with @offset, compensating for the size // of the LEA instruction itself - void lea_rip_inject_prev(std::uint8_t * opcode, std::int32_t offset); + static 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 @@ -356,16 +365,19 @@ namespace pslang::jit::linux_x86_64 // 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); + static 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); + static 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); + // + static void resolve_offset(std::uint8_t * opcode, std::int32_t offset); + private: template void do_push(Args ... values); diff --git a/libs/jit/include/pslang/jit/arch/linux_x86_64/resolver.hpp b/libs/jit/include/pslang/jit/arch/linux_x86_64/resolver.hpp new file mode 100644 index 0000000..fcb465b --- /dev/null +++ b/libs/jit/include/pslang/jit/arch/linux_x86_64/resolver.hpp @@ -0,0 +1,10 @@ +#pragma once + +#include + +namespace pslang::jit::linux_x86_64 +{ + + std::unique_ptr make_resolver(); + +} diff --git a/libs/jit/include/pslang/jit/arch/macos_aarch64/compiler.hpp b/libs/jit/include/pslang/jit/arch/macos_aarch64/compiler.hpp index 3e74832..dc7124f 100644 --- a/libs/jit/include/pslang/jit/arch/macos_aarch64/compiler.hpp +++ b/libs/jit/include/pslang/jit/arch/macos_aarch64/compiler.hpp @@ -1,11 +1,11 @@ #pragma once -#include +#include #include namespace pslang::jit::macos_aarch64 { - void compile(program_context & pcontext, ir::compiled_module const & module_in); + compiled_module compile(ir::compiled_module const & module_in); } diff --git a/libs/jit/include/pslang/jit/arch/macos_aarch64/instruction_builder.hpp b/libs/jit/include/pslang/jit/arch/macos_aarch64/instruction_builder.hpp index 7880605..1c34480 100644 --- a/libs/jit/include/pslang/jit/arch/macos_aarch64/instruction_builder.hpp +++ b/libs/jit/include/pslang/jit/arch/macos_aarch64/instruction_builder.hpp @@ -104,6 +104,8 @@ namespace pslang::jit::macos_aarch64 // plus a signed 19-bit @offset multiplied by 4, and store it into register @reg_dst void ldr_pc(std::uint8_t reg_dst, std::int32_t offset); + static void ldr_pc_inject(std::uint8_t * opcode, std::int32_t offset); + // Add a 12-bit @value to the register @reg_src and store the result in @reg_dst void add_imm(std::uint8_t reg_src, std::uint8_t reg_dst, std::uint16_t value); @@ -178,7 +180,7 @@ namespace pslang::jit::macos_aarch64 // Inject the 19-bit signed @offset into the opcode of a cbz or cbnz instruction // starting at @opcode - void cb_inject(std::uint8_t * opcode, std::int32_t offset); + static void cb_inject(std::uint8_t * opcode, std::int32_t offset); // Unconditionally move the program counter to the value of // 26-bit signed @offset multiplied by 4 @@ -198,12 +200,12 @@ namespace pslang::jit::macos_aarch64 // Inject the 26-bit signed @offset into the opcode of b or bl instruction // starting at @opcode - void b_inject(std::uint8_t * opcode, std::int32_t offset); + static void b_inject(std::uint8_t * opcode, std::int32_t offset); // Load the current program count plus a signed 21-bit offset into register @reg_dst void adr(std::uint8_t reg_dst, std::int32_t offset); - void adr_inject(std::uint8_t * opcode, std::int32_t offset); + static void adr_inject(std::uint8_t * opcode, std::int32_t offset); // Load a floating-point value from current program counter plus a // 19-bit signed @offset multiplied by 4, and store it in floating-point diff --git a/libs/jit/include/pslang/jit/arch/macos_aarch64/resolver.hpp b/libs/jit/include/pslang/jit/arch/macos_aarch64/resolver.hpp new file mode 100644 index 0000000..e76a310 --- /dev/null +++ b/libs/jit/include/pslang/jit/arch/macos_aarch64/resolver.hpp @@ -0,0 +1,10 @@ +#pragma once + +#include + +namespace pslang::jit::macos_aarch64 +{ + + std::unique_ptr make_resolver(); + +} diff --git a/libs/jit/include/pslang/jit/compiler.hpp b/libs/jit/include/pslang/jit/compiler.hpp new file mode 100644 index 0000000..6b5161e --- /dev/null +++ b/libs/jit/include/pslang/jit/compiler.hpp @@ -0,0 +1,12 @@ +#pragma once + +#include +#include +#include + +namespace pslang::jit +{ + + compiled_module compile(ir::compiled_module const & module_in, abi abi); + +} diff --git a/libs/jit/include/pslang/jit/executable.hpp b/libs/jit/include/pslang/jit/executable.hpp deleted file mode 100644 index 9e0af99..0000000 --- a/libs/jit/include/pslang/jit/executable.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include - -#include -#include - -namespace pslang::jit -{ - - std::shared_ptr make_host_executable(binary_storage const & storage); - -} diff --git a/libs/jit/include/pslang/jit/host_executable.hpp b/libs/jit/include/pslang/jit/host_executable.hpp new file mode 100644 index 0000000..f66b4b1 --- /dev/null +++ b/libs/jit/include/pslang/jit/host_executable.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include + +#include +#include +#include + +namespace pslang::jit +{ + + struct host_executable_program + { + std::shared_ptr mapping; + + // Offsets from the start of the mapping + std::unordered_map functions; + std::unordered_map globals; + + std::vector> entry_points; + }; + + host_executable_program make_host_executable(std::vector const & modules); + +} diff --git a/libs/jit/include/pslang/jit/jit.hpp b/libs/jit/include/pslang/jit/jit.hpp deleted file mode 100644 index eb1373e..0000000 --- a/libs/jit/include/pslang/jit/jit.hpp +++ /dev/null @@ -1,12 +0,0 @@ -#pragma once - -#include -#include -#include - -namespace pslang::jit -{ - - void compile(program_context & pcontext, ir::compiled_module const & module_in); - -} diff --git a/libs/jit/include/pslang/jit/linked_program.hpp b/libs/jit/include/pslang/jit/linked_program.hpp new file mode 100644 index 0000000..cd99a71 --- /dev/null +++ b/libs/jit/include/pslang/jit/linked_program.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include +#include + +namespace pslang::jit +{ + + struct linked_program_info + { + std::uint8_t storage; + + struct section + { + std::size_t begin; + std::size_t size; + }; + + section code; + section data; + section relocations; + }; + +} diff --git a/libs/jit/include/pslang/jit/linker.hpp b/libs/jit/include/pslang/jit/linker.hpp new file mode 100644 index 0000000..5efbb67 --- /dev/null +++ b/libs/jit/include/pslang/jit/linker.hpp @@ -0,0 +1,49 @@ +#pragma once + +#include +#include + +#include + +namespace pslang::jit +{ + + struct link_options + { + // If JIT-compiling into an in-memory host-executable binary, + // code and data sections must be aligned to OS page size. + std::size_t section_alignment = 1; + + // If AOT-compiling into an ELF executable, 1 GOT entry must + // be reserved to point to _DYNAMIC (the address of .dynamic segment) + std::size_t reserved_relocations = 0; + }; + + struct linked_program_info + { + struct section + { + std::size_t begin; + std::size_t size; + }; + + section code; + section data; + section relocations; + + std::size_t total_size; + + // Offsets from the beginning of combined storage + std::unordered_map functions; + std::unordered_map globals; + std::unordered_map relocation_offset; + + // One per each module + std::vector> entry_points; + }; + + linked_program_info prepare_link(std::vector const & modules, link_options const & options); + + void link(std::vector const & modules, linked_program_info const & info, std::uint8_t * storage); + +} diff --git a/libs/jit/include/pslang/jit/module.hpp b/libs/jit/include/pslang/jit/module.hpp new file mode 100644 index 0000000..9c488d1 --- /dev/null +++ b/libs/jit/include/pslang/jit/module.hpp @@ -0,0 +1,77 @@ +#pragma once + +#include +#include + +#include +#include +#include +#include +#include + +namespace pslang::ast +{ + + struct function_definition; + struct foreign_function_declaration; + struct variable_declaration; + +} + +namespace pslang::jit +{ + + struct compiled_module + { + struct internal_function_resolve_info + { + // Function whose address is to be resolved + ir::node_ref function; + // Offset to call/jump/lea/mov instruction that needs to be patched + std::int32_t instruction_offset; + }; + + struct internal_global_resolve_info + { + // Global variable whose address is to be resolved + ir::node_ref global; + // Offset to call/jump/lea/mov instruction that needs to be patched + std::int32_t instruction_offset; + }; + + struct external_function_resolve_info + { + // Name of external (foreign) function whose address needs to be resolved + // On all platforms, this uses GOT (Global Offset Table): an array of 8-byte + // pointers to functions resolved dynamically before executing the compiled code + // It is the address of this table entry that needs to be resolved + std::string name; + // Offset to call/jump/lea/mov instruction that needs to be patched + // (though this should always be mov, as there's no reason to jump/call the GOT entry) + std::int32_t instruction_offset; + }; + + jit::abi abi; + + // Executable code + std::vector code; + + // Module entry point + // Can be null if the module's entry point is a nop + std::optional entry_point; + + // Global variables + std::vector data; + + // Offset to function code start in `code` buffer + std::unordered_map functions; + + // Offset to global variable memory start in `data` buffer + std::unordered_map globals; + + std::vector internal_function_resolve; + std::vector internal_global_resolve; + std::vector external_function_resolve; + }; + +} diff --git a/libs/jit/include/pslang/jit/program_context.hpp b/libs/jit/include/pslang/jit/program_context.hpp deleted file mode 100644 index c8f1121..0000000 --- a/libs/jit/include/pslang/jit/program_context.hpp +++ /dev/null @@ -1,40 +0,0 @@ -#pragma once - -#include -#include - -#include -#include -#include -#include - -namespace pslang::ast -{ - - struct function_definition; - -} - -namespace pslang::jit -{ - - struct program_context - { - struct foreign_resolve_info - { - std::string name; - // Offset in bytes to the place in code blob - // containing the 64-bit address of the foreign symbol - std::int32_t offset; - }; - - jit::abi abi; - - binary_storage storage = {}; - - std::unordered_map symbols = {}; - std::int32_t entry_point = 0; - std::vector foreign_resolve = {}; - }; - -} diff --git a/libs/jit/include/pslang/jit/resolver.hpp b/libs/jit/include/pslang/jit/resolver.hpp new file mode 100644 index 0000000..4517a58 --- /dev/null +++ b/libs/jit/include/pslang/jit/resolver.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include + +#include +#include + +namespace pslang::jit +{ + + struct resolver + { + // Given a jump/call/lea/mov instruction at @opcode that uses an + // immediate signed 32-bit offset, add the @offset to the existing offset + virtual void resolve(std::uint8_t * opcode, std::int32_t offset) = 0; + + virtual ~resolver() {} + }; + + std::unique_ptr make_resolver(isa isa); + +} diff --git a/libs/jit/include/pslang/jit/storage.hpp b/libs/jit/include/pslang/jit/storage.hpp deleted file mode 100644 index 914be36..0000000 --- a/libs/jit/include/pslang/jit/storage.hpp +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include -#include -#include - -namespace pslang::jit -{ - - struct binary_storage - { - struct range - { - std::size_t begin; - std::size_t end; - }; - - std::vector storage; - std::vector data; - std::vector code; - - std::size_t size() const { return storage.size(); } - std::size_t align(); - }; - - std::size_t native_page_size(); - -} diff --git a/libs/jit/source/arch/linux_x86_64/compiler.cpp b/libs/jit/source/arch/linux_x86_64/compiler.cpp index 0482363..b47fccb 100644 --- a/libs/jit/source/arch/linux_x86_64/compiler.cpp +++ b/libs/jit/source/arch/linux_x86_64/compiler.cpp @@ -52,18 +52,21 @@ namespace pslang::jit::linux_x86_64 struct value_address { - // None means RIP-based addressing - std::optional base; + std::optional base = std::nullopt; + ir::node_ref global = {}; std::int32_t offset; }; + bool is_global(ir::node_ref it) + { + return std::holds_alternative(it->instruction); + } + struct local_context { bool use_frame_pointer = true; - std::unordered_map extern_symbols; - std::unordered_map small_structs; std::unordered_map nodes; @@ -116,7 +119,7 @@ namespace pslang::jit::linux_x86_64 struct populate_globals_visitor { - std::vector & storage; + compiled_module & result; local_context & lcontext; template @@ -130,36 +133,16 @@ namespace pslang::jit::linux_x86_64 throw std::runtime_error("global IR node with initializer larger than type size"); auto alignment = ast::type_alignment(*type); - auto offset = storage.size(); + auto offset = result.data.size(); offset = ((offset + (alignment - 1)) / alignment) * alignment; - storage.resize(offset + size); - std::copy(node.initializer.begin(), node.initializer.end(), storage.begin() + offset); - lcontext.nodes[it] = offset; - } - }; - - struct populate_const_data_visitor - { - program_context & pcontext; - local_context & lcontext; - - template - void apply(Node const & node, types::type_ptr const &) - {} - - void apply(ir::extern_symbol const & node, types::type_ptr const &) - { - 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); + result.data.resize(offset + size); + std::copy(node.initializer.begin(), node.initializer.end(), result.data.begin() + offset); + result.globals[it] = offset; } }; struct literal_visitor { - program_context & pcontext; - local_context & lcontext; instruction_builder & builder; void operator()(ast::bool_literal const & node) @@ -205,8 +188,7 @@ namespace pslang::jit::linux_x86_64 struct compile_visitor { - program_context & pcontext; - ir::compiled_module const & module; + compiled_module & module; local_context & lcontext; instruction_builder & builder; @@ -220,7 +202,7 @@ namespace pslang::jit::linux_x86_64 void apply(ir::node_ref it, ir::literal const & node, types::type_ptr const & type) { - std::visit(literal_visitor{pcontext, lcontext, builder}, node.value); + std::visit(literal_visitor{builder}, node.value); if (types::is_integer_like_type(*type)) store(it, reg::rax); else if (types::is_floating_point_type(*type)) @@ -323,9 +305,14 @@ namespace pslang::jit::linux_x86_64 { auto const address = node_address(node.arg1); if (address.base) + { builder.lea(*address.base, address.offset, reg::rax); + } else - builder.lea_rip_prev(address.offset, reg::rax); + { + push_resolve_global(address.global); + builder.lea_rip_prev(0, reg::rax); + } store(it, reg::rax); } break; @@ -606,11 +593,16 @@ namespace pslang::jit::linux_x86_64 { if (types::equal(*array_type->element_type, *pointer_type->referenced_type)) { - auto arg1_address = node_address(node.arg1); + auto const arg1_address = node_address(node.arg1); if (arg1_address.base) + { builder.lea(*arg1_address.base, arg1_address.offset, reg::rax); + } else - builder.lea_rip_prev(arg1_address.offset, reg::rax); + { + push_resolve_global(arg1_address.global); + builder.lea_rip_prev(0, reg::rax); + } store(it, reg::rax); return; } @@ -729,14 +721,15 @@ namespace pslang::jit::linux_x86_64 void apply(ir::node_ref it, ir::instruction_address const & node, types::type_ptr const & type) { - lcontext.node_resolve.emplace_back(pcontext.storage.size(), node.target); + lcontext.node_resolve.emplace_back(builder.code.size(), node.target); builder.lea_rip(0, reg::rax); store(it, reg::rax); } void apply(ir::node_ref it, ir::extern_symbol const & node, types::type_ptr const & type) { - builder.mov_read_rip_prev(lcontext.extern_symbols[node.name] - (std::int32_t)pcontext.storage.size(), reg::rax); + module.external_function_resolve.push_back({.name = node.name, .instruction_offset = (std::int32_t)builder.code.size()}); + builder.mov_read_rip_prev(0, reg::rax); store(it, reg::rax); } @@ -768,7 +761,7 @@ namespace pslang::jit::linux_x86_64 void apply(ir::node_ref, ir::jump const & node, types::type_ptr const & type) { - lcontext.jump_resolve.push_back({(std::int32_t)pcontext.storage.size(), node.target}); + lcontext.jump_resolve.push_back({(std::int32_t)builder.code.size(), node.target}); builder.jump(0); } @@ -777,7 +770,7 @@ namespace pslang::jit::linux_x86_64 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}); + lcontext.cjump_resolve.push_back({(std::int32_t)builder.code.size(), node.target}); builder.jump_if_zero(0); } @@ -786,7 +779,7 @@ namespace pslang::jit::linux_x86_64 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}); + lcontext.cjump_resolve.push_back({(std::int32_t)builder.code.size(), node.target}); builder.jump_if_nonzero(0); } @@ -839,8 +832,6 @@ namespace pslang::jit::linux_x86_64 { 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]; @@ -869,8 +860,8 @@ namespace pslang::jit::linux_x86_64 if (return_value_is_large_struct) { - auto address = node_address(it); // Function call node cannot have RIP-relative address + auto const address = node_address(it); builder.lea(address.base.value(), address.offset, reg::rdi); } @@ -931,8 +922,9 @@ namespace pslang::jit::linux_x86_64 void apply(ir::node_ref it, ir::call const & node, types::type_ptr const & type) { + // TODO: call function from a different module? apply_call(it, node, type, [&]{ - lcontext.call_resolve.emplace_back(pcontext.storage.size(), node.target); + lcontext.call_resolve.emplace_back(builder.code.size(), node.target); builder.call_imm(0); }); } @@ -959,8 +951,6 @@ namespace pslang::jit::linux_x86_64 { if (auto small_struct = classify_small_struct(lcontext, type)) { - auto address = node_address(*node.value); - std::uint8_t reg_index = 0; std::uint8_t fp_reg = 0; for (int o : {0, 1}) @@ -976,7 +966,7 @@ namespace pslang::jit::linux_x86_64 } else { - copy_memory(node_address(*node.value), {reg::rdi, 0}, size); + copy_memory(node_address(*node.value), {reg::rdi, {}, 0}, size); } } else if (types::is_integer_like_type(*type)) @@ -1019,7 +1009,7 @@ namespace pslang::jit::linux_x86_64 { stack_position[it] = argument_position[argument->index]; } - else if (std::holds_alternative(it->instruction)) + else if (is_global(it)) { // stack position doesn't make sense for globals } @@ -1037,7 +1027,7 @@ namespace pslang::jit::linux_x86_64 auto it = begin; - lcontext.nodes[it] = pcontext.storage.size(); + lcontext.nodes[it] = builder.code.size(); if (lcontext.use_frame_pointer) { @@ -1116,7 +1106,7 @@ namespace pslang::jit::linux_x86_64 // Uncomment to debug per-node instruction generation: builder.nop(); - lcontext.nodes[it] = pcontext.storage.size(); + lcontext.nodes[it] = builder.code.size(); std::visit([&](auto const & instruction){ apply(it, instruction, it->inferred_type); }, it->instruction); } } @@ -1154,10 +1144,11 @@ namespace pslang::jit::linux_x86_64 } } + // Must not be called for global nodes - they use RIP-based addressing with future relocation value_address node_address(ir::node_ref it) { - if (std::holds_alternative(it->instruction)) - return {.base = std::nullopt, .offset = lcontext.nodes.at(it) - static_cast(builder.code.size())}; + if (is_global(it)) + return {.global = it, .offset = 0}; else return {.base = reg::rsp, .offset = stack_size - stack_position.at(it)}; } @@ -1166,18 +1157,29 @@ namespace pslang::jit::linux_x86_64 { auto const address = node_address(it); if (address.base) + { builder.mov_read(*address.base, address.offset + offset, reg_dst); + } else - builder.mov_read_rip_prev(address.offset + offset, reg_dst); + { + push_resolve_global(address.global); + builder.mov_read_rip_prev(offset, reg_dst); + } } void store(ir::node_ref it, reg reg_src, std::int32_t offset = 0) { auto const address = node_address(it); if (address.base) + { + auto const address = node_address(it); builder.mov_write(reg_src, *address.base, address.offset + offset); + } else - builder.mov_write_rip_prev(reg_src, address.offset + offset); + { + push_resolve_global(address.global); + builder.mov_write_rip_prev(reg_src, offset); + } } void load_xmm(ir::node_ref it, reg reg_dst, std::uint8_t size, std::int32_t offset = 0) @@ -1188,13 +1190,17 @@ namespace pslang::jit::linux_x86_64 auto const address = node_address(it); if (address.base) { + auto const address = node_address(it); if (size == 4) builder.mov_read_xmm_32(*address.base, address.offset + offset, reg_dst); else builder.mov_read_xmm(*address.base, address.offset + offset, reg_dst); } else + { + // TODO throw std::runtime_error("RIP-relative XMM read is not supported"); + } } void store_xmm(ir::node_ref it, reg reg_src, std::uint8_t size, std::int32_t offset = 0) @@ -1211,12 +1217,15 @@ namespace pslang::jit::linux_x86_64 builder.mov_write_xmm(reg_src, *address.base, address.offset + offset); } else + { + // TODO throw std::runtime_error("RIP-relative XMM write is not supported"); + } } void copy_memory(value_address src, value_address dst, std::size_t size) { - auto const storage_size_at_start = (std::int32_t)pcontext.storage.size(); + auto const storage_size_at_start = (std::int32_t)builder.code.size(); reg reg_src; if (src.base) @@ -1224,6 +1233,7 @@ namespace pslang::jit::linux_x86_64 else { reg_src = find_free_reg(reg::r10, dst.base); + push_resolve_global(src.global); builder.lea_rip_prev(0, reg_src); } @@ -1233,7 +1243,8 @@ namespace pslang::jit::linux_x86_64 else { reg_dst = find_free_reg(reg::r10, reg_src); - builder.lea_rip_prev(storage_size_at_start - (std::int32_t)pcontext.storage.size(), reg_dst); + push_resolve_global(dst.global); + builder.lea_rip_prev(0, reg_dst); } std::int32_t offset = 0; @@ -1277,66 +1288,52 @@ namespace pslang::jit::linux_x86_64 } } } + + void push_resolve_global(ir::node_ref global) + { + module.internal_global_resolve.push_back({.global = global, .instruction_offset = (std::int32_t)builder.code.size()}); + } }; } - void compile(program_context & pcontext, ir::compiled_module const & module_in) + compiled_module compile(ir::compiled_module const & module_in) { + compiled_module result; + local_context lcontext; - auto data_begin = pcontext.storage.align(); { - populate_globals_visitor visitor{.storage = pcontext.storage.storage, .lcontext = lcontext}; + populate_globals_visitor visitor{result, lcontext}; for (auto it = module_in.nodes->begin(); it != module_in.nodes->end(); ++it) std::visit([&](auto const & instruction){ visitor.apply(it, instruction, it->inferred_type); }, it->instruction); } -#ifndef NDEBUG - // Force the data page to be allocated - // Helps with debugging in QtCreator (which pokes ~16 bytes before the code - // page and can hit unmapped memory) - if (data_begin == pcontext.storage.size()) - pcontext.storage.storage.push_back(0); -#endif - - auto data_end = pcontext.storage.align(); - if (data_begin != data_end) - pcontext.storage.data.push_back({data_begin, data_end}); - - auto code_begin = data_end; - - { - populate_const_data_visitor visitor{pcontext, lcontext}; - for (auto it = module_in.nodes->begin(); it != module_in.nodes->end(); ++it) - std::visit([&](auto const & instruction){ visitor.apply(instruction, it->inferred_type); }, it->instruction); - } - - instruction_builder builder{pcontext.storage.storage}; + instruction_builder builder{result.code}; for (auto const & function : module_in.functions) { - pcontext.symbols[function.first] = pcontext.storage.size(); - compile_visitor visitor{pcontext, module_in, lcontext, builder}; + result.functions[function.second.begin] = result.code.size(); + compile_visitor visitor{result, lcontext, builder}; visitor.compile(function.first, function.second.begin, function.second.end); } - pcontext.entry_point = lcontext.nodes.at(module_in.entry_point); + if (module_in.entry_point) + result.entry_point = module_in.functions.at(module_in.entry_point).begin; for (auto const & resolve : lcontext.jump_resolve) - builder.jump_inject_prev(pcontext.storage.storage.data() + resolve.offset, lcontext.nodes.at(resolve.target) - resolve.offset); + builder.jump_inject_prev(result.code.data() + resolve.offset, lcontext.nodes.at(resolve.target) - resolve.offset); for (auto const & resolve : lcontext.cjump_resolve) - builder.cjump_inject_prev(pcontext.storage.storage.data() + resolve.offset, lcontext.nodes.at(resolve.target) - resolve.offset); + builder.cjump_inject_prev(result.code.data() + resolve.offset, lcontext.nodes.at(resolve.target) - resolve.offset); 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); + builder.lea_rip_inject_prev(result.code.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); + builder.call_imm_inject_prev(result.code.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}); + return result; } } 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 dac2c8c..0044b69 100644 --- a/libs/jit/source/arch/linux_x86_64/instruction_builder.cpp +++ b/libs/jit/source/arch/linux_x86_64/instruction_builder.cpp @@ -28,6 +28,11 @@ namespace pslang::jit::linux_x86_64 return rex(W, std::uint8_t(R) >> 3, std::uint8_t(X) >> 3, std::uint8_t(B) >> 3); } + bool is_rex(std::uint8_t value) + { + return (value & 0xf0) == 0x40; + } + // Memory static constexpr std::uint8_t MOD_MEM = 0; // Memory + 8-bit offset @@ -609,6 +614,17 @@ namespace pslang::jit::linux_x86_64 cjump_inject(opcode, offset - 6); } + void instruction_builder::mov_rip_inject(std::uint8_t * opcode, std::int32_t offset) + { + auto src = (std::uint8_t const *)(&offset); + std::copy(src, src + 4, opcode + 3); + } + + void instruction_builder::mov_rip_inject_prev(std::uint8_t * opcode, std::int32_t offset) + { + mov_rip_inject(opcode, offset - 7); + } + void instruction_builder::lea_rip_inject(std::uint8_t * opcode, std::int32_t offset) { auto src = (std::uint8_t const *)(&offset); @@ -644,6 +660,35 @@ namespace pslang::jit::linux_x86_64 do_push(rex(0, {}, {}, reg_addr), 0xff_ub, modrm(MOD_REG, 0b010, reg_addr)); } + void instruction_builder::resolve_offset(std::uint8_t * opcode, std::int32_t offset) + { + if (is_rex(*opcode)) ++opcode; + + if (*opcode == 0x8d) + { + // lea rip prev + opcode += 2; + } + else if (*opcode == 0x8b) + { + // mov read rip prev + opcode += 2; + } + else if (*opcode == 0x89) + { + // mov write rip prev + opcode += 2; + } + else + throw std::runtime_error("Internal error: unknown opcode to patch"); + + std::int32_t current_offset; + std::copy(opcode, opcode + 4, (std::uint8_t *)¤t_offset); + current_offset += offset; + auto begin = (std::uint8_t const *)¤t_offset; + std::copy(begin, begin + 4, opcode); + } + template void instruction_builder::do_push(Args ... values) { diff --git a/libs/jit/source/arch/linux_x86_64/resolver.cpp b/libs/jit/source/arch/linux_x86_64/resolver.cpp new file mode 100644 index 0000000..ce94efe --- /dev/null +++ b/libs/jit/source/arch/linux_x86_64/resolver.cpp @@ -0,0 +1,26 @@ +#include +#include + +namespace pslang::jit::linux_x86_64 +{ + + namespace + { + + struct resolver_impl + : resolver + { + void resolve(std::uint8_t * opcode, std::int32_t offset) override + { + instruction_builder::resolve_offset(opcode, offset); + } + }; + + } + + std::unique_ptr make_resolver() + { + return std::make_unique(); + } + +} diff --git a/libs/jit/source/arch/macos_aarch64/compiler.cpp b/libs/jit/source/arch/macos_aarch64/compiler.cpp index d10d1c6..f54269e 100644 --- a/libs/jit/source/arch/macos_aarch64/compiler.cpp +++ b/libs/jit/source/arch/macos_aarch64/compiler.cpp @@ -27,7 +27,6 @@ namespace pslang::jit::macos_aarch64 std::unordered_map> struct_hfa; - std::unordered_map extern_symbols; std::unordered_map nodes; std::unordered_map f16_constants; @@ -43,6 +42,8 @@ namespace pslang::jit::macos_aarch64 std::vector branch_resolve; std::vector cbranch_resolve; std::vector adr_resolve; + + std::vector external_function_resolve; }; std::uint8_t fp_mode_for(types::type const & type) @@ -115,7 +116,7 @@ namespace pslang::jit::macos_aarch64 struct populate_globals_visitor { - std::vector & storage; + compiled_module & result; local_context & lcontext; template @@ -129,18 +130,18 @@ namespace pslang::jit::macos_aarch64 throw std::runtime_error("global IR node with initializer larger than type size"); auto alignment = ast::type_alignment(*type); - auto offset = storage.size(); + auto offset = result.data.size(); offset = ((offset + (alignment - 1)) / alignment) * alignment; - storage.resize(offset + size); - std::copy(node.initializer.begin(), node.initializer.end(), storage.begin() + offset); - lcontext.nodes[it] = offset; + result.data.resize(offset + size); + std::copy(node.initializer.begin(), node.initializer.end(), result.data.begin() + offset); + result.globals[it] = offset; } }; struct populate_const_data_visitor { - program_context & pcontext; local_context & lcontext; + instruction_builder & builder; template void apply(Node const & node, types::type_ptr const &) @@ -150,38 +151,30 @@ namespace pslang::jit::macos_aarch64 { if (auto f16_literal = std::get_if(&node.value)) { - lcontext.f16_constants[f16_literal->value.repr] = pcontext.storage.size(); + lcontext.f16_constants[f16_literal->value.repr] = builder.code.size(); push_bytes(f16_literal->value.repr); // Ensure 4-byte alignment push_bytes(std::uint16_t{0}); } else if (auto f32_literal = std::get_if(&node.value)) { - lcontext.f32_constants[f32_literal->value] = pcontext.storage.size(); + lcontext.f32_constants[f32_literal->value] = builder.code.size(); push_bytes(f32_literal->value); } else if (auto f64_literal = std::get_if(&node.value)) { - lcontext.f32_constants[f64_literal->value] = pcontext.storage.size(); + lcontext.f32_constants[f64_literal->value] = builder.code.size(); push_bytes(f64_literal->value); } } - void apply(ir::extern_symbol const & node, types::type_ptr const &) - { - std::int32_t offset = pcontext.storage.size(); - lcontext.extern_symbols[node.name] = offset; - pcontext.foreign_resolve.push_back({node.name, offset}); - push_bytes(nullptr); - } - private: template void push_bytes(T const & value) { auto begin = (std::uint8_t const *)(&value); auto end = begin + sizeof(value); - pcontext.storage.storage.insert(pcontext.storage.storage.end(), begin, end); + builder.code.insert(builder.code.end(), begin, end); } }; @@ -193,7 +186,6 @@ namespace pslang::jit::macos_aarch64 struct literal_visitor { - program_context & pcontext; local_context & lcontext; instruction_builder & builder; @@ -236,7 +228,7 @@ namespace pslang::jit::macos_aarch64 void operator()(ast::f16_literal const & node) { auto offset = lcontext.f16_constants.at(node.value.repr); - std::int32_t current = pcontext.storage.size(); + std::int32_t current = builder.code.size(); builder.adr(0, (offset - current) / 4); builder.ldr_fp(0, 1, 0, 0); } @@ -244,14 +236,14 @@ namespace pslang::jit::macos_aarch64 void operator()(ast::f32_literal const & node) { auto offset = lcontext.f32_constants.at(node.value); - std::int32_t current = pcontext.storage.size(); + std::int32_t current = builder.code.size(); builder.ldr_fp_pc(0, 0, (offset - current) / 4); } void operator()(ast::f64_literal const & node) { auto offset = lcontext.f64_constants.at(node.value); - std::int32_t current = pcontext.storage.size(); + std::int32_t current = builder.code.size(); builder.ldr_fp_pc(0, 1, (offset - current) / 4); } }; @@ -312,8 +304,6 @@ namespace pslang::jit::macos_aarch64 struct compile_visitor { - program_context & pcontext; - ir::compiled_module const & module_in; local_context & lcontext; instruction_builder & builder; @@ -327,7 +317,7 @@ namespace pslang::jit::macos_aarch64 void apply(ir::node_ref it, ir::literal const & node, types::type_ptr const & type) { - std::visit(literal_visitor{pcontext, lcontext, builder}, node.value); + std::visit(literal_visitor{lcontext, builder}, node.value); if (types::is_integer_like_type(*type)) store(it, 0); else if (types::is_floating_point_type(*type)) @@ -734,14 +724,15 @@ namespace pslang::jit::macos_aarch64 void apply(ir::node_ref it, ir::instruction_address const & node, types::type_ptr const &) { - lcontext.adr_resolve.emplace_back(pcontext.storage.size(), node.target); + lcontext.adr_resolve.emplace_back(builder.code.size(), node.target); builder.adr(0, 0); store(it, 0); } void apply(ir::node_ref it, ir::extern_symbol const & node, types::type_ptr const &) { - builder.ldr_pc(0, (lcontext.extern_symbols[node.name] - (std::int32_t)pcontext.storage.size()) / 4); + lcontext.external_function_resolve.push_back({.name = node.name, .instruction_offset = (std::int32_t)builder.code.size()}); + builder.ldr_pc(0, 0); store(it, 0); } @@ -774,7 +765,7 @@ namespace pslang::jit::macos_aarch64 void apply(ir::node_ref, ir::jump const & node, types::type_ptr const &) { - lcontext.branch_resolve.emplace_back(pcontext.storage.size(), node.target); + lcontext.branch_resolve.emplace_back(builder.code.size(), node.target); builder.b(0); } @@ -782,7 +773,7 @@ namespace pslang::jit::macos_aarch64 { load(node.condition, 0); extend(0, node.condition->inferred_type); - lcontext.cbranch_resolve.emplace_back(pcontext.storage.size(), node.target); + lcontext.cbranch_resolve.emplace_back(builder.code.size(), node.target); builder.cbz(0, 0); } @@ -790,7 +781,7 @@ namespace pslang::jit::macos_aarch64 { load(node.condition, 0); extend(0, node.condition->inferred_type); - lcontext.cbranch_resolve.emplace_back(pcontext.storage.size(), node.target); + lcontext.cbranch_resolve.emplace_back(builder.code.size(), node.target); builder.cbnz(0, 0); } @@ -921,7 +912,7 @@ namespace pslang::jit::macos_aarch64 void apply(ir::node_ref it, ir::call const & node, types::type_ptr const & type) { apply_call(it, node, type, [&]{ - lcontext.branch_resolve.emplace_back(pcontext.storage.size(), node.target); + lcontext.branch_resolve.emplace_back(builder.code.size(), node.target); builder.bl(0); }); } @@ -1039,7 +1030,7 @@ namespace pslang::jit::macos_aarch64 auto it = begin; - lcontext.nodes[it] = pcontext.storage.size(); + lcontext.nodes[it] = builder.code.size(); if (stack_size > 0) builder.sub_imm(31, 31, stack_size); if (lcontext.use_frame_pointer) @@ -1109,7 +1100,7 @@ namespace pslang::jit::macos_aarch64 // Uncomment to debug per-node instruction generation: // builder.nop(); - lcontext.nodes[it] = pcontext.storage.size(); + lcontext.nodes[it] = builder.code.size(); std::visit([&](auto const & instruction){ apply(it, instruction, it->inferred_type); }, it->instruction); } } @@ -1217,50 +1208,49 @@ namespace pslang::jit::macos_aarch64 } - void compile(program_context & pcontext, ir::compiled_module const & module_in) + compiled_module compile(ir::compiled_module const & module_in) { + compiled_module result; + local_context lcontext; - auto data_begin = pcontext.storage.align(); { - populate_globals_visitor visitor{.storage = pcontext.storage.storage, .lcontext = lcontext}; + populate_globals_visitor visitor{result, lcontext}; for (auto it = module_in.nodes->begin(); it != module_in.nodes->end(); ++it) std::visit([&](auto const & instruction){ visitor.apply(it, instruction, it->inferred_type); }, it->instruction); } - auto data_end = pcontext.storage.align(); - if (data_begin != data_end) - pcontext.storage.data.push_back({data_begin, data_end}); - auto code_begin = data_end; + instruction_builder builder{result.data}; { - populate_const_data_visitor visitor{pcontext, lcontext}; + populate_const_data_visitor visitor{lcontext, builder}; for (auto it = module_in.nodes->begin(); it != module_in.nodes->end(); ++it) std::visit([&](auto const & instruction){ visitor.apply(instruction, it->inferred_type); }, it->instruction); } - instruction_builder builder{pcontext.storage.storage}; for (auto const & function : module_in.functions) { - pcontext.symbols[function.first] = pcontext.storage.size(); - compile_visitor visitor{pcontext, module_in, lcontext, builder}; + result.functions[function.second.begin] = result.code.size(); + compile_visitor visitor{lcontext, builder}; visitor.compile(function.first, function.second.begin, function.second.end); } - pcontext.entry_point = lcontext.nodes.at(module_in.entry_point); + if (module_in.entry_point) + result.entry_point = module_in.functions.at(module_in.entry_point).begin; for (auto const & resolve : lcontext.branch_resolve) - builder.b_inject(pcontext.storage.storage.data() + resolve.offset, (lcontext.nodes.at(resolve.target) - resolve.offset) / 4); + builder.b_inject(result.code.data() + resolve.offset, (lcontext.nodes.at(resolve.target) - resolve.offset) / 4); for (auto const & resolve : lcontext.cbranch_resolve) - builder.cb_inject(pcontext.storage.storage.data() + resolve.offset, (lcontext.nodes.at(resolve.target) - resolve.offset) / 4); + builder.cb_inject(result.code.data() + resolve.offset, (lcontext.nodes.at(resolve.target) - resolve.offset) / 4); for (auto const & resolve : lcontext.adr_resolve) - builder.adr_inject(pcontext.storage.storage.data() + resolve.offset, lcontext.nodes.at(resolve.target) - resolve.offset); + builder.adr_inject(result.code.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}); + result.external_function_resolve = std::move(lcontext.external_function_resolve); + + return result; } } diff --git a/libs/jit/source/arch/macos_aarch64/instruction_builder.cpp b/libs/jit/source/arch/macos_aarch64/instruction_builder.cpp index 9a56e72..622fa02 100644 --- a/libs/jit/source/arch/macos_aarch64/instruction_builder.cpp +++ b/libs/jit/source/arch/macos_aarch64/instruction_builder.cpp @@ -143,6 +143,15 @@ namespace pslang::jit::macos_aarch64 do_push(0x58000000u | (reg_dst & REG_MASK) | ((((std::uint32_t)offset) & 0x7ffffu) << 5)); } + void instruction_builder::ldr_pc_inject(std::uint8_t * opcode, std::int32_t offset) + { + check_bits(offset, 19, "Bad ldr_pc_inject offset value"); + auto offset_val = std::uint32_t(offset) & 0x7ffffu; + auto dst = (std::uint32_t *)opcode; + *dst &= 0xff00001fu; + *dst |= offset_val << 5; + } + void instruction_builder::add_imm(std::uint8_t reg_src, std::uint8_t reg_dst, std::uint16_t value) { check_bits(value, 12, "Bad add_imm value"); diff --git a/libs/jit/source/arch/macos_aarch64/resolver.cpp b/libs/jit/source/arch/macos_aarch64/resolver.cpp new file mode 100644 index 0000000..f72d6d5 --- /dev/null +++ b/libs/jit/source/arch/macos_aarch64/resolver.cpp @@ -0,0 +1,26 @@ +#include +#include + +namespace pslang::jit::macos_aarch64 +{ + + namespace + { + + struct resolver_impl + : resolver + { + void resolve(std::uint8_t * opcode, std::int32_t offset) override + { + instruction_builder::ldr_pc_inject(opcode, offset / 4); + } + }; + + } + + std::unique_ptr make_resolver() + { + return std::make_unique(); + } + +} diff --git a/libs/jit/source/jit.cpp b/libs/jit/source/compiler.cpp similarity index 65% rename from libs/jit/source/jit.cpp rename to libs/jit/source/compiler.cpp index 0135855..1e01d40 100644 --- a/libs/jit/source/jit.cpp +++ b/libs/jit/source/compiler.cpp @@ -1,4 +1,4 @@ -#include +#include #include #include @@ -7,16 +7,15 @@ namespace pslang::jit { - void compile(program_context & pcontext, ir::compiled_module const & module_in) + compiled_module compile(ir::compiled_module const & module_in, abi abi) { - switch (pcontext.abi.platform) + switch (abi.platform) { case platform::linux: - switch (pcontext.abi.isa) + switch (abi.isa) { case isa::x86_64: - linux_x86_64::compile(pcontext, module_in); - break; + return linux_x86_64::compile(module_in); case isa::aarch64: throw std::runtime_error("Linux aarch64 JIT compilation not supported"); } @@ -25,16 +24,17 @@ namespace pslang::jit case platform::windows: throw std::runtime_error("Windows JIT compilation not supported"); case platform::macos: - switch (pcontext.abi.isa) + switch (abi.isa) { case isa::x86_64: throw std::runtime_error("macOS x86_64 JIT compilation not supported"); case isa::aarch64: - macos_aarch64::compile(pcontext, module_in); - break; + return macos_aarch64::compile(module_in); } break; } + + throw std::runtime_error("Internal error: unknown isa+platform combination"); } } diff --git a/libs/jit/source/executable.cpp b/libs/jit/source/executable.cpp deleted file mode 100644 index 4614ef7..0000000 --- a/libs/jit/source/executable.cpp +++ /dev/null @@ -1,57 +0,0 @@ -#include - -#include -#include - -#ifdef __linux__ -#include -#endif - -#ifdef __APPLE__ -#include -#endif - -namespace pslang::jit -{ - - std::shared_ptr make_host_executable(binary_storage const & storage) - { -#if defined(__linux__) || defined(__APPLE__) - auto const page_size = native_page_size(); - - for (auto const & range : storage.data) - { - if ((range.begin % page_size) != 0) - throw std::runtime_error("Data range start is not aligned to page boundary"); - - if ((range.end % page_size) != 0) - throw std::runtime_error("Data range end is not aligned to page boundary"); - } - - for (auto const & range : storage.code) - { - if ((range.begin % page_size) != 0) - throw std::runtime_error("Code range start is not aligned to page boundary"); - - if ((range.end % page_size) != 0) - throw std::runtime_error("Code range end is not aligned to page boundary"); - } - - auto total_size = storage.storage.size(); - auto ptr = (std::uint8_t *)mmap(nullptr, total_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, 0, 0); - if (ptr == MAP_FAILED) - throw std::system_error(errno, std::generic_category()); - - std::copy(storage.storage.begin(), storage.storage.end(), ptr); - - for (auto const & range : storage.code) - if (mprotect(ptr + range.begin, range.end - range.begin, PROT_READ | PROT_EXEC) != 0) - throw std::system_error(errno, std::generic_category()); - - return std::shared_ptr(ptr, [total_size](void * ptr){ munmap(ptr, total_size); }); -#else - throw std::runtime_error("Host-executable modules are not supported for this platform"); -#endif - } - -} diff --git a/libs/jit/source/host_executable.cpp b/libs/jit/source/host_executable.cpp new file mode 100644 index 0000000..38bbbd6 --- /dev/null +++ b/libs/jit/source/host_executable.cpp @@ -0,0 +1,71 @@ +#include +#include +#include + +#include + +#if defined(__linux__) || defined(__APPLE__) +#include +#include +#endif + +#include + +namespace pslang::jit +{ + + namespace + { + + std::size_t native_page_size() + { +#if defined(__linux__) || defined(__APPLE__) + return getpagesize(); +#else + throw std::runtime_error("Native page size is not supported for this platform"); +#endif + } + + } + + host_executable_program make_host_executable(std::vector const & modules) + { +#if defined(__linux__) || defined(__APPLE__) + auto const page_size = native_page_size(); + auto const abi = host_abi(); + + for (auto const & module : modules) + if (module.abi != abi) + throw std::runtime_error("Internal error: wrong compiled module ABI for host-executable binary"); + + auto link_info = prepare_link(modules, {.section_alignment = page_size}); + + auto storage = (std::uint8_t *)mmap(nullptr, link_info.total_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, 0, 0); + if (storage == MAP_FAILED) + throw std::system_error(errno, std::generic_category()); + + host_executable_program result; + result.mapping.reset(storage, [total_size = link_info.total_size](void * ptr){ munmap(ptr, total_size); }); + + link(modules, link_info, storage); + + for (auto const & relocation : link_info.relocation_offset) + *(void **)(storage + relocation.second) = load_foreign(relocation.first); + + if (mprotect(storage + link_info.code.begin, link_info.code.size, PROT_READ | PROT_EXEC) != 0) + throw std::system_error(errno, std::generic_category()); + + if (mprotect(storage + link_info.relocations.begin, link_info.relocations.size, PROT_READ) != 0) + throw std::system_error(errno, std::generic_category()); + + result.functions = std::move(link_info.functions); + result.globals = std::move(link_info.globals); + result.entry_points = std::move(link_info.entry_points); + + return result; +#else + throw std::runtime_error("Host-executable modules are not supported for this platform"); +#endif + } + +} diff --git a/libs/jit/source/linker.cpp b/libs/jit/source/linker.cpp new file mode 100644 index 0000000..f484a43 --- /dev/null +++ b/libs/jit/source/linker.cpp @@ -0,0 +1,110 @@ +#include +#include + +namespace pslang::jit +{ + + namespace + { + + std::size_t align(std::size_t value, std::size_t alignment) + { + return ((value + alignment - 1) / alignment) * alignment; + } + + } + + linked_program_info prepare_link(std::vector const & modules, link_options const & options) + { + linked_program_info result; + + std::size_t total_code_size = 0; + std::size_t total_data_size = 0; + std::size_t total_relo_size = options.reserved_relocations * 8; + + for (auto const & module : modules) + { + for (auto const & function : module.functions) + result.functions[function.first] = total_code_size + function.second; + + for (auto const & global : module.globals) + result.globals[global.first] = total_data_size + global.second; + + for (auto const & relocation : module.external_function_resolve) + if (!result.relocation_offset.contains(relocation.name)) + { + result.relocation_offset[relocation.name] = total_relo_size; + total_relo_size += 8; + } + + if (module.entry_point) + result.entry_points.push_back(total_code_size + module.functions.at(*module.entry_point)); + else + result.entry_points.push_back(std::nullopt); + + total_code_size += module.code.size(); + total_data_size += module.data.size(); + } + + total_code_size = align(total_code_size, options.section_alignment); + total_data_size = align(total_data_size, options.section_alignment); + total_relo_size = align(total_relo_size, options.section_alignment); + + std::size_t const data_offset = total_code_size; + std::size_t const relo_offset = total_code_size + total_data_size; + result.total_size = total_code_size + total_data_size + total_relo_size; + + for (auto & global : result.globals) + global.second += data_offset; + + for (auto & relocation : result.relocation_offset) + relocation.second += relo_offset; + + result.code = {0, total_code_size}; + result.data = {total_code_size, total_data_size}; + result.relocations = {total_code_size + total_data_size, total_relo_size}; + + return result; + } + + void link(std::vector const & modules, linked_program_info const & info, std::uint8_t * storage) + { + // TODO: resolve internal symbols (across different modules) + + if (modules.empty()) + return; + + auto resolver = make_resolver(modules.front().abi.isa); + + { + auto dst = storage + info.code.begin; + + for (auto const & module : modules) + { + auto end = std::copy(module.code.data(), module.code.data() + module.code.size(), dst); + + for (auto const & resolve : module.internal_global_resolve) + { + std::int32_t offset = info.globals.at(resolve.global) - (dst - storage + resolve.instruction_offset); + resolver->resolve(dst + resolve.instruction_offset, offset); + } + + for (auto const & resolve : module.external_function_resolve) + { + std::int32_t offset = info.relocation_offset.at(resolve.name) - (dst - storage + resolve.instruction_offset); + resolver->resolve(dst + resolve.instruction_offset, offset); + } + + dst = end; + } + } + + { + auto dst = storage + info.data.begin; + + for (auto const & module : modules) + dst = std::copy(module.data.data(), module.data.data() + module.data.size(), dst); + } + } + +} diff --git a/libs/jit/source/resolver.cpp b/libs/jit/source/resolver.cpp new file mode 100644 index 0000000..7321ab7 --- /dev/null +++ b/libs/jit/source/resolver.cpp @@ -0,0 +1,21 @@ +#include +#include +#include + +namespace pslang::jit +{ + + std::unique_ptr make_resolver(isa isa) + { + switch (isa) + { + case isa::x86_64: + return linux_x86_64::make_resolver(); + case isa::aarch64: + return macos_aarch64::make_resolver(); + } + + throw std::runtime_error("Unknown ISA for resolver"); + } + +} diff --git a/libs/jit/source/storage.cpp b/libs/jit/source/storage.cpp deleted file mode 100644 index f213ee9..0000000 --- a/libs/jit/source/storage.cpp +++ /dev/null @@ -1,32 +0,0 @@ -#include - -#ifdef __linux__ -#include -#endif - -#ifdef __APPLE__ -#include -#endif - -namespace pslang::jit -{ - - std::size_t binary_storage::align() - { - auto const page_size = native_page_size(); - auto aligned_size = ((storage.size() + (page_size - 1)) / page_size) * page_size; - storage.resize(aligned_size); - return storage.size(); - } - - std::size_t native_page_size() - { -#if defined(__linux__) || defined(__APPLE__) - return getpagesize(); -#else - throw std::runtime_error("Native page size is not supported for this platform"); -#endif - } - -} -