Как я ускорила парсинг строк в serde_json на 20%

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

serde — основной фреймворк для сериализации и десериализации в Rust. Его используют как крейт по умолчанию во всей экосистеме. serde_json — это официальный serde-миксин для JSON, так что каждый раз, когда нужно что-то парсить, люди обращаются именно к нему. Конечно, есть и другие библиотеки, специализующиеся на парсинге JSON, например simd-json, но популярность у них, мягко говоря, удручающая. serde_json значительно популярнее: на момент написания от него зависят аж целых 26916 крейта, а от simd-json — всего 66.

Это делает serde_json хорошей мишенью (не как у Jia Tan) для оптимизаций. Велик шанс того, что многим из тысяч пользователей переход на simd-json позволил бы добиться ускорения, но, пока они этого не делают, более мелкие оптимизации — лучше, чем совсем ничего, и такие улучшения — глобальный выигрыш для экосистемы.

С чего бы начать...?

Недавно я работала над #[iex]. В качестве мерила производительности я использовала serde и serde_json, и во время адаптации их к #[iex] я обнаружила некоторые вызывающие вопросы решения в важном для быстродействия коде.

Задача #[iex] — обработка ошибок, так что error-путь оказался первой целью бенчмарков. К моему удивлению, error-путь serde_json оказался более чем в два раза медленнее success-пути на тех же данных:

Почему? Проброс ошибок не может быть настолько медленным. Запуск под perf указал на бутылочное горлышко в виде невинной функции:

fn position_of_index(&self, i: usize) -> Position {
    let mut position = Position { line: 1, column: 0 };
    for ch in &self.slice[..i] {
        match *ch {
            b'\n' => {
                position.line += 1;
                position.column = 0;
            }
            _ => {
                position.column += 1;
            }
        }
    }
    position
}

...которая вызывается для форматирования ошибки в position(), задокументированной как:

/// Position of the most recent call to next().
///
/// ...
///
/// Only called in case of an error, so performance is not important.

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

Можно ли что-то с этим сделать? position_of_index() переводит индекс в строке в пару "строка, колонка". Разобьём задачу на две попроще:

• Посчитать число \n в self.slice[..i] — это номер строки в 0-нумерации, и

• Посчитать расстояние между i и позицией последнего \n в self.slice[..i] — это номер колонки в 1-нумерации.

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

В обоих случаях memchr использует SIMD, что позволяет ему обгонять простой проход циклом. В самом деле, замена реализации на

fn position_of_index(&self, i: usize) -> Position {
    let start_of_line = match memchr::memrchr(b'\n', &self.slice[..i]) {
        Some(position) => position + 1,
        None => 0,
    };
    Position {
        line: 1 + memchr::memchr_iter(b'\n', &self.slice[..start_of_line]).count(),
        column: i - start_of_line,
    }
}

...привносит значительное ускорение:

Error-путь всё ещё медленнее, чем success-путь, но разница теперь менее значительная.

Я отправила пулл-реквест, содержащий эту оптимизацию и гадала, вольют его или нет. В serde_json всё же зависимостей кот наплакал, а разработчик крейта заботится о скорости компиляции, так что не факт, что PR с дополнительной зависимостью не завернут.

На удивление, PR быстро вмержили. Неплохой первый вклад в проект.

А дальше?

dtolnay посоветовал мне порыскать по коду в поисках мест, которые можно ускорить похожим образом. (Не могу переоценить его помощь в этой истории.)

Первым делом мне бросился в глаза цикл в парсинге строк:

while self.index < self.slice.len() && !ESCAPE[self.slice[self.index] as usize] {
    self.index += 1;
}

Мы хотим найти первый не-экранирующий символ. Экранирующими символами являются не только \ (по очевидным причинам) и " (обозначает конец строки), но ещё и все ASCII-символы с кодом, меньшим 0x20 — спецификация JSON запрещает управляющим символам находиться в строках: "line 1\nline2" — корректный JSON, а замена \n на буквальный перевод строки делает его некорректным.

Если бы всё, что нужно было тут — найти первый \ или ", можно было бы достать функцию memchr2 из того же крейта memchr, и дело в шляпе. А тут малину портит третье условие, и нужно доставать более тяжёлую артиллерию. Откуда?

В поисках экранирования

Первая попытка

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

Идея заключалась в

• Использовании memchr2 для поиска первого \ или ", и последующем

