Attempt to fix JIT icache coherency on Arm64

The JIT compiler was experiencing intermittent failures on Arm64/Apple
Silicon due to missing instruction cache invalidation after patching
branch instructions. When update_branch_imm() modified branch targets
in JIT-compiled code, the CPU's icache wasn't being invalidated, causing
it to execute stale cached instructions instead of the newly patched
ones.

This manifested as non-deterministic test failures, particularly in
compute-intensive benchmarks like the pi calculation test.

The fix adds sys_icache_invalidate() after memcpy() in update_branch_imm
to ensure the icache is synchronized with the data cache after code
modification. This is critical on Arm64 cores which have separate L1
instruction and data caches.

Author: jserv

src/jit.c

@@ -593,6 +593,10 @@ static void update_branch_imm(struct jit_state *state,
     assert((imm & 3) == 0);
     uint32_t insn;
     imm >>= 2;
+#if defined(__APPLE__) && defined(__aarch64__)
+    /* Must be in write mode to read/write MAP_JIT memory on Apple ARM64 */
+    pthread_jit_write_protect_np(false);
+#endif
     memcpy(&insn, state->buf + offset, sizeof(uint32_t));
     if ((insn & 0xfe000000U) == 0x54000000U /* Conditional branch immediate. */
         || (insn & 0x7e000000U) ==
@@ -607,10 +611,8 @@ static void update_branch_imm(struct jit_state *state,
         assert(false);
         insn = BAD_OPCODE;
     }
-#if defined(__APPLE__) && defined(__aarch64__)
-    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
@@ -2160,13 +2162,31 @@ static const void *dispatch_table[] = {
 void clear_hot(block_t *block)
 {
     block->hot = false;
+    /* Note: Don't clear offset here - it causes issues with F extension disabled.
+     * The offset will be overwritten when the block is re-translated anyway.
+     */
 }
 
 static void code_cache_flush(struct jit_state *state, riscv_t *rv)
 {
     should_flush = false;
+    /* Invalidate the entire code cache area that will be reused */
+    if (state->offset > state->org_size) {
+#if defined(__APPLE__) && defined(__aarch64__)
+        pthread_jit_write_protect_np(false);
+#endif
+        /* Clear the code cache area to prevent stale code execution */
+        memset(state->buf + state->org_size, 0,
+               state->offset - state->org_size);
+        sys_icache_invalidate(state->buf + state->org_size,
+                              state->offset - state->org_size);
+#if defined(__APPLE__) && defined(__aarch64__)
+        pthread_jit_write_protect_np(true);
+#endif
+    }
     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)
@@ -2229,6 +2249,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);
@@ -2318,12 +2339,39 @@ void jit_translate(riscv_t *rv, block_t *block)
     memset(state->jumps, 0, MAX_JUMPS * sizeof(struct jump));
     state->n_jumps = 0;
     block->offset = state->offset;
+    uint32_t translation_start = state->offset;
     translate_chained_block(state, rv, block);
     if (unlikely(should_flush)) {
+        /* Reset block offset since translation was incomplete */
+        block->offset = 0;
+        block->hot = false;
         code_cache_flush(state, rv);
         goto restart;
     }
     resolve_jumps(state);
+    /* Ensure all instruction cache is synchronized after translation */
+    if (state->offset > translation_start) {
+#if defined(__APPLE__) && defined(__aarch64__)
+        /* Must be in write mode to invalidate cache on Apple ARM64 */
+        pthread_jit_write_protect_np(false);
+#endif
+#if defined(__aarch64__)
+        /* ARM64 needs data synchronization before cache invalidation */
+        __asm__ __volatile__("dsb sy" ::: "memory");
+#endif
+        sys_icache_invalidate(state->buf + translation_start,
+                              state->offset - translation_start);
+#if defined(__aarch64__)
+        /* ARM64 needs instruction synchronization after cache invalidation */
+        __asm__ __volatile__("isb" ::: "memory");
+#endif
+#if defined(__APPLE__) && defined(__aarch64__)
+        /* Re-enable execution mode */
+        pthread_jit_write_protect_np(true);
+#endif
+    }
+    /* Memory barrier to ensure all writes complete before marking hot */
+    __asm__ __volatile__("" ::: "memory");
     block->hot = true;
 }