SIMD: векторные дорожки для кодеков и константновременного кода

SIMD (single instruction, multiple data) — это когда одна инструкция процессора обрабатывает сразу несколько значений. 128-битный SIMD-регистр вмещает шестнадцать байтов, восемь 16-битных или четыре 32-битных целых, и одна инструкция складывает их все попарно. Каждая такая ячейка называется дорожкой (lane), а тип вроде u8x16 называет тип дорожки и их число: шестнадцать дорожек u8. Видеокодеки, фильтры изображений и шифры проводят бо́льшую часть времени в циклах, которые SIMD ускоряет в несколько раз. На x86 наборы таких инструкций называются SSE и AVX2 (256-битные регистры), на ARM — NEON.

В kola два семейства целочисленных SIMD-типов. В месте вызова они выглядят одинаково, а подчиняются противоположным правилам. Дорожки ct отказываются почти от всех операций, чтобы ими можно было пользоваться в коде с константным временем. Медиадорожки, наоборот, умеют всё зависящее от данных, что нужно кодеку, и поэтому в код с константным временем не допускаются. Поменяться местами им не даёт чекер. В главе разобраны оба семейства, точно зафиксированные части их семантики, небольшое семейство с плавающей точкой и то, как их переводят в машинный код бэкенды. Код с константным временем — это код, время работы и обращения к памяти которого не зависят от секретных значений, так что замеры времени ничего не выдают; ct fn — функция, у которой это свойство проверяет компилятор (Константное время (ct)).

Два семейства, одна форма

Векторный тип — это значение. Он копируется, как массив фиксированного размера, и никогда не бывает дескриптором. Любой векторный тип создаётся и читается четырьмя одинаковыми способами: T::splat(x) заполняет все дорожки одним скаляром, T::of(a, b, ...) принимает по скаляру на дорожку, а T::load(slice, off) и T::store(slice, off, v) переносят вектор из среза его типа элементов и обратно (&[u8] для u8x16, &[u32] для u32x4 и так далее). Новой грамматики здесь нет. Векторный тип — обычный тип, а T::op(...) — обычный вызов, поэтому чекер и может применять к двум семействам разные правила без отдельного синтаксиса.

Дорожки ct — это u32x4 и его 256-битный восьмидорожечный близнец u32x8. Они поддерживают побитовые & | ^, арифметику с переносом по модулю +% -% *% и сдвиги <</>> на литеральную величину — и больше ничего: ни сдвига на величину, известную только во время выполнения, ни ==, ни арифметики с трапом. Ни одна операция из этого набора не может ветвиться или выбирать адрес памяти по значению дорожки. Набор константновременной по построению, и именно поэтому эти два типа разрешены внутри ct fn.

Медиадорожки — это i16x8 (видео и DCT), u8x16 (пиксели) и i32x4 (аккумуляторы преобразований), а также их 256-битные близнецы i16x16, u8x32 и i32x8. Они дают сложение и вычитание с насыщением (результат, который переполнился бы, останавливается на наибольшем или наименьшем значении вместо того, чтобы перейти через границу, — ровно то, что нужно для яркости пикселя), минимум и максимум, сравнения, возвращающие маску (вектор, дорожки которого состоят из одних единиц там, где сравнение выполнилось, и из нулей в остальных), select, перестановку байтов shuffle, упаковку, распаковку, расширение, попарное умножение со сложением madd и единый для всего вектора сдвиг на величину времени выполнения. Большинство этих операций зависят от данных по самой своей сути (перестановка байтов — это поиск по таблице), поэтому всё семейство не-ct. Побитовых операций и операторов сдвига у медиадорожек нет: они принадлежат дорожкам ct.

Вот обе формы рядом: сложение пикселей с насыщением на медиадорожках и сдвиг на литерал на дорожках ct.

fn media_construct(a: &[u8], b: &[u8], out: &mut [u8]) !{} {
    let va = u8x16::splat(a[0]);
    let vb = u8x16::load(b, 0);
    u8x16::store(out, 0, u8x16::sat_add(va, vb));
}

fn ct_construct(a: &[u32], out: &mut [u32]) !{} {
    let v = u32x4::of(a[0], a[1], a[2], a[3]);
    u32x4::store(out, 0, v << 1u32);
}

