JIT Compilation — bytecode থেকে native, চলতে চলতে

সম্পূর্ণ, চলমান (~110 লাইন)। Linux-এ যথার্থ; macOS-এ নিচের note।

import ctypes, struct, time

# ---- ১. Stack bytecode (আমাদের IR) ------------------------------
PUSH, LOAD, STORE, ADD, SUB, JZ, JMP, RET = range(8)
HOT_THRESHOLD = 1000

def make_program(N):
    return [
        (PUSH, 0),        # 0 : acc = 0
        (STORE, 0),       # 1 : (slot 0 = acc)
        (PUSH, N),        # 2 : i = N
        (STORE, 1),       # 3 : (slot 1 = i)
        (LOAD, 1),        # 4 : ← loop head (backward target)
        (JZ, 15),         # 5 : i == 0 হলে end-এ
        (LOAD, 0),        # 6 : acc
        (LOAD, 1),        # 7 : i
        (ADD,),           # 8 : acc + i
        (STORE, 0),       # 9 :
        (LOAD, 1),        # 10:
        (PUSH, 1),        # 11:
        (SUB,),           # 12: i - 1
        (STORE, 1),       # 13:
        (JMP, 4),         # 14: → loop head (backward!)
        (LOAD, 0),        # 15: ← end
        (RET,),           # 16: return acc
    ]

# ---- ২. Interpreter (tier 0) + profiler --------------------------
def run_interpreted(prog):
    st, var, pc, backjumps = [], {}, 0, 0
    while True:
        op = prog[pc]
        if op[0] == PUSH:    st.append(op[1]); pc += 1
        elif op[0] == LOAD:  st.append(var[op[1]]); pc += 1
        elif op[0] == STORE: var[op[1]] = st.pop(); pc += 1
        elif op[0] == ADD:   b, a = st.pop(), st.pop(); st.append(a+b); pc += 1
        elif op[0] == SUB:   b, a = st.pop(), st.pop(); st.append(a-b); pc += 1
        elif op[0] == JZ:    pc = op[1] if st.pop() == 0 else pc + 1
        elif op[0] == JMP:
            if op[1] < pc:
                backjumps += 1
                if backjumps == HOT_THRESHOLD:
                    print(f"[profiler] backward jump #{backjumps} "
                          f"@ bytecode {pc} -> hot loop (head {op[1]})")
            pc = op[1]
        elif op[0] == RET:   return st.pop(), backjumps

# ---- ৩. JIT compiler (two-pass) ----------------------------------
SIZE = {PUSH: 11, LOAD: 8, STORE: 8, ADD: 6, SUB: 6, JZ: 10, JMP: 5, RET: 2}
NSLOT = 2

def compile_native(prog):
    # pass 1: প্রতিটা bytecode-index-এর native offset
    off, o = [], 0
    for ins in prog:
        off.append(o); o += SIZE[ins[0]]
    code_end = o
    buf = bytearray(code_end + 8 * NSLOT)

    def put(pos, bs):
        buf[pos:pos+len(bs)] = bs

    for pc, ins in enumerate(prog):
        op, p = ins[0], off[pc]
        if op == PUSH:                        # mov rax,imm64 ; push rax
            put(p, b'\x48\xb8' + struct.pack('q', ins[1]) + b'\x50')
        elif op == LOAD:                      # mov rax,[rip+d] ; push rax
            d = code_end + 8*ins[1] - (p + 7)
            put(p, b'\x48\x8b\x05' + struct.pack('i', d) + b'\x50')
        elif op == STORE:                     # pop rax ; mov [rip+d],rax
            d = code_end + 8*ins[1] - (p + 8)
            put(p, b'\x58' + b'\x48\x89\x05' + struct.pack('i', d))
        elif op == ADD:                       # pop rcx; pop rax; add; push
            put(p, b'\x59\x58\x48\x01\xc8\x50')
        elif op == SUB:
            put(p, b'\x59\x58\x48\x29\xc8\x50')
        elif op == JZ:                        # pop; test rax,rax; je far
            put(p, b'\x58\x48\x85\xc0\x0f\x84'
                 + struct.pack('i', off[ins[1]] - (p + 10)))
        elif op == JMP:
            put(p, b'\xe9' + struct.pack('i', off[ins[1]] - (p + 5)))
        elif op == RET:                       # pop rax; ret
            put(p, b'\x58\xc3')
    return bytes(buf)

# ---- ৪. mmap + execute -------------------------------------------
import sys
if sys.platform == 'darwin':
    PROT_RWX, MAP_PRIV_ANON = 0x07, 0x02 | 0x1000 | 0x800   # MAP_JIT সহ
else:
    PROT_RWX, MAP_PRIV_ANON = 0x07, 0x02 | 0x20             # MAP_ANONYMOUS

libc = ctypes.CDLL(None)
libc.mmap.restype = ctypes.c_void_p
libc.mmap.argtypes = [ctypes.c_void_p]*2 + [ctypes.c_int]*3 + [ctypes.c_long]
page = libc.mmap(None, 4096, PROT_RWX, MAP_PRIV_ANON, -1, 0)
assert page not in (-1, None), "mmap failed"

FTYPE = ctypes.CFUNCTYPE(ctypes.c_int64)

if __name__ == '__main__':
    N = 200_000
    prog = make_program(N)

    t0 = time.perf_counter()
    r1, backjumps = run_interpreted(prog)
    t_interp = time.perf_counter() - t0

    machine_code = compile_native(prog)
    print(f"[JIT] hot loop confirmed ({backjumps:,} backward jumps) "
          f"-> native emitted, {len(machine_code)} bytes")
    ctypes.memmove(page, machine_code, len(machine_code))

    fn = FTYPE(page)
    t0 = time.perf_counter()
    r2 = fn()
    t_jit = time.perf_counter() - t0

    assert r1 == r2, (r1, r2)
    print(f"interpreted : {t_interp*1000:8.2f} ms")
    print(f"jitted      : {t_jit*1000:8.2f} ms")
    print(f"speedup     : x{t_interp/t_jit:,.0f}   (sum={r1})")

প্রত্যাশিত output (Linux/macOS, x86-64; সময় machine-ভেদে বদলাবে):

[profiler] backward jump #1000 @ bytecode 14 -> hot loop (head 4)
[JIT] hot loop confirmed (200,000 backward jumps) -> native emitted, 150 bytes
interpreted :   915.87 ms
jitted      :     1.11 ms
speedup     : x825   (sum=20000100000)

একই গাণিতিক ফল, দুই path — কিন্তু native path শত-সহস্র গুণ কম dispatch। আর এই 108-byte-এর ভেতরে আগের তিন লেসনের সবকিছু: template-based code generation, jump-offset “relocation”, আর permission-nach।

নিজে বাড়ান:

1. MUL opcode যোগ করুন (48 0F AF C8 = imul rax,rcx — অবশিষ্ট bytes একই) — square-sum loop compile করে যাচাই করুন
2. Hot detection কে পুরো প্রোগ্রাম না বানিয়ে trace-ভিত্তিক করুন: backward jump থেকে linear trace ধরে শুধু সেটাই compile, শেষে trace-এর শুরুতে patch-in (TraceMonkey-র মূল ভাবনা)
3. W^X সঠিকভাবে মানুন: mmap-কে RW রাখুন, emit-এর পরে mprotect(page, size, PROT_READ|PROT_EXEC) ডাকুন — তারপর call