Add Arm64 TSAN support and fix JIT cache coherency

This commit adds ThreadSanitizer (TSAN) support for ARM64/Apple Silicon
and fixes critical JIT instruction cache coherency issues.

ARM64 TSAN Support:
- Extended TSAN-compatible memory allocation to ARM64 architecture
- Main memory allocated at fixed address 0x150000000000 (21TB)
- JIT buffer allocated at 0x151000000000 with MAP_JIT for Apple Silicon
- Both allocations avoid TSAN shadow memory and enable race detection
- Note: Requires ASLR disabled on macOS (SIP restrictions may apply)

JIT Cache Coherency Fixes:
1. Fixed pthread_jit_write_protect_np() ordering in update_branch_imm
2. Added sys_icache_invalidate() in update_branch_imm
3. Added cache invalidation in resolve_jumps() for x86_64

Author: jserv

src/emulate.c

@@ -24,6 +24,7 @@ extern struct target_ops gdbstub_ops;
 #endif
 
 #include "decode.h"
+#include "log.h"
 #include "mpool.h"
 #include "riscv.h"
 #include "riscv_private.h"
@@ -1205,10 +1206,20 @@ void rv_step(void *arg)
 #endif
         ) {
             jit_translate(rv, block);
-            ((exec_block_func_t) state->buf)(
-                rv, (uintptr_t) (state->buf + block->offset));
-            prev = NULL;
-            continue;
+            /* Only execute if translation succeeded (block is hot) */
+            if (block->hot) {
+                rv_log_debug("JIT: Executing block pc=0x%08x, offset=%u",
+                             block->pc_start, block->offset);
+                ((exec_block_func_t) state->buf)(
+                    rv, (uintptr_t) (state->buf + block->offset));
+                prev = NULL;
+                continue;
+            }
+            /* Fall through to interpreter if translation failed */
+            rv_log_debug(
+                "JIT: Translation failed for block pc=0x%08x, using "
+                "interpreter",
+                block->pc_start);
         }
         set_reset(&pc_set);
         has_loops = false;

src/io.c

@@ -27,18 +27,32 @@ memory_t *memory_new(uint32_t size)
         return NULL;
     assert(mem);
 #if HAVE_MMAP
-#if defined(TSAN_ENABLED) && defined(__x86_64__)
+#if defined(TSAN_ENABLED)
     /* ThreadSanitizer compatibility: Use MAP_FIXED to allocate at a specific
-     * address within TSAN's app range (0x7cf000000000 - 0x7ffffffff000).
+     * address to avoid conflicts with TSAN's shadow memory.
+     */
+#if defined(__x86_64__)
+    /* x86_64: Allocate within TSAN's range (0x7cf000000000 - 0x7ffffffff000).
      *
      * Fixed address: 0x7d0000000000
      * Size: up to 4GB (0x100000000)
      * End: 0x7d0100000000 (well within app range)
-     *
-     * This guarantees the allocation won't land in TSAN's shadow memory,
-     * preventing "unexpected memory mapping" errors.
      */
     void *fixed_addr = (void *) 0x7d0000000000UL;
