Use a custom executable memory allocator.

Author: Hans-Kristian Arntzen

CMakeLists.txt

@@ -12,6 +12,7 @@ endif()
 add_library(parallel-rsp STATIC
 		rsp/vfunctions.cpp
 		rsp_jit.cpp rsp_jit.hpp
+		jit_allocator.cpp jit_allocator.hpp
 		rsp_disasm.cpp rsp_disasm.hpp
 		rsp/ls.cpp rsp/pipeline.h
 		rsp/reciprocal.cpp rsp/reciprocal.h

jit_allocator.cpp

@@ -0,0 +1,100 @@
+#ifdef _WIN32
+#define WIN32_LEAN_AND_MEAN
+#include <windows.h>
+#else
+#include <sys/mman.h>
+#endif
+#include <limits>
+#include <algorithm>
+
+#include "jit_allocator.hpp"
+
+namespace RSP
+{
+namespace JIT
+{
+static constexpr bool huge_va = std::numeric_limits<size_t>::max() > 0x100000000ull;
+// On 64-bit systems, we will never allocate more than one block, this is important since we must ensure that
+// relative jumps are reachable in 32-bits.
+static constexpr size_t block_size = huge_va ? (256 * 1024 * 1024) : (2 * 1024 * 1024);
+Allocator::~Allocator()
+{
+#ifdef _WIN32
+	for (auto &block : blocks)
+		VirtualFree(block.code, 0, MEM_RELEASE);
+#else
+#endif
+}
+
+static size_t align_page(size_t offset)
+{
+	return (offset + 4095) & ~size_t(4095);
+}
+
+static bool commit_read_write(void *ptr, size_t size)
+{
+#ifdef _WIN32
+	return VirtualAlloc(ptr, size, MEM_COMMIT, PAGE_READWRITE) == ptr;
+#else
+#endif
+}
+
+static bool commit_execute(void *ptr, size_t size)
+{
+#ifdef _WIN32
+	DWORD old_protect;
+	return VirtualProtect(ptr, align_page(size), PAGE_EXECUTE, &old_protect) != 0;
+#else
+#endif
+}
+
+bool Allocator::commit_code(void *code, size_t size)
+{
+	return commit_execute(code, size);
+}
+
+void *Allocator::allocate_code(size_t size)
+{
+	size = align_page(size);
+	if (blocks.empty())
+		blocks.push_back(reserve_block(std::max(size, block_size)));
+
+	auto *block = &blocks.back();
+	if (!block->code)
+		return nullptr;
+
+	block->offset = align_page(block->offset);
+	if (block->offset + size > block->size)
+		block = nullptr;
+
+	if (!block)
+	{
+		if (huge_va)
+			abort();
+		blocks.push_back(reserve_block(std::max(size, block_size)));
+		block = &blocks.back();
+	}
+
+	if (!block)
+		return nullptr;
+
+	void *ret = block->code + block->offset;
+	block->offset += size;
+
+	if (!commit_read_write(ret, size))
+		return nullptr;
+	return ret;
+}
+
+Allocator::Block Allocator::reserve_block(size_t size)
+{
+	Block block;
+#ifdef _WIN32
+	block.code = static_cast<uint8_t *>(VirtualAlloc(nullptr, size, MEM_RESERVE, PAGE_READWRITE));
+	block.size = size;
+	return block;
+#else
+#endif
+}
+}
+}

jit_allocator.hpp

@@ -0,0 +1,33 @@
+#pragma once
+
+#include <vector>
+#include <stdint.h>
+
+namespace RSP
+{
+namespace JIT
+{
+class Allocator
+{
+public:
+	Allocator() = default;
+	~Allocator();
+	void operator=(const Allocator &) = delete;
+	Allocator(const Allocator &) = delete;
+
+	void *allocate_code(size_t size);
+	static bool commit_code(void *code, size_t size);
+
+private:
+	struct Block
+	{
+		uint8_t *code = nullptr;
+		size_t size = 0;
+		size_t offset = 0;
+	};
+	std::vector<Block> blocks;
+
+	static Block reserve_block(size_t size);
+};
+}
+}

rsp_jit.cpp

@@ -32,15 +32,12 @@ namespace JIT
 {
 CPU::CPU()
 {
-	cleanup_jit_states.reserve(16 * 1024);
 	init_jit("RSP");
 	init_jit_thunks();
 }
 
 CPU::~CPU()
 {
-	for (auto *_jit : cleanup_jit_states)
-		jit_destroy_state();
 	finish_jit();
 }
 
@@ -426,6 +423,14 @@ void CPU::init_jit_thunks()
 	// Return status. This register is considered common for all thunks.
 	jit_retr(JIT_REGISTER_MODE);
 
+	jit_realize();
+	jit_word_t code_size;
+	jit_get_code(&code_size);
+	void *thunk_code = allocator.allocate_code(code_size);
+	if (!thunk_code)
+		abort();
+	jit_set_code(thunk_code, code_size);
+
 	thunks.enter_frame = reinterpret_cast<int (*)(void *)>(jit_emit());
 	thunks.enter_thunk = jit_address(entry_label);
 	thunks.return_thunk = jit_address(return_label);
@@ -434,7 +439,10 @@ void CPU::init_jit_thunks()
 	//jit_disassemble();
 	jit_clear_state();
 	//printf(" === END DISASM ===\n");
-	cleanup_jit_states.push_back(_jit);
+	jit_destroy_state();
+
+	if (!Allocator::commit_code(thunk_code, code_size))
+		abort();
 }
 
 Func CPU::get_jit_block(uint32_t pc)
@@ -1882,6 +1890,14 @@ Func CPU::jit_region(uint64_t hash, unsigned pc_word, unsigned instruction_count
 	for (auto &b : local_branches)
 		jit_patch_at(b.node, branch_targets[b.local_index]);
 
+	jit_realize();
+	jit_word_t code_size;
+	jit_get_code(&code_size);
+	auto *block_code = allocator.allocate_code(code_size);
+	if (!block_code)
+		abort();
+	jit_set_code(block_code, code_size);
+
 	auto ret = reinterpret_cast<Func>(jit_emit());
 
 #ifdef TRACE_DISASM
@@ -1891,7 +1907,10 @@ Func CPU::jit_region(uint64_t hash, unsigned pc_word, unsigned instruction_count
 	printf(" === DISASM END ===\n\n");
 #endif
 	jit_clear_state();
-	cleanup_jit_states.push_back(_jit);
+	jit_destroy_state();
+
+	if (!Allocator::commit_code(block_code, code_size))
+		abort();
 	return ret;
 }
 

rsp_jit.hpp

@@ -10,6 +10,7 @@
 
 #include "rsp_op.hpp"
 #include "state.hpp"
+#include "jit_allocator.hpp"
 
 extern "C"
 {
@@ -135,8 +136,6 @@ class alignas(64) CPU
 
 	int enter(uint32_t pc);
 
-	std::vector<jit_state_t *> cleanup_jit_states;
-
 	void init_jit_thunks();
 
 	struct
@@ -198,6 +197,7 @@ class alignas(64) CPU
 	std::vector<Link> local_branches;
 
 	RegisterCache regs;
+	Allocator allocator;
 };
 } // namespace JIT
 } // namespace RSP