Применение JIT-компилятора в Python

В пререлизной версии Python 3.13 добавлен экспериментальный JIT-компилятор. То есть, он может войти в состав полноценной версии осенью. Говорится о том, что он делает код намного быстрее в работе.

Однако, если я правильно понял, то реализовать JIT-компилятор можно по разному, т.к. существует несколько неофициальных компиляторов для Python: pypy / numba. Следовательно, ускорять работу приложения они могут по разному или в разных местах.

Возникает вопрос - что именно оптимизируют JIT-компиляторы в Python и, соответственно, в каких типах задач их можно использовать?


Ответы (1 шт):

Автор решения: mrgervant

Ниже представлена найденная информация о JIT в Python

Но другие возможные ответы также востребованы:

  • с объяснением иных механизмов оптимизации
  • с опытом реального использования какого-либо из JIT-компиляторов

В ближайшее время значительного ускорения работы Python-приложений от официального JIT ждать не стоит (как указал @insolor). Это следует из PEP-744 - "JIT-компиляция", в котором говорится пока что только о 5% ускорения для одной из платформ.

Однако, в дальнейшем прирост производительности будет больше, когда разработчики добавят больше механизмов оптимизации.

Некоторые варианты ускорения работы JIT-компиляторами

Замена медленных инструкций

Хотя такой способ оптимизации относят к JIT-компиляторам, в Python оптимизацию внедрили в интерпретатор в версии 3.11 (PEP-659 - "Специализированный адаптивный интерпретатор") - с его помощью оптимизируется байткод.

В языках с динамической типизацией (как Python) функция заранее не знает, какие данные могут в неё прийти:

def add(x, y):
    return x + y

Такая функция спокойно обработает как числа, так и строки - главное, чтобы типы совпадали:

add(1, 2)      # 3 - число
add('1', '2')  # 12 - строка

По умолчанию будет использоваться опкод BINARY_OP - с ним будет тратиться время на выбор подходящей команды по переданному аргументу ("+"):

import dis

def add(x, y):
    return x + y

dis.dis(add, adaptive=True)

# RESUME                   0
# LOAD_FAST__LOAD_FAST     0 (x)
# LOAD_FAST                1 (y)
# BINARY_OP                0 (+)
# RETURN_VALUE

Если же мы используем цикл, в котором вызываем функцию add() с int-ами, то на основе информации о предыдущих вызовах мы можем заменить медленную общую инструкцию на более быструю специализированную BINARY_OP_ADD_INT:

for i in range(100):
    add(i, 1)

dis.dis(add, adaptive=True)

# RESUME                   0
# LOAD_FAST__LOAD_FAST     0 (x)
# LOAD_FAST                1 (y)
# BINARY_OP_ADD_INT        0 (+)
# RETURN_VALUE

Если всё же в какой-то вызов были переданы не числа, то происходит деоптимизация - возврат к базовому выполнению инструкций.

Оптимизация нескольких инструкций, работающих вместе

Продолжая идею оптимизации одной команды, мы можем отследить использование сразу нескольких байткодов и создать более быстрый путь для их выполнения. Для этого сначала производится трассировка (поиск часто используемого участка кода для оптимизации), после которой можно убрать из байткодов некоторые микрооперации (micro-ops) и использовать только оставшиеся из них, вместо исходных полных байткодов.

Рассмотрим тот же пример с циклом в диапазоне range(100), в котором складываются 2 целых числа.

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

Далее, в коде команды цикла FOR_ITER_RANGE есть _ITER_CHECK_RANGE. Она гарантирует, что объект, по которому выполняется итерация, по-прежнему является диапазоном. Но в данном конкретном случае мы знаем, что это всегда будет диапазон, так что защитная микрооперация также может быть удалена.

Аналогично, микрооперации для BINARY_OP_ADD_INT начинаются с _GUARD_BOTH_INT - микрооперация, которая защищает от изменения типов операндов и выбирает обычный, медленный путь, если это произойдет. Однако типы не меняются, поэтому защиту можно снять.

В результате можно использовать не байткоды, а сокращенные микрооперации (по примеру с циклом сложения они сокращаются с 19 до 10 штук - практически в 2 раза).

Однако, само по себе выполнение такой оптимизации только замедлит выполнение программы, поскольку накладные расходы на создание и использование 10 микроопераций выше, чем на 7 полных байткодов. Поэтому её необходимо использовать совместно с JIT-компилятором, который сможет использовать при компиляции оптимизированные пути выполнения.

Copy-and-Patch компиляция

Данный вариант JIT-компиляции планируется использовать в официальном Python, начиная 3.13 (как и метод оптимизации нескольких инструкций). Здесь можно ознакомиться со статьей, описывающей общую концепцию данного варианта JIT-компиляции.

Основная идея: при помощи LLVM предварительно собирается объектный файл .o (шаблон) в формате ELF, в котором остаются "дыры" - места для вставки в них данных во время выполнения программы. JIT заменяет сгенерированные в ходе интерпретации программы инструкции байткода (оптимизированные способами, указанными выше) на представления в машинном коде, попутно подставляя необходимые для вычислений данные в оставленные "дыры".

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

Какой-либо специфики использования у Copy-and-Patch компиляции нет. Так, в указанной статье исследователи оптимизировали работу как с вебом, так и с БД.


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

Дальнейшие пути оптимизации пока что не озвучивались.

Для примера: у IBM перечислено 19 способов оптимизации в JIT-компиляторе для Java.


Источники информации по способам оптимизации и JIT-компиляции в официальном Python:


Примеры других JIT-ов для Python и их путей оптимизации

PyPy

PyPy использует для ускорения работы не просто трассировку, а метатрассировку. Приставка "мета" означает, что трассировка выполняется при исполнении интерпретатора, а не программы. Интерпретатор PyPy может отслеживать свои собственные операции и оптимизировать пути выполнения кода. Хорошо оптимизированный интерпретатор идёт по специфическому пути из машинных кодов - эти коды уже могут использоваться в окончательной компиляции.

Однако для этого требуется довольно много вызовов функции (нужно около 3000 исполнений функции, чтобы она начала работать быстрее).

Кроме того, из-за того, что PyPy сам по сути написан на Python (точнее RPython), то ему для исполнения функций, написанных на C, требуется создавать дополнительные прослойки, что в итоге замедляет работу приложения.

Использование PyPy сможет ускорить ваш код, если:

  • процессы работает хотя бы несколько секунд, а желательно длительное время (у JIT-компилятора будет достаточно времени для сбора информации)
  • выполняется чистый Python код, а не функции из C-библиотек

Статьи с информацией про PyPy:


Numba

Numba добавляет JIT в снабженный примечаниями код на Python - для этого в коде добавляются декораторы (интерпретатор менять не нужно).

Numba считывает байткод Python для оформленной функции и объединяет его с информацией о типах входных аргументов функции. Далее происходит анализ и оптимизация кода, после чего с помощью библиотеки компилятора LLVM генерируется машинный код функции, адаптированной к возможностям используемого процессора. Эта скомпилированная версия затем используется при каждом вызове функции.

Однако, указанные оптимизации ограничены математическими операциями и действиями с библиотекой NumPy. Есть отдельное расширение для SciPy. Другие библиотеки (в частности, pandas) Numba уже не понимает.

Если ваш код ориентирован на числовые вычисления (много математики), часто использует NumPy и/или имеет много циклов, то Numba будет хорошим выбором.


Статьи с информацией про Numba:


→ Ссылка