Как мы ускорили расчёт факторов ранжирования в поиске Ozon с помощью динамической компиляции

Всем привет! Меня зовут Петя Портнов, я работаю в Ozon ведущим разработчиком в команде среднего поиска — слоя, который ранжирует поисковую выдачу.

Представьте, что вы вводите запрос в поисковую строку маркетплейса. За этим простым действием скрывается сложный поисковый пайплайн: миллионы товаров фильтруются, ранжируются и сортируются по релевантности. Но как именно система решает, что показать первым? В основе этого решения лежат вычисления, среди которых — сотни разнообразных формул, учитывающих цену, рейтинг, популярность, персонализацию и другие факторы. По мере развития системы таких формул становится всё больше, а сами они усложняются. В какой-то момент вычисления превращаются в узкое место: начинают потреблять значительную долю CPU, создают множество промежуточных объектов — и так для каждого поискового запроса. Возникает вопрос: как снизить стоимость таких вычислений в JVM?

В этой статье я расскажу, что сделали мы, чтобы снизить нагрузку на систему: как заменили интерпретирующий движок формул на динамический компилятор, выполняющий построение эффективного байт-кода, отлично векторизующегося JIT-компилятором. Это текстовая версия доклада с Joker 2025 с дополнениями, которые не вошли в выступление или появились в проекте уже после конференции.

Максимально кратко: поиск в e-commerce

Поиск в e-commerce — это не просто сравнение строк, это сложная многослойная система. Ранее мы уже рассказывали об устройстве нашего поиска.

В этой схеме базовый поиск обрабатывает миллионы товаров, используя полнотекстовый поиск и дешёвые формулы релевантности для первичного отбора. Это позволяет отобрать из огромного набора товаров несколько тысяч наиболее подходящих под запрос.

Средний поиск работает с этой, уже отфильтрованной выборкой, решая проблему точного ранжирования: какой товар должен быть выше, а какой — ниже.

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

Среди них можно выделить calculated-фичи — произвольные формулы, определяющие вычисления над набором других фич-входов. Например, релевантность товара может вычисляться как некая формула:

где  — фича цены товара,  — фича рейтинга товара, — фича срока доставки, а — специфические веса.

Движок вычисления calculated-фич получает определения фич в рантайме раз в пару часов. Поближе с архитектурой можно познакомиться в статье моего коллеги про факторы ранжирования в runtime поиска.

Нам же предстоит понять, почему calculated-фичи — это огромный объём вычислений и как мы смогли решить проблему оптимизации этих вычислений.

Проблема: дорогие вычисления и лишние аллокации

Формулы для факторов описываются на предметно-ориентированном языке (DSL). Для простоты давайте дальше рассматривать упрощённый пример формулы расчёта фактора :

где , ,  — другие факторы ранжирования, в том числе константы. На практике формулы гораздо сложнее: нетривиальные функции, векторные операции и даже условные выражения. В дополнение к этому встречаются и формулы с вложенностью, превышающей сотню.

Все эти выражения описаны в нашем DSL (записанном в Protobuf-форме):

