Improve "JALR" execution using JIT-cache

vacantron · vacantron · commit c8f98768f01a · 2024-08-02T03:00:42.000+08:00
Currently, the "JALR" indirect jump instruction turns the mode of
rv32emu from T2C back to the interpreter. This commit introduces a
"JIT-cache" table lookup to make it redirect to the JIT-ed code entry
and avoids the mode change.

There are several scenarios benefitting from this approach, e.g.
function pointer invocation and far-way function call. The former
like "qsort" can be speeded up two times, and the latter like
"Fibonacci", which compiles from the hand-written assembly, can even
reach x4.3 performance enhencement.
diff --git a/Makefile b/Makefile
@@ -141,7 +141,7 @@ ifneq ("$(LLVM_CONFIG)", "")
 ifneq ("$(findstring -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS, "$(shell $(LLVM_CONFIG) --cflags)")", "")
 ENABLE_T2C := 1
 $(call set-feature, T2C)
-OBJS_EXT += t2c.o
+OBJS_EXT += jit-cache.o t2c.o
 CFLAGS += -g $(shell $(LLVM_CONFIG) --cflags)
 LDFLAGS += $(shell $(LLVM_CONFIG) --libs)
 else
diff --git a/build/Fibonacci.elf b/build/Fibonacci.elf
diff --git a/src/jit-cache.c b/src/jit-cache.c
@@ -0,0 +1,36 @@
+/*
+ * rv32emu is freely redistributable under the MIT License. See the file
+ * "LICENSE" for information on usage and redistribution of this file.
+ */
+
+#include <stdlib.h>
+#include <string.h>
+
+#include "jit-cache.h"
+
+jit_cache_t *jit_cache_init()
+{
+    return calloc(MAX_JIT_CACHE_TABLE_ENTRIES, sizeof(jit_cache_t));
+}
+
+void jit_cache_exit(jit_cache_t *cache)
+{
+    free(cache);
+}
+
+void jit_cache_update(jit_cache_t *cache,
+                      uint32_t pc,
+                      uint32_t from,
+                      void *entry)
+{
+    uint32_t pos = pc & (MAX_JIT_CACHE_TABLE_ENTRIES - 1);
+
+    cache[pos].pc = pc;
+    cache[pos].from = from;
+    cache[pos].entry = entry;
+}
+
+void jit_cache_clear(jit_cache_t *cache)
+{
+    memset(cache, 0, MAX_JIT_CACHE_TABLE_ENTRIES * sizeof(jit_cache_t));
+}
diff --git a/src/jit-cache.h b/src/jit-cache.h
@@ -0,0 +1,37 @@
+/*
+ * rv32emu is freely redistributable under the MIT License. See the file
+ * "LICENSE" for information on usage and redistribution of this file.
+ */
+
+#pragma once
+
+#if !defined(__x86_64__) && !defined(__aarch64__)
+/* make sure LLVM can access the correct address of structure member */
+#error "The JIT-cache only supports 64-bit machine (LP64)."
+#endif
+
+#include <stdint.h>
+
+#define MAX_JIT_CACHE_TABLE_ENTRIES (1 << 12)
+
+/* clang-format off */
+typedef struct jit_cache {
+    __attribute__((aligned(4))) uint32_t pc;   /* program counter of ELF */
+    __attribute__((aligned(4))) uint32_t from; /* from whether t1c or t2c,
+                                                  easily to build LLVM IR 
+                                                  because of the alignment
+                                                */
+    void *entry; /* entry of JIT-ed code */
+} jit_cache_t;
+/* clang-format on */
+
+jit_cache_t *jit_cache_init();
+
+void jit_cache_exit(jit_cache_t *cache);
+
+void jit_cache_update(jit_cache_t *cache,
+                      uint32_t pc,
+                      uint32_t from,
+                      void *entry);
+
+void jit_cache_clear(jit_cache_t *cache);
diff --git a/src/jit.c b/src/jit.c
@@ -37,6 +37,7 @@
 #include "cache.h"
 #include "decode.h"
 #include "io.h"
+#include "jit-cache.h"
 #include "jit.h"
 #include "riscv.h"
 #include "riscv_private.h"
@@ -1864,6 +1865,7 @@ static void code_cache_flush(struct jit_state *state, riscv_t *rv)
     state->offset = state->org_size;
     state->n_blocks = 0;
     set_reset(&state->set);
+    jit_cache_clear(rv->jit_cache);
     clear_cache_hot(rv->block_cache, (clear_func_t) clear_hot);
     return;
 }
