Может ли студент второго курса написать JIT-компилятор Питона, конкурирующий по производительности с промышленным решением? С учётом того, что он это сделает за две недели за зачёт по программированию. Как оказалось, может, но с нюансами. Обучаясь в РТУ МИРЭА, на специальности "Программная инженерия" я попал на семестровый курс программирования на Питоне. Питон я знал до этого, поэтому не хотелось много с ним возиться. Благо творчество студентов поощряется, иногда даже "автоматами". Собственно, стимулируемый этим "автоматом" и тягой к написанию системных модулей я написал JIT-компилятор, который назвал MetaStruct. С кодом проекта можно ознакомиться в репозитории. Предыдущий мой опыт в написании низкоуровневых программ оказался нежизнеспособным и весьма поучительным. Но об этом сегодня речь не пойдёт. Стандартная реализация Python - CPython - достаточно медленная. В сравнении с C++ называют замедление в 20-30 раз. Но целое сообщество программистов на Питоне готовы заплатить эту цену ради удобства синтаксиса, быстроты написания, изящности и выразительности кода. На этой почве появляются разнообразные способы оптимизации выполнения программ на Питоне. Такие диалекты как Cython и RPython пытаются решить проблему "разгона" Питона за счёт статической типизации и компиляции модулей. В области JIT-компиляции промышленным решением является проект Numba, спонсируемый такими технологическими гигантами как Intel, AMD и NVIDIA. Именно с этим пакетом мне предложили и посоревноваться, написав миниатюрный JIT-компилятор программ на Питоне. В этой статье я хочу рассказать, с какими трудностями я, как программист достаточно прикладной, столкнулся при написании такой довольно низкоуровневой вещи, как миниатюрный JIT-компилятор. На схеме выше показано, какие этапы проходит функция на Питоне, становясь скомпилированным модулем на С++: Функция, которую мы хотим оптимизировать, помечается аннотацией @jit, примерно так: Аннотация, получая объект функции, с помощью inspect.getsource(func_object) получает текст функции в виде строки. С помощью функции ast.parse(func_py_text) текст функции превращается в абстрактное синтаксическое дерево (AST) языка Питон Моя программа проходится по дереву через метод visit(), наследуясь от ast.NodeVisitor, и получает на выходе текст программы на C++, который записывается в файл. Для примера выше, он будет примерно таким: Через subprocess.run() происходит вызов компилятора g++, который выдаёт динамически подключаемую библиотеку (в зависимости от платформы файлом .dll или .so) При помощи вызова ctypes.LibraryLoader(CDLL).LoadLibrary(dll_filename) Динамическая библиотека подключается к среде выполнения и даёт доступ к скомпилированному варианту исходной функции. Конечный пользователь, добавивший над функцией аннотацию @jit, пользуется совершенно другим вариантом своего кода, ничего не подозревая. Процесс достаточно трудоёмкий для функции сложения из примера, но при частых вызовах и большом количестве вычислений внутри функции время компиляции окупается. Если бы это был не Питон, а какой-нибудь предметно-ориентированный язык, то пришлось бы писать парсер и обход получившегося абстрактного дерева, и решение не было бы уже таким коротким. Но в моём случае, инфраструктура Питона и его гибкость сыграли мне на руку. Наверное, стоит от технической части переходить к части визуализации и маркетинга. Созданный алгоритм JIT-компиляции был протестирован на нескольких простых алгоритмических задачах: Сумма двух чисел. Хеш-функция для целых чисел. Вычисление экспоненты через ряд Тейлора. Числа Фибоначчи. С расчётами и графиками можно подробнее ознакомиться в Jupyter-блокноте Для оценки времени выполнения использованы функции timeit() и repeat() модуля timeit. Для отрисовки графиков - модуль matplotlib В примерах будут сравниваться четыре реализации функций: Просто функции питона. Оптимизированные аннотацией @jit. Оптимизированные аннотацией @numba.jit. Запущенные на интерпретаторе PyPy Примечание: Я добавил к рассмотрению PyPy по просьбе одного из читателей. И этот метод оптимизации Питона действительно иногда является очень эффективным, что подтверждают графики ниже. Но при его использовании есть много нюансов.  Первый, это совершенно другая среда запуска, количество поддерживаемых библиотек которой значительно меньше, встроить в большой проект с уймой сторонних модулей просто так не получится.  Второй, версии PyPy выходят позже версий самого Питона. На момент написания статьи в PyPy не было конструкции match/case, поэтому пришлось использовать более простую и длинную реализацию обхода дерева. Поэтому я призываю не сильно разочаровываться в jit-компиляторах, которые в отличие от PyPy не зависят от конкретной версии Питона и списка библиотек вашего проекта. На задаче сложения двух целых чисел никакой оптимизации не видно, даже наоборот. Накладные расходы на вызов функции из dll-файла и обработка результата занимает много времени по сравнению с самими расчётами. Numba обставила моего "питомца" в 3 раза на этом примере. Модуль PyPy отработал в 30 раз (0.003 секунды против 0.1 секунды) быстрее, чем CPython. Обычно, для чисел из небольшого диапазона в качестве хеша используют их самих. Однако на просторах Интернета я нашёл такую хеш-функцию: Автором сообщения утверждается, что значение параметра 0x45d9f3b позволяет достичь наибольшей "случайности" бит внутри числа. По крайней мере, для хеш-функций такого вида. Numba оказалась хорошо оптимизированной под битовые операции. Не совсем понятно, откуда она взялась. Оставим этот вопрос открытым, но мне кажется, спонсорство главных производителей процессоров и видеокарт не прошло даром. Мой же вариант оказался слегка быстрее простого Питона, и то не всегда. PyPy и тут обставил оптимизаторы, выполнив прогоны за 0.002 секунды, то есть в 100 раз быстрее, чем Numba. Странное большое время для маленького x, выяснилось, обосновано тем, что Numba делает какие-то отложенные шаги компиляции при первом запуске. На общей её производительности это почти никак не сказывается (на втором графике с методами оптимизации аномалия исчезла, потому что был произведён повторный прогон). Питон явно показал себя неважно, поэтому посмотрим на двух оптимизаторов и PyPy отдельно. Кому интересен матан, экспонента считается по формуле соответствующего ряда Тейлора: Алгоритм прекращается, когда разница между дельтами двух итераций становится меньше порога, либо превращается в машинный ноль. Суммирование происходит от меньших членов к большим для уменьшения потерь точности. Наконец-то моё творение начало соперничать с Numba. На больших объёмах вычислений однозначного лидера нет. PyPy уже потерял преимущество в два порядка и выполняется с такой же скоростью, как и jit-оптимизаторы. Несмотря на то, что аннотация позволяет компилировать функции по одной, в ней всё ещё можно использовать рекурсию. На рекурсии Питон вообще перестал за себя отвечать. Что там с оптимизаторами? Внезапно, реализованная в проекте компиляция начала работать в 4 раза быстрее, чем Numba и в 8 раз быстрее PyPy. Получается, что с задачами разветвлённой рекурсии мой JIT-компилятор неплохо справляется. Это одно из самых интересных мест всего исследования, которое можно было бы продолжить. На самом деле, такой подход к оптимизации не нов в мире программирования. Чем-то похожим занимался Владимир Макаров, оптимизируя Ruby до уровня языка передачи регистров RTL в своём проекте MJIT. Существуют даже оптимизации сделанные поверх решения Макарова, о которых можно почитать здесь. В частности, в исследованиях отмечается, что выбор компилятора, будь то GCC или LLVM, существенно не сказывается на производительности. В моём решении использован g++ из-за большей портируемости скомпилированного кода. Для ускорения вычислений в проекте Ruby используются также предкомплированные заголовки. Однако, для студенческой работы такой уровень оптимизации не требуется. Конечно же, всё заработало не с первого раза. Вероятно, даже не с десятого. Поэтому хотелось бы привести здесь небольшую "работу над ошибками" Питон медленный во многом из-за динамической типизации, так как довольно много времени уходит на определение типа переменной перед её использованием. Также, идеология "всё есть объект" раздувает примитивные типы данных до размера остальных объектов и классов. Чтобы ускорить вычисления, нужно использовать именно примитивы, а не объекты. Проблема в том, что из кода на Питоне не всегда очевидно, какого типа будет переменная. Продвинутые оптимизаторы умеют определять тип переменной "на лету" из контекста. Я решил не усложнять жизнь , а усложнить код, и использовать аннотации типов. Про использование аннотаций типов есть хорошие статьи в официальной документации или на Хабре. Пусть нужно скомпилировать ту самую функцию сложения: Какой-нибудь компилируемый язык, C++ например, за такую программу спасибо не скажет. Добавим аннотации: В этом примере явного объявления типов требуют три вещи: Аргументов функции. Возвращаемое из функции значение. Локальные переменные. В базовой реализации будет только три типа данных: Тип данных Python Тип данных C++ bool bool int int float double Использование более сложных типов данных выходит за задачу миниатюрности компилятора. Но стоит отметить, что строки, коллекции и объекты Питона могут быть поддержаны с использованием Python C API С++ видоизменяет названия функций согласно их сигнатуре и аргументам. Например, функцию int f(int x) компилятор может преобразовать в _Z1fi. Подробнее о соглашении именования функции при компиляции можно узнать, например, здесь После переименования к функциям уже нельзя обратиться по первоначальному названию. Конечно, можно было бы написать свой алгоритм, который делает те же преобразования, что и компилятор. Но на самом деле, существует более простое решение, к которому мне пришлось в итоге прийти. При добавлении к объявлению функции префикса extend "C" имена не будут кодироваться: Так происходит, потому что мы явно указываем, что имена функций должны кодироваться по соглашению языка C, то есть, никак. Как программисту, плохо знакомому с чем-то ниже C++, мне было трудно понять, как подключить собранную dll-библиотеку к Питону. Изначально была идея использовать rundll32.exe для запуска. Почитав Википедию и ещё одно обсуждение, я немного разочаровался в прикладной применимости и портируемости этого решения. Потом я нашёл статью с говорящим названием: What’s the guidance on when to use rundll32? Easy: Don’t use it Только после этого я был направлен на путь истинный, а точнее, на использование модуля ctypes. В этот момент мои скомпилированные функции впервые начали возвращать мне значения прямо в Питоне. Всё шло гладко, пока я оперировал целыми числами. При попытке считать дробные... возвращались целые очень странного формата. Как оказалось, по умолчанию, все функции, импортируемые из dll работают с аргументами как с int и возвращают тоже int. Тут до меня дошло, что возвращать просто строку программы на C++ из транслятора недостаточно. Из функций необходимо ещё получить сигнатуру, чтобы потом применить нужные типы из модуля ctypes. Метод посещения объявления функции начал выглядеть вот так: Появился метод встраивания этих типов в сигнатуру dll: Для того чтобы функции вызывались так же хорошо, как с целыми числами, достаточно было поменять у функций внутри загруженного dll атрибуты argtypes и restype. Сами типы для подстановки были взяты из модуля ctypes: После всех этих накручиваний костылей махинаций ответ экспоненты начал совпадать с Питоном. Подробнее про использование функций, загруженных из DLL, можно почитать в этой статье Хороший код не пишется сразу. Например, в моём проекте основные блоки пришлось переписать два раза. В крупных проектах борьба с тех. долгом вообще может уходить в бесконечность, но я пока что, к сожалению или к счастью, с этим не сталкивался. Я хотел показать на примере функции синтаксического разбора бинарной операции ast.BinOp как можно по-разному писать код, который будет в разной степени сложно поддерживать. Это самая первая реализация: По сравнению с блоком if-elif-elif-elif-...-else такой код кажется проще. Но тут происходит явное обращение к типу через type(), что не очень хорошо. Тут мне посоветовали начать уже использовать плюшки версии Питона 3.10 на полную катушку, а именно, применить сравнение по шаблону и оператор match/case Код стал чуточку короче и выразительнее. Шаблоны после слова case можно всячески усложнять, выбирая всё более специфичные случаи. Какой-нибудь switch/case в таких языках программирования как C++ или Java такого себе не может позволить. Подробнее про возможности оператора match/case можно узнать тут Вот ещё маленький пример, решающий проблему занижения регистра для констант True и False: Выглядит покрасивее, чем if node.value == True:. Но тут уже на вкус и цвет. Вот я и пришёл к тому, что давно написали за меня, но я об этом никогда не слышал, поэтому ещё не применил. Как ни странно, модуль, описывающий абстрактное дерево в виде структуры, предоставляет также и методы для его обхода. Этот метод называется ast.NodeVisitor.visit() Как следует из названия, NodeVisitor реализует шаблон проектирования Посетитель, позволяющий создавать новую внешнюю функциональность с минимальным изменением уже написанного кода. Для написания своего посетителя необходимо объявить класс-наследник класса ast.NodeVisitor: Код остальных методов посетителя можно изучить в репозитории. Код стал более модульным, названия и сигнатуры методов уже придуманы за нас, выбор нужного метода происходит без нашего участия. Вот на этом варианте я и решил остановиться. Во время написания кода на Питоне у меня возникает чувство ощущения красоты, краткости и мощи собственного кода. В этом есть какая-то магия. Дабы не потерять доверие читателя, я провёл подсчёт строк кода транслятора: Код Объём Аннотация и процесс компиляции 36 Транслятор на NodeVisitor 221 Итого 257 Остальные файлы, как оказалось, к процессу исполнения напрямую не причастны. Я даже не стал исключать того большого числа пустых строк, которого требует PEP8. Получается, что даже немного наврал читателю насчёт числа строк. Пусть он меня простит. Думаю, такой небольшой проект можно было бы легко поддерживать, если бы в этом была бы необходимость. Хочу сказать спасибо Петру Николаевичу Советову за наставничество при написании этого проекта и этой статьи. Как оказывается, писать статьи труднее, чем писать код в ящик. Проекту в плане функциональности есть куда расти. Поддержка строк, коллекций, объектов, классов. Правда с учётом полученного зачёта, предлагаю энтузиастам, взяв моё решение за основу, добиться большей производительности и функциональности. Предел у этого совершенства всё равно есть. На поддержке библиотек такие решения оптимизации обычно отказываются работать либо поддерживают самые популярные и базовые, такие как Numpy или PIL. И всё же, если очень захотеть, можно заставить Питон работать быстрее, уничтожая один из извечных аргументов программистов на C++ и Java против использования Python.  Студент направления «Программная инженерия»