message Expression {  oneof kind {    // Константное значение    float constant = 1;    // Значение именованной фичи    string feature = 2;    // Унарный оператор    Unary unary = 3;    // Бинарный оператор    Binary binary = 4;    // Условный оператор    Conditional conditional = 5;    // Векторный оператор    Vector vector = 6;    ...  }}

В свою очередь, выражения могут ссылаться на другие выражения:

// Бинарное выражениеmessage Binary {  // Конкретная выполняемая операция  Operator operator = 1;  // Левая часть выражения  Expression lhs = 2;  // Правая часть выражения  Expression rhs = 3;   enum Operator {    // lhs + rhs    OPERATOR_SUM = 1;    // lhs - rhs    OPERATOR_SUB = 2;    // lhs * rhs    OPERATOR_MUL = 3;    // lhs / rhs    OPERATOR_DIV = 4;    // lhs^pow    OPERATOR_POW = 5;    // min(lhs, rhs)    OPERATOR_MIN = 6;    // max(lhs, rhs)     OPERATOR_MAX = 7;    ...  }}

Для примера выше получим такое представление:

{  "binary": {    "operator": "OPERATOR_DIV",    "lhs": {      "feature": "f_foo",    },    "rhs": {      "binary": {        "operator": "OPERATOR_SUM",        "lhs": {          "feature": "f_bar"        },        "rhs": {          "feature": "f_baz"        }      }    }  }}

По сути, мы имеем некое подобие AST — абстрактного синтаксического дерева. Ввиду его рекурсивности напрашивается довольно простая реализация вычисления: задача узла Expression в формуле — вычислить на данных входах значение, получающееся в его результате. Отсюда имеем интерфейс:

public sealed interface Calculation {    /// Вычисляет значение выражения    ///    /// @param length длина результирующего массива    /// @param inputs входные фичи (каждая длины `length`)    /// @return результат вычисления выражения    float[] apply(int length, ReadOnlyArray.AsFloats<?>... inputs);    ...}

Реализации интерфейса выполняют вычисления соответствующих им операторов:

record BinaryOperator(Calculation lhs, Calculation rhs, FloatBinaryOperator operator) implements Calculation {    @Override    public float[] apply(int length, ReadOnlyArray.AsFloats<?>... inputs) {        assert length > 0 : "length must be positive but is %d".formatted(length);        // Рекурсивно считаем левое и правое выражения        var lhsValue = lhs.apply(length, inputs);        assert lhsValue.length == length;        var rhsValue = rhs.apply(length, inputs);        assert rhsValue.length == length;        // Выполняем эту операцию, собирая результат в `lhsValue`,        // чтобы не создавать дополнительные массивы        for (var index = 0; index < lhsValue.length; index++) {            lhsValue[index] = operator.apply(lhsValue[index], rhsValue[index]);        }        return lhsValue;    }}

Это простое и понятное представление. Оно легко реализуется для разных операций, легко расширяется, и его очень просто построить из приходящего нам Protobuf-представления.

Для нашего примера на этапе парсинга построилось бы подобное представление:

new BinaryOperator(    new Feature(0 /* вычисленный на этапе парсинга индекс `f_foo` в `inputs` */)    new BinaryOperator(        new Feature(1 /* вычисленный на этапе парсинга индекс `f_bar` в `inputs` */),        new Feature(2 /* вычисленный на этапе парсинга индекс `f_baz` в `inputs` */),        (lhs, rhs) -> lhs * rhs    ),    (lhs, rhs) -> lhs + rhs)

Поэтому долгое время мы действительно использовали именно его. Ключевое слово: использовали.

Недостатки такого подхода:

• Высокая стоимость по CPU. Вычисление формул в среднем составляло более 15% от всей работы сервиса. Причём из-за рекурсивности подхода вместе с тем, что всякий вызов виртуальный (см. красные фреймы на профиле), формулы не инлайнились и были крайне неблагоприятны для нетривиальных оптимизаций JIT-компилятора.

  

• Лишние аллокации. На вычисление формул приходилось около 18% от всех аллокаций нашего сервиса. Каждая промежуточная операция, поскольку работает в изоляции (не смотрит внутрь дочерних операций), создаёт промежуточные массивы, которые всё так же неохотно вырезаются JIT-компилятором.

  

Ручной подход к вычислениям

Ручная формула

Нам показалось логичным предположить, что «ручная» запись конкретной формулы могла бы дать иной результат на профилях. В конце концов, JIT-компилятор заточен именно под то, как код пишут люди.

Например, расчёт нашей f_qux мог бы выглядеть так:

// Попытка ручной реализации нашей формулыenum ManualCalculation implements Calculation {    INSTANCE;    @Override    public float[] apply(int length, ReadOnlyArray.AsFloats<?>... inputs) {        // Сразу вычисляем значения всех входов        var foo = inputs[0].asFloats();        var bar = inputs[1].asFloats();        var baz = inputs[2].asFloats();        // Массив для агрегации результата, пока пишем достаточно наивно без особой экономии массивов        var result = new float[length];        for (var index = 0; index < length; index++) {            // In-line применение вложенных операций            result[index] = foo[index] / (bar[index] + baz[index]);        }        return result;    }}

Все виртуальные вызовы в одном месте, единственный цикл, никакой рекурсии, понятный паттерн использования массивов. Давайте узнаем, поможет ли это JIT-компилятору? Очевидно, что проверить эту гипотезу на всём многообразии формул и при реальной нагрузке было практически невозможно, поэтому решили написать реалистичные бенчмарки для проверки базовых гипотез.

Бенчмарки

Параметры бенчмарков:

Устройство бенчмарка:

/// Варианты реализации:///   - `MANUAL`: написанная вручную,///   - `EVALUATED`: композиция из Calculation (она же рекурсивная реализация, baseline)@Param({"MANUAL", "EVALUATED"})private EvaluationMode mode;/// Количество товаров, по сути, длина векторов@Param("2000")private int length;private Calculation calculation;private ReadOnlyFloatArray[] inputs;@Setuppublic void setUp() {    // Выбор реализации    calculation = switch (mode) {        case MANUAL -> ManualCalculation.INSTANCE;        case EVALUATED -> EvaluatedCalculation.INSTANCE;        case COMPILED -> null;    };    // Детерминированная генерация входных данных    var random = new Random(0x1234_0001);    var foo = FormulaBenchmarks.generateInputs(random, length, 1f, 100f);    var bar = FormulaBenchmarks.generateInputs(random, length, 1f, 10f);    var baz = FormulaBenchmarks.generateInputs(random, length, 20f, 30f);    inputs = new ReadOnlyFloatArray[]{foo, bar, baz};}

Собственно, сам бенчмарк:

@Benchmarkpublic void apply(Blackhole blackhole) {    // Вычисляем формулу    blackhole.consume(calculation.apply(length, inputs));}

Вот какие результаты мы получили:

Обращу внимание, что тут и далее результаты бенчмарков приводятся в операциях в секунду, то есть чем больше значение, тем лучше.

Да-да, рекурсивная интерпретация неожиданно выиграла у наивной «ручной» реализации. Если бы мы остановились на этом этапе, можно было бы сделать ошибочный вывод, что рекурсивная интерпретация и так хороша, оптимизировать нечего. Но у нас появилось подозрение, что пример с простой формулой не очень показателен. Этому могло быть несколько причин, например, достаточно малое разнообразие возможных формул во время бенчмарка (по сути, одна) и малая вложенность (всего лишь Calculation в Calculation), не вставляющие палки в колёса JIT-компилятору.

Пример посложнее

Мы решили взять менее упрощённый (но всё ещё простой) синтетический пример, отражающий сложно вложенную формулу. Выбрали последовательную сумму из 10 фич:

public enum ManualCalculation implements Calculation.Manual {    INSTANCE;    @Override    public float[] apply(int length, ReadOnlyArray.AsFloats<?>... inputs) {        var result = new float[length];        // Перебор слагаемых.        for (var input : inputs) {            var inputValues = input.asFloats();            for (var index = 0; index < length; index++) {                result[index] += inputValues[index];            }        }        return result;    }}

Здесь стоит оговориться, что перед нами не стояло цели написать идеальное ручное вычисление. Была задача оценить, насколько быстрее может быть вычисление в сравнении с тем, что даёт нам парсер:

// Честно записанное вычисление, как его собрал бы парсерprivate static final class EvaluatedCalculation {    public static final Calculation INSTANCE = new Calculation.BinaryOperator(        new Calculation.BinaryOperator(            new Calculation.BinaryOperator(                new Calculation.BinaryOperator(                    new Calculation.BinaryOperator(                        new Calculation.BinaryOperator(                            new Calculation.BinaryOperator(                                new Calculation.BinaryOperator(                                    new Calculation.BinaryOperator(                                        new Calculation.Input(0),                                        new Calculation.Input(1),                                        Float::sum                                    ),                                    new Calculation.Input(2),                                    Float::sum                                ),                                new Calculation.Input(3),                                Float::sum                            ),                            new Calculation.Input(4),                            Float::sum                        ),                        new Calculation.Input(5),                        Float::sum                    ),                    new Calculation.Input(6),                    Float::sum                ),                new Calculation.Input(7),                Float::sum            ),            new Calculation.Input(8),            Float::sum        ),        new Calculation.Input(9),        Float::sum    );}

Вот здесь ожидание совпало с реальностью: рекурсивное вычисление оказалось на порядок медленнее теоретически возможного:

Бенчмарки подтвердили гипотезу: для сложных формул мы действительно принципиально отстаём по производительности, и надо что-то с этим делать.

Решение: компиляция формул в байт-код JVM

Базовая идея

Нам не подходил вариант писать каждую формулу вручную. Их тысячи, они постоянно меняются, и составляют их специалисты Data Science. Более того, с ростом формулы и вложенности в ней когнитивная нагрузка от её перевода в эффективную Java-форму быстро растёт. Значит, нужен автоматический способ превращать декларативное описание формулы в оптимальный императивный код.

Мы решили динамически компилировать выражения в нерекурсивную форму.

Общая идея: в рантайме нужно сгенерировать и загрузить уже знакомое нам вычисление. По сути, динамически собрать реализацию Calculation для конкретной формулы.

// Всякая специализированная реализация должна расширять этот класс,// чтобы было проще их отличать на профилях и при отладкеabstract non-sealed class Compiled implements Calculation {}

Сразу возник вопрос, что именно генерировать в рантайме: Java-код или байт-код?

Вариант с генерацией Java-кода и моментальной компиляцией его с помощью javac отвергли из-за сложности поддержки вложенных конструкций, необходимости включать компилятор Java в окружение выполнения (javax.tools.JavaCompiler не включён в JRE) и избыточности промежуточных представлений (AST-формулы → Java-код → внутренние представления javac, включая AST → байт-код JVM).

Мы выбрали вариант с генерацией байт-кода JVM. Его преимущества:

• Производительность генерации: сразу генерировать байт-код быстрее, чем генерировать Java-код и затем компилировать его.

• Контроль и гибкость: позволяет тонко управлять использованием стека и локальных переменных.

• Композируемость: рекурсивная структура AST идеально ложится на рекурсивную процедуру генерации (но не исполнения) вычислений.

• Job security Интересный технический челлендж: написание компилятора — это и вызов, и нетривиальная инженерная задача.

При этом мы понимали, что у этого подхода есть и явные минусы: устройство этой системы трудно объяснить без вводной в устройство абстрактной JVM, сгенерированный байт-код сложнее отлаживать (в конечном счёте мы смогли решить обе проблемы).

Байт-код JVM

JVM — стековая машина с локальными переменными. Большинство инструкций оперирует стеком операндов. Для исследования байт-кода можно воспользоваться javap из JDK, который также доступен через Compiler Explorer, куда я добавил пример с расчётом f_qux.

Давайте для начала рассмотрим простой пример того, во что компилируется вычисление для одинарных float-значений:

float f_qux(float foo, float bar, float baz) {    var sum = bar + baz;    return foo / sum;}

Compiler explorer покажет следующее (префикс f используется для операций, оперирующих float):

    // Изначально стек пустой:    // []    // В локальных переменных лежат this и аргументы функции    // 0: this, 1: foo, 2: bar, 3: baz// var sum = bar + baz;    // Кладём на стек 2-ю локальную переменную:    // -> [bar]    fload_2    // Кладём на стек 3-ю локальную переменную:    // -> [bar, baz]    fload_3    // Складываем два элемента наверху стека и кладём результат на стек:    // [bar, baz] -> [bar + baz]    fadd    // Сохраняем результат в 4-ю локальную переменную    // -> []    // -> 4: bar + baz    fstore 4// return foo / sum;     // Кладём на стек 1-ю локальную переменную:    // -> [foo]    fload_1    // Кладём на стек 4-ю локальную переменную:    // -> [foo, bar + baz]    fload 4    // Делим элементы на стеке и кладём результат на стек    // -> [foo / (bar + baz)]    fdiv    // Забираем значение со стека и возвращаем из функции    // -> []    freturn

Что здесь уместно отметить:

• вычисления выполняются над элементами на стеке;

• локальные переменные определяются только их индексами.

По тому же принципу строятся и вызовы функций. Простой код:

void printMax(float x, float y) {    System.out.println(Math.max(x, y));}

Во что он превращается в байт-коде:

    // []    // 0: this, 1: x, 2: y     // Читаем поле java.lang.System.out и кладём на стек:    // -> [System.out]    getstatic #7 // Field java/lang/System.out:Ljava/io/PrintStream;     // Кладём на стек 1-ю локальную переменную:    // -> [System.out, x]    fload_1     // Кладём на стек 2-ю локальную переменную:    // -> [System.out, x, y]     fload_2     // Вызываем статический метод `float java.lang.Math.max(float, float)`    // с двумя аргументами со стека и кладём результат на стек:    // -> [System.out, max(x, y)]    invokestatic  #13 // Method java/lang/Math.max:(FF)F     // Вызываем виртуальный метод `void java.io.PrintStream.println(float)`    // на аргументе со стека с аргументом со стека:    // -> []    invokevirtual #19 // Method java/io/PrintStream.println:(F)V    // Выходим из void-метода    return

Теперь, с этим знанием, дело остаётся за малым: за реализацией.

Простой подход к генерации байт-кода

В действительности взаимодействие с нашими векторами практически не отличается. В дополнение к примитивам выше добавляются:

• взаимодействия с массивами по индексам;

• простые циклы по счётчикам.

Но до этого мы ещё дойдём, а пока стоит ключевой вопрос: как нам композировать реализацию компилятора, чтобы его было просто разрабатывать и развивать?

На самом деле, эта стековая модель очень похожа на то, как мы поступали в интерпретирующей реализации. Каждый вложенный Expression соответствует отдельному вызову с фиксированными входами и выходами, далее результат таким же образом комбинируются. Поэтому давайте пойдём по такому пути:

1. Каждый сегмент байт-кода, соответствующего применению операции, должен изменять стек следующим образом: […] → […, float[]], то есть порождать на стеке результат своего применения (и ничего больше).

2. Сам сегмент волен выполнять любые операции, необходимые для своего расчёта, включая вложенные вычисления (также соответствующие п. 1).

На практике оказалось, что так довольно просто писать реализацию вычисления.

Генерация константы

Например, так выглядит код, выполняющий генерацию байт-кода для константного выражения (вектор длины length, заполненный константой):

// […]// Создаём на стеке пустой массив длины из локальной переменной// -> […, float[length] values]pushNewResultArray();// Дублируем ссылку на массив на стеке:// -> […, float[length] values, float[length] values]visitInsn(DUP);// Генерируем вызов `Arrays.fill(float[], float)` с заданной в формуле константой:// -> […, float[length] values]AccArrays.fillFloat(expression.getNumber());// […, float[]]

Для генерации байт-кода используется стандартный для этого инструмент — OW2 ASM. С недавних пор в JDK существует Classfile API, также предназначенное для манипуляций с байт-кодом. Однако на момент разработки проекта оно ещё не было стабилизировано (хотя сейчас оно нам доступно, ведь мы используем Java 25).

Для константы 42 сгенерированный байт-код будет выглядеть так:

// […]// Кладём `length` на стек:// -> […, int length]iload_1// Конструируем массив длины `length`:// -> […, float[length]]newarray float// Дублируем ссылку на массив на стеке:// -> […, float[length], float[length]]dup// Кладём константу 42f на стек// -> […, float[length], float[length], float 42f]ldc #7 // float 42.0f// Вызываем метод `java.util.Arrays.fill(float[], float)`// -> […, float 42f]invokestatic #8 // Method java/util/Arrays.fill:([FF)V// […, float[]]

Примечательно, что обычным Java-кодом такой байт-код не получить: чтобы передать массив в Arrays.fill и при этом не потерять ссылку на него, его нужно куда-то сохранить, например, в локальную переменную. Здесь же мы оперируем сугубо стеком. В реальности это не даёт никакой разницы с точки зрения производительности, но позволяет минимизировать избыточность в байт-коде, а также экономит размер тела функции (строго ограниченный 65 535 байтами).

Генерация бинарного оператора

Как уже говорили, это позволяет легко композировать вложенные выражения. Далее привели пример кодогенератора для уже знакомой ноды Binary:

// […]try (    // Наш вспомогательный метод для присвоения локальной переменной    // названия и выделения свободного номера    // Информация о названиях локальных переменных не является обязательной,    // но помогает в отладке генерируемого байт-кода    var lhs = locals.<float[]>var("lhs", FLOAT_ARRAY_DESCRIPTOR);     var rhs = locals.<float[]>var("rhs", FLOAT_ARRAY_DESCRIPTOR)) {    // Вставляем байт-код вложенной операции расчёта левого операнда:    // -> […, float[length] lhs]    pushExpression(binary.getLhs());    // Для удобства сразу сохраняем результат в локальную переменную:    // -> […]    lhs.visitInsn(ASTORE);    // Вставляем байт-код вложенной операции расчёта правого операнда:    // -> […, float[length] rhs]    pushExpression(binary.getRhs());    // Для удобства сразу сохраняем результат в локальную переменную:    // -> […]    rhs.visitInsn(ASTORE);     // Вставляем байт-код выполнения бинарной операции над локальными переменными lhs и rhs:    // -> […]    applyBinaryOperator(operator, lhs, rhs);    // По контракту `applyBinaryOperator` сохраняет результат в `lhs`    // Достаём это значение на стек по контракту для сегмента:    // […, float[length] result]    lhs.visitInsn(ALOAD);}// […, float[]]

Как видно, байт-кодогенератор обходит AST рекурсивно, но при этом генерирует последовательный байт-код.

Отладка

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

Далее привели пример декомпилированного представления для уже знакомой f_qux:

// Названия локальных переменных уточнены для простоты понимания примера,// но декомпилированное тело выражения является результатом работы Fernflowerpublic final class CalculationCompiler$CompiledCalculation$123    extends Calculation.Compiled {    @Override    public float[] apply(int length, ReadOnlyArray.AsFloats<?>... inputs) {        if (length == 0) {            return new float[0];        }        var foo = inputs[0].asFloats();        var bar = inputs[1].asFloats();        var baz = inputs[2].asFloats();        for (var index = 0; index < length; index++) {            bar[index] += baz[index];        }        for (var index = 0; index < length; index++) {            foo[index] /= bar[index];        }        return foo;    }}

Результаты: цифры и профили

Бенчмарки

Итак, что же показали наши бенчмарки на динамически скомпилированных реализациях? Скомпилированное вычисление суммы из 10 фич (COMPILED в таблицах) не только достигло уровня ручной реализации, но и заметно обогнало её:

Более того, на синтетическом примере с foo—bar—baz оно тоже оказалось самым быстрым, обогнав прежде побеждавший там EVALUATED:

Хотя этот пример и признан синтетическим, результат всё равно приятный.

Запуск в production

Бенчмарки полезны, но главное — результат на реальной нагрузке. Что же мы получили там?

Запуск Compiled-вычислений вместо старых Evaluated дал нам ожидаемый результат:

• Доля CPU, потребляемая вычислением факторов ранжирования, снизилась с >15% до ~5–6%.

  

• На профиле исчезли глубокие цепочки рекурсивных вызовов. Код стал отлично инлайниться и автоматически векторизоваться JIT-компилятором HotSpot (об этом далее).

• Allocation rate снизился с >18% до ~15–16%. В первой версии это не было главной целью, но оказалось приятным бонусом.

  

• Общее время обработки поискового запроса сократилось на 10 миллисекунд — что хорошо для пользовательского опыта.

Тестирование

Модульные тесты

Для вычисления calculated-фич уже был выстроен пайплайн, использующий методы property-based тестирования для обнаружения проблем. Если коротко, есть обычные тесты, покрывающие фиксированные сценарии вида «для функции, записанной в AST как <…>, собери её Evaluated (а теперь и Compiled) форму и проверь, что она даёт результат <…> на входе <…>». Помимо этого, есть тесты вида «придумай функцию, варьирующую по фрагментам <…>, собери её Evaluated (а теперь и Compiled) форму и проверь, что она даёт результат <…> на входе <…>».

При этом генерация фрагментов выражений происходит случайно, но обязательно покрывает вырожденные случаи (экстремумы, нули и так далее) и остаётся воспроизводимой. В случае провала теста, test engine автоматически запоминает условия падения и начинает редуцировать их до минимального воспроизводящегося примера

Здесь нам удалось найти специфическую функцию, sigmoid, расчёт которой расходился со старой реализацией на некоторых специфических входах, как и всегда в таких ситуациях, из-за float-арифметики.

Здесь хорошо подходит классический комикс xkcd про особенности float-арифметики

— Эй, смотри: . Странно.
— Да. Именно так меня в колледже выгнали из ACM.
— … Что?
— Во время соревнования я сказал программистам нашей команды, что — это стандартный тест для обработчиков чисел с плавающей точкой: мол, должно получиться , если нет ошибок округления.
— Это ужасно.
— Да, они перелопатили половину своих алгоритмов в поисках серьёзной ошибки, прежде чем поняли, в чём дело.
— Ещё я слышал, что корень четвёртой степени из — это π.

Сравнение с референтной реализацией

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

1. Включили логику компиляции формул в продакшене: для приходящих формул создавали их скомпилированное представление Calculation.Compiled, загружали его, но не использовали для расчётов. Это позволило отловить проблемы в генерации, которые не покрыли наши автоматические тесты. Спойлер: тут всё сразу оказалось успешно.

2. Включили расчёт с использованием Compiled-представления, но при этом продолжали использовать результат расчёта старой реализации и сравнивали с ним результат Compiled-представления. Здесь валидация также прошла успешно, проблем не выявили.

3. Начали использовать расчёт с использованием Compiled-представления, отключили старую реализацию. На выходе получили те самые графики ожидаемого улучшения.

Снапшотные тесты

Со временем у нас возникла потребность вносить доработки в компилятор. Мы поняли, что было бы удобно наблюдать diff не только в коде кодогенератора, но и в результате его работы на тестовых данных. Например, внося оптимизацию, исключающую лишние константы, я ожидал бы увидеть, что исчезли лишние new float[length] и зависящие от них вычисления, при этом прочие части не изменились.

У разработчиков компиляторов (и не только) есть устоявшаяся для этого методология — снапшотное тестирование (snapshot testing, golden testing). Идея довольно простая: по сути, это привычный нам assertEquals(expected, actual), у которого expected не задан явно, а генерируется автоматически при первом вызове (изначально он, очевидно, совпадает с actual) и сохраняется в файл-снапшот. При будущих вызовах actual сравнивается со значением в файле-снапшоте. В случае несовпадения разработчик может принять изменение, заменив expected на новое значение, или распознать ошибку.

Мы же в таких тестах смотрим на результирующий байт-код и результат его декомпиляции для определённых тестовых формул.

Исследование границ: какие оптимизации ещё возможны

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

Подход к проверке гипотез

Возник вопрос: как проверять гипотезы, если они требуют доработок компилятора (порой нетривиальных), а их эффект заранее неизвестен?

На помощь снова пришли бенчмарки и то, что у нас был доступ к исходным формулам. Вместо того чтобы переделывать компилятор под эксперимент, нам было достаточно:

1. Найти примеры формул, которые могли бы согласно гипотезе выиграть от предлагаемой оптимизации.

2. Получить для них дамп сгенерированного байт-кода.

3. Получить с помощью декомпилятора приближённую Java-форму.

4. Сравнить в бенчмарке этот вариант и вариант с внесёнными вручную изменениями, которые мы хотим реализовать в компиляторе.

Дальше рассмотрим несколько наших экспериментов.

Гипотеза 1: Явные проверки перед циклами — range check

Одно из сильных свойств экосистемы JVM — это мощный JIT-компилятор HotSpot. Он способен применять ряд оптимизаций, характерных для вашего окружения, включая конкретный используемый CPU. Это позволяет ему, в частности, выполнять автоматическую векторизацию кода, заменяя поэлементные операции их векторными аналогами (например, инструкциями из AVX-512).

Другое важное свойство экосистемы JVM — это достаточно строгие гарантии целостности с поправкой на цитату Вани Углянского про FFM API:

Так что теперь в Java официально есть UB, с чем всех и поздравляю!

Это значит, что в случае выхода за границы массива поведение строго определено: JVM должна гарантировать, что в таком случае случится IndexOutOfBoundsException. Причём его stack-trace должен вести к месту, где это случилось.

Отсюда возникло предположение: в скомпилированном коде мы индексируем массивы в цикле и предварительно никак не проверяем диапазон, из-за этого JIT-компилятор не может доказать, что мы всегда попадаем в диапазон массива, а значит, вынужден по-честному перебирать элементы вместо использования широких векторных инструкций. Значит, вставка явной проверки range check могла бы помочь ему в автовекторизации циклов.

Для нашего примера изменение выглядит так:

// Массивы длины `length`, но JIT-компилятор может не знать об этой гарантии var foo = inputs[0].asFloats();var bar = inputs[1].asFloats();var baz = inputs[2].asFloats();// Добавляем явную проверку длин, чтобы дать JIT-компилятору гарантию отсутствия выхода за границы нижеif (foo.length != length || bar.length != length || baz.length != length) {    throw new AssertionError();}// Ожидаем, что теперь эти циклы будут лучше векторизоватьсяfor (var index = 0; index < length; index++) {    bar[index] += baz[index];}for (var index = 0; index < length; index++) {    foo[index] /= bar[index];}

Результат бенчмарков оказался неожиданным: никакой разницы.

Однако, раз уж мы выдвинули гипотезу, следует в ней разобраться. Для этого есть отличный инструмент: hsdis — плагин к HotSpot, позволяющий печатать в человекочитаемом виде дизассемблер для сгенерированного им кода (бинарные сборки от Алексея Шипилёва).

Итак, давайте посмотрим, что сгенерировал JIT-компилятор для наших формул. Для этого добавим hsdis в библиотеки, доступные JVM, а в бенчмарках в качестве дополнительных флагов виртуальной машины укажем:

-XX:+UnlockDiagnosticVMOptions \-XX:CompileCommand=print,<путь до класса>/calculated/compiled/CalculationCompiler$CompiledCalculation*.apply`

Что же мы увидим?

Во-первых, проверки range check действительно появились (внимание на *if_icmpne в листинге):

0x00007fe2043e7c84: vzeroupper0x00007fe2043e7c87: call 0x00007fe203cb6580 ; ImmutableOopMap {};*if_icmpne {reexecute=1 rethrow=0 return_oop=0}; - (reexecute)ru.ozon.o2.midway.experimental.plugins.api.plugins.calculated.compiled.OptBoundsCheckingBenchmarks$CalculationMode$2::apply@40 (line 71); {runtime_call UncommonTrapBlob}0x00007fe2043e7c8c: nopl 0xb000c7c(%rax,%rax,1) ; {other}0x00007fe2043e7c94: mov $0xffffff45,%esi0x00007fe2043e7c99: data16 xchg %ax,%ax0x00007fe2043e7c9c: vzeroupper0x00007fe2043e7c9f: call 0x00007fe203cb6580 ; ImmutableOopMap {};*if_icmpne {reexecute=1 rethrow=0 return_oop=0}; - (reexecute)ru.ozon.o2.midway.experimental.plugins.api.plugins.calculated.compiled.OptBoundsCheckingBenchmarks$CalculationMode$2::apply@47 (line 71); {runtime_call UncommonTrapBlob}0x00007fe2043e7ca4: nopl 0xc000c94(%rax,%rax,1) ; {other}0x00007fe2043e7cac: mov $0xffffff45,%esi0x00007fe2043e7cb1: mov %edx,%ebp0x00007fe2043e7cb3: nop0x00007fe2043e7cb4: vzeroupper0x00007fe2043e7cb7: call 0x00007fe203cb6580 ; ImmutableOopMap {};*ifne {reexecute=1 rethrow=0 return_oop=0}; - (reexecute)ru.ozon.o2.midway.experimental.plugins.api.plugins.calculated.compiled.OptBoundsCheckingBenchmarks$CalculationMode$2::apply@1 (line 64); {runtime_call UncommonTrapBlob}0x00007fe2043e7cbc: nopl 0xd000cac(%rax,%rax,1) ; {other}0x00007fe2043e7cc4: mov $0xffffff45,%esi0x00007fe2043e7cc9: mov %ebx,(%rsp)0x00007fe2043e7ccc: vzeroupper0x00007fe2043e7ccf: call 0x00007fe203cb6580 ; ImmutableOopMap {};*if_icmpeq {reexecute=1 rethrow=0 return_oop=0}; - (reexecute)ru.ozon.o2.midway.experimental.plugins.api.plugins.calculated.compiled.OptBoundsCheckingBenchmarks$CalculationMode$2::apply@54 (line 71); {runtime_call UncommonTrapBlob}0x00007fe2043e7cd4: nopl 0xe000cc4(%rax,%rax,1) ; {other}

То есть как минимум мы смотрим действительно на то, что нас интересует.

Во-вторых, и это самое главное, векторизация успешно сработала в обоих случаях. Мы получили идентичный ACM с точностью до регистров:

0x00007fdec03e70d3: vmovdqu32 0x10(%rax,%r10,4),%zmm0 ; {no_reloc}0x00007fdec03e70de: vaddps 0x10(%r13,%r10,4),%zmm0,%zmm00x00007fdec03e70e9: vmovdqu32 %zmm0,0x10(%rax,%r10,4)0x00007fdec03e70f4: vmovdqu32 0x50(%rax,%r10,4),%zmm00x00007fdec03e70ff: vaddps 0x50(%r13,%r10,4),%zmm0,%zmm00x00007fdec03e710a: vmovdqu32 %zmm0,0x50(%rax,%r10,4)0x00007fdec03e7115: vmovdqu32 0x90(%rax,%r10,4),%zmm00x00007fdec03e7120: vaddps 0x90(%r13,%r10,4),%zmm0,%zmm00x00007fdec03e712b: vmovdqu32 %zmm0,0x90(%rax,%r10,4)0x00007fdec03e7136: vmovdqu32 0xd0(%rax,%r10,4),%zmm00x00007fdec03e7141: vaddps 0xd0(%r13,%r10,4),%zmm0,%zmm00x00007fdec03e714c: vmovdqu32 %zmm0,0xd0(%rax,%r10,4)0x00007fdec03e7157: vmovdqu32 0x110(%rax,%r10,4),%zmm00x00007fdec03e7162: vaddps 0x110(%r13,%r10,4),%zmm0,%zmm00x00007fdec03e716d: vmovdqu32 %zmm0,0x110(%rax,%r10,4)0x00007fdec03e7178: vmovdqu32 0x150(%rax,%r10,4),%zmm00x00007fdec03e7183: vaddps 0x150(%r13,%r10,4),%zmm0,%zmm00x00007fdec03e718e: vmovdqu32 %zmm0,0x150(%rax,%r10,4)0x00007fdec03e7199: vmovdqu32 0x190(%rax,%r10,4),%zmm00x00007fdec03e71a4: vaddps 0x190(%r13,%r10,4),%zmm0,%zmm00x00007fdec03e71af: vmovdqu32 %zmm0,0x190(%rax,%r10,4)0x00007fdec03e71ba: vmovdqu32 0x1d0(%rax,%r10,4),%zmm00x00007fdec03e71c5: vaddps 0x1d0(%r13,%r10,4),%zmm0,%zmm00x00007fdec03e71d0: vmovdqu32 %zmm0,0x1d0(%rax,%r10,4) ; {no_reloc}

0x00007fe2043e75e3: vmovdqu32 0x10(%r14,%r9,4),%zmm00x00007fe2043e75ee: vaddps 0x10(%rbp,%r9,4),%zmm0,%zmm00x00007fe2043e75f9: vmovdqu32 %zmm0,0x10(%r14,%r9,4)0x00007fe2043e7604: vmovdqu32 0x50(%r14,%r9,4),%zmm00x00007fe2043e760f: vaddps 0x50(%rbp,%r9,4),%zmm0,%zmm00x00007fe2043e761a: vmovdqu32 %zmm0,0x50(%r14,%r9,4)0x00007fe2043e7625: vmovdqu32 0x90(%r14,%r9,4),%zmm00x00007fe2043e7630: vaddps 0x90(%rbp,%r9,4),%zmm0,%zmm00x00007fe2043e763b: vmovdqu32 %zmm0,0x90(%r14,%r9,4)0x00007fe2043e7646: vmovdqu32 0xd0(%r14,%r9,4),%zmm00x00007fe2043e7651: vaddps 0xd0(%rbp,%r9,4),%zmm0,%zmm00x00007fe2043e765c: vmovdqu32 %zmm0,0xd0(%r14,%r9,4)0x00007fe2043e7667: vmovdqu32 0x110(%r14,%r9,4),%zmm00x00007fe2043e7672: vaddps 0x110(%rbp,%r9,4),%zmm0,%zmm00x00007fe2043e767d: vmovdqu32 %zmm0,0x110(%r14,%r9,4)0x00007fe2043e7688: vmovdqu32 0x150(%r14,%r9,4),%zmm00x00007fe2043e7693: vaddps 0x150(%rbp,%r9,4),%zmm0,%zmm00x00007fe2043e769e: vmovdqu32 %zmm0,0x150(%r14,%r9,4)0x00007fe2043e76a9: vmovdqu32 0x190(%r14,%r9,4),%zmm00x00007fe2043e76b4: vaddps 0x190(%rbp,%r9,4),%zmm0,%zmm00x00007fe2043e76bf: vmovdqu32 %zmm0,0x190(%r14,%r9,4) ; {no_reloc}0x00007fe2043e76ca: vmovdqu32 0x1d0(%r14,%r9,4),%zmm00x00007fe2043e76d5: vaddps 0x1d0(%rbp,%r9,4),%zmm0,%zmm00x00007fe2043e76e0: vmovdqu32 %zmm0,0x1d0(%r14,%r9,4)

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

Итог: отказались от добавления дополнительных проверок range check в компиляторе из-за отсутствия эффекта.

Бонус: подтвердили тезис о том, что Compiled-выражения быстрые, в том числе за счёт успешной автовекторизации.

Гипотеза 2: Inlining констант

Другая оптимизация, которую захотелось попробовать, — это подстановка констант. Как можно было заметить выше, константы мы также приводим к векторам, то есть для общности заводим массивы, заполненные одной и той же константой для всех тысяч позиций. Это действительно упрощает реализацию компилятора за счёт того самого правила […] → […, float[]], однако в контексте оптимизаций и улучшения производительности звучит расточительно. Что, если попробовать честно подставлять константы на место?

Например, если мы заменим в нашем выражении baz на константу 100, то его представление без этой оптимизации будет выглядеть так:

var baz = new float[length];Arrays.fill(baz, 100);for (var index = 0; index < length; index++) {    bar[index] += baz[index];}

Хотя вместо этого мы могли бы заинлайнить константу и получить просто:

for (var index = 0; index < length; index++) {    bar[index] += 100;}

Для точности эксперимента бенчмарки проводились с включённым профилированием аллокаций (-prof gc в JMH).

Бенчмарк показал положительный эффект. Интересно, что главный выигрыш оказался в наиболее ценной для нас метрике ops/s — около 50%(!):

Итог: оптимизация требует статического анализа AST (то есть сложна в реализации), но этого того стоит.

Гипотеза 3: Использование Vector API

Также из любопытства мы попробовали Vector API для ручной (явной) векторизации. У нас не было особых ожиданий, потому что JIT-компилятор в HotSpot и так отлично справляется с векторизацией.

Действительно, чуда не произошло — Vector API проиграл:

С важной оговоркой, что бенчмарк проводился на Java 21. В качестве «бонуса» — буду рад, если кто-то попробует самостоятельно эффективно записать наш foo-bar-baz через Vector API в комментариях.

Итог: сейчас у нас нет потребности переписывать компилятор на генерацию кода, использующего Vector API. Маловероятно, что мы от этого что-то выиграем, а скорее, даже проиграем.

Гипотеза 4: переиспользование массивов

Может показаться, что все гипотезы об улучшении производительности рождаются сугубо теоретически. Однако некоторые из них — это результат анализа профилей приложения, который мы регулярно проводим. Так, при очередном анализе профиля аллокаций мы заметили, что промежуточные аллокации превышают 27% и всё ещё составляют большую долю от общего числа.

Отвечая на вопрос, почему доля увеличилась с момента запуска проекта: за это время оптимизировали ряд других подсистем, и относительная доля этой части стала больше.

В результате выдвинули идею оптимизации аллокаций для одного распространённого сценария. Для определённых видов операций нам всегда известно, какие из операндов точно не изменяются. Например, для всё той же кодогенерации бинарной функции модифицируется левый операнд, в который накапливается результат, но никогда не изменяется правый операнд. Значит, для значения правого операнда можно генерировать код, не копирующий данные из источника. У нашего ReadOnlyFloatArray есть метод float[] asFloatsUnsafe() — он не копирует внутренний массив, но требует от вызывающей стороны не изменять его:

try (    var lhs = locals.<float[]>var("lhs", FLOAT_ARRAY_DESCRIPTOR);     var rhs = locals.<float[]>var("rhs", FLOAT_ARRAY_DESCRIPTOR)) {    pushExpression(binary.getLhs(), false /* не read-only; сюда может записываться результат */);    lhs.visitInsn(ASTORE);    pushExpression(binary.getRhs(), /* read-only; можно использовать ReadOnlyFloatArray#asFloatsUnsafe() */);    rhs.visitInsn(ASTORE);     applyBinaryOperator(operator, lhs, rhs);    lhs.visitInsn(ALOAD);}

И, в свою очередь, ниже по иерархии нужен подобный метод для генерации чтения значения фичи:

loadInputFeature(index);visitMethodInsn(    INVOKEINTERFACE,    READ_ONLY_ARRAY_AS_FLOATS_TYPE_NAME,    readOnly ? "asFloatsUnsafe" : "asFloats",    FLOAT_ARRAY_METHOD_DESCRIPTOR,    true);

За пару часов реализовали эту функциональность в компиляторе, провалидировали изменения с помощью снапшотных тестов, проверили на локальных бенчмарках и запустили в продакшен. Оптимизация оказалась полезной: относительная доля аллокаций опустилась до ~17,5%, заодно снизив общее потребление CPU на вычисление calculated-фич: ~6,7–7,9% → ~6,2–6,7%.

И хотя эта оптимизация менее мощная, чем потенциальная реализация Гипотезы 2, но она проста в реализации и хорошо показала себя. К тому же она никак не мешает реализации инлайнинга констант и, скорее даже, предшествует ему.

Выводы и рекомендации

Универсальные выводы

Здесь собраны выводы, которые могут оказаться полезными для всех, независимо от необходимости разрабатывать подобный сложный инструмент и даже независимо от используемого языка программирования:

• Собирайте профили, а не гадайте. Именно профилирование позволило нам найти реальную проблему в продакшене и прийти к её успешному решению. Тут не можем не посоветовать async-profiler. Кстати, для Java-сервисов в Ozon у нас есть механизм автоматического профилирования по расписанию и по кнопке, позволяющий разработчикам детально исследовать работу их приложения.

• Понимайте ваше окружение. В случае с экосистемой JVM — это собственно VM и, в частности, понимание основ работы JIT-компилятора: инлайнинга, векторизации, обработки исключения. Эти знания могут оказаться критически важными при написании высокопроизводительного кода.

• Пишите бенчмарки для гипотез. Прежде чем внедрять сложную оптимизацию в код, попробуйте смоделировать её с помощью изолированного бенчмарка. Это сэкономит время и силы. Само собой, для Java классическим инструментом является JMH.

• Оптимизируйте сверху вниз. Сначала ищите архитектурные улучшения, такие как замена интерпретации на компиляцию, и только после этого занимайтесь микрооптимизациями. Вклад последних может быть огромен в изоляции, но ничтожен на общем профиле. И наоборот, не столь радикальные оптимизации могут давать заметный практический эффект, если они попадают на горячий путь.

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

Специальные выводы

Здесь несколько более узкоспециализированных советов, которые также могут пригодиться:

• Генерация байт-кода — мощный инструмент. Для задач, связанных с вычислениями DSL, динамическая генерация байт-кода может дать феноменальный прирост, хотя и требует высокой квалификации.

• Изучайте листинги. Если закапываетесь во внутреннее устройство JVM (а порой это необходимо), не забывайте о специализированных инструментах вроде -XX:+PrintAssembly с hsdis или Compiler Explorer. Они дают бесценную информацию о том, во что действительно компилируется ваш код.

Заключение

Путь от «умной», но неэффективной рекурсивной интерпретации к специализированному компилятору байт-кода оказался крайне результативным. Мы получили гибкую систему, способную эффективно вычислять даже самые причудливые формулы от команды Data Science. Этот опыт подтвердил, что для сложных вычислительных задач иногда стоит выйти за рамки стандартных паттернов и взять на себя ответственность за то, как код будет исполняться на виртуальной машине. Иногда самый прямой путь к эффективности — это дать возможность JIT-компилятору увидеть простой и ясный цикл вместо горы абстракций.

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