В этой статье я покажу, как написать рудиментарный, нативный x86-64 just-in-time компилятор (JIT) на CPython, используя только встроенные модули.

Код предназначен для UNIX-систем, таких как macOS и Linux, но его должно быть легко транслировать на другие системы, типа Windows. Весь код опубликован на github.com/cslarsen/minijit.

Цель — сгенерировать в рантайме новые версии нижеприведённого ассемблерного кода и выполнить их.

48 b8 ed ef be ad de  movabs $0xdeadbeefed, %rax
00 00 00
48 0f af c7           imul   %rdi,%rax
c3                    retq

В основном, мы будем иметь дело с левой частью кода — байтовой последовательностью 48 b8 ed ... и так далее. Эти 15 байтов в машинном коде составляют функцию x86-64, которая умножает свой аргумент на константу 0xdeadbeefed. На этапе JIT будут созданы функции с разными такими константами. Такая надуманная форма специализации должна продемонстрировать базовую механику JIT-компиляции.

Основная стратегия заключается в том, чтобы с помощью встроенного питоновского модуля ctypes загрузить стандартную библиотеку C. Оттуда мы получим доступ к системным функциям для взаимодействия c менеджером виртуальной памяти. Используем mmap для получения блока памяти, выровненного по границе страницы. Выравнивание необходимо для исполнения кода. По этой причине мы не берём обычную функцию malloc, поскольку она может вернуть память, выходящую за границы страницы.

Используем функцию mprotect для пометки блока памяти как доступного только для чтения и исполняемого. После этого должна появиться возможность вызова нашего свежескомпилированного блока кода посредством ctypes.

Шаблонная часть

Прежде чем что-либо сделать, нужно загрузить стандартную библиотеку C.

import ctypes
import sys

if sys.platform.startswith("darwin"):
    libc = ctypes.cdll.LoadLibrary("libc.dylib")
    # ...
elif sys.platform.startswith("linux"):
    libc = ctypes.cdll.LoadLibrary("libc.so.6")
    # ...
else:
    raise RuntimeError("Unsupported platform")

Есть и другие способы сделать это, например

>>> import ctypes
>>> import ctypes.util
>>> libc = ctypes.CDLL(ctypes.util.find_library("c"))
>>> libc
<CDLL '/usr/lib/libc.dylib', handle 110d466f0 at 103725ad0>

Чтобы определить размер страницы, вызовем sysconf(_SC_PAGESIZE). Константа _SC_PAGESIZE равняется 29 на macOS, но 30 на Linux. Мы просто жёстко закодируем их в нашей программе. Вы можете найти размер страницы, изучив системные заголовочные файлы или написав простую программу на C для вывода. Более надёжное и элегантное решение — использовать модуль cffi вместо ctypes, потому что он умеет автоматически парсить заголовочные файлы. Однако, поскольку мы поставили цель использовать стандартный дистрибутив CPython, то продолжим работать с ctypes.

Нам нужны несколько дополнительных констант для mmap и проч. Они записаны ниже. Может быть, вам придётся поискать их для других вариантов UNIX.

import ctypes
import sys

if sys.platform.startswith("darwin"):
    libc = ctypes.cdll.LoadLibrary("libc.dylib")
    _SC_PAGESIZE = 29
    MAP_ANONYMOUS = 0x1000
    MAP_PRIVATE = 0x0002
    PROT_EXEC = 0x04
    PROT_NONE = 0x00
    PROT_READ = 0x01
    PROT_WRITE = 0x02
    MAP_FAILED = -1 # voidptr actually
elif sys.platform.startswith("linux"):
    libc = ctypes.cdll.LoadLibrary("libc.so.6")
    _SC_PAGESIZE = 30
    MAP_ANONYMOUS = 0x20
    MAP_PRIVATE = 0x0002
    PROT_EXEC = 0x04
    PROT_NONE = 0x00
    PROT_READ = 0x01
    PROT_WRITE = 0x02
    MAP_FAILED = -1 # voidptr actually
else:
    raise RuntimeError("Unsupported platform")

Хотя это не является строгим требованием, но очень полезно передать ctypes типы функций, которые мы собираемся использовать. Таким образом, будут подняты исключения в случае смешивания недопустимых типов. Например:

# Set up sysconf
sysconf = libc.sysconf
sysconf.argtypes = [ctypes.c_int]
sysconf.restype = ctypes.c_long

Это говорит ctypes, что функция sysconf принимает четырёхбайтовое целое число, а выдаёт длинное целое. После этого можно узнать текущий размер страницы следующей командой:

pagesize = sysconf(_SC_PAGESIZE)

Машинный код, который мы собираемся сгенерировать, будет интерпретироваться как беззнаковые 8-битные байты, поэтому зарегистрируем беззнаковый тип указателя на эти байты:

