Как я ускорила парсинг строк в 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 1nline2" — корректный 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 не , и

• Двухпроходная: вычисляем 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] != ''; ++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] = '';
    }
    return len;
}

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

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

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

 

Мораль сей басни такова: векторизация — мощная штука, но применение векторизации к простой части при оставлении скалярной более сложную часть не привносит никаких улучшений. И если мы хотим улучшить код, придётся использовать 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 становился медленее, чем скалярный код. Мы решили, что регрессии на таких коротких строках, что в реальных данных встречаются редковато — это цена, которую мы готовы заплатить за скорость в других случаях.

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

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

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

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

Юникод

Кстати, а как там с юникодом? serde_json умеет парсить его как в неэкранированном, так и в экранированном виде, например "🥺" и "ud83eudd7a". Неэкранированный юникод парсить тривиально, а вот u-экраны добавляют боли и всё усложняют.

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

Для парсинга шестнадцатеричной цифры в число нужно перевести непересекающиеся интервалы '0'..='9', 'A'..='F', 'a'..='f' в 0..16. Можно написать с условиями:

match c {
    b'0'..=b'9' => c - b'0',
    b'A'..=b'F' => c - b'A' + 10,
    b'a'..=b'f' => c - b'a' + 10,
    _ => return Err(..),
}

...или без них, как в стандартной библиотеке Rust.

Но ничто не обойдёт LUT.

static HEX: [u8; 256] = {
    const __: u8 = 255; // not a hex digit
    [
        //   1   2   3   4   5   6   7   8   9   A   B   C   D   E   F
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 0
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 1
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 2
        00, 01, 02, 03, 04, 05, 06, 07, 08, 09, __, __, __, __, __, __, // 3
        __, 10, 11, 12, 13, 14, 15, __, __, __, __, __, __, __, __, __, // 4
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 5
        __, 10, 11, 12, 13, 14, 15, __, __, __, __, __, __, __, __, __, // 6
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 7
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 8
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // 9
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // A
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // B
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // C
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // D
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // E
        __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, __, // F
    ]
};

fn decode_hex_val(val: u8) -> Option<u16> {
    let n = HEX[val as usize] as u16;
    if n == 255 {
        None
    } else {
        Some(n)
    }
}

Такой подход использовался в serde_json раньше и работал достаточно неплохо (лучше, чем std: код в std должен уметь работать с разными основаниями, а serde_json нет). Эту функцию использовали примерно так:

let mut n = 0;
for _ in 0..4 {
    n = (n << 4) + decode_hex_val(self.slice[self.index])?;
    self.index += 1;
}

Как минимум 3 shl, 3 add, щепотка mov и cmp с 255 и заправлено условными прыжками. Можно лучше.

Начнем с убирания ? из цикла. Это вполне просто: вместо того, чтобы хранить HEX как массив [u8; 256], можно хранить его как [u32; 256], заменяя __ на u32::MAX. Ни в каком валидном экранированном символе нету старших битов, так что можно проверку на ошибку провести после цикла:

let mut n = 0;
for _ in 0..4 {
    n = (n << 4) + HEX[self.slice[self.index] as usize];
    self.index += 1;
}
ensure!(n >= 65536, "Invalid Unicode escape");
let n = n as u16;

Выжимаем максикум

Сэкономить на памяти и кеше заменой u32 на u16 на первый взгляд кажется невозможной задачей, так как u16::MAX = 0xFFFF в первой цифре перейдет в 0xFxxx, и отличить от настоящего символа будет невозможно.

Оказывается, это всё же достижимо. Юля (моя девушка, переводчица) придумала фокус. Заменим __ на u16::MAX, n << 4 на n.rotate_left(4), а сложение на побитовое ИЛИ:

let mut n = 0;
for _ in 0..4 {
    n = n.rotate_left(4) | HEX[self.slice[self.index] as usize];
    self.index += 1;
}

Если все цифры были корректными, то ничего не сломается, и n — верный кодпоинт. rol на x86 настолько же эффективен, как и shl, так что проблем с производительностью возникнуть не должно. В случае, если какой-то из символов оказался плохим, он "отравит" n и каждая последующая итерация будет удерживать значение 0xFFFF. U+FFFF — кодпоинт, не отвечающий ни за какой символ, а значит в реальных данных он будет редок, и мы можем позволить себе после цикла проверить n == 0xFFFF и перепроверить вычисления. Прикольно?

clueless.jpg