+#elif defined(__aarch64__)
+    /* ARM64 (macOS/Apple Silicon): Use higher address range.
+     *
+     * Fixed address: 0x150000000000 (21TB)
+     * Size: up to 4GB (0x100000000)
+     * End: 0x150100000000
+     *
+     * This avoids TSAN's shadow memory and typical process allocations.
+     * Requires ASLR disabled via: setarch $(uname -m) -R
+     */
+    void *fixed_addr = (void *) 0x150000000000UL;
+#else
+#error "TSAN is only supported on x86_64 and aarch64"
+#endif
     data_memory_base = mmap(fixed_addr, size, PROT_READ | PROT_WRITE,
                             MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
     if (data_memory_base == MAP_FAILED) {

src/jit.c

@@ -42,6 +42,7 @@
 #include "decode.h"
 #include "io.h"
 #include "jit.h"
+#include "log.h"
 #include "riscv.h"
 #include "riscv_private.h"
 #include "utils.h"
@@ -593,24 +594,30 @@ static void update_branch_imm(struct jit_state *state,
     assert((imm & 3) == 0);
     uint32_t insn;
     imm >>= 2;
+    rv_log_debug("JIT: Patching branch at offset=%u, imm=%d", offset, imm * 4);
+    /* Read instruction while in execute mode (MAP_JIT requirement) */
     memcpy(&insn, state->buf + offset, sizeof(uint32_t));
     if ((insn & 0xfe000000U) == 0x54000000U /* Conditional branch immediate. */
         || (insn & 0x7e000000U) ==
                0x34000000U) { /* Compare and branch immediate. */
         assert((imm >> 19) == INT64_C(-1) || (imm >> 19) == 0);
+        insn &= ~(0x7ffffU << 5); /* Clear old offset bits */
         insn |= (imm & 0x7ffff) << 5;
     } else if ((insn & 0x7c000000U) == 0x14000000U) {
         /* Unconditional branch immediate.  */
         assert((imm >> 26) == INT64_C(-1) || (imm >> 26) == 0);
+        insn &= ~0x03ffffffU; /* Clear old offset bits */
         insn |= (imm & 0x03ffffffU) << 0;
     } else {
         assert(false);
         insn = BAD_OPCODE;
     }
 #if defined(__APPLE__) && defined(__aarch64__)
+    /* Switch to write mode only for writing */
     pthread_jit_write_protect_np(false);
 #endif
     memcpy(state->buf + offset, &insn, sizeof(uint32_t));
+    sys_icache_invalidate(state->buf + offset, sizeof(uint32_t));
 #if defined(__APPLE__) && defined(__aarch64__)
     pthread_jit_write_protect_np(true);
 #endif
@@ -2164,9 +2171,12 @@ void clear_hot(block_t *block)
 
 static void code_cache_flush(struct jit_state *state, riscv_t *rv)
 {
+    rv_log_info("JIT: Flushing code cache (n_blocks=%d, n_jumps=%d, offset=%u)",
+                state->n_blocks, state->n_jumps, state->offset);
     should_flush = false;
     state->offset = state->org_size;
     state->n_blocks = 0;
+    state->n_jumps = 0; /* Reset jump count when flushing */
     set_reset(&state->set);
     clear_cache_hot(rv->block_cache, (clear_func_t) clear_hot);
 #if RV32_HAS(T2C)
@@ -2196,6 +2206,7 @@ static void translate(struct jit_state *state, riscv_t *rv, block_t *block)
 
 static void resolve_jumps(struct jit_state *state)
 {
+    rv_log_debug("JIT: Resolving %d jumps", state->n_jumps);
     for (int i = 0; i < state->n_jumps; i++) {
         struct jump jump = state->jumps[i];
         int target_loc;
@@ -2218,6 +2229,10 @@ static void resolve_jumps(struct jit_state *state)
                     (if (jump.target_satp == state->offset_map[i].satp), )
                     {
                         target_loc = state->offset_map[i].offset;
+                        rv_log_debug(
+                            "JIT: Jump %d resolved to block pc=0x%08x, "
+                            "offset=%d",
+                            i, jump.target_pc, target_loc);
                         break;
                     }
                 }
@@ -2229,6 +2244,7 @@ static void resolve_jumps(struct jit_state *state)
 
         uint8_t *offset_ptr = &state->buf[jump.offset_loc];
         memcpy(offset_ptr, &rel, sizeof(uint32_t));
+        sys_icache_invalidate(offset_ptr, sizeof(uint32_t));
 #elif defined(__aarch64__)
         int32_t rel = target_loc - jump.offset_loc;
         update_branch_imm(state, jump.offset_loc, rel);
@@ -2308,23 +2324,35 @@ void jit_translate(riscv_t *rv, block_t *block)
             ) {
                 block->offset = state->offset_map[i].offset;
                 block->hot = true;
+                rv_log_debug("JIT: Cache hit for block pc=0x%08x, offset=%u",
+                             block->pc_start, block->offset);
                 return;
             }
         }
         assert(NULL);
         __UNREACHABLE;
     }
+    rv_log_debug("JIT: Starting translation for block pc=0x%08x",
+                 block->pc_start);
 restart:
     memset(state->jumps, 0, MAX_JUMPS * sizeof(struct jump));
     state->n_jumps = 0;
     block->offset = state->offset;
     translate_chained_block(state, rv, block);
     if (unlikely(should_flush)) {
+        /* Mark block as not translated since translation was incomplete */
+        block->hot = false;
+        /* Don't reset offset - it will be set correctly on restart */
+        rv_log_debug("JIT: Translation triggered flush for block pc=0x%08x",
+                     block->pc_start);
         code_cache_flush(state, rv);
         goto restart;
     }
     resolve_jumps(state);
     block->hot = true;
+    rv_log_debug(
+        "JIT: Translation completed for block pc=0x%08x, offset=%u, size=%u",
+        block->pc_start, block->offset, state->offset - block->offset);
 }
 
 struct jit_state *jit_state_init(size_t size)
@@ -2336,10 +2364,12 @@ struct jit_state *jit_state_init(size_t size)
 
     state->offset = 0;
     state->size = size;
-#if defined(TSAN_ENABLED) && defined(__x86_64__)
+#if defined(TSAN_ENABLED)
     /* ThreadSanitizer compatibility: Allocate JIT code buffer at a fixed
      * address above the main memory region to avoid conflicts.
-     *
+     */
+#if defined(__x86_64__)
+    /* x86_64 memory layout:
      * Main memory: 0x7d0000000000 - 0x7d0100000000 (4GB for FULL4G)
      * JIT buffer:  0x7d1000000000 + size
      *
@@ -2348,7 +2378,32 @@ struct jit_state *jit_state_init(size_t size)
      */
     void *jit_addr = (void *) 0x7d1000000000UL;
     state->buf = mmap(jit_addr, size, PROT_READ | PROT_WRITE | PROT_EXEC,
-                      MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
+                      MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED
+#if defined(__APPLE__)
+                          | MAP_JIT
+#endif
+                      ,
+                      -1, 0);
+#elif defined(__aarch64__)
+    /* ARM64 memory layout (macOS/Apple Silicon):
+     * Main memory: 0x150000000000 - 0x150100000000 (4GB for FULL4G)
+     * JIT buffer:  0x151000000000 + size
+     *
+     * Apple Silicon requires MAP_JIT for executable memory. The fixed
+     * address is chosen to avoid TSAN's shadow memory and typical process
+     * allocations. Requires ASLR disabled via: setarch $(uname -m) -R
+     */
+    void *jit_addr = (void *) 0x151000000000UL;
+    state->buf = mmap(jit_addr, size, PROT_READ | PROT_WRITE | PROT_EXEC,
+                      MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED
+#if defined(__APPLE__)
+                          | MAP_JIT
+#endif
+                      ,
+                      -1, 0);
+#else
+#error "TSAN is only supported on x86_64 and aarch64"
+#endif
     if (state->buf == MAP_FAILED) {
         free(state);
         return NULL;

src/main.c

@@ -304,7 +304,7 @@ int main(int argc, char **args)
         .args_offset_size = ARGS_OFFSET_SIZE,
         .argc = prog_argc,
         .argv = prog_args,
-        .log_level = LOG_TRACE,
+        .log_level = LOG_INFO,
         .run_flag = run_flag,
         .profile_output_file = prof_out_file,
         .cycle_per_step = CYCLE_PER_STEP,