Существует три уровня понимания того, как работает SIMD (ну, по крайней мере, на данный момент я нахожусь на 3-м уровне):

1. Компиляторы умны! Они автоматически векторизуют весь код!
2. Компиляторы тупы, автоматическая векторизация хрупка, ее очень легко нарушить несвязанными изменениями в коде. Всегда лучше вручную написать конкретные инструкции SIMD.
3. Написать SIMD вручную действительно сложно — для каждой архитектуры процессора придется писать разный код. Кроме того, вы, вероятно, понимаете, что компилятор напишет на ассемблере скалярный код лучше вас. Что заставляет вас думать, что вы превзойдете компилятор в SIMD, где еще больше странных инструкций и запретов? Компиляторы — это инструменты. Они могут надежно векторизовать код, если он написан в форме, поддающейся векторизации.

Недавно я перешел со второго уровня на третий, и я заметил, как модель, используемая компилятором, щелкнула у меня в голове. В этом посте я хочу объяснить общую структуру компиляторов, пригодную для оптимизации статических языков, таких как Rust или C++. После этого я применю эту структуру к автоматической векторизации.

Я не работал над бэкендами компиляторов, оптимизирующих производство, поэтому нижеследующее не будет абсолютно идеальным, но эти модели определенно полезны, по крайней мере, для меня!

❯ Глазами компилятора

Первая часть нашей головоломки – понимание того, как компилятор видит код. Некоторые данные взяты из книги “The SSA Book” или книги авторства LLVM — “Language Reference”.

Помимо компиляторов рассмотрим “WebAssembly Specification”. Хоть WASM и плохой выбор для компилятора, он имеет много структурных сходств, а спецификация ядра исключительно пригодна для чтения.

Единицей оптимизации является функция. Давайте возьмем простую функцию, подобную следующей:

fn sum(xs: &[i32]) -> i32 {
  let mut total = 0;
  for i in 0..xs.len() {
    total = total.wrapping_add(xs[i]);
  }
  total
}

В каком-нибудь псевдо-компиляторе это выглядело бы примерно так:

fn sum return i32 {
  param xs_ptr: ptr
  param xs_len: size

  local total: i32 = 0
  local i: size = 0
  local x: i32

loop:
  branch_if i >= xs_len :ret
  load x base=xs_ptr offset=i
  add total x
  add i 1
  goto :loop

ret:
  return total
}

Наиболее важной характеристикой здесь является то, что существует два вида «сущностей»:

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

Вторая — локальные переменные. Локальные переменные — это не байты — это целые числа, они подчиняются математическим операциям, которые компилятор уже понимает.

Например, если компилятор видит цикл, подобный:

param n: u32
local i: u32 = 0
local total: u32
local tmp

loop:
  branch_if i >= n :ret
  set tmp i
  mul tmp 4
  add t tmp
  goto :loop

ret:
  return total

Он понимает, что в каждом цикле tmp становится i*4 и оптимизирует код до:

param n: u32
local i: u32 = 0
local total: u32
local tmp = 0

loop:
  branch_if i >= n :ret
  add t tmp
  add tmp 4  # replace multiplication with addition
  goto :loop

ret:
  return total

Это работает, потому что все здесь, — это просто цифры. Если бы мы выполнили то же самое вычисление, но все числа были расположены в памяти, компилятору было бы значительно сложнее понять, что преобразование на самом деле правильное. Что, если хранилище для n и total на самом деле перезаписывается? Что, если tmp пересекается с чем-то, чего даже нет в текущей функции?

Однако существует мост между миром математических локальных переменных и миром байтов памяти — инструкции load и store. Команда load берет диапазон байтов в памяти, интерпретирует байты как целое число и сохраняет это целое число в локальной переменной. Инструкция store делает все то же, но наоборот. Загружая что-либо из памяти в локальный файл, компилятор получает возможность понимать это. Таким образом, компилятору не нужно отслеживать общее содержимое памяти. Ему нужно лишь проверить, надо ли и стоит ли загружать что-то в определенный момент.

Таким образом, компилятор на самом деле понимает все не так и хорошо, он может по-настоящему рассуждать только об одной функции за раз и только о локальных переменных в этой функции.