@@ -1934,6 +1936,9 @@ static void translate_chained_block(struct jit_state *state,
 
     set_add(&state->set, block->pc_start);
     offset_map_insert(state, block->pc_start);
+    block->offset = state->offset;
+    jit_cache_update(rv->jit_cache, block->pc_start, 1,
+                     state->buf + block->offset);
     translate(state, rv, block);
     if (unlikely(should_flush))
         return;
diff --git a/src/jit.h b/src/jit.h
@@ -51,6 +51,6 @@ void jit_translate(riscv_t *rv, block_t *block);
 typedef void (*exec_block_func_t)(riscv_t *rv, uintptr_t);
 
 #if RV32_HAS(T2C)
-void t2c_compile(block_t *block, uint64_t mem_base);
+void t2c_compile(riscv_t *, block_t *);
 typedef void (*exec_t2c_func_t)(riscv_t *);
 #endif
diff --git a/src/riscv.c b/src/riscv.c
@@ -32,6 +32,7 @@
 #include <pthread.h>
 #endif
 #include "cache.h"
+#include "jit-cache.h"
 #include "jit.h"
 #define CODE_CACHE_SIZE (4 * 1024 * 1024)
 #endif
@@ -199,8 +200,7 @@ static void *t2c_runloop(void *arg)
             pthread_mutex_lock(&rv->wait_queue_lock);
             list_del_init(&entry->list);
             pthread_mutex_unlock(&rv->wait_queue_lock);
-            t2c_compile(entry->block,
-                        (uint64_t) ((memory_t *) PRIV(rv)->mem)->mem_base);
+            t2c_compile(rv, entry->block);
             free(entry);
         }
     }
@@ -291,6 +291,7 @@ riscv_t *rv_create(riscv_user_t rv_attr)
         mpool_create(sizeof(chain_entry_t) << BLOCK_IR_MAP_CAPACITY_BITS,
                      sizeof(chain_entry_t));
     rv->jit_state = jit_state_init(CODE_CACHE_SIZE);
+    rv->jit_cache = jit_cache_init();
     rv->block_cache = cache_create(BLOCK_MAP_CAPACITY_BITS);
     assert(rv->block_cache);
 #if RV32_HAS(T2C)
@@ -392,6 +393,7 @@ void rv_delete(riscv_t *rv)
 #endif
     mpool_destroy(rv->chain_entry_mp);
     jit_state_exit(rv->jit_state);
+    jit_cache_exit(rv->jit_cache);
     cache_free(rv->block_cache);
 #endif
     free(rv);
diff --git a/src/riscv_private.h b/src/riscv_private.h
@@ -155,6 +155,7 @@ struct riscv_internal {
     struct mpool *block_mp, *block_ir_mp;
 
     void *jit_state;
+    void *jit_cache;
 #if RV32_HAS(GDBSTUB)
     /* gdbstub instance */
     gdbstub_t gdbstub;
diff --git a/src/t2c.c b/src/t2c.c
@@ -11,6 +11,7 @@
 #include <llvm-c/Transforms/PassBuilder.h>
 #include <stdlib.h>
 
+#include "jit-cache.h"
 #include "jit.h"
 #include "riscv_private.h"
 
@@ -49,15 +50,15 @@ FORCE_INLINE LLVMBasicBlockRef t2c_block_map_search(struct LLVM_block_map *map,
     return NULL;
 }
 
-#define T2C_OP(inst, code)                                                \
-    static void t2c_##inst(                                               \
-        LLVMBuilderRef *builder UNUSED, LLVMTypeRef *param_types UNUSED,  \
-        LLVMValueRef start UNUSED, LLVMBasicBlockRef *entry UNUSED,       \
-        LLVMBuilderRef *taken_builder UNUSED,                             \
-        LLVMBuilderRef *untaken_builder UNUSED, uint64_t mem_base UNUSED, \
-        rv_insn_t *ir UNUSED)                                             \
-    {                                                                     \
-        code;                                                             \
+#define T2C_OP(inst, code)                                               \
+    static void t2c_##inst(                                              \
+        LLVMBuilderRef *builder UNUSED, LLVMTypeRef *param_types UNUSED, \
+        LLVMValueRef start UNUSED, LLVMBasicBlockRef *entry UNUSED,      \
+        LLVMBuilderRef *taken_builder UNUSED,                            \
+        LLVMBuilderRef *untaken_builder UNUSED, riscv_t *rv UNUSED,      \
+        uint64_t mem_base UNUSED, rv_insn_t *ir UNUSED)                  \
+    {                                                                    \
+        code;                                                            \
     }
 
 #define T2C_LLVM_GEN_ADDR(reg, rv_member, ir_member)                          \
@@ -135,6 +136,12 @@ FORCE_INLINE void t2c_gen_call_io_func(LLVMValueRef start,
                    &io_param, 1, "");
 }
 
+static bool t2c_jit_cache_triggered;
+static LLVMTypeRef t2c_t1c_cache_func_proto;
+static LLVMValueRef t2c_t1c_cache_func;
+static LLVMTypeRef t2c_t2c_cache_func_proto;
+static LLVMTypeRef jit_cache_ty;
+
 #include "t2c_template.c"
 #undef T2C_OP
 
@@ -174,14 +181,15 @@ typedef void (*t2c_codegen_block_func_t)(LLVMBuilderRef *builder UNUSED,
                                          LLVMBasicBlockRef *entry UNUSED,
                                          LLVMBuilderRef *taken_builder UNUSED,
                                          LLVMBuilderRef *untaken_builder UNUSED,
+                                         riscv_t *rv UNUSED,
                                          uint64_t mem_base UNUSED,
                                          rv_insn_t *ir UNUSED);
 
 static void t2c_trace_ebb(LLVMBuilderRef *builder,
                           LLVMTypeRef *param_types UNUSED,
                           LLVMValueRef start,
                           LLVMBasicBlockRef *entry,
-                          uint64_t mem_base,
+                          riscv_t *rv,
                           rv_insn_t *ir,
                           set_t *set,
                           struct LLVM_block_map *map)