• Посимвольном проходе по строке, чтобы убедиться, что управляющих символов в ней нет.

Смысл в том, что перекладывание поиска \ и " на более быстрый алгоритм должно улучшить общую производительность.

В реальности всё пошло не так: код замедлился, потому что дважды пройтись по строке (сначала быстро, а затем медленно) всегда хуже, чем пройтись по этой же строке единожды, настолько же медленно. Да, в отдельности сравнение ch < 0x20 немного быстрее, чем доступ к памяти (ESCAPE[...]), но это улучшение настолько мизерное, что меркнет перед двухпроходностью алгоритма и удвоенному требованию к пропускной способности RAM.

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

strlcpy(char *dst, const char *src, size_t size) копирует строку из src в dst, срезая её по size - 1 символов. Один байт зарезервирован под необходимый нуль-терминатор; в рамках этого обсуждения его проигнорируем. Сравнивались две реализации:

• Однопроходная: делаем *dst++ = *src++ не более size - 1 раз, пока *src не \0, и

• Двухпроходная: вычисляем len = strlen(src) и вызываем memcpy(dst, src, min(len, size - 1)).

Двухпроходный алгоритм был быстрее, потому что strlen и memcpy были glibc-шными фунциями, оптимизированными через SIMD, а цикл однопроходного алгоритма был скалярным. Автор понял это и написал свою реализацию strlen и memcpy, чтобы сделать сравнение более честным:

size_t bespoke_strlcpy(char *dst, const char *src, size_t size) {
    size_t len = 0;
    for (; src[len] != '\0'; ++len) {} // strlen()

    if (size > 0) {
        size_t to_copy = len < size ? len : size - 1;
        for (size_t i = 0; i < to_copy; ++i) // memcpy()
            dst[i] = src[i];
        dst[to_copy] = '\0';
    }
    return len;
}

GCC не дурак и умеет находить и заменять такие циклы на вызовы соответствующих функций из glibc, поэтому автор явно отключил такие фокусы через -fno-builtin. И даже так, двухпроходный алгоритм оказался быстрее однопроходного.

Вот только есть одно "но". -fno-builtin не режет все оптимизации, связанные с memcpy — цикл, похожий на memcpy, всё ещё может быть векторизован. Так и поступил GCC с bespoke_strlcpy, и вышло, что автор сравнивал скалярный strlen (проверить \0, зациклиться) + векторизованный memcpy (проверить размер, зациклиться) и скалярный strlcpy (проверить \0, проверить размер, зациклиться).

Отключение автовекторизации через -fno-tree-vectorize расставляет всё на свои места и делает двухпроходный алгоритм медленнее однопроходного, каким он и должен быть, потому что теперь идет сравнение двух циклов (проверить \0, зациклиться; проверить размер, зациклиться) с одним (проверить \0, проверить размер, зациклиться), и второй вариант быстрее из-за меньшей нагрузки на предсказатель ветвлений и меньшего количества обращений к памяти.

Мораль сей басни такова: векторизация — мощная штука, но применение векторизации к простой части при оставлении скалярной более сложную часть не привносит никаких улучшений. И если мы хотим улучшить код, придётся использовать SIMD в обоих частях.

Вторая попытка

Исходный подход теперь немного преобразовался:

• Берём memchr2, чтобы найти первый \ или ", и

• Используем SIMD, написанный своими руками, чтобы удостовериться в отсутствии управляющих символов.

Придётся переизобрести колесо и слепить свой велосипед, но если задуматься, выходит неплохо. В success-пути мы ищем позиции \ и ", но для управляющих символов лишь проверяем отсутствие. Можно воспользоваться этим фактом и вытащить из-под цикла условный прыжок, заменяя это:

for simd_word in to_simd_words(data):
    if any(simd_word < 0x20):
        ...

...на это:

mask = False
for simd_word in to_simd_words(data):
    mask |= simd_word < 0x20
if any(mask):
    ...

Достаточно быстро мы поняли, что этот метод тоже проигрышный. Короткие строки обрабатываются медленно из-за непрямого вызова (выбираемой в рантайме) функции для поиска \ и ", а на длинных мы упираемся в два прохода. Плоховато, нужно что-нибудь получше.

Принимаем судьбу

Придется за один проход искать \, " и управляющие символы.

Я же пыталась оставить serde_json простым, каким он был до меня. В своих проектах сложность кода для меня не преграда, но что-то, что потом придется поддерживать другим людям, лучше писать максимально просто. Так что писать SIMD-код под каждую архитектуру — не вариант.

