Как можно компилировать типизированный Python

Прошло уже целых 9 лет с тех пор, как состоялся документ PEP 484, в котором сообществу Python были ниспосланы типы. Многих это сильно разозлило, и в широких массах этот ход осуждался (1). С тех пор жители Интернета неоднократно заявляли, что стремятся выяснить: в самом ли деле это означает, что теперь можно компилировать Python в нативный код и таким образом его ускорять? Вопрос совершенно оправданный. Он возник у меня на самом раннем этапе моих разработок, касающихся Python-компиляторов. Итак, осуществимо ли это?

Нет. Но в каком-то роде и «да», с оговорками. Сейчас объясню. Разберём этот вопрос на примере «компиляции перед исполнением» (AOT) в коде на CPython или в смежном с ним коде. В настоящее время CPython – основная подобная реализация в коде на Python. Средства динамической (JIT) компиляции – уже другая категория, и они также будут подробнее описаны ниже. Совершенно новой информации в этом посте нет, я всего лишь постараюсь помочь вам разобраться в ворохе известных академических и отраслевых знаний.

Ключевой тезис статьи таков: типы – это очень развёрнутые подсказки, и иногда они врут.

Не то, о чём вы подумали

Многим нравится вступать в дискуссии и заявлять вещи из разряда «У нас на 100% типизированная база кода с применением Mypy/Pyright/Pytype/Pyre. Я просто буду использовать типы при компиляции – и получится генерировать более качественный код». Кстати, исходно я так и хотел назвать эту статью: «Я просто буду пользоваться типами в компиляторе». Но, на самом деле, так не получится (2). Ниже я приведу пару примеров, демонстрирующих, почему. Моя цель – не загасить вас градом примеров, я просто хочу добиться, чтобы люди отучились говорить «просто».

Например, давайте рассмотрим эту маленькую и миленькую функцию на Python. Она краткая, типизированная и, как видите, она просто складывает целые числа. Тут всё ясно: после компиляции от неё останется две машиненые инструкции. Верно?

def add(x: int, y: int) -> int:
    return x + y

add(3, 4)  # => 7

К сожалению, нет. Аннотации типов из разряда x: T означают: «x – это экземпляр T или экземпляр подкласса T» (3). Такие случаи совершенно обычны в языках программирования и, спасибо Барбаре Лисков, они, как правило, несут осмысленную семантику. Но в данном случае это не способствует производительности. Известно, насколько непроста диспетчеризация двоичных операторов в Python – всё дело в субклассах. Это значит, что придётся выполнять очень много кода, если на месте x и y могут быть любые субклассы int.

class C(int):
    def __add__(self, other):
        send_a_nasty_email()  # эх
        return 42

Если это кажется знакомым, то, возможно, вы читали статью Бен-Эшера и Ротема, статью Фолкона или статью Бэрани. Во всех этих статьях обсуждается, как практически все аспекты языка Python сводятся к вызову функции, и почему из-за этого Python плохо поддаётся статическому анализу. Этот пост построен немного иначе, поскольку теперь у нас есть типы, а, как принято считать, типы упрощают анализ.

Если не считать выполнения произвольного кода, то также не очевидно, как именно при диспетчеризации операторов выбирается подходящая функция __add__. Вам не обязательно понимать или даже читать простыню текста, размещённую ниже и объясняющую, почему так. Можете просто сказать «ах», «ох» и «почему так много if-операторов».

_MISSING = object()

