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!