# 8-bit unsigned pointer type
c_uint8_p = ctypes.POINTER(ctypes.c_uint8)

Ниже просто выложены остальные типы функций, которые будем использовать. Для отчётов об ошибках хорошо, чтобы была доступной функция strerror. Используем munmap для уничтожения блока машинного кода, когда закончим с ним. Так операционная система сможет повторно использовать эту память.

strerror = libc.strerror
strerror.argtypes = [ctypes.c_int]
strerror.restype = ctypes.c_char_p

mmap = libc.mmap
mmap.argtypes = [ctypes.c_void_p,
                 ctypes.c_size_t,
                 ctypes.c_int,
                 ctypes.c_int,
                 ctypes.c_int,
                 # Below is actually off_t, which is 64-bit on macOS
                 ctypes.c_int64]
mmap.restype = c_uint8_p

munmap = libc.munmap
munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
munmap.restype = ctypes.c_int

mprotect = libc.mprotect
mprotect.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
mprotect.restype = ctypes.c_int

На данном этапе сложно оправдать использование Python вместо C. В случае C нам не нужен вышеперечисленный шаблонный код. Но Python даст гораздо больше свободы в дальнейших экспериментах.

Теперь мы готовы написать обёртку mmap.

def create_block(size):
    ptr = mmap(0, size, PROT_WRITE | PROT_READ,
            MAP_PRIVATE | MAP_ANONYMOUS, 0, 0)

    if ptr == MAP_FAILED:
        raise RuntimeError(strerror(ctypes.get_errno()))

    return ptr

Эта функция использует mmap для выделения нам памяти, выровненной по границам страницы. Мы помечаем участок памяти флагами PROT как доступный для чтения и записи, а также помечаем его как приватный и анонимный. Последнее означает, что другие процессы не смогут увидеть этот участок памяти и что он не имеет файловой поддержки. Руководство по mmap в Linux более подробно раскрывает эту тему (только убедитесь, что открыли руководство именно для своей системы). Если вызов mmap неудачный, мы вызываем ошибку Python.

Чтобы пометить память как исполняемую:

def make_executable(block, size):
    if mprotect(block, size, PROT_READ | PROT_EXEC) != 0:
        raise RuntimeError(strerror(ctypes.get_errno()))

Таким вызовом mprotect мы помечаем участок памяти как доступный для чтения и исполняемый. Если хотим, можем также сделать его доступным для записи, но некоторые системы откажутся исполнять код из памяти, открытой для записи. Иногда это называют функцией безопасности W^X.

Для уничтожения блока памяти делаем следующее:

def destroy_block(block, size):
    if munmap(block, size) == -1:
        raise RuntimeError(strerror(ctypes.get_errno()))

Весёлая часть

Теперь мы наконец-то готовы написать безумно простой код JIT!

Вспомните листинг ассемблера из начала статьи: это небольшая функция — без локального стекового кадра — которая умножает число на входе на константу. На Python мы бы записали её так:

def create_multiplication_function(constant):
    return lambda n: n * constant

Это действительно надуманный пример, но соответствует определению JIT. В конце концов, мы действительно создаём нативный код в рантайме и исполняем его. Легко представить более продвинутые примеры, такие как JIT-компиляция Brainfuck в машинный код x86-64. Или использование инструкций AVX для молниеносных математических операций по векторизации.

Разборка машинного кода в начале статьи в реальности была выполнена путём компиляции и дизассемблирования следующего кода C:

#include <stdint.h>

uint64_t multiply(uint64_t n)
{
  return n*0xdeadbeefedULL;
}
If you want to compile it yourself, use something like

$ gcc -Os -fPIC -shared -fomit-frame-pointer     -march=native multiply.c -olibmultiply.so

Здесь я оптимизировал по размеру (-Os), чтобы сгенерировать минимальный машинный код, позиционно-независимый (-fPIC) для предотвращения использования переходов в абсолютных адресах, без каких-либо указателей фреймов (-fomit-frame-pointer) для удаления излишнего кода установки стека (но он может понадобиться для более продвинутых функций) и с использованием нативного набора инструкций имеющегося процессора (-march=native).

Мы могли бы передать -S и получить листинг дизассемблера, но нам интересен именно машинный код, так что вместо этого используем инструмент вроде objdump:

$ objdump -d libmultiply.so
...
0000000000000f71 <_multiply>:
 f71:   48 b8 ed ef be ad de    movabs $0xdeadbeefed,%rax
 f78:   00 00 00 
 f7b:   48 0f af c7             imul   %rdi,%rax
 f7f:   c3                      retq