def sub(lhs: Any, rhs: Any, /) -> Any:
        # lhs.__sub__
        lhs_type = type(lhs)
        try:
            lhs_method = debuiltins._mro_getattr(lhs_type, "__sub__")
        except AttributeError:
            lhs_method = _MISSING

        # lhs.__rsub__ (чтобы знать, должен ли rhs.__rsub__ вызываться первым)
        try:
            lhs_rmethod = debuiltins._mro_getattr(lhs_type, "__rsub__")
        except AttributeError:
            lhs_rmethod = _MISSING

        # rhs.__rsub__
        rhs_type = type(rhs)
        try:
            rhs_method = debuiltins._mro_getattr(rhs_type, "__rsub__")
        except AttributeError:
            rhs_method = _MISSING

        call_lhs = lhs, lhs_method, rhs
        call_rhs = rhs, rhs_method, lhs

        if (
            rhs_type is not _MISSING  # Нас это волнует?
            and rhs_type is not lhs_type  # Может ли RHS быть субклассом?
            and issubclass(rhs_type, lhs_type)  # RHS – это субкласс!
            and lhs_rmethod is not rhs_method  # а отличается ли __r*__ в самом деле?
        ):
            calls = call_rhs, call_lhs
        elif lhs_type is not rhs_type:
            calls = call_lhs, call_rhs
        else:
            calls = (call_lhs,)

        for first_obj, meth, second_obj in calls:
            if meth is _MISSING:
                continue
            value = meth(first_obj, second_obj)
            if value is not NotImplemented:
                return value
        else:
            raise TypeError(
                f"unsupported operand type(s) for -: {lhs_type!r} and {rhs_type!r}"
            )

Этот огромный листинг взят из поста Бретта Кэннона, ссылка на который дана выше. Здесь на «псевдокоде» Python показано, что происходит под капотом на C, когда на Python выполняется lhs — rhs.

Но давайте проигнорируем эту проблему и предположим, что наследовать int невозможно. Проблема решена, верно? Остаётся посчитать производительность?

К сожалению, нет. Притом, что мы можем обеспечить быструю диспетчеризацию бинарных операторов и других методов, целые числа в Python – это большие объекты, выделяемые в куче. Таким образом, каждая операция над ними – это вызов функции к PyLong_Add или т.п. Конечно, эти функции годами оптимизировались и со временем стали гораздо быстрее, работа с ними пока всё равно идёт медленнее, чем с машинными целыми числами. Но давайте предположим, что достаточно умный компилятор способен автоматически распаковывать «сравнительно мелкие» большие числа, превращая их в машинные слова в самом начале функции. Тогда с ними можно делать быстрые математические операции. Если очень захотим, то можем реализовать даже быструю математику над числами с плавающей точкой. Надеюсь, теперь проблема решена?

# Можно хорошо повеселиться, делая очень большие числа! Это часть языка Python
def big_pow(x: int) -> int:
    # Даже если распаковать здесь `x`, результат всё равно может получиться огромным.
    return 2**x

К сожалению, нет. Большинство математических ядер не просто используют встроенные функции и операторы. Они вызывают внешние библиотеки C, например, NumPy или SciPy, а функции из этих библиотек ожидают получить в работу объекты PyLongObject, выделенные в куче. Этот C-API просто не готов открывать доступ к низкоуровневым функциям, оперирующим машинными целыми числами, а ещё не готов к работе с мечеными указателями. Такое изменение в API и ABI можно было бы считать тектоническим. Но ладно, давайте в рамках этого поста предположим, что команда разработчиков компилятора владеет волшебной палочкой, и ей под силу всё это сделать.

// Будем хранить 63-битные целые числа внутри самого указателя.
static const long kSmallIntTagBits = 1;
static const long kBits = 64 - kSmallIntTagBits;
static const long kMaxValue = (long{1} << (kBits - 1)) - 1;

PyObject* PyLong_FromLong(long x) {
    if (x < kMaxValue) {
        return (PyObject*)((unsigned long)value << kSmallIntTagBits);
    }
    return MakeABigInt(x);
}

PyObject* PyLong_Add(PyObject* left, PyObject* right) {
    if (IsSmallInt(left) && IsSmallInt(right)) {
        long result = Untag(left) + Untag(right);
        return PyLong_FromLong(result);
    }
    return BigIntAdd(left, right);
}

Значит, на этом можно остановиться, правильно?