test both_families_construct_and_store {
    let a = [250u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8,
             0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8];
    let b = [10u8, 0u8, 30u8, 40u8, 0u8, 0u8, 0u8, 0u8,
             0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8];
    let mut mo = [0u8; 16];
    media_construct(&a, &b, &mut mo);
    assert mo[0] == 255u8;   // 250 + 10 насыщается до 255
    assert mo[1] == 250u8;   // 250 + 0

    let au = [1u32, 2u32, 3u32, 4u32];
    let mut co = [0u32; 4];
    ct_construct(&au, &mut co);
    assert co[0] == 2u32 && co[3] == 8u32;
}

Узкий набор операций ct входит в сам тип и действует везде, а не только внутри ct fn. Сдвиг u32x4 на величину времени выполнения запрещён и в обычной функции:

fn bad_shift(v: &[u32], n: u32, out: &mut [u32]) !{} {
    let va = u32x4::load(v, 0);
    u32x4::store(out, 0, va << n);
}

error[card §9]: vector shift amount must be a literal (shift-by-immediate only)
  --> bs.kl:3:32
   |
 3 |     u32x4::store(out, 0, va << n);
   |                                ^
   = help: build a variable shift from fixed-amount shifts if needed

В обратную сторону: медиаоперация внутри ct fn запрещена по §6b. Чекер сообщает, что ct fn вызывает функцию без ct (load, sat_add, store), и называет каждый такой вызов.

Константновременное ядро и ядро кодека хотят от одних и тех же 128 бит регистра противоположного. Первому нужен набор операций, которым данные атакующего не могут управлять. Второму нужны как раз управляемые данными операции: переставить строку по вектору движения, выбрать дорожку по результату сравнения. Будь тип один на обе задачи, пришлось бы выбирать между ct-типом, который больше не ct, и типом для кодека без перестановки. Когда типов два, правка, которая тянется к медиаоперации в бывшем константновременном коде, останавливается прямо на этой строке.

Правило среза для load и store

T::load и T::store принимают срез, а локальный массив фиксированного размера в самом вызове интринсика в срез не приводится. u8x16::load(&a, 0) для локального [u8; 16] компилятор отвергнет с сообщением u8x16::load needs a &[u8] slice. Массив нужно передать через обычный параметр функции типа &[u8] или &mut [u8], где работает привычное приведение массива к срезу. Так написаны все ядра этой главы: маленькая функция над срезами, которую вызывает test или main, владеющий массивами.

Медиаоперации

У 128-битных медиадорожек общие создание, доступ к памяти и арифметика по модулю +% -% *%. Кроме того, у них есть:

Сдвиг на величину времени выполнения принимает беззнаковое число (usize, u64 и так далее); i32 будет отвергнут, потому что у отрицательной величины сдвига нет смысла. Если величина не меньше ширины дорожки, shl даёт 0, а shr заполняет дорожку знаковым битом. Сдвига на величину времени выполнения для u8x16 нет намеренно (в x86 нет сдвига упакованных байтов), как нет и сдвига на свою величину в каждой дорожке. u8x16::store_low32(slice, off, v) записывает только младшие четыре байта.

Зафиксированная семантика

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

Нагляднее всего это видно на u8x16::sad, сумме абсолютных разностей, которой на x86 соответствует psadbw. Видеокодер экономит место, описывая блок пикселей как «похожий вон на тот блок из предыдущего кадра, чуть сдвинутый», а чтобы найти лучшее совпадение, сравнивает блоки-кандидаты по сумме абсолютных разностей их пикселей. Поэтому sad — самая часто выполняемая операция кодера. Можно ожидать, что она вернёт одно число. Она возвращает i32x4: сумму разностей младших восьми байтов в дорожке 0, старших восьми — в дорожке 2 и нули в дорожках 1 и 3. kola сохраняет ровно эту раскладку. Вектор вместо скаляра тоже выбран намеренно: кодер накапливает много строк простыми сложениями векторов и сводит результат к скаляру один раз в конце, а не платит за перенос из вектора в скаляр на каждой строке. 256-битный u8x32::sad — это две независимые 128-битные половины, так что группа g всегда попадает в дорожку 2g.

fn sad_row16(a: &[u8], b: &[u8], lanes: &mut [i32]) !{} {
    i32x4::store(lanes, 0, u8x16::sad(u8x16::load(a, 0), u8x16::load(b, 0)));
}

