Какое-то время назад memorysafety.org объявил о конкурсе по повышению производительности rav1d — порта AV1-декодера dav1d на Rust.

Моя фамилия Равид, совсем как название декодера, поэтому я решил, что будет забавно попробовать (хоть я и, вероятно, не смогу участвовать в конкурсе).

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

Предыстория и методики

rav1d — это порт dav1d, созданный (1) обработкой dav1d при помощи c2rust, (2) внедрением в dav1d ассемблерно-оптимизированных функций и (3) модификацией кода, чтобы он больше соответствовал стандартам Rust и был безопаснее.

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

Чуть позже был объявлен конкурс со следующей отправной точкой:

Наш написанный на Rust декодер rav1d пока примерно на 5% медленнее, чем написанный на C декодер dav1d.

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

От нас не ожидают существенных улучшений, и может оказаться, что с некоторыми регрессиями трудно справиться (например, для LLVM функцию на Rust оптимизировать сложнее, чем версию на C), но попробовать стоит, в том числе и потому, что версия для aarch64 (моя среда), вероятно, менее оптимизирована, чем для x86_64.

Моя методика заключалась в следующем:

1. Использовать сэмплирующий профилировщик для создания снэпшотов двух прогонов с одними и теми же данными.

2. Использовать оптимизированные ассемблерные вызовы в качестве «якорей», потому что они должны полностью совпадать.

3. Сравнить версии функций на Rust и на C, и в случае достаточно больших различий в функции исследовать её.

Отправная точка

Для начала нам нужно выполнить сборку и сравнить производительность локально (при помощи hyperfine и файлов образцов, указанных в правилах конкурса и в CI декодера rav1d).

Чтобы не усложнять, мы будем использовать однопоточную версию (--threads 1).

Для rav1d:

$ git clone git@github.com:memorysafety/rav1d.git && cd rav1d && git log -n1
commit a654c1e82adb2d9a33ae50d2a82a7a747102cbb6
$ rustc --version --verbose # set by rust-toolchain.toml
rustc 1.88.0-nightly (b45dd71d1 2025-04-30)
...
LLVM version: 20.1.2
$ cargo build --release
    Finished `release` profile [optimized] target(s) in ..
$ hyperfine --warmup 2 "target/release/dav1d -q -i Chimera-AV1-8bit-1920x1080-6736kbps.ivf -o /dev/null --threads 1"
Benchmark 1: target/release/dav1d -q -i Chimera-AV1-8bit-1920x1080-6736kbps.ivf -o /dev/null --threads 1
  Time (mean ± σ):     73.914 s ±  0.151 s    [User: 73.295 s, System: 0.279 s]
  Range (min … max):   73.770 s … 74.132 s    10 runs

Для dav1d:

$ git clone https://code.videolan.org/videolan/dav1d.git && cd dav1d && git checkout 1.5.1
$ brew install llvm@20 && export CC=clang; $CC --version
Homebrew clang version 20.1.4
$ meson setup build "-Dbitdepths=['8','16']"
$ bear -- ninja -C build tools/dav1d
...
[88/88] Linking target tools/dav1d
$ hyperfine --warmup 2 "build/tools/dav1d -q -i Chimera-AV1-8bit-1920x1080-6736kbps.ivf -o /dev/null --threads 1"
Benchmark 1: build/tools/dav1d -q -i Chimera-AV1-8bit-1920x1080-6736kbps.ivf -o /dev/null --threads 1
  Time (mean ± σ):     67.912 s ±  0.541 s    [User: 67.208 s, System: 0.282 s]
  Range (min … max):   66.933 s … 68.948 s    10 runs

То есть в случае обработки файла-образца rav1d примерно на 9% (6 секунд) медленнее, чем dav1d  (по крайней мере на чипе M3).

(В идеале clang и rustc должны использовать одну и ту же версию LLVM, но, вероятно, разница в версиях патчей приемлема.)
(Измерения проведены на MacBook Air M3 с восьмью ядрами.)

Профилирование

Я использовал сэмплирующий профилировщик samply, с которым сейчас обычно работаю:

./dav1d $ sudo samply record ./build/tools/dav1d -q -i /Chimera-AV1-8bit-1920x1080-6736kbps.ivf -o /dev/null --threads 1
./rav1d $ sudo samply record ./target/release/dav1d -q -i /Chimera-AV1-8bit-1920x1080-6736kbps.ivf -o /dev/null --threads 1

(Двоичный файл Rust тоже называется dav1d, что немного странно.)