К сожалению, осталось ещё несколько концов, которые необходимо зачистить. Возможно, числовое ядро кода Python у вас красивое и аккуратно типизированное, но ему придётся взаимодействовать с внешним миром. А внешний мир зачастую не так хорошо типизирован. Спасибо Джереми Сику и Валиду Тахе за то, что дали нам постепенную типизацию — только благодаря ей в Python вообще что-то типизируется – но вы не сможете компилировать на основе типов нетипизированный код и при этом рассчитывать, что он получится быстрым.

Это значит, что на входе в ваши типизированные функции вам придётся проверять типы тех объектов, что в них поступают. Может быть, вы сможете спроектировать систему так, что у функции будет несколько входных точек – одна для нетипизированных вызовов, одна для типизированных и, возможно, даже ещё одна для неупакованных вызовов. Но сложность такой гипотетической системы будет расти, причём, быстро. А есть ещё целая куча других осложнений и элементов динамического поведения, о которых я даже не упоминаю.

def typed_function_unboxed(x: int64, y: int64) -> int64:
    return int64_add(x, y)

def typed_function(x: int, y: int) -> int:
    return typed_function_unboxed(unbox(x), unbox(y))

def typed_function_shell(x, y):
    if not isinstance(x, int):
        raise TypeError("...")
    if not isinstance(y, int):
        raise TypeError("...")
    return typed_function(cast(x, int), cast(y, int))

f = make_typed_function(typed_function_shell, typed_function, typed_function_unboxed)
f(3, 4)  # Здесь диспетчеризация теряет управляемость 

Кстати, дело здесь далеко не ограничивается типами! Например, при связывании переменных, в особенности, глобальных, страдает производительность. Глобальные значения, даже встроенные, почти всегда могут быть затёрты любым случайным кодом Python.

«Постой, Макс», — скажете вы, — «но ведь библиотеки компилятора Python, например, Numba, явно работают нормально. В динамических компиляторах такая практика отработана годами. Так в чём же дело?»

Динамические компиляторы могут достигать такой эффективности, поскольку могут спекулятивно предполагать некоторые вещи, не являющиеся константами во время компиляции. Тогда, если допущение перестаёт выполняться, всегда в качестве запасного варианта можно прибегнуть к (более медленному) интерпретатору. Именно так PyPy, наследница Psyco, показывает чудеса техники в деле специализации структур данных; динамическому компилятору не приходится обрабатывать каждый отдельный случай. Такая деоптимизация с применением интерпретатора – одна из причин, по которым сложно понять JIT, и эта фича не слишком помогает заранее компилировать код.

Вот как может выглядеть трассировка PyPy; этот пример взят из статьи 2011 года Runtime Feedback in a Meta-Tracing JIT for Efficient Dynamic Languages. Каждый guard потенциально может послужить выходом из динамически скомпилированного кода в интерпретатор:

# inst1.getattr("a")
map1 = inst1.map
guard(map1 == 0xb74af4a8)
index1 = Map.getindex(map1, "a")
guard(index1 != -1)
storage1 = inst1.storage
result1 = storage1[index1]

Кстати: то, что компилирует Numba – это не совсем Python…

Диалекты

Numba отличная! Слов не подобрать, насколько я стараюсь не придираться к Numba в этой статье. Команда ударно поработала над Numba и как как следует её спроектировала – получился компилятор, производящий очень быстрый числовой код.

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

Это также верно для многих других компиляторов, спроектированных на основе работы с типами – таких, которые, якобы, оптимизируют «код на Python»:

• Cython/Pyrex
• MicroPython Viper
• Mojo
• Mypyc
• Shed Skin
• Starkiller (PDF)
• Static Python
• Typed Python

И в особенности для тех, что оптимизируют числа:

• Codon
• Numba
• lpython
• Pyccel
• Pythran
• Taichi
• TensorFlow JIT/XLA
• TorchScript

Возможно, какие-то подобные инструменты я просто забыл или не нашёл на них ссылки.
Правда, здесь возникает интересный вопрос: а что, если намеренно и явно попытаться обойтись без наиболее динамических возможностей? Можно ли от этого выиграть в производительности? Оказывается, что да. Ещё как.

Усиленные гарантии