test sad_matches_hand_computed_value {
    let a = [10u8, 10u8, 10u8, 10u8, 10u8, 10u8, 10u8, 10u8,
             20u8, 20u8, 20u8, 20u8, 20u8, 20u8, 20u8, 20u8];
    let b = [7u8, 7u8, 7u8, 7u8, 7u8, 7u8, 7u8, 7u8,
             15u8, 15u8, 15u8, 15u8, 15u8, 15u8, 15u8, 15u8];
    let mut lanes = [0i32; 4];
    sad_row16(&a, &b, &mut lanes);
    assert lanes[0] == 24;   // 8 дорожек * |10 - 7|
    assert lanes[1] == 0;    // всегда ноль
    assert lanes[2] == 40;   // 8 дорожек * |20 - 15|
    assert lanes[3] == 0;    // всегда ноль
    sad_row16(&a, &a, &mut lanes);
    assert lanes[0] == 0 && lanes[2] == 0;
}

Ещё три фиксации закрывают остальное ядро. abs от минимального значения дорожки (i16 -32768 или i32 -2147483648) по модулю возвращает само это значение, без трапа и без насыщения, как pabsw/pabsd и vabsq в NEON. avg на u8x16 и u8x32 — это среднее с округлением (a + b + 1) >> 1, вычисленное в расширенной разрядности, так что переполниться оно не может: из 255 и 255 получается 255, а не 0. Именно так округляет среднее бипредсказание в H.264, а простое (a + b) / 2 округляет вниз и восстанавливает другую картинку. hadd на i16x8 и i32x4 складывает соседние пары с переносом по модулю: пары a попадают в младшую половину, пары b — в старшую.

fn abs16(v: &[i16], out: &mut [i16]) !{} {
    i16x8::store(out, 0, i16x8::abs(i16x8::load(v, 0)));
}

fn avg8(a: &[u8], b: &[u8], out: &mut [u8]) !{} {
    u8x16::store(out, 0, u8x16::avg(u8x16::load(a, 0), u8x16::load(b, 0)));
}

test abs_of_int_min_wraps {
    let mut src = [0i16; 8];
    src[0] = -32768i16;
    let mut out = [0i16; 8];
    abs16(&src, &mut out);
    assert out[0] == -32768i16;   // по модулю; ни трапа, ни насыщения
}

test avg_rounds_up {
    let mut a = [0u8; 16];
    let mut b = [1u8; 16];
    a[15] = 255u8;
    b[15] = 255u8;
    let mut out = [0u8; 16];
    avg8(&a, &b, &mut out);
    assert out[0] == 1u8;      // (0 + 1 + 1) >> 1 округляется вверх
    assert out[15] == 255u8;   // 255 и 255 не переполняются
}

256-битные близнецы и правило половин

u8x32, i16x16 и i32x8 предлагают операции тех же форм над 32, 16 и 8 дорожками, а сверх того i16x16::madd, i32x8::mul_hi, упаковку и T::join(lo, hi), который собирает 256-битный вектор из двух 128-битных половин соответствующего типа.

Одно правило здесь удивляет тех, кто считает, что 256 бит — это «вдвое больше дорожек при том же поведении». unpack_lo, unpack_hi и упаковка с насыщением работают по 128-битным половинам: каждая половина 32-байтного вектора распаковывается или упаковывается независимо, как это делают vpunpck* и vpack* на настоящем железе, а не как одно линейное чередование 32 байтов. Расширение устроено наоборот, линейно: widen_lo_i16 берёт дорожки источника с 0 по 15 по порядку, а widen_hi_i16 — с 16 по 31, как vpmovzxbw. Если вы переносите ядро на AVX2 вручную, важны оба правила.

fn unpack32(a: &[u8], b: &[u8], lo: &mut [u8], hi: &mut [u8]) !{} {
    let va = u8x32::load(a, 0);
    let vb = u8x32::load(b, 0);
    u8x32::store(lo, 0, u8x32::unpack_lo(va, vb));
    u8x32::store(hi, 0, u8x32::unpack_hi(va, vb));
}

fn widen32(a: &[u8], lo: &mut [i16], hi: &mut [i16]) !{} {
    let v = u8x32::load(a, 0);
    i16x16::store(lo, 0, u8x32::widen_lo_i16(v));
    i16x16::store(hi, 0, u8x32::widen_hi_i16(v));
}

test unpack_is_per_128_bit_half {
    let mut a = [0u8; 32];
    let mut b = [0u8; 32];
    let mut i: usize = 0;
    while i < 32 {
        a[i] = narrow[u8](i).unwrap();
        b[i] = narrow[u8](100 + i).unwrap();
        i = i + 1;
    }
    let mut lo = [0u8; 32];
    let mut hi = [0u8; 32];
    unpack32(&a, &b, &mut lo, &mut hi);
    assert lo[0] == 0u8 && lo[1] == 100u8;
    assert lo[14] == 7u8 && lo[15] == 107u8;
    // вторая половина начинается заново с a[16] и b[16]; линейное
    // чередование продолжилось бы с a[8] и b[8]
    assert lo[16] == 16u8 && lo[17] == 116u8;
}

