Tail-calling: разбираемся в новом интерпретаторе в CPython

В последнее время в моём инфополе появилось много шума вокруг нового типа интерпретатора в Python: tail-calling. Я посмотрел PR на Github, из которого понял, что [[clang::musttail]] должен ускорить рантайм на 5%. 

Ещё я почитал Соболева, но понял только то, что эта инструкция генерирует вызов метода в asm-коде как jmp, а не call, то есть экономит один стэк-фрейм — посмотреть можно тут. Но почему эти инструкции в данном случае эквивалентны и сработают в CPython — непонятно. Так что давайте разбираться вместе!

Интерпретаторы в CPython: switch/case и computed-goto

Сейчас в CPython есть два интерпретатора. Первый — старый-добрый switch/case, который выглядит так:

while (true) {   instruction i = get_next_instruction();   switch (i) {     case POP:       // Code for a POP instruction       ...       break;     case ADD:       // Code for an ADD instruction       ...       break;     ...   } }

Он примитивен, но работает хорошо: берёт следующую инструкцию, смотрит её опкод и переходит в нужный блок case. Код в CPython — так было раньше, но сейчас его перенесли в generated_cases.c.h. Просто, надёжно, стабильно… но медленно! Да, есть техники оптимизации switch/case: например, таблицы переходов, бинарный поиск. Но всё равно приходится жонглировать аргументами и делать вызовы функций.

Второй тип интерпретатора — относительно новый computed-goto. Насколько я понял, он не полагается на оптимизацию switch/case компилятором, а сам преобразовывает большой switch в таблицу. Ключом в нём будет байт-код инструкции, а значением — метка, куда надо прыгнуть:

static void* dispatch_table[] = { &&POP_TOP, &&POP_JUMP_IF_TRUE, ... };  #define DISPATCH() goto *dispatch_table[get_next_instruction()]  DISPATCH();  POP_TOP:   // Code for POP_TOP   ...   DISPATCH(); POP_JUMP_IF_TRUE:   // Code for POP_JUMP_IF_TRUE   ...   DISPATCH();

В реальном коде это выглядит чуть сложнее — метка обёрнута в макрос TARGET:

TARGET(POP_JUMP_IF_TRUE) {   ...   DISPATCH(); } TARGET(POP_TOP) {   ...   DISPATCH(); }

Обратите внимание на макрос DISPATCH — он пригодится, чтобы легко подменять тип интерпретатора на этапе компиляции. 

В общем, оно работает, но сама функция диспетчеризации получается настолько огромной, что компилятор отказывается её далее оптимизировать. Зато лучше работает CPU branch predictor. Раньше он пытался выполнить сравнение опкода и предсказать переход, но это было зачастую бессмысленно. Сейчас же сравнения не происходит: сразу берётся адрес кода, который отвечает за обработку той или иной инструкции. В итоге производительность именно этого блока кода, который отвечает за выбор обработчика следующей инструкции (а это основной цикл!),  выросла на 20%.

Время радоваться? Ну… нет.

Текущие проблемы и идея решения

Как упоминал выше, функция диспетчеризации получается размером примерно с файл generated_cases.c.h — 5700 строк. Уверен, что компилятор даже не попытается её оптимизировать. В итоге на реальных замерах получили прирост даже не 10%, как ожидалось, а так, на уровне погрешности.

Надо срочно чинить этот момент, чтобы компилятор мог снова оптимизировать код внутри обработчиков. Самое простое — опять обернуть все TARGET в функции, но это приведёт к потере предыдущей computed-goto оптимизации.

Тут на помощь приходит оптимизация хвостового вызова. Если последнее, что делает первая функция — это вызов другой функции, то она может просто передать управление во вторую без использования инструкции вызова (call), а заодно и повторно использовать кадр стека. Это будет безопасно, потому что регистры тоже под нашим контролем. Так что мы сможем нагенерировать по функции на каждый опкод (они получатся маленькими) и применить хвостовую оптимизацию. В этом случае компилятор сможет оптимизировать каждую из них и мы будем скакать по коду через jmp, а не через call.

Реализация в CPython

Самый важный кусок кода, который реализует эту идею, находится в ceval_macros.h. Приведу небольшой фрагмент:

#ifdef Py_TAIL_CALL_INTERP     // Note: [[clang::musttail]] works for GCC 15, but not __attribute__((musttail)) at the moment. #   define Py_MUSTTAIL [[clang::musttail]] #   define Py_PRESERVE_NONE_CC __attribute__((preserve_none))     Py_PRESERVE_NONE_CC typedef PyObject* (*py_tail_call_funcptr)(TAIL_CALL_PARAMS);  #   define TARGET(op) Py_PRESERVE_NONE_CC PyObject *_TAIL_CALL_##op(TAIL_CALL_PARAMS) #   define DISPATCH_GOTO() \         do { \             Py_MUSTTAIL return (INSTRUCTION_TABLE[opcode])(TAIL_CALL_ARGS); \         } while (0) #   define JUMP_TO_LABEL(name) \         do { \             Py_MUSTTAIL return (_TAIL_CALL_##name)(TAIL_CALL_ARGS); \         } while (0) #   define JUMP_TO_PREDICTED(name) \         do { \             Py_MUSTTAIL return (_TAIL_CALL_##name)(frame, stack_pointer, tstate, this_instr, oparg); \         } while (0) #    define LABEL(name) TARGET(name) #elif USE_COMPUTED_GOTOS #  define TARGET(op) TARGET_##op: #  define DISPATCH_GOTO() goto *opcode_targets[opcode] #  define JUMP_TO_LABEL(name) goto name; #  define JUMP_TO_PREDICTED(name) goto PREDICTED_##name; #  define LABEL(name) name: #else #  define TARGET(op) case op: TARGET_##op: #  define DISPATCH_GOTO() goto dispatch_opcode #  define JUMP_TO_LABEL(name) goto name; #  define JUMP_TO_PREDICTED(name) goto PREDICTED_##name; #  define LABEL(name) name: #endif

Тут происходит определение макросов в зависимости от указанного на этапе компиляции интерпретатора. По умолчанию TARGET разворачивается в case, при computed-goto — в метку, а в случае нового tail-calling — в страшную конструкцию, к которой мы ещё вернёмся:

__attribute__((preserve_none)) PyObject *_TAIL_CALL_##op(TAIL_CALL_PARAMS)

За саму оптимизацию отвечает макрос Py_MUSTTAIL, который содержит директиву [[clang::musttail]]. Именно она подсказывает компилятору, что вместо call тут вполне уместно подставить jmp и переиспользовать код. 

Вернёмся к godbolt, который предоставил Никита Соболев в свой заметке. Вот так описаны функции:

int foo(int a);  int bar(int a) {   a += 2;   return foo(a); }  int bar_tail(int a) {   a += 2;   [[clang::musttail]] return foo(a); }

А вот во что они компилируются:

bar(int):         push    rbp         mov     rbp, rsp     sub     rsp, 16         mov     dword ptr [rbp - 4], edi         mov     eax, dword ptr [rbp - 4]         add     eax, 2         mov     dword ptr [rbp - 4], eax         mov     edi, dword ptr [rbp - 4]     call    foo(int)@PLT     add     rsp, 16         pop     rbp         ret  bar_tail(int):         push    rbp         mov     rbp, rsp         mov     dword ptr [rbp - 4], edi         mov     eax, dword ptr [rbp - 4]         add     eax, 2         mov     dword ptr [rbp - 4], eax         mov     edi, dword ptr [rbp - 4]         pop     rbp     jmp     foo(int)@PLT 

Различающиеся строки я сдвинул влево для удобства чтения. В итоге мы видим два момента: call действительно заменён на jmp, и не происходит выделения 16 байт на стеке. Мне кажется, это хорошая заявка на победу.

Опять проблемы

А почему только сейчас до этого додумались? Ну, директива компилятора [[clang::musttail]] появилась в clang только в 2021 году, а в gcc её добавили только в прошлом году. Но это лишь половина проблемы…

Вторая половина — ABI, который CPython должен поддерживать. Он требует, чтобы функции перед выполнением сохраняли состояния регистров («callee-saved» registers) на стеке, а по окончанию работы они должны быть обратно восстановлены. Насколько я понял, эти операции выполняются ещё до выбора, как вызывать функцию — через jmp или call, так что это соглашение портит всю оптимизацию. Каждая из функций-обработчиков опкодов постоянно бы что-то записывала на стек, хотя толку в этом ноль.

Решение нашлось в директиве компилятора preserve_none, которая меняет соглашение о вызове функций. Вот тут-то и пригодился макрос Py_PRESERVE_NONE_CC. Теперь за сохранение и восстановление регистров ответственна вызывающая сторона, что позволяет просто скипнуть эти действия в случае вызова через [[clang::musttail]]. Но есть нюанс: эта директива не стандартизована, и может быть удалена в будущем 🙂

Фух, выдыхаем!

И насколько сложно было это сделать? После прочтения, наверно, выглядит довольно просто. Однако, не было подходящего инструментария: [[clang::musttail]] появился не так давно в сравнении с CPython, а preserve_none разработчики и вовсе использовали на свой страх и риск. Всё гениальное просто!

Но не на всех архитектурах и компиляторах эта оптимизация доступна: насколько я понимаю, пока только на AArch64 и x86_64 и только на clang-19. Хотя есть надежды на скором портировании и под gcc.

А что же в итоге получилось по скорости? Довольно внушительные 10% на типичных питонячьих тестах! Так что ждём python 3.14, так как новый интерпретатор уже попал в main. Viva la faster-cpython!