Implement function calls

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 offset = pcontext.symbols.at("sqr");
-			auto fptr = (float(*)())(executable.data.get() + offset);
+			auto fptr = (unsigned(*)(unsigned))(executable.data.get() + offset);
-			auto x = fptr();
+			auto x = fptr(30);
 			std::cout << "Result: " << std::boolalpha << x << std::endl;
 		}
 	}

examples/jit_test.psl

@@ -1,3 +1,5 @@
-func test() -> f32:
+func sqr(x : f32) -> f32:
-	let pi : f16 = 3.141592h
+	return x * x
-	return (pi * pi) as f32
+
+func smoothstep(x : f32) -> f32:
+	return sqr(x) * (3.0 - 2.0 * x)

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

libs/jit/source/arch/aarch64/compiler.cpp

@@ -10,6 +10,8 @@
 #include <unordered_map>
 #include <vector>
 
+#include <iostream>
+
 namespace pslang::jit::aarch64
 {
 
@@ -21,6 +23,15 @@ namespace pslang::jit::aarch64
 			std::unordered_map<float, std::int32_t> f16_constants;
 			std::unordered_map<float, std::int32_t> f32_constants;
 			std::unordered_map<double, std::int32_t> f64_constants;
+
+			struct resolve_info
+			{
+				std::string name;
+				// Must be 'adr' instruction
+				std::int32_t instruction_offset;
+			};
+
+			std::vector<resolve_info> 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<ast::identifier>(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);
+		}
 	}
 
 }

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));