@@ -194,7 +202,8 @@ static void t2c_trace_ebb(LLVMBuilderRef *builder,
 
     while (1) {
         ((t2c_codegen_block_func_t) dispatch_table[ir->opcode])(
-            builder, param_types, start, entry, &tk, &utk, mem_base, ir);
+            builder, param_types, start, entry, &tk, &utk, rv,
+            (uint64_t) ((memory_t *) PRIV(rv)->mem)->mem_base, ir);
         if (!ir->next)
             break;
         ir = ir->next;
@@ -214,8 +223,7 @@ static void t2c_trace_ebb(LLVMBuilderRef *builder,
                 LLVMPositionBuilderAtEnd(untaken_builder, untaken_entry);
                 LLVMBuildBr(utk, untaken_entry);
                 t2c_trace_ebb(&untaken_builder, param_types, start,
-                              &untaken_entry, mem_base, ir->branch_untaken, set,
-                              map);
+                              &untaken_entry, rv, ir->branch_untaken, set, map);
             }
         }
         if (ir->branch_taken) {
@@ -230,14 +238,16 @@ static void t2c_trace_ebb(LLVMBuilderRef *builder,
                 LLVMPositionBuilderAtEnd(taken_builder, taken_entry);
                 LLVMBuildBr(tk, taken_entry);
                 t2c_trace_ebb(&taken_builder, param_types, start, &taken_entry,
-                              mem_base, ir->branch_taken, set, map);
+                              rv, ir->branch_taken, set, map);
             }
         }
     }
 }
 
-void t2c_compile(block_t *block, uint64_t mem_base)
+void t2c_compile(riscv_t *rv, block_t *block)
 {
+    t2c_jit_cache_triggered = false;
+
     LLVMModuleRef module = LLVMModuleCreateWithName("my_module");
     LLVMTypeRef io_members[] = {
         LLVMPointerType(LLVMVoidType(), 0), LLVMPointerType(LLVMVoidType(), 0),
@@ -254,6 +264,22 @@ void t2c_compile(block_t *block, uint64_t mem_base)
     LLVMTypeRef param_types[] = {LLVMPointerType(struct_rv, 0)};
     LLVMValueRef start = LLVMAddFunction(
         module, "start", LLVMFunctionType(LLVMVoidType(), param_types, 1, 0));
+
+    LLVMTypeRef t1c_args[2] = {LLVMInt64Type(), LLVMInt64Type()};
+    t2c_t1c_cache_func_proto =
+        LLVMFunctionType(LLVMVoidType(), t1c_args, 2, false);
+    t2c_t1c_cache_func =
+        LLVMAddFunction(module, "t2c_invoke_cache", t2c_t1c_cache_func_proto);
+
+    LLVMTypeRef t2c_args[1] = {LLVMInt64Type()};
+    t2c_t2c_cache_func_proto =
+        LLVMFunctionType(LLVMVoidType(), t2c_args, 1, false);
+
+    /* Notice to the alignment */
+    LLVMTypeRef jit_cache_memb[3] = {LLVMInt32Type(), LLVMInt32Type(),
+                                     LLVMPointerType(LLVMVoidType(), 0)};
+    jit_cache_ty = LLVMStructType(jit_cache_memb, 3, false);
+
     LLVMBasicBlockRef first_block = LLVMAppendBasicBlock(start, "first_block");
     LLVMBuilderRef first_builder = LLVMCreateBuilder();
     LLVMPositionBuilderAtEnd(first_builder, first_block);
@@ -266,8 +292,8 @@ void t2c_compile(block_t *block, uint64_t mem_base)
     struct LLVM_block_map map;
     map.count = 0;
     /* Translate custon IR into LLVM IR */
-    t2c_trace_ebb(&builder, param_types, start, &entry, mem_base,
-                  block->ir_head, &set, &map);
+    t2c_trace_ebb(&builder, param_types, start, &entry, rv, block->ir_head,
+                  &set, &map);
     /* Offload LLVM IR to LLVM backend */
     char *error = NULL, *triple = LLVMGetDefaultTargetTriple();
     LLVMExecutionEngineRef engine;
@@ -296,7 +322,13 @@ void t2c_compile(block_t *block, uint64_t mem_base)
         abort();
     }
 
+    if (t2c_jit_cache_triggered) {
+        LLVMAddGlobalMapping(engine, t2c_t1c_cache_func,
+                             ((struct jit_state *) rv->jit_state)->buf);
+    }
+
     /* Return the function pointer of T2C generated machine code */
     block->func = (exec_t2c_func_t) LLVMGetPointerToGlobal(engine, start);
+    jit_cache_update(rv->jit_cache, block->pc_start, 2, block->func);
     block->hot2 = true;
 }
diff --git a/src/t2c_template.c b/src/t2c_template.c