По умолчанию samply использует частоту 1000 Гц, то есть (например), на любой diff пятисот сэмплов в функции придётся примерно 0,5 секунд разницы времени выполнения.

Обычно сузить поиск до интересных вариантов можно, начав с режиме «инвертированного стека», но в этот раз мы хотим сосредоточиться на якорях, которые должны совпадать: на ассемблерных функциях.

Мы можем просматривать полные снэпшоты профилировщика онлайн в Firefox Profiler (dav1d, rav1d); ниже приведены релевантные отфильтрованные блоки.

Вот версия dav1d (на C) (общее количество сэмплов: около 69500):

А вот версия rav1d (на Rust) (общее количество сэмплов: около 75150):

Посмотрите на выделенные функции dav1d_cdef_brow_8bpc и rav1d_cdef_brow.
Total — это количество сэмплов, при котором функция находилась «в любом месте стека», что подразумевает и все вызываемые ею «дочерние» функции. Self — это количество сэмплов, в которых она была исполняемой функций, то есть без учёта количества дочерних сэмплов.

Между dav1d и rav1d есть небольшое различие: хотя расширение _neon обозначает относящиеся к Arm ассемблерные функции, общие для обоих двоичных файлов, мы видим, что:

1. dav1d вызывает cdef_filter_8x8_neon и cdef_filter_4x4_neon, и каждая из них размещает релевантные ассемблерные функции (соответственно, версию 8 или 4).

2. rav1d вызывает cdef_filter_neon_erased, обрабатывающую размещение всех ассемблерных функций.

Также мы видим, что в обоих снэпшотах количество сэмплов cdef_filter8_pri_sec_edged_8bpc_neon почти идентично, и это значит, что мы находимся на верном пути.

Пока закроем глаза на функцию cdef_filter4_pri_edged_8bpc_neon, сэмплы которой не совпадают.

Это значит, что (А) количество сэмплов Self для dav1d_cdef_brow_8bpc должно соответствовать количеству сэмплов rav1d_cdef_brow, и (Б) сумма количества сэмплов Self  функций cdef_filter_{8x8,4x4}_neon должна совпадать с количеством сэмплов Self для cdef_filter_neon_erased.

Теперь мы видим кое-что любопытное: во второй части суммированное количество сэмплов Self cdef_filter_{8x8,4x4}_neon равно примерно 400 сэмплам, а у rav1d’s cdef_filter_neon_erased — почти 670 сэмплов. Также мы видим, что dav1d_cdef_brow_8bpc имеет 1790 сэмплов, а rav1d_cdef_brow — 2350 сэмплов.

В целом эта разница составляет примерно 1% общего времени выполнения rav1d!

Если перейти к реализации cdef_filter_neon_erased, не обращая внимания на преобразование указателей при помощи .cast(), то можно найти лишь одну «большую работу», не относящуюся к механизмам вызова ассемблерного кода:

#[deny(unsafe_op_in_unsafe_fn)]
pub unsafe extern "C" fn cdef_filter_neon_erased<
    BD: BitDepth,
    const W: usize,
    const H: usize,
    const TMP_STRIDE: usize,
    const TMP_LEN: usize,
>(
    // .. вырезано ..
) {
    use crate::src::align::Align16;

    // .. вырезано ..

    let mut tmp_buf = Align16([0u16; TMP_LEN]);
    let tmp = &mut tmp_buf.0[2 * TMP_STRIDE + 8..];
    
    padding::Fn::neon::<BD, W>().call::<BD>(tmp, dst, stride, left, top, bottom, H, edges);
    filter::Fn::neon::<BD, W>().call(dst, stride, tmp, pri_strength, sec_strength, dir, damping, H, edges, bd);
}

TMP_LEN может быть 12 * 16 + 8 = 200 или 12 * 8 + 8 = 104, так что в худшем случае tmp_buf = [u16; 200]. Большой объём обнуляемой памяти для временного буфера!

Что здесь делает dav1d?

#define DEFINE_FILTER(w, h, tmp_stride)                                      \
static void                                                                  \
cdef_filter_##w##x##h##_neon(/* .. вырезано .. */)                               \
{                                                                            \
    ALIGN_STK_16(uint16_t, tmp_buf, 12 * tmp_stride + 8,);                   \
    uint16_t *tmp = tmp_buf + 2 * tmp_stride + 8;                            \
    BF(dav1d_cdef_padding##w, neon)(tmp, dst, stride,                        \
                                    left, top, bottom, h, edges);            \
    BF(dav1d_cdef_filter##w, neon)(dst, stride, tmp, pri_strength,           \
                                   sec_strength, dir, damping, h, edges      \
                                   HIGHBD_TAIL_SUFFIX);                      \
}