❯ Тыкать компилятор лицом в код

Компиляторы «близоруки». Это можно исправить, предоставив компилятору больше контекста, что является задачей двух основных оптимизаций.

Первая основная оптимизация — это “inlining” (встраивание). Это замена “тела” при вызове. Преимущество здесь не в том, что мы устраняем оверхед на вызов функции, это относительно незначительно. Важно то, что локальные данные как вызывающего, так и вызываемого объекта теперь находятся в одной области, и компилятор может оптимизировать их вместе.

Взглянем снова на этот код на Rust:

fn sum(xs: &[i32]) -> i32 {
  let mut total = 0;
  for i in 0..xs.len() {
    total = total.wrapping_add(xs[i]);
  }
  total
}

Выражение xs[i] на самом деле является вызовом функции. Индексирование проверяет границы перед доступом к элементу массива. После вставки его в sum компилятор может увидеть, что это мертвый код, и устранить его.

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

Вторая основная оптимизация — это SRA (скалярная замена агрегатов). Это обобщение идеи “давайте использовать load, чтобы избежать рассуждений о памяти и вместо этого рассуждать о локальных числах”, которую мы уже видели.

Если у вас есть функция, подобная:

fn permute(xs: &mut Vec<i32>) {
  ...
}

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

fn permute(xs: &mut Vec<i32>) {
  local ptr: ptr
  local len: usize
  local cap: usize

  load ptr xs.ptr
  load len xs.len
  load cap xs.cap

  ...

  store xs.ptr ptr
  store xs.len len
  store xs.cap cap
}

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

❯ Возможности и Невозможности

Используя эту ментальную модель компилятора, который:

1. Оптимизирует каждую функцию;
2. Может использовать встраивание;
3. Отлично замечает взаимосвязи между локальными переменными и перестраивает код на основе этого;
4. Способен ограниченно рассуждать о памяти (а именно, решать, когда безопасно загружать или сохранять);

мы можем описать, какой код можно легко оптимизировать, а какой код предотвращает оптимизацию, рассуждая об абстракциях с нулевыми затратами (zero-cost abstractions).

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

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

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

// Compiler is guaranteed to be able to inline call to `f`.
fn call1<F: Fn()>(f: F) {
  f()
}

// Compiler _might_ be able to inline call to `f`.
fn call2(f: fn()) {
  f()
}

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

Аналогично, косвенное обращение к памяти может вызвать проблемы у компилятора.

Для чего-то подобного:

struct Foo {
  bar: Bar,
  baz: Baz,
}

Структура Foo полностью прозрачна для компилятора.

Однако здесь:

struct Foo {
  bar: Box<Bar>,
  baz: Baz,
}

Все не очень ясно. То, что предоставлено памяти, занимаемой Foo, в общем случае не переносится в память, занимаемую Bar. Опять же, во многих случаях компилятор может рассуждать с помощью блоков благодаря уникальности, но это не гарантируется.
Хорошей домашней работой сейчас будет взглянуть на итераторы Rust и понять, почему они выглядят так, как они выглядят.

Почему сигнатура и значение map таковы?

#[inline]
fn map<B, F>(self, f: F) -> Map<Self, F>
where
  Self: Sized,
  F: FnMut(Self::Item) -> B,
{
  Map::new(self, f)
}

Еще одним важным моментом, касающимся памяти, является то, что, как правило, компилятор не может изменить общую компоновку. SROA может загрузить некоторые данные в набор локальных переменных, которые затем могут, например, заменить представление “указатель и индекс” на “пару указателей”. Но в конце концов SROA пришлось бы упаковать “указатель и индекс” обратно и сохранить это представление обратно в памяти. Это связано с тем, что расположение памяти является общим для всех функций, поэтому функция не может в одностороннем порядке диктовать свои правила.

В совокупности эти наблюдения дают базовое представление о том, как должен работать код.

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

❯ SIMD

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

Прямая реализация выглядела бы следующим образом:

use std::iter::zip;

// 650 milliseconds
fn common_prefix(xs: &[u8], ys: &[u8]) -> usize {
  let mut result = 0;
  for (x, y) in zip(xs, ys) {
    if x != y { break; }
    result += 1
  }
  result
}

