diff --git a/apps/interpreter/source/main.cpp b/apps/interpreter/source/main.cpp index e138bbd..2dfcec9 100644 --- a/apps/interpreter/source/main.cpp +++ b/apps/interpreter/source/main.cpp @@ -173,9 +173,9 @@ int main(int argc, char ** argv) { // TODO: remove, testing-only code; should execute entry point instead - auto offset = pcontext.symbols.at("test"); - auto fptr = (float(*)())(executable.data.get() + offset); - auto x = fptr(); + auto offset = pcontext.symbols.at("sqr"); + auto fptr = (unsigned(*)(unsigned))(executable.data.get() + offset); + auto x = fptr(30); std::cout << "Result: " << std::boolalpha << x << std::endl; } } diff --git a/examples/jit_test.psl b/examples/jit_test.psl index 6620a08..a644ef6 100644 --- a/examples/jit_test.psl +++ b/examples/jit_test.psl @@ -1,3 +1,5 @@ -func test() -> f32: - let pi : f16 = 3.141592h - return (pi * pi) as f32 +func sqr(x : f32) -> f32: + return x * x + +func smoothstep(x : f32) -> f32: + return sqr(x) * (3.0 - 2.0 * x) diff --git a/libs/jit/include/pslang/jit/arch/aarch64/instruction_builder.hpp b/libs/jit/include/pslang/jit/arch/aarch64/instruction_builder.hpp index 12f2e15..982559f 100644 --- a/libs/jit/include/pslang/jit/arch/aarch64/instruction_builder.hpp +++ b/libs/jit/include/pslang/jit/arch/aarch64/instruction_builder.hpp @@ -113,10 +113,18 @@ namespace pslang::jit::aarch64 // 26-bit signed @offset multiplied by 4 void b(std::int32_t offset); + // Unconditionally move the program counter to the value of the register @reg + void b_reg(std::uint8_t reg); + // Inject the 26-bit signed @offset into the opcode of b instruction // starting at @opcode 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); + // Load a floating-point value from current program counter plus a // 19-bit signed @offset multiplied by 4, and store it in floating-point // register @reg_dst. @mode should be 0 for 32-bit values and 1 for 64-bit diff --git a/libs/jit/source/arch/aarch64/compiler.cpp b/libs/jit/source/arch/aarch64/compiler.cpp index b581307..5f12fe5 100644 --- a/libs/jit/source/arch/aarch64/compiler.cpp +++ b/libs/jit/source/arch/aarch64/compiler.cpp @@ -10,6 +10,8 @@ #include #include +#include + namespace pslang::jit::aarch64 { @@ -21,6 +23,15 @@ namespace pslang::jit::aarch64 std::unordered_map f16_constants; std::unordered_map f32_constants; std::unordered_map f64_constants; + + struct resolve_info + { + std::string name; + // Must be 'adr' instruction + std::int32_t instruction_offset; + }; + + std::vector resolve; }; std::uint8_t fp_mode_for(types::type const & type) @@ -336,9 +347,13 @@ namespace pslang::jit::aarch64 builder.ldr_fp(0, fp_mode_for(*node.inferred_type), 31, (stack_offset - jt->second.frame_offset) / builtin_type_size(*node.inferred_type)); else builder.ldr(0, 31, (stack_offset - jt->second.frame_offset) / 8); - break; + return; } } + + // Not a variable - must be a function! + lcontext.resolve.push_back({node.name, (std::int32_t)pcontext.code.size()}); + builder.adr(0, 0); } void apply(ast::unary_operation const & node) @@ -586,6 +601,49 @@ namespace pslang::jit::aarch64 } } + void apply(ast::function_call const & node) + { + apply(*node.function); + push(0); + + for (std::size_t i = node.arguments.size(); i --> 0;) + { + auto const & arg = node.arguments[i]; + apply(*arg); + auto type = ast::get_type(*arg); + if (types::is_bool_type(*type) || types::is_integer_type(*type)) + { + push(0); + } + else if (types::is_floating_point_type(*type)) + { + push_fp(0, fp_mode_for(*type)); + } + } + + std::uint8_t reg = 0; + std::uint8_t fp_reg = 0; + for (auto const & arg : node.arguments) + { + auto type = ast::get_type(*arg); + if (types::is_bool_type(*type) || types::is_integer_type(*type)) + { + pop(reg); + ++reg; + } + else if (types::is_floating_point_type(*type)) + { + pop_fp(fp_reg, fp_mode_for(*type)); + ++fp_reg; + } + } + + pop(reg); + push(30); + builder.b_reg(reg); + pop(30); + } + void apply(ast::assignment const & node) { auto identifier = std::get_if(node.lhs.get()); @@ -701,7 +759,7 @@ namespace pslang::jit::aarch64 void do_apply(ast::function_definition const & node) { - // TODO: floating-point / struct arguments + // TODO: struct arguments scopes.emplace_back(); std::uint8_t reg = 0; @@ -815,10 +873,14 @@ namespace pslang::jit::aarch64 void compile(program_context & pcontext, ast::statement_list_ptr const & statements) { local_context lcontext; - populate_constants_visitor{{}, {}, pcontext, lcontext}.apply(*statements); - compile_visitor{{}, pcontext, lcontext}.apply(*statements); + + instruction_builder builder{pcontext.code}; + for (auto const & resolve : lcontext.resolve) + { + builder.adr_inject(pcontext.code.data() + resolve.instruction_offset, pcontext.symbols.at(resolve.name) - resolve.instruction_offset); + } } } \ No newline at end of file diff --git a/libs/jit/source/arch/aarch64/instruction_builder.cpp b/libs/jit/source/arch/aarch64/instruction_builder.cpp index b3e61d8..fa5a582 100644 --- a/libs/jit/source/arch/aarch64/instruction_builder.cpp +++ b/libs/jit/source/arch/aarch64/instruction_builder.cpp @@ -152,6 +152,11 @@ namespace pslang::jit::aarch64 do_push(0x14000000u | (std::uint32_t(offset) & 0x3ffffffu)); } + void instruction_builder::b_reg(std::uint8_t reg) + { + do_push(0xd63f0000u | ((reg & REG_MASK) << 5)); + } + void instruction_builder::b_inject(std::uint8_t * opcode, std::int32_t offset) { auto dst = (std::uint32_t *)opcode; @@ -159,6 +164,21 @@ namespace pslang::jit::aarch64 *dst |= (std::uint32_t(offset) & 0x3ffffffu); } + void instruction_builder::adr(std::uint8_t reg_dst, std::int32_t offset) + { + auto offset_val = std::uint32_t(offset) & 0x1fffffu; + do_push(0x10000000u | (reg_dst & REG_MASK) | ((offset_val >> 2) << 5) | ((offset_val & 0x3u) << 29)); + } + + void instruction_builder::adr_inject(std::uint8_t * opcode, std::int32_t offset) + { + auto offset_val = std::uint32_t(offset) & 0x1fffffu; + auto dst = (std::uint32_t *)opcode; + *dst &= 0x9f00001fu; + *dst |= (offset_val >> 2) << 5; + *dst |= (offset_val & 0x3u) << 29; + } + void instruction_builder::ldr_fp_pc(std::uint8_t reg_dst, std::uint8_t mode, std::int32_t offset) { do_push(0x1c000000u | ((mode & 0x1u) << 30) | (reg_dst & REG_MASK) | ((std::uint32_t(offset) & 0x7ffffu) << 5));