Whole-Stage Code Generation

Whole-Stage Code Generation (WSCG) — это техника компиляции в средах выполнения языков программирования, при которой множество операторов или узлов графа выполнения (например, в системе обработки данных) объединяются в единую, скомпилированную на лету функцию, исключающую промежуточные накладные расходы на виртуальные вызовы, материализацию промежуточных данных и передачу управления между отдельными этапами. Основная цель WSCG — повышение производительности за счёт уменьшения числа абстракций времени выполнения и более эффективного использования кэша процессора.

История

Идея объединения нескольких этапов обработки в один компилируемый фрагмент восходит к ранним работам в области компиляции запросов к базам данных и систем потоковой обработки. В 1990-х годах в проектах, таких как Volcano и Gamma, применялись техники «push-based» и «pull-based» конвейеризации, однако они всё ещё требовали передачи данных между операторами через виртуальные функции.

Прорыв произошёл в 2010-х годах с развитием just-in-time (JIT) компиляции в виртуальных машинах. В 2014 году инженеры компании Databricks представили проект Tungsten, в рамках которого для Apache Spark была разработана техника Whole-Stage Code Generation. Она позволила генерировать на лету Java-байт-код, объединяющий несколько операторов запроса (например, фильтрацию, проекцию и агрегацию) в один цикл. Это привело к значительному ускорению обработки данных, особенно в сценариях с большими объёмами.

Позднее WSCG была адаптирована в других системах, включая Apache Flink (начиная с версии 1.9), а также в некоторых компиляторах машинного обучения, таких как TVM.

Принцип работы

WSCG применяется в системах, где данные проходят через цепочку преобразований (операторов). В традиционной реализации каждый оператор является отдельным объектом, который получает входные данные, обрабатывает их и передаёт следующему оператору. Это приводит к следующим накладным расходам:

• Виртуальные вызовы: каждый вызов метода next() или process() требует динамической диспетчеризации.
• Материализация промежуточных данных: каждый оператор может создавать временные объекты (например, строки или кортежи), которые затем передаются дальше.
• Промахи кэша: из-за разрозненного кода и данных кэш процессора используется неэффективно.

WSCG решает эти проблемы следующим образом:

1. Анализ графа операторов: система определяет последовательность операторов, которые могут быть объединены (например, фильтр, затем проекция, затем агрегация).
2. Генерация кода: на основе этого графа генерируется единая функция на языке программирования (например, Java, Scala, C++), которая выполняет все преобразования за один проход по данным.
3. Компиляция: сгенерированный код компилируется в машинный код с помощью JIT-компилятора (например, в JVM) или компилируется заранее (AOT).
4. Выполнение: скомпилированная функция применяется к каждому элементу данных (например, строке таблицы или записи потока) без промежуточных вызовов.

Пример

Рассмотрим запрос: SELECT name, age FROM users WHERE age > 18. В традиционной реализации:

• Оператор Scan читает строки.
• Оператор Filter проверяет условие age > 18.
• Оператор Project выбирает поля name и age.

При WSCG генерируется код, аналогичный следующему (псевдокод на Java):

``java void processRow(Row row) { if (row.getInt("age") > 18) { emit(row.getString("name"), row.getInt("age")); } } ``

Этот код выполняется для каждой строки без создания промежуточных объектов и без вызовов виртуальных методов.

Применение

Apache Spark (Tungsten)

Apache Spark — одна из самых известных систем, использующих WSCG. В рамках проекта Tungsten (начиная с Spark 2.0) WSCG применяется для оптимизации выполнения запросов SQL и DataFrame. Система генерирует Java-код, который затем компилируется с помощью JIT-компилятора HotSpot. Это позволяет Spark обрабатывать данные со скоростью, близкой к скорости нативных приложений, написанных на C++.

Apache Flink

В Apache Flink WSCG была внедрена в версии 1.9 для оптимизации выполнения потоковых и пакетных запросов. Flink генерирует код на Java, который объединяет несколько операторов (например, map, filter, keyBy) в один. Это особенно эффективно для сценариев с высокой пропускной способностью, таких как обработка событий в реальном времени.

Системы машинного обучения

В компиляторах машинного обучения, таких как Apache TVM и XLA (Accelerated Linear Algebra), WSCG применяется для объединения нескольких операций нейронной сети (например, свёртки, активации и пулинга) в один оптимизированный вычислительный граф. Это позволяет уменьшить накладные расходы на передачу данных между ядрами GPU и повысить производительность вывода.

Базы данных

Некоторые современные in-memory базы данных, такие как HyPer (ныне часть SAP HANA) и Umbra, используют WSCG для компиляции запросов в нативный машинный код. Это позволяет достичь производительности, сравнимой с ручной оптимизацией на C++.

Преимущества и недостатки

Преимущества

• Высокая производительность: WSCG может ускорить выполнение запросов в 2–10 раз по сравнению с традиционными интерпретируемыми или конвейерными подходами, особенно на больших объёмах данных.
• Снижение накладных расходов: устранение виртуальных вызовов и промежуточных объектов уменьшает нагрузку на сборщик мусора (в JVM-средах) и улучшает использование кэша.
• Улучшенная предсказуемость: сгенерированный код часто лучше поддаётся оптимизации со стороны JIT-компилятора, так как он представляет собой компактный и горячий участок.

Недостатки

• Сложность отладки: сгенерированный код трудно анализировать и отлаживать, так как он создаётся динамически и не имеет прямого соответствия исходному коду пользователя.
• Увеличение времени компиляции: генерация и компиляция кода на лету могут добавить задержку при первом выполнении запроса (холодный старт).
• Ограниченная применимость: WSCG эффективна только для операций, которые могут быть объединены в один линейный проход (например, фильтрация, проекция, простые агрегации). Для сложных операций, таких как соединения с несколькими источниками, требуется более сложная оптимизация.
• Зависимость от JIT-компилятора: производительность WSCG сильно зависит от возможностей JIT-компилятора (например, HotSpot C2). В средах без JIT (например, в AOT-компиляции) выгода может быть меньше.

Сравнение с альтернативами

Интересные факты

• В Apache Spark WSCG может генерировать код длиной в несколько тысяч строк для одного запроса, но JIT-компилятор HotSpot способен эффективно его оптимизировать.
• В некоторых тестах производительности Spark с WSCG обгоняет нативные реализации на C++ (например, в задачах агрегации) благодаря более эффективному использованию кэша.
• WSCG используется не только в базах данных, но и в компиляторах для квантовых вычислений, где объединение операций уменьшает количество ошибок.

Источники

• Armbrust, M., et al. "Spark SQL: Relational Data Processing in Spark." Proceedings of the 2015 ACM SIGMOD International Conference on Management of Data, 2015.
• Neumann, T. "Efficiently Compiling Efficient Query Plans for Modern Hardware." Proceedings of the VLDB Endowment, 2011.
• Apache Spark Documentation: "Tungsten Execution Engine."
• Apache Flink Documentation: "Code Generation."
• Chen, T., et al. "TVM: An Automated End-to-End Optimizing Compiler for Deep Learning." Proceedings of the 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2018.