Если у вас уже есть ментальная модель автоматической векторизации, или если вы посмотрите на выходные данные ассемблера, вы можете понять, что функция в записанном виде работает по одному байту за раз, что намного медленнее, чем должно быть. Давайте это исправим!

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

// 450 milliseconds
fn common_prefix(xs: &[u8], ys: &[u8]) -> usize {
  let chunk_size = 16;

  let mut result = 0;

  'outer: for (xs_chunk, ys_chunk) in
    zip(xs.chunks_exact(chunk_size), ys.chunks_exact(chunk_size))
  {
    for (x, y) in zip(xs_chunk, ys_chunk) {
      if x != y { break 'outer; }
      result += 1
    }
  }

  for (x, y) in zip(&xs[result..], &ys[result..]) {
    if x != y { break; }
    result += 1
  }

  result
}

Забавно, что это уже немного быстрее, но еще не совсем то. В частности, SIMD должен обрабатывать все значения в блоке параллельно и одинаково. В нашем приведенном выше коде у нас есть break, который означает, что обработка nой пары байтов зависит от n-1-й пары. Давайте исправим это, отключив использование короткой схемы. Мы проверим, совпадает ли весь фрагмент байтов или нет, но нам будет все равно, какой конкретный байт является несоответствующим:

// 80 milliseconds
fn common_prefix3(xs: &[u8], ys: &[u8]) -> usize {
  let chunk_size = 16;

  let mut result = 0;
  for (xs_chunk, ys_chunk) in
    zip(xs.chunks_exact(chunk_size), ys.chunks_exact(chunk_size))
  {
    let mut chunk_equal: bool = true;
    for (x, y) in zip(xs_chunk, ys_chunk) {
      // NB: &, unlike &&, doesn't short-circuit.
      chunk_equal = chunk_equal & (x == y);
    }

    if !chunk_equal { break; }
    result += chunk_size;
  }

  for (x, y) in zip(&xs[result..], &ys[result..]) {
    if x != y { break; }
    result += 1
  }

  result
}

И эта версия, наконец, позволяет “впустить” векторизацию, сокращая время выполнения почти на порядок. Теперь мы можем сжать эту версию с помощью итераторов.

// 80 milliseconds
fn common_prefix5(xs: &[u8], ys: &[u8]) -> usize {
  let chunk_size = 16;

  let off =
    zip(xs.chunks_exact(chunk_size), ys.chunks_exact(chunk_size))
      .take_while(|(xs_chunk, ys_chunk)| xs_chunk == ys_chunk)
      .count() * chunk_size;

  off + zip(&xs[off..], &ys[off..])
    .take_while(|(x, y)| x == y)
    .count()
}

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

В частности, для SIMD:

• Мы изменяем алгоритм для работы с блоками элементов.
• Внутри каждого блока, мы убеждаемся, что ничего не разветвляется и все элементы обрабатываются одинаково.

❯ Заключение

Компиляторы — это инструменты. Несмотря на то, что иногда происходит изрядная доля “оптимистичных” преобразований, основная часть воздействия оптимизирующего компилятора достигается за счет гарантированной оптимизации с определенными предварительными условиями. Компиляторы близоруки — им трудно рассуждать о коде вне текущей функции и значениях, не хранящихся в локальных переменных. Встраивание и скалярная замена агрегатов — это две оптимизации, позволяющие исправить ситуацию. Абстракции с нулевой стоимостью работают за счёт выражения возможностей для гарантированных оптимизаций.

Если вам понравился этот пост, я настоятельно рекомендую “A Catalogue of Optimizing Transformations” Франсиса Аллена.

---

Возможно, захочется почитать и это:

• ➤ Пример использования WebAssembly-модуля, скомпилированного из Rust, в React-приложении
• ➤ TypeScript: разрабатываем WebAssembly-компилятор
• ➤ Компилируем быстрые консольные .exe приложения на PHP 8.1 в 2023 году, а почему бы и нет? (upd)
• ➤ Откат операционной системы, с отечественной аппаратной картой «ЛанСил»
• ➤ 96 лет со дня рождения Джона Маккарти