Годы идут, а люди снова и снова для себя открывают, насколько туго Python поддаётся статической оптимизации. Функции наподобие следующей, которые «только лишь» загружают атрибуты совершенно не поддаются оптимизации вне контекста:

def get_an_attr(o):
    return o.value

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

Следовательно, загрузка атрибутов эквивалентна выполнению непрозрачных включений пользовательского кода.

Даже если функция типизирована, описанная выше проблема с созданием субклассов не исчезает. Сохраняются проблемы с атрибутами экземпляров, удалением атрибутов, затенением дескрипторов и так далее.

class C:
    def __init__(self):
        self.value = 5


def get_an_attr(o: C):
    return o.value

Правда, если целенаправленно отказаться от многих из этих динамических возможностей, всё начинает становиться интереснее. Например, компилятор Static Python входит в проект Cinder и позволяет пожертвовать динамизмом ради скорости. Если модуль помечен как статический при помощи import __static__, то Cinder будет компилировать байт-код не стандартным компилятором, а Static Python, который компилирует не Python, а другой язык!

Приведу пример. Компилятор SP автоматически разобьёт на слоты (расслотит) все классы в модуле SP и запретит слот __dict__. Таким образом, перестают работать те фичи, что раньше работали — например, динамическое создание и удаление атрибутов и т.д. Но, с другой стороны, теперь операции загрузки атрибутов из статических классов укладываются всего в три машинные инструкции. Такой компромисс, преобразующий компилятор, разумен. Большинство готово на него пойти.

import __static__

class C:
    def __init__(self):
        self.value: int = 5

def get_an_attr(o: C):
    return o.value
# ...
# mov rax,QWORD PTR [rsi+0x10]  # Загрузить поле
# test rax,rax                  # Проверить, не равно ли оно null
# je 0x7f823ba17283             # Может быть, выбросить исключение
# ...

Static Python, чтобы это сделать, требуются всего лишь имеющиеся аннотации Python. При работе с ним действует несколько больше ограничений, чем с Mypy (например, невозможно игнорировать ignore ошибки в ваших типах), но синтаксис Python при этом не меняется. Но в данном случае важно знать, что это напрямую не следует из трансляции, направляемой через типы. Требуется специально выбрать более строгий режим проверки и изменить поведение для всего типизированного ядра базы кода. Лучше это делать постепенно (SP может вызывать нетипизированный Python и наоборот). В таком случае потребуется изменить представление объектов во время выполнения: от пары заголовок+словарь переходим к заголовок+массив. Именно поэтому он (сейчас) реализован как пользовательский компилятор, пользовательский интерпретатор байт-кода, причём, с поддержкой Cinder JIT. Подробнее в этом разобраться поможет статья команды Static Python, написанная в коллаборации с Brown. В ней изложены некоторые подробности постепенной типизации и объяснено, почему этот подход разумен.

В данном случае было бы упущением не упомянуть Mypyc (оптимизатор и генератор кода для Mypy). Mypyc очень похож на Static Python в том, что берёт код, в чём-то напоминающий Python с типами и генерирует типизированно-оптимизированный код. Отличается он в том, что генерирует расширения C и по умолчанию делает меченые указатели для целых чисел. В зависимости от конкретного практического случая – в частности, от истории развёртывания – может случиться так, что именно этот компилятор подойдёт вам лучше всего! Например, инструмент форматирования Black имел большой успех в работе с Mypyc.

Дополнительно к имеющимся в распоряжении нормальным типам, и Static Python, и Mypyc позволяют типизировать параметры и другие переменные как примитивные целые числа, например, int8. Поэтому у вас может получиться арифметика распаковки, которую многие при первом чтении ожидают встретить в таких сниппетах, как первый приведённый в этом посте.

# Следующий сниппет допускается в Mypy, но не в Static Python, потому что
# это было бы опасно
def foo(x: int) -> int:
    return x

foo("hello")  # тип: игнорируем

В данном примере

