JIT Compilation

Just-in-time compilation for runtime code generation and execution.

Quick Start

import ctypes

# Generate machine code for simple function
# x86-64 machine code for: int add(int a, int b) { return a + b; }
machine_code = bytes([
    0x55,                    # push rbp
    0x48, 0x89, 0xe5,        # mov rbp, rsp
    0x89, 0x7d, 0xfc,        # mov [rbp-4], edi (a)
    0x89, 0x75, 0xf8,        # mov [rbp-8], esi (b)
    0x8b, 0x45, 0xfc,        # mov eax, [rbp-4]
    0x03, 0x45, 0xf8,        # add eax, [rbp-8]
    0x5d,                    # pop rbp
    0xc3                     # ret
])

# Allocate executable memory
size = len(machine_code)
buf = ctypes.create_string_buffer(machine_code, size)

# Make memory executable
libc = ctypes.CDLL('libc.so.6')
libc.mprotect.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
addr = ctypes.addressof(buf)
page_size = 4096
page_addr = addr & ~(page_size - 1)
libc.mprotect(page_addr, page_size, 7)  # PROT_READ|PROT_WRITE|PROT_EXEC

# Create function pointer and call
func_type = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_int)
func = func_type(addr)
result = func(5, 3)
print(f"5 + 3 = {result}")

Using PyPy's RPython JIT

from rpython.rlib.jit import JitDriver, promote

# Define JIT driver
jitdriver = JitDriver(
    greens=['pc', 'bytecode'],
    reds=['stack', 'variables']
)

def interpret(bytecode):
    """Simple bytecode interpreter with JIT."""
    pc = 0
    stack = []
    variables = {}

    while pc < len(bytecode):
        jitdriver.jit_merge_point(
            pc=pc,
            bytecode=bytecode,
            stack=stack,
            variables=variables
        )

        opcode = bytecode[pc]
        pc = promote(pc)  # Hint to JIT

        if opcode == 'LOAD_CONST':
            value = bytecode[pc + 1]
            stack.append(value)
            pc += 2

        elif opcode == 'ADD':
            b = stack.pop()
            a = stack.pop()
            stack.append(a + b)
            pc += 1

        elif opcode == 'RETURN':
            return stack.pop()

    return None

Dynamic Code Generation with LLVM

from llvmlite import ir, binding as llvm

class JITCompiler:
    """JIT compiler using LLVM."""

    def __init__(self):
        llvm.initialize()
        llvm.initialize_native_target()
        llvm.initialize_native_asmprinter()

        self.target = llvm.Target.from_default_triple()
        self.target_machine = self.target.create_target_machine()
        self.backing_mod = llvm.parse_assembly("")
        self.engine = llvm.create_mcjit_compiler(self.backing_mod, self.target_machine)

    def compile_function(self, func_ir):
        """Compile LLVM IR to machine code and return function pointer."""
        mod = llvm.parse_assembly(str(func_ir))
        mod.verify()

        self.engine.add_module(mod)
        self.engine.finalize_object()
        self.engine.run_static_constructors()

        func_ptr = self.engine.get_function_address("jit_func")
        return func_ptr

    def create_add_function(self):
        """Generate add function dynamically."""
        module = ir.Module(name="jit_module")
        i32 = ir.IntType(32)

        fnty = ir.FunctionType(i32, [i32, i32])
        func = ir.Function(module, fnty, name="jit_func")

        block = func.append_basic_block(name="entry")
        builder = ir.IRBuilder(block)

        result = builder.add(func.args[0], func.args[1])
        builder.ret(result)

        return module

# Usage
jit = JITCompiler()
module = jit.create_add_function()
func_ptr = jit.compile_function(module)

# Call compiled function
import ctypes
cfunc = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_int)(func_ptr)
result = cfunc(10, 20)
print(f"Result: {result}")

Numba JIT Decorator

from numba import jit, njit
import numpy as np