Трюк Майкрофта

Есть приём, который был в ходу ещё до того, как в процессоры подвезли нативный SIMD. Вместо использования больших (например, 128-битных) слов, можно использовать машинные, которые почти везде 64-битные — и использовать побитовые операции для симулирования поэлементных операций. Этот подход называют SIMD Within A Register (SWAR, SIMD внутри регистра). Простой пример: перевести восемь латинских символов в нижний регистр можно через x | 0x2020202020202020.

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

Трюк применяется так: для c: i8 условие "это управляющий символ" выглядит как c >= 0 && c < 0x20, что можно переписать как c >= 0 && c - 0x20 < 0. Это условие проверяет, что у c знаковый бит — 0, а у c - 0x20 — 1, а значит, условие можно записать как !c & (c - 0x20) & 0x80 != 0.

Так, для восьми упакованных байтов, нужно вычислять !c & (c - 0x2020202020202020) & 0x8080808080808080. Если получаем 0, то классно, управляющих символов нет, иначе ищем в маске младший ненулевой байт — это и будет первым вхождением управляющего символа.

Но есть нюанс. c - 0x20 в условии выше — вычитание с переносом, и если упаковывать эти операции, то могут возникнуть займы битов между разными байтами. К счастью, это не проблема: ломаются только старшие байты и только начиная с того, который был меньше 0x20. Так как мы ищем первый байт, то проблемы со старшими байтами нас не волнуют.

Конечно, это работает только на little-endian архитектурах. На big-endian придется предварительно развернуть байты в c.

А как насчёт \ и "? Условие для \ выглядит как c ^ b'\\' >= 0 && c ^ b'\\' < 1 — формула выше, в которой 0x20 заменили на 0x01. Вишенка на торте: у b'\\' не выставлен знаковый бит, поэтому c ^ b'\\' >= 0 можно упростить до c >= 0.

Собираем всю кучку формул, тыкаем пару минут и легким движением руки формула преображается в красивую:

!c
& (
    (c - 0x2020202020202020)
    | ((c ^ (b'\\' * 0x0101010101010101)) - 0x0101010101010101)
    | ((c ^ (b'"' * 0x0101010101010101)) - 0x0101010101010101)
)
& 0x8080808080808080

Всего девять битовых операций (считая a & !b за одну). Для сравнения, на x86 с помощью SIMD можно сделать за семь — не такая большая разница помимо раза в два-четыре меньшей пропускной способности (в зависимости от длины машинного слова и векторного машинного слова).

Для коротких строк throughput не имеет значения, в отличие от latency. Из-за этого эффекта "настоящий" SIMD проигрывал SWAR-у по скорости на json-benchmark. (Впрочем, вполне возможно, что я допустила где-то ошибку.) По итогу мы сошлись на коде с использованием SWAR.

Палка о двух концах

Каждый раз, когда вы оптимизируете что-то разворотом цикла или с помощью SIMD, неплохо было бы подумать о том, стоит ли игра свеч и достаточной ли длины данные, чтобы вообще хоть что-то выиграть. Например, алгоритма поиска длины строки длиной от 0 до 16 без ветвлений -- это, конечно, хорошо, но простейшая реализация strlen циклом запросто уделает её по скорости, если строки обычно трехбайтовые.

Вот на это я и напоролась. На строках длиной примерно 5 символов подход с использованием SWAR становился медленее, чем скалярный код. Мы решили, что регрессии на таких коротких строках, что в реальных данных встречаются редковато — это цена, которую мы готовы заплатить за скорость в других случаях.

Но есть одна крайне частая короткая строка — пустая. Похожая регрессия проявилась на строках с идущими подряд экранированными символами, например? \r\n или \uD801\uDC37, что достаточно часто встречаются в данных с Юникодом. Такие регрессии непозволительны, но, к счастью, это легко исправить: не запускать SWAR-цикл, если искомый символ — первый.

Получаем хорошее ускорение:

citm_catalog DOM фликерит, так что регрессий по сути нет вообще. Ну, кроме одной: пустые строки парсятся чуть дольше, но на специально подобранном бенчмарке замедление было в пределах всего 2%.

Когда лексинг оказывается сложным

Юникод

Кстати, а как там с юникодом? serde_json умеет парсить его как в неэкранированном, так и в экранированном виде, например "