• CPython полностью игнорирует аннотации
• В Mypy проскальзывает несоответствие типов type: ignore
  
  • Довольно интересно, чем это оборачивается: Mypyc откладывает эту проблему на время исполнения, делая инъекцию unbox, которая позволяет чисто выбросить TypeError
• Static Python не позволит скомпилироваться этому коду
  
  • Обратите внимание, что присутствующий в этом же проекте код, не относящийся к Static Python, не будет проверяться во время компиляции и, следовательно, даст только ошибку времени выполнения TypeError на границе, если будет вызван foo с не-int

Все эти варианты поведения имеют право на существование, ведь у каждого проекта иные цели.

В других проектах это выражено ещё сильнее. Например, проект Mojo нацелен на создание гораздо более обширного и явственно иного нового языка, который является полноценным надмножеством Python (4). Выполняемый в Mojo код Python должен и далее работать именно как заявлено, но на уровне модулей может делаться выбор в пользу меньшего динамизма. Итерация за итерацией код Python превращается в нечто существенно иное. Кроме того, здесь делается ещё немало приятных вещей, но разговор о них выходит за рамки этого поста.

Рассмотрим, например, следующий фрагмент, в котором определяется структура (новая фича), которая читается как класс Python, но даёт более серьёзные гарантии по изменяемости и связыванию:

@value
struct MyPair:
    var first: Int
    var second: Int
    def __lt__(self, rhs: MyPair) -> Bool:
        return self.first < rhs.first or
              (self.first == rhs.first and
               self.second < rhs.second)

Даже создаётся впечатление, что декоратор value даёт семантику значения.

Заглянем в документацию Mojo.

Структуры в Mojo статические: они связываются во время компиляции (во время выполнения добавлять методы невозможно). Работая со структурами, можно пожертвовать гибкостью ради производительности, притом, что на практике структуры безопасны и просты в использовании.

[…]

В Mojo компоновка и содержимое «структуры» задаются заранее и не поддаются изменению, когда программа уже работает. В отличие от Python, где можно добавлять, удалять или менять атрибуты объекта на лету, Mojo не позволяет делать такое со структурами. Таким образом, нельзя удалить метод при помощи del или изменить его значение прямо в процессе выполнения программы.

Выглядит красиво. Присмотримся к нему получше, когда откроют исходный код.

Наконец, даже если вы не пытаетесь оптимизировать саму генерацию кода, правильно построенный процесс связывания всего кода в бандл приложения позволяет значительно сэкономить время при запуске. Если вам не требуется сбрасывать данные на диск, как минимум, по разу на каждый импортируемый модуль, то можно серьёзно сэкономить время и выиграть в размещении кода. Пользователь ngoldbaum на lobste.rs отмечает, что PyOxidizer позволяет связать часть исходного кода, сделав из него сегмент данных для исполняемого файла.