@jit(nopython=True)
def fibonacci(n):
    """JIT-compiled fibonacci function."""
    if n < 2:
        return n
    return fibonacci(n - 1) + fibonacci(n - 2)

@njit
def matrix_multiply(A, B):
    """JIT-compiled matrix multiplication."""
    m, n = A.shape
    n, p = B.shape
    C = np.zeros((m, p))

    for i in range(m):
        for j in range(p):
            for k in range(n):
                C[i, j] += A[i, k] * B[k, j]

    return C

# Functions are compiled on first call
result = fibonacci(10)

Bytecode Generation and Execution

import dis
import types

def generate_bytecode_function(operations):
    """Generate Python bytecode dynamically."""
    # Create code object
    code_string = compile(operations, '<generated>', 'exec')

    # Extract function from code
    for const in code_string.co_consts:
        if isinstance(const, types.CodeType):
            func = types.FunctionType(const, globals())
            return func

    return None

# Example: generate add function
operations = """
def add(a, b):
    return a + b
"""

func = generate_bytecode_function(operations)
print(func(5, 3))

Tracing JIT

class TracingJIT:
    """Simple tracing JIT implementation."""

    def __init__(self):
        self.traces = {}
        self.trace_threshold = 10
        self.counters = {}

    def should_trace(self, func_id):
        """Check if function should be traced."""
        count = self.counters.get(func_id, 0)
        count += 1
        self.counters[func_id] = count
        return count >= self.trace_threshold

    def record_trace(self, func_id, operations):
        """Record execution trace."""
        if func_id not in self.traces:
            self.traces[func_id] = []
        self.traces[func_id].append(operations)

    def optimize_trace(self, trace):
        """Optimize recorded trace."""
        optimized = []

        # Example: constant folding
        i = 0
        while i < len(trace):
            op = trace[i]

            if (i + 2 < len(trace) and
                op[0] == 'LOAD_CONST' and
                trace[i+1][0] == 'LOAD_CONST' and
                trace[i+2][0] == 'ADD'):

                # Fold constants
                value = op[1] + trace[i+1][1]
                optimized.append(('LOAD_CONST', value))
                i += 3
            else:
                optimized.append(op)
                i += 1

        return optimized

    def compile_trace(self, trace):
        """Compile optimized trace to native code."""
        # Generate machine code from trace
        # (simplified example)
        pass

Method JIT with Specialization

class MethodJIT:
    """JIT with method specialization based on types."""

    def __init__(self):
        self.specialized_versions = {}

    def specialize(self, func, arg_types):
        """Create specialized version for specific argument types."""
        key = (func.__name__, arg_types)

        if key not in self.specialized_versions:
            # Generate specialized code
            specialized = self.generate_specialized(func, arg_types)
            self.specialized_versions[key] = specialized

        return self.specialized_versions[key]

    def generate_specialized(self, func, arg_types):
        """Generate type-specialized version of function."""
        # Example: if all args are int, use integer operations
        if all(t == int for t in arg_types):
            # Generate int-specific version
            def specialized(*args):
                # Fast integer path
                return func(*args)
            return specialized
        else:
            return func

# Usage
jit = MethodJIT()

def add(a, b):
    return a + b

int_add = jit.specialize(add, (int, int))
result = int_add(5, 3)

Inline Caching

class InlineCache:
    """Simple inline cache for method lookups."""

    def __init__(self):
        self.cache = {}

    def lookup(self, obj, method_name):
        """Cached method lookup."""
        obj_type = type(obj)
        cache_key = (obj_type, method_name)

        if cache_key in self.cache:
            return self.cache[cache_key]

        method = getattr(obj, method_name)
        self.cache[cache_key] = method
        return method

# Usage
cache = InlineCache()

class Point:
    def distance(self):
        return (self.x ** 2 + self.y ** 2) ** 0.5

p = Point()
distance_method = cache.lookup(p, 'distance')