DEFINE_FILTER(8, 8, 16)
DEFINE_FILTER(4, 8, 8)
DEFINE_FILTER(4, 4, 8)

Спустя несколько расширений макросов мы находим uint16_t tmp_buf[200] __attribute__((aligned(16)));

Это означает, что tmp_buf не инициализируется функциями cdef_filter_{8x8,4x4}_neon: вместо этого он используется как место для записи ассемблерной функции padding, а позже в том же виде для ассемблерной функции filter. Похоже, компилятор не знает, как устранить эту инициализацию; чтобы подробнее изучить это, можно использовать --emit=llvm-ir:

$ RUSTFLAGS="--emit=llvm-ir" cargo build --release --target aarch64-apple-darwin

; rav1d::src::cdef::neon::cdef_filter_neon_erased
; Function Attrs: nounwind
define internal void @_ZN5rav1d3src4cdef4neon23cdef_filter_neon_erased17h7e4dbe8ecff68724E(ptr noundef %dst, i64 noundef %stride, ptr noundef %left, ptr noundef %top, ptr noundef %bottom, i32 noundef %pri_strength, i32 noundef %sec_strength, i32 noundef %dir, i32 noundef %damping, i32 noundef %edges, i32 noundef %bitdepth_max, ptr nocapture readnone %_dst, ptr nocapture readnone %_top, ptr nocapture readnone %_bottom) unnamed_addr #1 {
start:
  %tmp_buf = alloca [400 x i8], align 16
  call void @llvm.lifetime.start.p0(i64 400, ptr nonnull %tmp_buf)
  call void @llvm.memset.p0.i64(ptr noundef nonnull align 16 dereferenceable(400) %tmp_buf, i8 0, i64 400, i1 false)
  %_37 = getelementptr inbounds nuw i8, ptr %tmp_buf, i64 80
  call void @dav1d_cdef_padding8_16bpc_neon(ptr noundef nonnull %_37, ptr noundef %dst, i64 noundef %stride, ptr noundef %left, ptr noundef %top, ptr noundef %bottom, i32 noundef 8, i32 noundef %edges) #121
  %edges2.i = zext i32 %edges to i64
  %_0.i.i.i.i = and i32 %bitdepth_max, 65535
  call void @dav1d_cdef_filter8_16bpc_neon(ptr noundef %dst, i64 noundef %stride, ptr noundef nonnull readonly align 2 %_37, i32 noundef %pri_strength, i32 noundef %sec_strength, i32 noundef %dir, i32 noundef %damping, i32 noundef 8, i64 noundef %edges2.i, i32 noundef %_0.i.i.i.i) #121
  call void @llvm.lifetime.end.p0(i64 400, ptr nonnull %tmp_buf)
  ret void
}

Устраняем необязательное обнуление буферов при помощи MaybeUninit

На самом деле, это должно быть достаточно просто! Как раз на такой случай в Rust есть std::mem::MaybeUninit:

-let mut tmp_buf = Align16([0u16; TMP_LEN])
+let mut tmp_buf = Align16([MaybeUninit::<u16>::uninit(); TMP_LEN]);

Мы по-прежнему можем безопасно брать sub-slice (&mut tmp_buf.0[2 * TMP_STRIDE + 8..]), но нам нужно будет обновить сигнатуры внутренних функций, чтобы они использовали новый тип (tmp: *mut MaybeUninit<u16>, tmp: &[MaybeUninit<u16>]).

Так как использовавший их код всё равно был небезопасным, нам не нужно добавлять новые блоки unsafe; достаточно убедиться, что старый код не изменился (относительно dav1d), тогда можно быть уверенными, что буфер обнуляется.

Ранее cdef_filter_neon_erased имела 670 сэмплов Self. Запустив профилировщик повторно, мы получим новый снэпшот:

Всего 274 сэмплов! Чуть меньше, чем количество сэмплов Self  dav1d’s cdef_filter_{8x8,4x4}_neon.

Возможно, это не единственное место, где время впустую тратится на обнуление буферов? Поискав другие большие буферы Align16, я обнаружил удачную находку:

pub(crate) fn rav1d_cdef_brow<BD: BitDepth>(/* .. вырезано ..*/)
{
    // .. вырезано ..

    for by in (by_start..by_end).step_by(2) {
        // .. вырезано ..
        let mut lr_bak =
            Align16([[[[0.into(); 2 /* x */]; 8 /* y */]; 3 /* plane */ ]; 2 /* idx */]);
        
        // .. вырезано ..
    }
}

Соответствующий код из dav1d тоже не инициализирует этот буфер. Здесь переключиться на MaybeUninit будет сложнее, но мы всё равно можем обеспечить скромное улучшение: если перенести инициализацию lr_bak на верхний уровень, то инициализацию достаточно будет выполнять только один раз!

pub(crate) fn rav1d_cdef_brow<BD: BitDepth>(/* .. вырезано ..*/)
{
    // .. вырезано ..
+   let mut lr_bak =
+       Align16([[[[0.into(); 2 /* x */]; 8 /* y */]; 3 /* plane */ ]; 2 /* idx */]);
        
    for by in (by_start..by_end).step_by(2) {
        // .. вырезано ..
-       let mut lr_bak =
-           Align16([[[[0.into(); 2 /* x */]; 8 /* y */]; 3 /* plane */ ]; 2 /* idx */]);
        
        // .. вырезано ..
    }
}

Так как dav1d всё равно никогда его не инициализирует, логично, что любые данные, читаемые из этого буфера, были заранее записаны с валидным значением. Экономия здесь очень мала, но ценна каждая копейка!

Запустив полный бенчмарк, мы получили хорошее увеличение скорости по сравнению с исходными 73.914 s ± 0.151 s:

$ hyperfine --warmup 2 "target/release/dav1d -q -i Chimera-AV1-8bit-1920x1080-6736kbps.ivf -o /dev/null --threads 1"
Benchmark 1: target/release/dav1d -q -i Chimera-AV1-8bit-1920x1080-6736kbps.ivf -o /dev/null --threads 1
  Time (mean ± σ):     72.644 s ±  0.250 s    [User: 72.023 s, System: 0.239 s]
  Range (min … max):   72.281 s … 73.098 s    10 runs

До dav1d с его 67.912 s ± 0.541 s нам всё ещё далеко, но улучшение на 1,2 секунды (1,5%) от общего времени выполнения — отличное начало, покрывающее примерно 20% разницы в производительности между двумя версиями декодера.

Снова профилирование, но инвертированное

Давайте перезагрузим результаты работы профилировщика из начала статьи, но на этот раз используем режим «инвертированный стек».

Существует несколько вариантов исследования способов оптимизации, но моё внимание привлекла функция add_temporal_candidate: разница между версиями на Rust и C достаточно существенна (примерно 400 сэмплов, около 0,5 секунды), а сама функция выглядит безобидно: примерно пятьдесят строк if и for с несколькими вызовами коротких вспомогательных функций.

Чтобы разобраться, куда девается производительность, можно попробовать перекомпилировать rav1d с отладочными символами. Удобно, что в Cargo.toml rav1d определён [profile.release-with-debug]; это позволяет нам выполнить следующее:

$ cargo build --profile=release-with-debug
$ sudo samply record target/release-with-debug/dav1d ...

Получаемое нами немного отличается от предыдущего: профиль release-with-debug будет не столь оптимизированным, а вызовы маленьких функций кажутся больше, чем на самом деле; зато мы получим построчную разбивку сэмплов функции, что должно направить нас по нужному пути.

