Just‐in‐time compilation - retrotruestory/M1DEV GitHub Wiki

Here's how to implement Just-In-Time (JIT) compilation for Magic-1:

        .cseg
        .global _init_jit
        .global _compile_block
        .global _execute_jit_code

// Initialize JIT compiler
_init_jit:
        enter   8
        push    a
        push    b
        
        // Setup executable memory region
        ld.16   b,#JIT_CODE_BASE
        ld.16   c,#JIT_CODE_SIZE
        
        // Map pages as executable
        wcpte   #0x9000,(b)    ; Present + Executable
        
        pop     b
        pop     a
        leave
        ret

// Compile code block
// void* compile_block(uint8_t* bytecode, uint16_t size)
_compile_block:
        enter   8
        push    a
        push    b
        push    c
        
        ld.16   b,12(sp)       ; Load bytecode ptr
        ld.16   c,14(sp)       ; Load size
        ld.16   a,_next_code   ; Get next free code slot
        
.compile_loop:
        ld.8    b,(b)          ; Load bytecode
        call    _translate_op   ; Translate to native code
        br.eq   .done          ; Exit if translation failed
        
        sub.16  c,#1           ; Decrement counter
        br.ne   .compile_loop  ; Continue if more bytecode
        
.done:
        st.16   _next_code,a   ; Update next free slot
        pop     c
        pop     b
        pop     a
        leave
        ret

Create the JIT interface:

#ifndef JIT_H
#define JIT_H

#include <stdint.h>

// JIT memory layout
#define JIT_CODE_BASE   0x20000
#define JIT_CODE_SIZE   0x10000
#define MAX_CODE_BLOCKS 256

// Block descriptor
typedef struct {
    void*    code_ptr;     // Pointer to compiled code
    uint16_t code_size;    // Size of compiled code
    uint16_t entry_point;  // Entry point offset
} jit_block_t;

// Function declarations
extern void init_jit(void);
extern void* compile_block(uint8_t* bytecode, uint16_t size);
extern int execute_jit_code(void* code_ptr);
extern void invalidate_block(jit_block_t* block);

#endif

Implement the JIT compiler core:

#include "jit.h"
#include "vm_memory.h"

// Translation table entry
typedef struct {
    uint8_t  bytecode;     // Original bytecode
    uint16_t native_size;  // Size of native code
    uint8_t* native_code;  // Template for native code
} translation_t;

// Track compiled blocks
static jit_block_t blocks[MAX_CODE_BLOCKS];
static uint16_t num_blocks = 0;
static uint8_t* next_code = (uint8_t*)JIT_CODE_BASE;

// Simple translation table
static translation_t translations[] = {
    // ADD instruction
    { 
        0x01, // ADD bytecode
        4,    // 4 bytes native code
        (uint8_t[]){
            0x80,       // ld.16 a,(sp)
            0x81,       // add.16 a,b
            0x82,       // st.16 (sp),a
            0xC3        // ret
        }
    },
    // More translations...
};

// Translate single bytecode instruction
static int translate_op(uint8_t bytecode, uint8_t** target) {
    // Find translation
    for(int i = 0; i < sizeof(translations)/sizeof(translation_t); i++) {
        if(translations[i].bytecode == bytecode) {
            // Copy native code template
            memcpy(*target, translations[i].native_code, 
                   translations[i].native_size);
            *target += translations[i].native_size;
            return 1;
        }
    }
    return 0; // No translation found
}

// Compile bytecode block
void* compile_block(uint8_t* bytecode, uint16_t size) {
    if(num_blocks >= MAX_CODE_BLOCKS) {
        return NULL;    // No more slots
    }
    
    uint8_t* code_start = next_code;
    uint8_t* code_ptr = code_start;
    
    // Translate each instruction
    for(int i = 0; i < size; i++) {
        if(!translate_op(bytecode[i], &code_ptr)) {
            return NULL;    // Translation failed
        }
    }
    
    // Register new block
    blocks[num_blocks].code_ptr = code_start;
    blocks[num_blocks].code_size = code_ptr - code_start;
    blocks[num_blocks].entry_point = 0;
    
    next_code = code_ptr;
    num_blocks++;
    
    return code_start;
}

Add optimization passes:

#include "jit.h"

// Peephole optimization pass
static void peephole_optimize(uint8_t* code, uint16_t size) {
    for(int i = 0; i < size-1; i++) {
        // Example: combine consecutive loads/stores
        if(code[i] == 0x80 && code[i+1] == 0x82) {
            // Replace with single operation
            code[i] = 0x84;    // Combined load/store
            // Remove second instruction
            memmove(&code[i+1], &code[i+2], size-i-2);
            size--;
            i--;
        }
    }
}

// Register allocation pass
static void allocate_registers(uint8_t* code, uint16_t size) {
    // Track register usage
    uint8_t reg_map[8] = {0};
    
    for(int i = 0; i < size; i++) {
        // Analyze register usage
        uint8_t reg = code[i] & 0x07;
        if(reg_map[reg] == 0) {
            // Allocate new register
            reg_map[reg] = 1;
        }
    }
}

Example usage:

#include <stdio.h>
#include "jit.h"

int main() {
    // Initialize JIT compiler
    init_jit();
    
    // Simple bytecode program
    uint8_t program[] = {
        0x01,   // ADD
        0x02,   // MUL
        0x03,   // DIV
        0x04    // RET
    };
    
    // Compile program
    void* code = compile_block(program, sizeof(program));
    if(!code) {
        printf("Compilation failed\n");
        return 1;
    }
    
    // Execute compiled code
    int result = execute_jit_code(code);
    printf("Result: %d\n", result);
    
    return 0;
}

Update Makefile:

CC = clcc
AS = m1_as
CFLAGS = -O2
OBJECTS = jit_compiler.o jit_optimizer.o test_jit.o

test_jit: $(OBJECTS)
    $(CC) $(CFLAGS) -o $@ $(OBJECTS)

%.o: %.s
    $(AS) -o $@ $<

%.o: %.c
    $(CC) $(CFLAGS) -c $<

clean:
    rm -f test_jit *.o

Key features:

  1. Dynamic code generation
  2. Basic block translation
  3. Peephole optimization
  4. Register allocation
  5. Executable memory management

The JIT compiler provides:

This can be extended with:

  • More sophisticated optimizations
  • Instruction combining
  • Loop optimization
  • Profiling-guided compilation
⚠️ **GitHub.com Fallback** ⚠️