Это уже делается с зафиксированными встроенными модулями (до версии 3.11 был вариант только с importlib 3.11), а ранее такие попытки предпринимались с целыми приложениями (раз, два, три (выведено в продакшен здесь). Может быть, имеются и другие примеры с эскалацией разной степени успешности.

Другие подходы

Nuitka (выше не упоминался) – это компилятор, позволяющий транслировать целые программы с Python на C. Насколько могу судить, в процессе компиляции он не пользуется вашими аннотациями типов. Вместо этого он задействует для обнаружения типов собственный конвейер, использующий и втягивание функций, и другие приёмы. Поправьте меня пожалуйста, если я ошибаюсь!

В проекте Oils есть написанная на Python оболочка, ускоряемая компилятором mycpp, переделывающим Python в C++. Этот компилятор не является инструментом общего назначения. Вероятно, он может предназначаться только для сборки и эксплуатации одной базы кода. Таким образом удаётся полностью игнорировать очень динамическую семантику Python (которая в оболочке Oil не используется), поскольку у программы, рассчитанной на одного пользователя, видимое поведение всегда одинаковое.

На других языках

Если вам не нравится такая идея с цельным «типизированным ядром», есть и другие компиляторы, например, SubstrateVM от Graal, позволяющие при помощи встроенной в них магии проанализировать всю вашу программу.

SubstrateVM – это AOT-компилятор для Java, рассматривающий всю вашу базу кода как один модуль. Он выполняет некоторый интенсивный межпроцедурный статический анализ, проверяя и доказывая некоторые факты о вашем коде и тем самым повышая производительность. Также он аккуратно меняет язык. Чтобы выполнить такой анализ, он запрещает загрузку произвольных классов во время выполнения. Также он ограничивает объём рефлексии некоторым известным подмножеством фич.

Что всё это значит для среднего Python-программиста

Если вы занимаетесь научным проектом или проектом, касающимся машинного обучения, то, возможно, вам придётся написать много склеивающего кода вокруг быстрого ядра, выполняющего хардкорную математику. Если добавите типы и воспользуетесь одним из тех отличных компиляторов, что перечислены выше, то, вероятно, можете рассчитывать на выигрыш в производительности.

Но, если вы работаете над каким-нибудь другим проектом, то можно и не найти такого явного узла, влияющего на производительность. Более того, у тех объектов, с которыми вы работаете, может не быть быстрых примитивных представлений, таких как int. Вероятно, у вас есть набор структур данных и пр. В таком случае вам могут пригодиться некоторые из вышеупомянутых проектов, например, Mypyc, Static Python и Nuitka. Но вам также потребуется добавить типы, исправить какие-либо связанные с ними ошибки, а затем изменить порядок сборки и запуска вашего приложения.

К сожалению (это касается всех, кто пишет на Python) огромная часть того кода, который должен работать быстро, создаётся в виде расширений на С – так исторически сложилось. Таким образом, вы можете как угодно оптимизировать интересующую вас часть приложения на Python, но эти вставки на С будут непрозрачными (если только ваш компилятор не понимает их так, как Numba понимает NumPy). Это значит, что ради достижения максимальной производительности вам может потребоваться переписать на Python тот код, что сейчас имеется у вас на C. Или, как минимум, типизировать границу между языками.

Что это значит для других динамических языков

Эти вопросы, касающиеся прямолинейной компиляции на основе типов, а также связанные с ней проблемы, актуальны не только для Python. Команде Sorbet пришлось хорошо поработать, чтобы реализовать такие возможности для Ruby. Такие же вопросы ставятся по поводу типизированного JavaScript (вида TypeScript) и, возможно, по поводу других языков. Приходится разрабатывать похожие решения.

Также есть созданный Мануэлем Серрано AOT-компилятор (PDF), рассчитанный на обработку целых программ на JS, он называется Hopc. Он проигнорирует ваши аннотации типов и проведёт собственный анализ. Серрано отмечает:

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

[…]

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

Кажется, он прав.

Заключение

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

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

Дополненительное чтение

Уже после публикации оригинала этого поста автор нашёл статью Optimizing and Evaluating Transient Gradual Typing, в которой рассказано, как добавлять нудные проверки типов в CPython и убирать избыточные, а затем заметил, что этот источник цитируется и в статье Брауна. Вся серия статей по-настоящему интересна.

________________________________________
1. Если что, это слова Дугласа Адамса. ↩
2. Я здесь даже не стану затрагивать тонкие материи, связанные с метапрограммированием, так как не думаю, что кому-то может прийти в голову воспользоваться этими возможностями для повышения производительности. Есть множество вариантов использования аннотаций типов в Python. Мы рассмотрим только простейшие стандартные аннотации x: SomeTypeName. ↩
3. Можно сказать, что в проекте Static Python есть своего рода внутренний тип Exact, которым можно пользоваться в духе x: Exact[int], чтобы запрещать создание субклассов, но напрямую он не раскрывается. По словам Карла Майера, дело в том, что в Mypy, Pyre и других проверочных инструментах не предусмотрено простого способа проверить этот тип. В системе типов Python такая возможность также не поддерживается. ↩
4. Автор благодарит Криса Грегори за то, что он посоветовал подробнее в этом разобраться! ↩