Если немного проскроллить, можно заметить, что строки if cand.mv.mv[0] == mv { и if cand.mv == mvp { суммарно покрывают 600 сэмплов!

Давайте поднимем определение mv: Mv наверх:

#[derive(Clone, Copy, PartialEq, Eq, Default, FromZeroes, FromBytes, AsBytes)]
#[repr(C)]
pub struct Mv {
    pub y: i16,
    pub x: i16,
}

Хм. Как этот код может быть медленным? Это просто #[derive(PartialEq)].

Ещё подозрительнее то, что версия dav1d немного отличается, в ней для тех же сравнений используется mvstack[n].mv.n == mvp.n. Но что такое n? Взглянув на определение mv в dav1d, мы обнаружим следующее:

typedef union mv {
    struct {
        int16_t y, x;
    };
    uint32_t n;
} mv;

Похоже, авторы dav1d знали, что сравнение двух i16 может быть медленным, поэтому при сравнении двух mv они обращаются с ними, как с u32s.

Заменяем равенство полей побайтовым равенством, что лучше оптимизируется

Может ли это быть проблемой?

Определение Mv в качестве union имеет в Rust очень большой недостаток: из-за этого доступ к любым полям union становится unsafe, что «заразит» все случаи использования Mv, а это противоположно тому, к чему мы обычно стремимся в Rust (пытаемся инкапсулировать небезопасность в безопасный API).

К счастью, есть и другой способ: мы можем использовать transmute для повторной интерпретации Mv в качестве u32, и применить его для реализации PartialEq.

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

Очевидно, версия с transmute лучше, но можем ли мы избежать блока unsafe?¹

Оказывается, крейт zerocopy может статически верифицировать требования безопасности struct для представления в виде &[u8], что позволяет нам написать следующее:

use zerocopy::{AsBytes, FromBytes, FromZeroes};

#[derive(Clone, Copy, Eq, Default, FromZeroes, FromBytes, AsBytes)]
#[repr(C)]
pub struct Mv {
    pub y: i16,
    pub x: i16,
}

impl PartialEq for Mv {
    #[inline(always)]
    fn eq(&self, other: &Self) -> bool {
        self.as_bytes() == other.as_bytes()
    }
}

При этом мы получим тот же (оптимизированный) ассемблерный код, который мы видели при использовании transmute.

Реализовав подобные оптимизации для RefMvs{Mv,Ref}Pair, можно ещё раз запустить бенчмарк:

$ hyperfine --warmup 2 "target/release/dav1d -q -i Chimera-AV1-8bit-1920x1080-6736kbps.ivf -o /dev/null --threads 1"
Benchmark 1: target/release/dav1d -q -i Chimera-AV1-8bit-1920x1080-6736kbps.ivf -o /dev/null --threads 1
  Time (mean ± σ):     72.182 s ±  0.289 s    [User: 71.501 s, System: 0.242 s]
  Range (min … max):   71.850 s … 72.722 s    10 runs

Получаем улучшение ещё на 0,5 секунды по сравнению с предыдущим результатом (72.644 s ± 0.250 s), или на 2,3% по сравнению с отправной точкой (73.914 s ± 0.151 s).

Мы всего на 4,2 секунды отстаём от dav1d с его 67.912 s ± 0.541 s, то есть мы покрыли около 30% разницы в производительности.

Вы можете задаться вопросом, почему стандартная реализация PartialEq приводит к генерации плохого кода: комментарий в PR, добавляющем эти impl, указывает на Rust issue 140167, который связан именно с этим типом проблем.

Если рассмотреть случай C, при использовании struct { int16_t y, x; } можно инициализировать только y, оставив x неинициализированным. Если равенство проверяется при помощи this.y == other.y && this.x == other.x и все y разные, никаких UB мы не получим.

Следовательно, нельзя оптимизировать это до единственной загрузки в память и сравнения, если только код не может гарантировать, что все поля всегда инициализированы. Однако процитируем комментарий @hanna-kruppe про этот issue:

Это не просто упущенная возможность оптимизации. Хотя загрузка второго поля не может загрузить poison/undef, это свойство имеет control dependency. ..
Устранить эту проблему сложно: не думаю, что в LLVM есть способ выразить указание «загрузка через этот указатель всегда считывает инициализированные байты».

Подведём итог

При помощи нескольких снэпшотов профилировщика samply мы сравнили работу rav1d и dav1d с одним и тем же файлом данных, отметили разницу в 6 секунд (9%) и нашли две простые возможности для оптимизации:

1. Устранив затратную инициализацию нулями на горячем пути выполнения кода, специфичного для Arm (PR), снизив таким образом время выполнения на 1,2 секунды (-1,6%).

2. Заменив стандартные impl PartialEq небольших числовых struct оптимизированной версией, которая реинтерпретирует их, как байты (PR), снизив время выполнения на 0,5 секунды (-0,7%).

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

Мейнтейнеры проекта rav1d быстро отреагировали на мои PR, помогли сделать их более корректными и оптимальными (огромная благодарность @kkysen).

Между двумя реализациями по-прежнему существует разница примерно в 6%, поэтому можно найти ещё множество оптимизаций. Подозреваю, такая методика сравнения снэпшотов профилировщика dav1d и rav1d позволит обнаружить хотя бы некоторые из них.

Можете попробовать сделать это сами! Возможно, когда-то rav1d станет быстрее, чем dav1d...

---

1. Примечание о безопасности: хотя функция use_transmute безопасна, здесь есть тонкость: так как mem::align_of::<Mv> != mem::align_of::<u32>(), мы обязаны разыменовать &Mv заранее. Попробуйте запустить Miri в Playground.