Если вы не очень знакомы с ассемблером, объясню, как работает эта функция. Сначала функция movabs просто помещает непосредственное число (immediate number) в регистр RAX. Непосредственное — это такой жаргон на ассемблере для обозначения чего-то, указанного прямо в машинном коде. Другими словами, это встроенный аргумент для инструкции movabs. Так что теперь регистр RAX содержит константу 0xdeadbeefed.

Также — согласно конвенции AMD64 — первый целочисленный аргумент будет в регистре RDI, а возвращаемое значение в RAX. Так что регистр RDI содержит число-множитель. В этом суть команды imul, которая перемножает RAX и RDI, помещая результат в RAX. Наконец, мы извлекаем со стека 64-битный адрес возврата и переходим к нему командой RETQ. На этом уровне легко представить, как можно реализовать программирование в стиле передачи продолжений.

Обратите внимание, что константа 0xdeadbeefed указана в формате с обратным порядком байтов (little-endian). Нужно не забыть сделать то же самое в коде.

Теперь мы готовы поместить всё в функцию Python.

def make_multiplier(block, multiplier):
    # Encoding of: movabs <multiplier>, rax
    block[0] = 0x48
    block[1] = 0xb8

    # Little-endian encoding of multiplication constant
    block[2] = (multiplier & 0x00000000000000ff) >>  0
    block[3] = (multiplier & 0x000000000000ff00) >>  8
    block[4] = (multiplier & 0x0000000000ff0000) >> 16
    block[5] = (multiplier & 0x00000000ff000000) >> 24
    block[6] = (multiplier & 0x000000ff00000000) >> 32
    block[7] = (multiplier & 0x0000ff0000000000) >> 40
    block[8] = (multiplier & 0x00ff000000000000) >> 48
    block[9] = (multiplier & 0xff00000000000000) >> 56

    # Encoding of: imul rdi, rax
    block[10] = 0x48
    block[11] = 0x0f
    block[12] = 0xaf
    block[13] = 0xc7

    # Encoding of: retq
    block[14] = 0xc3

    # Return a ctypes function with the right prototype
    function = ctypes.CFUNCTYPE(ctypes.c_uint64)
    function.restype = ctypes.c_uint64
    return function

В нижней части мы возвращаем тип функции ctypes для использования в этом коде. Это несколько произвольное размещение, но я подумал, что хорошо будет поместить его рядом с машинным кодом.

Финальная часть

Теперь у нас есть основные части, которые можно совместить. Сначала — выделить одну страницу памяти:

pagesize = sysconf(_SC_PAGESIZE)
block = create_block(pagesize)

Затем сгенерировать машинный код. В качестве множителя выберем число 101.

mul101_signature = make_multiplier(block, 101)

Теперь помечаем участок памяти как исполняемый и доступный только для чтения:

make_executable(block, pagesize)

Берём адрес первого байта в блоке памяти подаём его в вызываемую функцию ctypes с правильным типом:

address = ctypes.cast(block, ctypes.c_void_p).value
mul101 = mul101_signature(address)

Чтобы получить адрес памяти блока, используем ctypes для его передачи пустому указателю и извлечения его значения. Наконец, инициализируем реальную функцию с этого адреса с помощью конструктора mul101_signature.

Вуаля! Теперь у нас кусочек нативного кода, который можно вызвать из Python. Если вы находитесь в среде REPL, то можно сделать это напрямую:

>>> print(mul101(8))
808

Заметьте, что эта маленькая функция умножения работает медленнее, чем нативные вычисления Python. Это главным образом из-за чужеродной библиотеки ctypes, использование которой несёт много накладных расходов: каждый раз при вызове функции нужно проверить, какие динамические типы вы ей передаёте, затем распаковать их и преобразовать, а потом сделать то же самое с возвращаемым значением. Так что есть смысл использовать ассемблер или если у вас есть доступ к каким-то новым инструкциям Intel, или для компиляции чего-то вроде Brainfuck в нативный код.

В конце, если хотите, можете позволить системе повторно использовать участок памяти, в котором находится функция. Имейте в виду, что после этого, вероятно, произойдёт сбой процесса, если попытаться снова обратиться к коду. Так что наверное лучше одновременно удалить вообще все упоминания Python:

destroy_block(block, pagesize)

del block
del mul101

Если запустите код в его полном виде из репозитория GitHub, то можете указать константу для умножения прямо в командной строке:

$ python mj.py 101
Pagesize: 4096
Allocating one page of memory
JIT-compiling a native mul-function w/arg 101
Making function block executable
Testing function
OK mul(0) = 0
OK mul(1) = 101
OK mul(2) = 202
OK mul(3) = 303
OK mul(4) = 404
OK mul(5) = 505
OK mul(6) = 606
OK mul(7) = 707
OK mul(8) = 808
OK mul(9) = 909
Deallocating function

Отладка JIT-кода

