Simplify JIT code buffer management

This commit is contained in:
Dmitry Stogov 2023-12-08 01:58:45 +03:00
parent 8b2a536d74
commit b7e7de5821
4 changed files with 107 additions and 111 deletions

13
ir.h
View File

@ -534,6 +534,12 @@ typedef const void *(*ir_get_veneer_t)(ir_ctx *ctx, const void *addr);
typedef bool (*ir_set_veneer_t)(ir_ctx *ctx, const void *addr, const void *veneer);
#endif
typedef struct _ir_code_buffer {
void *start;
void *end;
void *pos;
} ir_code_buffer;
struct _ir_ctx {
ir_insn *ir_base; /* two directional array - instructions grow down, constants grow up */
ir_ref insns_count; /* number of instructions stored in instructions buffer */
@ -593,12 +599,9 @@ struct _ir_ctx {
uint32_t entries_count;
uint32_t *entries; /* array of ENTRY blocks */
void *osr_entry_loads;
void *code_buffer;
size_t code_buffer_size;
ir_code_buffer *code_buffer;
#if defined(IR_TARGET_AARCH64)
int32_t deoptimization_exits;
int32_t veneers_size;
uint32_t code_size;
ir_get_exit_addr_t get_exit_addr;
ir_get_veneer_t get_veneer;
ir_set_veneer_t set_veneer;
@ -857,7 +860,7 @@ int ir_patch(const void *code, size_t size, uint32_t jmp_table_size, const void
uint32_t ir_cpuinfo(void);
/* Deoptimization helpers */
const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_per_group, const void *exit_addr, void *code_buffer, size_t code_buffer_size, size_t *size_ptr);
const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_per_group, const void *exit_addr, ir_code_buffer *code_buffer, size_t *size_ptr);
/* A reference IR JIT compiler */
IR_ALWAYS_INLINE void *ir_jit_compile(ir_ctx *ctx, int opt_level, size_t *size)

View File

@ -23,26 +23,27 @@
static bool aarch64_may_use_b(ir_ctx *ctx, const void *addr)
{
if (ctx->code_buffer) {
if (addr >= ctx->code_buffer && (char*)addr < (char*)ctx->code_buffer + ctx->code_buffer_size) {
return (ctx->code_buffer_size < B_IMM);
} else if ((char*)addr >= (char*)ctx->code_buffer + ctx->code_buffer_size) {
return (((char*)addr - (char*)ctx->code_buffer) < B_IMM);
} else if (addr < ctx->code_buffer) {
return (((char*)(ctx->code_buffer + ctx->code_buffer_size) - (char*)addr) < B_IMM);
if (addr >= ctx->code_buffer->start && (char*)addr < (char*)ctx->code_buffer->end) {
return (((char*)ctx->code_buffer->end - (char*)ctx->code_buffer->start) < B_IMM);
} else if ((char*)addr >= (char*)ctx->code_buffer->end) {
return (((char*)addr - (char*)ctx->code_buffer->start) < B_IMM);
} else if (addr < ctx->code_buffer->start) {
return (((char*)ctx->code_buffer->end - (char*)addr) < B_IMM);
}
}
return 1; //???
}
#if 0
static bool aarch64_may_use_adr(ir_ctx *ctx, const void *addr)
{
if (ctx->code_buffer) {
if (addr >= ctx->code_buffer && (char*)addr < (char*)ctx->code_buffer + ctx->code_buffer_size) {
return ( ctx->code_buffer_size < ADR_IMM);
} else if ((char*)addr >= (char*)ctx->code_buffer + ctx->code_buffer_size) {
return (((char*)addr - (char*)ctx->code_buffer) < ADR_IMM);
} else if (addr < ctx->code_buffer) {
return (((char*)(ctx->code_buffer + ctx->code_buffer_size) - (char*)addr) < ADR_IMM);
if (addr >= ctx->code_buffer->start && (char*)addr < (char*)ctx->code_buffer->end) {
return (((char*)ctx->code_buffer->end - (char*)ctx->code_buffer->start) < ADR_IMM);
} else if ((char*)addr >= (char*)ctx->code_buffer->end) {
return (((char*)addr - (char*)ctx->code_buffer->start) < ADR_IMM);
} else if (addr < ctx->code_buffer->start) {
return (((char*)ctx->code_buffer->end - (char*)addr) < ADR_IMM);
}
}
return 0;
@ -51,16 +52,17 @@ static bool aarch64_may_use_adr(ir_ctx *ctx, const void *addr)
static bool aarch64_may_use_adrp(ir_ctx *ctx, const void *addr)
{
if (ctx->code_buffer) {
if (addr >= ctx->code_buffer && (char*)addr < (char*)ctx->code_buffer + ctx->code_buffer_size) {
return ( ctx->code_buffer_size < ADRP_IMM);
} else if ((char*)addr >= (char*)ctx->code_buffer + ctx->code_buffer_size) {
return (((char*)addr - (char*)ctx->code_buffer) < ADRP_IMM);
} else if (addr < ctx->code_buffer) {
return (((char*)(ctx->code_buffer + ctx->code_buffer_size) - (char*)addr) < ADRP_IMM);
if (addr >= ctx->code_buffer->start && (char*)addr < (char*)ctx->code_buffer->end) {
return (((char*)ctx->code_buffer->end - (char*)ctx->code_buffer->start) < ADRP_IMM);
} else if ((char*)addr >= (char*)ctx->code_buffer->end) {
return (((char*)addr - (char*)ctx->code_buffer->start) < ADRP_IMM);
} else if (addr < ctx->code_buffer->start) {
return (((char*)ctx->code_buffer->end - (char*)addr) < ADRP_IMM);
}
}
return 0;
}
#endif
/* Determine whether "val" falls into two allowed ranges:
* Range 1: [0, 0xfff]
@ -5807,15 +5809,15 @@ void *ir_emit_code(ir_ctx *ctx, size_t *size_ptr)
}
size = *size_ptr;
if (ctx->code_buffer != NULL) {
if (IR_ALIGNED_SIZE(size, 16) > ctx->code_buffer_size) {
dasm_free(&data.dasm_state);
if (ctx->code_buffer) {
entry = ctx->code_buffer->pos;
entry = (void*)IR_ALIGNED_SIZE(((size_t)(entry)), 16);
if (size > (size_t)((char*)ctx->code_buffer->end - (char*)entry)) {
ctx->data = NULL;
ctx->status = IR_ERROR_CODE_MEM_OVERFLOW;
return NULL;
}
entry = ctx->code_buffer;
IR_ASSERT((uintptr_t)entry % 16 == 0);
ctx->code_buffer->pos = (char*)entry + size;
} else {
entry = ir_mem_mmap(size);
if (!entry) {
@ -5828,20 +5830,16 @@ void *ir_emit_code(ir_ctx *ctx, size_t *size_ptr)
}
ir_current_ctx = ctx;
ctx->veneers_size = 0;
if (data.jmp_table_label) {
ctx->code_size = dasm_getpclabel(&data.dasm_state, data.jmp_table_label);
} else if (data.rodata_label) {
ctx->code_size = dasm_getpclabel(&data.dasm_state, data.rodata_label);
} else {
ctx->code_size = size;
}
ret = dasm_encode(&data.dasm_state, entry);
if (ret != DASM_S_OK) {
IR_ASSERT(0);
dasm_free(&data.dasm_state);
if (ctx->code_buffer == NULL) {
if (ctx->code_buffer) {
if (ctx->code_buffer->pos == (char*)entry + size) {
/* rollback */
ctx->code_buffer->pos = (char*)entry - size;
}
} else {
ir_mem_unmap(entry, size);
}
ctx->data = NULL;
@ -5874,11 +5872,13 @@ void *ir_emit_code(ir_ctx *ctx, size_t *size_ptr)
dasm_free(&data.dasm_state);
*size_ptr += ctx->veneers_size;
if (ctx->code_buffer) {
size = (char*)ctx->code_buffer->pos - (char*)entry;
}
ir_mem_flush(entry, size);
if (ctx->code_buffer == NULL) {
if (!ctx->code_buffer) {
ir_mem_protect(entry, size);
}
@ -5886,7 +5886,7 @@ void *ir_emit_code(ir_ctx *ctx, size_t *size_ptr)
return entry;
}
const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_per_group, const void *exit_addr, void *code_buffer, size_t code_buffer_size, size_t *size_ptr)
const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_per_group, const void *exit_addr, ir_code_buffer *code_buffer, size_t *size_ptr)
{
void *entry;
size_t size;
@ -5896,12 +5896,12 @@ const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_pe
/* IR_ASSERT(aarch64_may_use_b(ctx, exit_addr)) */
IR_ASSERT(code_buffer);
if ((char*)exit_addr >= (char*)code_buffer && (char*)exit_addr < (char*)code_buffer + code_buffer_size) {
IR_ASSERT(code_buffer_size < B_IMM);
} else if ((char*)exit_addr >= (char*)code_buffer + code_buffer_size) {
IR_ASSERT(((char*)exit_addr - (char*)code_buffer) < B_IMM);
} else if ((char*)exit_addr < (char*)code_buffer) {
IR_ASSERT(((((char*)(code_buffer)) + code_buffer_size) - (char*)exit_addr) < B_IMM);
if ((char*)exit_addr >= (char*)code_buffer->start && (char*)exit_addr < (char*)code_buffer->end) {
IR_ASSERT(((char*)code_buffer->end - (char*)code_buffer->end) < B_IMM);
} else if ((char*)exit_addr >= (char*)code_buffer->end) {
IR_ASSERT(((char*)exit_addr - (char*)code_buffer->start) < B_IMM);
} else if ((char*)exit_addr < (char*)code_buffer->start) {
IR_ASSERT(((char*)code_buffer->end - (char*)exit_addr) < B_IMM);
} else {
IR_ASSERT(0);
}
@ -5933,25 +5933,21 @@ const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_pe
return NULL;
}
if (code_buffer != NULL) {
if (IR_ALIGNED_SIZE(size, 16) > code_buffer_size) {
dasm_free(&dasm_state);
entry = code_buffer->pos;
entry = (void*)IR_ALIGNED_SIZE(((size_t)(entry)), 16);
if (size > (size_t)((char*)code_buffer->end - (char*)entry)) {
return NULL;
}
entry = code_buffer;
IR_ASSERT((uintptr_t)entry % 16 == 0);
} else {
entry = ir_mem_mmap(size);
ir_mem_unprotect(entry, size);
}
code_buffer->pos = (char*)entry + size;
ir_current_ctx = NULL;
ret = dasm_encode(&dasm_state, entry);
if (ret != DASM_S_OK) {
IR_ASSERT(0);
dasm_free(&dasm_state);
if (code_buffer == NULL) {
ir_mem_unmap(entry, size);
if (code_buffer->pos == (char*)entry + size) {
/* rollback */
code_buffer->pos = (char*)entry - size;
}
return NULL;
}
@ -5960,10 +5956,6 @@ const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_pe
ir_mem_flush(entry, size);
if (code_buffer == NULL) {
ir_mem_protect(entry, size);
}
*size_ptr = size;
return entry;
}
@ -6013,8 +6005,8 @@ static int ir_add_veneer(dasm_State *Dst, void *buffer, uint32_t ins, int *b, ui
}
}
veneer = (char*)buffer + (Dst->codesize + ctx->veneers_size);
if (veneer > (void*)((char*)ctx->code_buffer + ctx->code_buffer_size)) {
veneer = ctx->code_buffer->pos;
if ((char*)ctx->code_buffer->end - (char*)veneer < 4 ) {
IR_ASSERT(0 && "too long jmp distance" && "jit buffer overflow");
return 0; /* jit_buffer_size overflow */
}
@ -6073,7 +6065,7 @@ static int ir_add_veneer(dasm_State *Dst, void *buffer, uint32_t ins, int *b, ui
/* generate B instruction */
*(uint32_t*)veneer = 0x14000000 | ((m >> 2) & 0x03ffffff);
ctx->veneers_size += 4;
ctx->code_buffer->pos = (char*)ctx->code_buffer->pos + 4;
return n;
}

View File

@ -252,8 +252,7 @@ typedef struct _ir_main_loader {
ir_sym *sym;
ir_ref sym_count;
void *data;
void *code_buffer;
size_t code_buffer_size;
ir_code_buffer code_buffer;
} ir_main_loader;
static bool ir_loader_add_sym(ir_loader *loader, const char *name, void *addr)
@ -641,32 +640,30 @@ static bool ir_loader_func_process(ir_loader *loader, ir_ctx *ctx, const char *n
size_t size;
void *entry;
if (l->code_buffer) {
ctx->code_buffer = (char*)l->code_buffer + l->size;
ctx->code_buffer_size = l->code_buffer_size - l->size;
ir_mem_unprotect(l->code_buffer, l->code_buffer_size);
if (l->code_buffer.start) {
ctx->code_buffer = &l->code_buffer;
ir_mem_unprotect(l->code_buffer.start, (char*)l->code_buffer.end - (char*)l->code_buffer.start);
}
entry = ir_emit_code(ctx, &size);
#ifndef _WIN32
if (l->run) {
if (!l->code_buffer) {
if (!l->code_buffer.start) {
ir_mem_unprotect(entry, size);
}
ir_gdb_register(name, entry, size, sizeof(void*), 0);
if (!l->code_buffer) {
if (!l->code_buffer.start) {
ir_mem_protect(entry, size);
}
}
#endif
if (l->code_buffer) {
ir_mem_protect(l->code_buffer, l->code_buffer_size);
if (l->code_buffer.start) {
ir_mem_protect(l->code_buffer.start, (char*)l->code_buffer.end - (char*)l->code_buffer.start);
}
if (entry) {
if (!l->code_buffer.start) {
l->size += size;
#if defined(IR_TARGET_AARCH64)
l->size += ctx->veneers_size;
#endif
l->size = IR_ALIGNED_SIZE(l->size, 16);
}
if (!ir_loader_add_sym(loader, name, entry)) {
fprintf(stderr, "\nERROR: Symbol redefinition: %s\n", name);
return 0;
@ -994,12 +991,15 @@ int main(int argc, char **argv)
#if defined(IR_TARGET_AARCH64)
if (dump_asm || dump_size || run) {
/* Preallocate 2MB JIT code buffer. On AArch64 it may be necessary to generate veneers. */
loader.code_buffer_size = 2 * 1024 * 1024;
loader.code_buffer = ir_mem_mmap(loader.code_buffer_size);
if (!loader.code_buffer) {
size_t size = 2 * 1024 * 1024;
loader.code_buffer.start = ir_mem_mmap(size);
if (!loader.code_buffer.start) {
fprintf(stderr, "ERROR: Cannot allocate JIT code buffer\n");
return 0;
}
loader.code_buffer.pos = loader.code_buffer.start;
loader.code_buffer.end = (char*)loader.code_buffer.start + size;
}
#endif
@ -1051,6 +1051,9 @@ finish:
}
if (dump_size) {
if (loader.code_buffer.start) {
loader.size = loader.code_buffer.pos - loader.code_buffer.start;
}
fprintf(stderr, "\ncode size = %lld\n", (long long int)loader.size);
}

View File

@ -21,8 +21,8 @@
#define IR_IS_FP_ZERO(insn) ((insn.type == IR_DOUBLE) ? (insn.val.u64 == 0) : (insn.val.u32 == 0))
#define IR_MAY_USE_32BIT_ADDR(addr) \
(ctx->code_buffer && \
IR_IS_SIGNED_32BIT((char*)(addr) - (char*)ctx->code_buffer) && \
IR_IS_SIGNED_32BIT((char*)(addr) - ((char*)ctx->code_buffer + ctx->code_buffer_size)))
IR_IS_SIGNED_32BIT((char*)(addr) - (char*)ctx->code_buffer->start) && \
IR_IS_SIGNED_32BIT((char*)(addr) - ((char*)ctx->code_buffer->end)))
#define IR_SPILL_POS_TO_OFFSET(offset) \
((ctx->flags & IR_USE_FRAME_POINTER) ? \
@ -9675,14 +9675,15 @@ next_block:;
}
size = *size_ptr;
if (ctx->code_buffer != NULL) {
if (IR_ALIGNED_SIZE(size, 16) > ctx->code_buffer_size) {
if (ctx->code_buffer) {
entry = ctx->code_buffer->pos;
entry = (void*)IR_ALIGNED_SIZE(((size_t)(entry)), 16);
if (size > (size_t)((char*)ctx->code_buffer->end - (char*)entry)) {
ctx->data = NULL;
ctx->status = IR_ERROR_CODE_MEM_OVERFLOW;
return NULL;
}
entry = ctx->code_buffer;
IR_ASSERT((uintptr_t)entry % 16 == 0);
ctx->code_buffer->pos = (char*)entry + size;
} else {
entry = ir_mem_mmap(size);
if (!entry) {
@ -9698,7 +9699,12 @@ next_block:;
if (ret != DASM_S_OK) {
IR_ASSERT(0);
dasm_free(&data.dasm_state);
if (ctx->code_buffer == NULL) {
if (ctx->code_buffer) {
if (ctx->code_buffer->pos == (char*)entry + size) {
/* rollback */
ctx->code_buffer->pos = (char*)entry - size;
}
} else {
ir_mem_unmap(entry, size);
}
ctx->data = NULL;
@ -9746,7 +9752,7 @@ next_block:;
}
#endif
if (ctx->code_buffer == NULL) {
if (!ctx->code_buffer) {
ir_mem_protect(entry, size);
}
@ -9754,7 +9760,7 @@ next_block:;
return entry;
}
const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_per_group, const void *exit_addr, void *code_buffer, size_t code_buffer_size, size_t *size_ptr)
const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_per_group, const void *exit_addr, ir_code_buffer *code_buffer, size_t *size_ptr)
{
void *entry;
size_t size;
@ -9763,8 +9769,8 @@ const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_pe
int ret;
IR_ASSERT(code_buffer);
IR_ASSERT(IR_IS_SIGNED_32BIT((char*)exit_addr - (char*)code_buffer));
IR_ASSERT(IR_IS_SIGNED_32BIT((char*)exit_addr - ((char*)code_buffer + code_buffer_size)));
IR_ASSERT(IR_IS_SIGNED_32BIT((char*)exit_addr - (char*)code_buffer->start));
IR_ASSERT(IR_IS_SIGNED_32BIT((char*)exit_addr - (char*)code_buffer->end));
Dst = &dasm_state;
dasm_state = NULL;
@ -9788,24 +9794,20 @@ const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_pe
return NULL;
}
if (code_buffer != NULL) {
if (IR_ALIGNED_SIZE(size, 16) > code_buffer_size) {
dasm_free(&dasm_state);
entry = code_buffer->pos;
entry = (void*)IR_ALIGNED_SIZE(((size_t)(entry)), 16);
if (size > (size_t)((char*)code_buffer->end - (char*)entry)) {
return NULL;
}
entry = code_buffer;
IR_ASSERT((uintptr_t)entry % 16 == 0);
} else {
entry = ir_mem_mmap(size);
ir_mem_unprotect(entry, size);
}
code_buffer->pos = (char*)entry + size;
ret = dasm_encode(&dasm_state, entry);
if (ret != DASM_S_OK) {
IR_ASSERT(0);
dasm_free(&dasm_state);
if (code_buffer == NULL) {
ir_mem_unmap(entry, size);
if (code_buffer->pos == (char*)entry + size) {
/* rollback */
code_buffer->pos = (char*)entry - size;
}
return NULL;
}
@ -9814,10 +9816,6 @@ const void *ir_emit_exitgroup(uint32_t first_exit_point, uint32_t exit_points_pe
ir_mem_flush(entry, size);
if (code_buffer == NULL) {
ir_mem_protect(entry, size);
}
*size_ptr = size;
return entry;
}