Пока я писала этот пост, я осознала, что это хрестоматийный пример оверинжиниринга. 0xFFFF превращается в 0xFxxx только если производить 16-битные вычисления. Но каменный век давно подошёл к концу, нынче можно использовать 32-битные числа. Будем хранить HEX как [i8; 256] и используем -1 для обозначения некорректной цифры. Тогда

let mut n = 0;
for _ in 0..4 {
    n = (n << 4) | HEX[self.slice[self.index] as usize] as i32;
    self.index += 1;
}

...вернёт неотрицательное число в случае успеха, и отрицательное — в случае ошибки. Может показаться, что as i32 занимает ресурсы процессора, но на x86 инструкция movsx умеет делать одновременно чтение и знаковое расширение, поэтому это не так.

В знаковых числах мне нравится то, что у большинства процессоров есть конкретная инструкция прыжка по знаковому биту. Вместо cmp r, imm; je label во многих случаях можно обойтись js label. На современных процессорах это не особо на производительность влияет, но, блин, красиво же!

Сдвиги

Сдвиги увеличивают latency. Плохо. Хочу без.

Эту проблему можно решить легко. Храним вместо одной таблицы две: HEX0 — HEX, но в виде [i16; 256], и HEX1 — HEX0 со значениями, сдвинутыми влево на 4 бита. Теперь и цикл можно с легкостью развернуть и получить итоговую реализацию.

fn decode_four_hex_digits(a: u8, b: u8, c: u8, d: u8) -> Option<u16> {
    let a = HEX1[a as usize] as i32;
    let b = HEX0[b as usize] as i32;
    let c = HEX1[c as usize] as i32;
    let d = HEX0[d as usize] as i32;

    let codepoint = ((a | b) << 8) | c | d;

    // A single sign bit check.
    if codepoint >= 0 {
        Some(codepoint as u16)
    } else {
        None
    }
}

Суммарно эти изменения ускорили парсинг "Войны и мира" Толстого на русском в виде JSON с 284 MB/s до 344 MB/s — на целых 21%.

Транскодинг

Препятствия

После последней оптимизации бутылочным горлышком парсинга юникод-строк оказалась перекодировка в UTF-8.

Весело, UTF-8 же задумывался как до боли простой. Напомню, что UTF-8 хранит кодпоинты одним из следующих способов:

• 1 байт: 0xxxxxxx

• 2 байта: 110xxxxx 10xxxxxx

• 3 байта: 1110xxxx 10xxxxxx 10xxxxxx

• 4 байта: 11110xxx 10xxxxxx 10xxxxxx 10xxxxxx

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

Стандартная библиотека Rust предоставляет энкодинг char в UTF-8 через, например, функцию char::encode_utf8, которая кладёт результат в предоставленный буфер. Есть, конечно, и ложка дёгтя. Сигнатура функции выглядит так:

fn encode_utf8(self, dst: &mut [u8]) -> &mut str;

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

Писалось это в предположении, что оптимизатор достаточно умный, чтобы выкинуть зануление. Это могло бы работать, но UTF-8 — кодировка с непостоянной длиной символа. Если занулено больше байт, чем encode_utf8 хочет, убирать зануление некорректно. Нужно занулить переменное число байт, но семантика такого действия оказывается для LLVM слишком сложной, и он это не оптимизирует.

 

В serde_json выкрутились:

scratch.extend_from_slice(c.encode_utf8(&mut [0u8; 4]).as_bytes());

[0u8; 4] — локальная переменная, так что занулять чуть больше байтов — не проблема, потому что aliasing analysis сильный. Ну, в теории.

На практике всё хуже. Vec::extend_from_slice нужно скопировать от одного до четырёх байт из локального буфера на кучу. LLVM с переменной длиной не справляется и тут, вызывая memcpy из libc. Зашибись!

 

Лучшим способом избежать memset и memcpy оказалась ручная генерация UTF-8. Алгоритмически там всё просто как два байта переслать, но придется добавить немного unsafe. Не лучшее решение, но пришлось смириться.

Вместе со всякими мелкими изменениями, производительность на "Войне и мире" увеличилась ещё на 9% до 374 MB/s.

Итоги

🥺

Мои действия ускорили serde_json на разных бенчмарках, напичканных строками, на 10%, 23% и 32%. Подобные JSON-ы часто встречаются на практике, поэтому такое улучшение принесёт пользу экосистеме в долгосрочной перспективе.