Если хотите продолжить обучение с помощью этой маленькой программы, то вскоре возникнет мысль дизассемблировать сгенерированный машинный код. Как вариант, можно просто использовать gdb или lldb, но нужно знать, откуда начать. Есть такой трюк: просто вывести hex-значение адреса блока с ожиданием нажатия клавиши:

print("address: 0x%x" % address)
print("Press ENTER to continue")
raw_input()

Затем просто запускаете программу в отладчике и во время паузы дизассемблируете область памяти. Конечно, также есть возможность пошаговой отладки ассемблерного кода, если хотите посмотреть, что происходит. Пример сессии lldb:

$ lldb python
...
(lldb) run mj.py 101
...
(lldb) c
Process 19329 resuming
...
address 0x1002fd000
Press ENTER to continue

В этом месте нажмите CTRL+C для возвращения в отладчик, затем дизассемблируйте код из области памяти:

(lldb) x/3i 0x1002fd000
    0x1002fd000: 48 b8 65 00 00 00 00 00 00 00  movabsq $0x65, %rax
    0x1002fd00a: 48 0f af c7                    imulq  %rdi, %rax
    0x1002fd00e: c3                             retq

Обратите внимание, что 65 в hex — это 101 в десятичной системе, то есть тот аргумент командной строки, который мы ввели выше.

Если нужен дизассемблер только в Python, рекомендую модуль Capstone.

Что дальше?

Хорошим упражнением была бы JIT-компиляция программ Brainfuck в нативный код. Если хотите сразу этим заняться, я открыл репозиторий GitHub по адресу github.com/cslarsen/brainfuck-jit. Я даже сделал презентацию Speaker Deck. Там показана JIT-компиляция и оптимизации, но вместо подхода из этой статьи для компиляции нативного кода используется GNU Lightning. Но должно быть исключительно просто использовать примеры без GNU Lightning или сгенерировать собственный код. Интересное наблюдение по проекту Brainfuck: если вы выполните просто JIT-компиляцию всех инструкций Brainfuck одна за другой, то не получите особого увеличения производительности даже в нативном коде. Весь прирост производительности происходит на этапе оптимизации кода, где вы забиваете целочисленными операциями одну или несколько инструкций x86. Ещё один кандидат на такую компиляцию — язык Forth.

Прежде чем серьёзно взяться за расширение этого JIT-компилятора, взгляните на проект PeachPy. Это гораздо более продвинутый проект, чем у нас, он включает в себя дизассемблер и поддерживает как будто весь набор инструкций x86-64, вплоть до AVX.

Как упомянуто выше, существует много накладных расходов при использовании ctypes для вызова функций. Чтобы устранить некоторые из них, можно использовать модуль cffi, но факт остаётся фактом: если вы хотите многократно вызывать очень маленькие JIT-функции, то обычно это быстрее сделать на чистом Python.

Какие ещё есть примечательные варианты использования? Мне встречались некоторые математические библиотеки на Python, которые переходят на векторные операции для повышения производительности. Но могу представить и другие вещи. Например, инструменты для сжатия и декомпрессии нативного кода, доступа к примитивам виртуализации и так далее. Знаю, что некоторые инструменты BPF и модули regex тоже выполняют JIT-компиляцию запросов для более быстрой обработки данных.

Что интересно в этом упражнении — так это то, что мы вступаем на территорию за пределами чистого ассемблера. Например, приходит в голову мысль, как различные инструкции дизассемблируются в одинаковые символы. Так, у инструкции RETQ иной opcode, чем у обычной инструкции RET, потому что она обрабатывает 64-битные значения. Может, это и не важно при программировании на ассемблере, потому что такие детали не всегда имеют значение, но стоит помнить о разнице. Я видел, как gcc, lldb и objdump выдавали слегка разный листинг дизассемблера для инструкций RETQ и MOVABSQ в одном и том же коде.

Ещё одно замечание. Я упоминал, что сделанный мной нативный компилятор Brainfuck изначально выдаёт не очень быстрый код. Его нужно оптимизировать. То есть код не становится быстрее автоматически от того факта, что у вас есть AVX, CUDA или что-то ещё. Жестокая правда в том, что в компилятор gcc зашита большая база оптимизаций, которые практически невозможно воспроизвести вручную. Для дополнительной информации я бы рекомендовал классическую лекцию Феликса фон Лейтнера об оптимизации исходного кода.

Что насчёт действительной компиляции?

Несколько человек написали в комментариях, что ожидали более подробного описания этапа, где по-настоящему выполняется компиляция. Это справедливое замечание. Как я говорил, это и вправду очень ограниченная форма компиляции, где мы практически ничего не делаем с кодом во время выполнения — всего лишь вставляем константу. Может быть, я напишу продолжение статьи, в котором рассмотрим чисто этап компиляции. Будьте на связи!