test widen_is_linear {
    let mut a = [0u8; 32];
    let mut i: usize = 0;
    while i < 32 {
        a[i] = narrow[u8](i).unwrap();
        i = i + 1;
    }
    let mut lo = [0i16; 16];
    let mut hi = [0i16; 16];
    widen32(&a, &mut lo, &mut hi);
    assert lo[0] == 0i16 && lo[15] == 15i16;
    assert hi[0] == 16i16 && hi[15] == 31i16;
}

Дорожки с плавающей точкой

f64x2 и f64x4 содержат две и четыре дорожки IEEE binary64. Как и скалярный f64, они вне ct. У них те же интринсики создания и доступа к памяти над &[f64], поэлементные + - * /, настоящий слитный T::fma(a, b, c) с одним округлением и T::hsum(v), который складывает в фиксированном порядке (v[0] + v[1] для двух дорожек, (v[0] + v[2]) + (v[1] + v[3]) для четырёх), так что результат на всех бэкендах одинаков. Как и со скалярным f64, операция, результатом которой становится NaN, вызывает трап, а загрузка и запись переносят уже хранящийся NaN без изменений.

fn square_plus_one(a: &[f64], out: &mut [f64]) !{} {
    let v = f64x2::load(a, 0);
    f64x2::store(out, 0, f64x2::fma(v, v, f64x2::splat(1.0)));
}

test fused_multiply_add {
    let a = [2.0, 3.0];
    let mut o = [0.0, 0.0];
    square_plus_one(&a, &mut o);
    assert o[0] == 5.0 && o[1] == 10.0;
    assert f64x2::hsum(f64x2::of(1.5, 2.5)) == 4.0;
}

Бэкенды

Для всех векторных типов эталон — интерпретатор. Каждый тест выше проходит под kola test, который запускает его в интерпретаторе и через C-бэкенд и сравнивает результаты файл за файлом; расхождение по определению считается ошибкой компилятора. Каждое зафиксированное правило этой главы сначала было записано как поведение интерпретатора, а затем воспроизведено на x86-64 SSE и AArch64 NEON.

• C-бэкенд (kola build --native, kola test --native) переводит в машинный код все векторные типы. Каждый из них — простая структура в kola.h, а каждая операция — вспомогательная функция, которая использует интринсики SSE или NEON, если они есть на целевой платформе, и переносимый скалярный путь в остальных случаях, совпадающий с интерпретатором до байта. Настоящими векторами AVX2 256-битные дорожки становятся только при сборке с --cpu x86-64-v3, и тогда полученный бинарник требует процессор этого уровня. На базовом уровне по умолчанию они идут запасным путём.

• Собственный бэкенд, которому не нужен компилятор C (kola test --own и kola build --link, который пробует его первым), тоже переводит все векторные типы и выполняет каждую 256-битную операцию как две 128-битные половины на SSE2 или NEON. Покрывает он пока не все операции: avg, например, отвергается по имени при сборке, и запасным вариантом служит kola test --native.

Векторные типы могут встречаться в сигнатурах функций — как параметры, возвращаемые значения, элементы кортежей и поля структур; ABI для них задаёт структура из kola.h.

Интринсики и аппаратные реализации

Ассемблерных вставок и FFI в kola нет. Быстрый код берётся из двух источников. Первый — именованные интринсики вроде описанных выше: у каждого точное определение, которое реализует интерпретатор и которому обязан соответствовать каждый бэкенд. Когда ядро кодека работает медленно, сначала стоит спросить, не отсутствует ли нужный интринсик: цикл SAD, написанный скалярным кодом, не скажет, нужен вам ассемблер или просто psadbw.

Второй источник — аппаратная реализация (hardware leg): написанная вручную нативная реализация одной функции, которая используется только на процессорах, где она поддерживается. Рядом с req и ens функция может нести предложение вида hw x86_64(aes, pclmul) = "symbol". Оно называет символ на C или ассемблере с ABI этой функции, и выполняется он только там, где есть все перечисленные возможности процессора. Определением при этом остаётся тело на kola: его выполняет интерпретатор, его свойства доказывает чекер, и с ним kola test --native сравнивает аппаратную реализацию; --no-hw отключает их все. std пользуется этим для AES-NI, PCLMUL, SHA-NI, CRC-32C и их аналогов на AArch64. Правила описаны в SPEC-CARD §8 и в главе Под капотом.