Skip to content

Commit

Permalink
Allow multiple JIT code pages. Issue #672
Browse files Browse the repository at this point in the history
  • Loading branch information
nickg committed May 1, 2023
1 parent 26d8f43 commit 9299dc6
Show file tree
Hide file tree
Showing 3 changed files with 119 additions and 25 deletions.
2 changes: 2 additions & 0 deletions src/diag.c
Original file line number Diff line number Diff line change
Expand Up @@ -915,6 +915,8 @@ void diag_femit(diag_t *d, FILE *f)
return;
else if (consumer != NULL && d->level > DIAG_DEBUG)
(*consumer)(d);
else if (d->level == DIAG_DEBUG && opt_get_int(OPT_UNIT_TEST))
return;
else {
SCOPED_LOCK(diag_lock);

Expand Down
100 changes: 75 additions & 25 deletions src/jit/jit-code.c
Original file line number Diff line number Diff line change
Expand Up @@ -67,8 +67,8 @@
#define SHT_X86_64_UNWIND 0x70000001
#endif

#define CODECACHE_ALIGN 4096
#define CODECACHE_SIZE 0x400000
#define CODE_PAGE_ALIGN 4096
#define CODE_PAGE_SIZE 0x400000
#define THREAD_CACHE_SIZE 0x10000
#define CODE_BLOB_ALIGN 256
#define MIN_BLOB_SIZE 0x4000
Expand All @@ -79,7 +79,9 @@

STATIC_ASSERT(MIN_BLOB_SIZE <= THREAD_CACHE_SIZE);
STATIC_ASSERT(MIN_BLOB_SIZE % CODE_BLOB_ALIGN == 0);
STATIC_ASSERT(CODECACHE_SIZE % THREAD_CACHE_SIZE == 0);
STATIC_ASSERT(CODE_PAGE_SIZE % THREAD_CACHE_SIZE == 0);

typedef struct _code_page code_page_t;

typedef struct _code_span {
code_cache_t *owner;
Expand All @@ -96,9 +98,15 @@ typedef struct _patch_list {
code_patch_fn_t fn;
} patch_list_t;

typedef struct _code_page {
code_cache_t *owner;
code_page_t *next;
uint8_t *mem;
} code_page_t;

typedef struct _code_cache {
nvc_lock_t lock;
uint8_t *mem;
code_page_t *pages;
code_span_t *spans;
code_span_t *freelist[MAX_THREADS];
code_span_t *globalfree;
Expand Down Expand Up @@ -132,10 +140,10 @@ static void code_cache_unwinder(uintptr_t addr, debug_frame_t *frame,
static void code_fault_handler(int sig, void *addr, struct cpu_state *cpu,
void *context)
{
code_cache_t *code = context;
code_page_t *page = context;

const uint8_t *pc = (uint8_t *)cpu->pc;
if (pc < code->mem || pc > code->mem + CODECACHE_SIZE)
if (pc < page->mem || pc > page->mem + CODE_PAGE_SIZE)
return;

uintptr_t mark = cpu->pc;
Expand All @@ -144,20 +152,33 @@ static void code_fault_handler(int sig, void *addr, struct cpu_state *cpu,
mark--; // Point to faulting instruction
#endif

for (code_span_t *span = code->spans; span; span = span->next) {
for (code_span_t *span = page->owner->spans; span; span = span->next) {
if (pc >= span->base && pc < span->base + span->size && span->name)
code_disassemble(span, mark, cpu);
}
}

#ifdef DEBUG
static bool code_cache_contains(code_cache_t *code, uint8_t *base, size_t size)
{
assert_lock_held(&code->lock);

for (code_page_t *p = code->pages; p; p = p->next) {
if (base >= p->mem && base + size <= p->mem + CODE_PAGE_SIZE)
return true;
}

return false;
}
#endif

static code_span_t *code_span_new(code_cache_t *code, ident_t name,
uint8_t *base, size_t size)
{
assert(base >= code->mem);
assert(base + size <= code->mem + CODECACHE_SIZE);

SCOPED_LOCK(code->lock);

assert(code_cache_contains(code, base, size));

code_span_t *span = xcalloc(sizeof(code_span_t));
span->name = name;
span->next = code->spans;
Expand All @@ -169,10 +190,37 @@ static code_span_t *code_span_new(code_cache_t *code, ident_t name,
return span;
}

static void code_page_new(code_cache_t *code)
{
assert_lock_held(&code->lock);

code_page_t *page = xcalloc(sizeof(code_page_t));
page->owner = code;
page->next = code->pages;
page->mem = map_jit_pages(CODE_PAGE_ALIGN, CODE_PAGE_SIZE);

add_fault_handler(code_fault_handler, page);
debug_add_unwinder(page->mem, CODE_PAGE_SIZE, code_cache_unwinder, code);

code->pages = page;

code_span_t *span = xcalloc(sizeof(code_span_t));
span->next = code->spans;
span->base = page->mem;
span->size = CODE_PAGE_SIZE;
span->owner = code;

code->globalfree = code->spans = span;
}

code_cache_t *code_cache_new(void)
{
code_cache_t *code = xcalloc(sizeof(code_cache_t));
code->mem = map_jit_pages(CODECACHE_ALIGN, CODECACHE_SIZE);

{
SCOPED_LOCK(code->lock);
code_page_new(code);
}

#ifdef HAVE_CAPSTONE
#if defined ARCH_X86_64
Expand All @@ -189,20 +237,20 @@ code_cache_t *code_cache_new(void)
fatal_trace("failed to set capstone detailed mode");
#endif

add_fault_handler(code_fault_handler, code);
debug_add_unwinder(code->mem, CODECACHE_SIZE, code_cache_unwinder, code);

code->globalfree = code_span_new(code, NULL, code->mem, CODECACHE_SIZE);

return code;
}

void code_cache_free(code_cache_t *code)
{
debug_remove_unwinder(code->mem);
remove_fault_handler(code_fault_handler, code);
for (code_page_t *it = code->pages, *tmp; it; it = tmp) {
debug_remove_unwinder(it->mem);
remove_fault_handler(code_fault_handler, it);

nvc_munmap(code->mem, CODECACHE_SIZE);
nvc_munmap(it->mem, CODE_PAGE_SIZE);

tmp = it->next;
free(it);
}

for (code_span_t *it = code->spans, *tmp; it; it = tmp) {
tmp = it->next;
Expand All @@ -214,8 +262,7 @@ void code_cache_free(code_cache_t *code)
#endif

#ifdef DEBUG
if (!opt_get_int(OPT_UNIT_TEST))
debugf("JIT code footprint: %zu bytes", code->used);
debugf("JIT code footprint: %zu bytes", code->used);
#endif

free(code);
Expand Down Expand Up @@ -345,7 +392,7 @@ code_blob_t *code_blob_new(code_cache_t *code, ident_t name, size_t hint)

code_span_t *free = relaxed_load(freeptr);
if (free == NULL) {
free = code_span_new(code, NULL, code->mem, 0);
free = code_span_new(code, NULL, code->pages->mem, 0);
relaxed_store(freeptr, free);
}

Expand Down Expand Up @@ -374,8 +421,11 @@ code_blob_t *code_blob_new(code_cache_t *code, ident_t name, size_t hint)
code->globalfree->base += take;
code->globalfree->size -= take;

if (code->globalfree->size == 0)
warnf("global JIT code buffer exhausted");
if (code->globalfree->size == 0) {
DEBUG_ONLY(debugf("requesting new %d byte code page", CODE_PAGE_SIZE));
code_page_new(code);
assert(code->globalfree->size == CODE_PAGE_SIZE);
}
}

assert(reqsz <= free->size);
Expand Down Expand Up @@ -445,7 +495,7 @@ void code_blob_emit(code_blob_t *blob, const uint8_t *bytes, size_t len)
{
if (unlikely(blob->overflow))
return;
else if (unlikely(blob->wptr + len >= blob->span->base + blob->span->size)) {
else if (unlikely(blob->wptr + len > blob->span->base + blob->span->size)) {
warnf("JIT code buffer for %s too small", istr(blob->span->name));
for (patch_list_t *it = blob->patches, *tmp; it; it = tmp) {
tmp = it->next;
Expand Down
42 changes: 42 additions & 0 deletions test/test_jit.c
Original file line number Diff line number Diff line change
Expand Up @@ -2011,6 +2011,47 @@ START_TEST(test_dce2)
}
END_TEST

START_TEST(test_code1)
{
const error_t expect[] = {
{ LINE_INVALID, "JIT code buffer for test too small" },
{ -1, NULL },
};
expect_errors(expect);

code_cache_t *code = code_cache_new();

code_blob_t *blob = code_blob_new(code, ident_new("test"), 100);
ck_assert_ptr_nonnull(blob);

void *mem LOCAL = xcalloc(1024 * 1024);

code_blob_emit(blob, mem, 128);
fail_if(blob->overflow);

code_blob_emit(blob, mem, 1024 * 1024);
fail_unless(blob->overflow);

jit_entry_fn_t entry = NULL;
code_blob_finalise(blob, &entry);
ck_assert_ptr_null(entry);

for (int i = 0; i < 200; i++) {
blob = code_blob_new(code, ident_new("loop"), 32 * 1024);
ck_assert_ptr_nonnull(blob);

code_blob_emit(blob, mem, 32 * 1024);
fail_if(blob->overflow);

code_blob_finalise(blob, &entry);
}

code_cache_free(code);

check_expected_errors();
}
END_TEST

Suite *get_jit_tests(void)
{
Suite *s = suite_create("jit");
Expand Down Expand Up @@ -2065,6 +2106,7 @@ Suite *get_jit_tests(void)
tcase_add_test(tc, test_lvn9);
tcase_add_test(tc, test_cfg3);
tcase_add_test(tc, test_dce2);
tcase_add_test(tc, test_code1);
suite_add_tcase(s, tc);

return s;
Expand Down

0 comments on commit 9299dc6

Please sign in to comment.