Девять правил для ускорения SIMD вашего кода на Rust (Часть 2)

Девять правил для повышения производительности SIMD в вашем коде на Rust (Часть 2)

Общие уроки повышения скорости ввода данных в крейт range-set-blaze в 7 раз

Большое спасибо Бену Лихтману (B3NNY) созванию Seattle Rust за правильное направление на SIMD.

Это часть 2 статьи о создании SIMD-кода на языке Rust. (См. Часть 1.) Мы рассмотрим правила с 7 по 9:

• 7. Используйте бенчмаркинг Criterion для выбора алгоритма и выясните, что LANES (почти всегда) должны быть 32 или 64.
• 8. Интегрируйте лучший SIMD-алгоритм в ваш проект с помощью as_simd, специального кода для i128/u128 и дополнительного контекстного бенчмаркинга.
• 9. Отделите лучший SIMD-алгоритм от вашего проекта (пока) с помощью опциональной функциональности в Cargo.

Вспомним правила с 1 по 6:

1. Используйте ночной Rust и core::simd, экспериментальный модуль SIMD стандарта Rust.
2. CCC: Проверьте, контролируйте и выберите возможности SIMD вашего компьютера.
3. Изучите core::simd, но выбирайте определенные части.
4. Найдите кандидатов на алгоритмы.
5. Используйте Godbolt и AI для понимания сборки вашего кода, даже если вы не знаете язык ассемблера.
6. Обобщите для всех типов и LANES с помощью встроенных обобщений, (и когда это не сработает) макросов и (когда это не сработает) трейтов.

Эти правила основаны на моем опыте попыток ускорения range-set-blaze, крейта для Rust, предназначенного для работы со “скомканными” целыми числами.

• Сегментация клиентов больше, чем кластеризация
• Оценка настроек компилятора Rust с помощью Criterion
• Сколько данных нам нужно? Балансировка машинного обучения и вопросов безопасности

Вспомним, что правило 6 из Части 1 показывает, как сделать SIMD-алгоритмы Rust полностью обобщенными в отношении типа и LANES. Теперь нам нужно выбрать наш алгоритм и настроить LANES.

Правило 7: Используйте бенчмаркинг Criterion для выбора алгоритма и выясните, что LANES (почти всегда) должны быть 32 или 64.

В этом правиле мы рассмотрим, как использовать популярный крейт criterion для бенчмаркинга и оценки наших алгоритмов и вариантов. В контексте range-set-blaze мы оценим:

• 5 алгоритмов — Regular, Splat0, Splat1, Splat2, Rotate
• 3 уровня SIMD-расширений — sse2 (128 бит), avx2 (256 бит), avx512f (512 бит)
• 10 типов элементов — i8, u8, i16, u16, i32, u32, i64, u64, isize, usize
• 5 чисел LANES — 4, 8, 16, 32, 64
• 4 длины входных данных — 1024, 10 240, 102 400, 1 024 000
• 2 процессора — AMD 7950X с avx512f, Intel i5–8250U с avx2

Бенчмарк измеряет среднее время выполнения каждой комбинации. Затем мы вычисляем пропускную способность в Мбайт/сек.

Ознакомьтесь с этой новой статьей о начале работы с Criterion. В этой статье также показано, как использовать (злоупотреблять?) Criterion для измерения влияния настроек компилятора, таких как уровень расширения SIMD.

Запуск бенчмарков приводит к созданию файла *.csv на 5000 строк, который начинается:

Group,Id,Parameter,Mean(ns),StdErr(ns)vector,regular,avx2,256,i16,16,16,1024,291.47,0.080141vector,regular,avx2,256,i16,16,16,10240,2821.6,3.3949vector,regular,avx2,256,i16,16,16,102400,28224,7.8341vector,regular,avx2,256,i16,16,16,1024000,287220,67.067vector,regular,avx2,256,i16,16,32,1024,285.89,0.59509...

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

Алгоритмы и Линии

Вот таблица сводных данных Excel, показывающая для каждого алгоритма пропускную способность (Мбайт/сек) по отношению к SIMD Линиям. Таблица усредняет пропускную способность по уровням расширения SIMD, типам элементов и длине ввода.

На моем настольном компьютере AMD:

На ноутбуке Intel:

Таблицы показывают, что Splat1 и Splat2 показывают лучшие результаты. Они также показывают, что большее количество Линий всегда лучше до 32 или 64.

Как, например, sse2 (шириной 128 бит) может обрабатывать 64 линии i64 (шириной 4096 бит)? Модуль Rust core::simd делает эту магию возможной, автоматически и эффективно разделяя 4096 бит на 32 куска по 128 бит. Обработка 32 128-битных кусков вместе (по-видимому) позволяет делать оптимизации, выходящие за рамки обработки независимых 128-битных блоков.

Уровни расширения SIMD

Установим LANES на 64 и сравним разные уровни расширения SIMD на компьютере AMD. Таблица усредняет пропускную способность по типу элемента и длине ввода.

На моем компьютере AMD, при использовании 64 линий, sse2 самый медленный. Сравнивая avx2 с avx512f, результаты разные. Опять же, алгоритмы Splat1 и Splat2 показывают лучшие результаты.

Типы элементов

Зададим уровень расширения SIMD равным avx512f и сравним разные типы элементов. Мы оставляем LANES на 64 и усредняем пропускную способность по длине ввода.

Мы видим, что наибольшую скорость обработки имеют бит-по-бит, 32-бит и 64-битные элементы. (Однако, для каждого элемента, более маленькие типы обрабатываются быстрее.) Splat1 и Splat2 являются самыми быстрыми алгоритмами, причем Splat1 немного лучше.

Длина ввода

Наконец, давайте установим тип элемента i32 и посмотрим, как меняется длина ввода в зависимости от производительности.

Мы видим, что все алгоритмы SIMD выполняют примерно одинаково при миллионе входных данных. Алгоритм Splat1, кажется, работает лучше других алгоритмов при коротких входных данных.

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

Выводы по тестам

Исходя из этих тестов, мы будем использовать алгоритм Splat1. Пока мы установим значения LANES от 32 до 64, но следующее правило может создать некоторые осложнения. Наконец, мы рекомендуем пользователям установить уровень расширения SIMD как минимум avx2.

Правило 8: Интегрируйте ваш лучший SIMD-алгоритм в ваш проект с помощью as_simd, специального кода для i128/u128 и дополнительных тестов в контексте.

as_simd

Перед добавлением поддержки SIMD основной конструктор RangeSetBlaze выглядел так:

let a = RangeSetBlaze::from_iter([1, 2, 3]);

Операции SIMD лучше всего работают с массивами, а не с итераторами. Более того, конструирование RangeSetBlaze из массива часто является естественным действием, поэтому я добавил новый конструктор from_slice:

#[inline]
pub fn from_slice(slice: impl AsRef<[T]>) -> Self {
    T::from_slice(slice)
}

Новый конструктор вызывает метод from_slice для каждого целого числа. Для всех типов целых чисел, кроме i128/u128, это следующее вызовы:

let (prefix, middle, suffix) = slice.as_simd();

Метод as_simd в версии Rust Nightly безопасно и быстро трансформирует массив следующим образом:

1. невыровненный prefix, который мы обрабатываем с помощью from_iter, как раньше.
2. middle, выровненный массив фрагментов структуры Simd
3. невыровненный suffix, который мы обрабатываем с помощью from_iter, как раньше.

Думайте о middle как о разбиении нашего входного массива целых чисел на куски размером 16 (или другой размер, определяемый значением LANES). Затем мы проходим по кускам с помощью функции is_consecutive, ищем последовательности true. Каждая последовательность становится отдельным диапазоном. Например, последовательность из 160 последовательных чисел от 1000 до 1159 (включительно) будет идентифицирована и заменена одним диапазоном Rust RangeInclusive 1000..=1159. Этот диапазон затем обрабатывается с помощью from_iter гораздо быстрее, чем обработка 160 отдельных чисел с помощью from_iter. Когда is_consecutive возвращает false, мы возвращаемся к обработке отдельных чисел в куске с помощью from_iter.

i128/u128

Как обрабатывать массивы типов, которые core::simd не поддерживает, такие как i128/u128? Пока что я просто обрабатываю их медленным способом с помощью from_iter.

Тесты в контексте

В заключение, проведите тестирование вашего SIMD-кода в контексте вашего основного кода, в идеале на представительных данных.

Пакет range-set-blaze уже включает бенчмарки. Один из них измеряет производительность при вводе 1 000 000 целых чисел с разными уровнями групировки. Средний размер группы варьируется от 1 (без группировки) до 100 000 групп. Давайте запустим этот бенчмарк с установленными значениями LANES: 4, 8, 16, 32 и 64. Мы будем использовать алгоритм Splat1 и расширение SIMD avx512f.

Для каждого размера группы показаны относительные скорости ввода 1 000 000 целых чисел. Для каждого размера группы, самое быстрое значение LANES равно 100%.

Мы видим, что для групп размером 10 и 100, лучший вариант – LANES=4. Однако, для групп размером 100 000, LANES=4 в 4 раза хуже лучшего варианта. С другой стороны, LANES=64 выглядит хорошо для групп размером 100 000, но для групп размером 100 и 1000 оно хуже лучшего варианта в 1.8 и 1.5 раза, соответственно.

Я решил установить значение LANES равным 16. Оно является оптимальным для групп размером 1000. Более того, оно никогда не хуже лучшего варианта более, чем в 1.25 раза.

С такими настройками мы можем запускать другие бенчмарки. График ниже показывает различные библиотеки для работы с диапазонами (включая range-set-blaze), выполняющие ту же задачу – ввод 1 000 000 целых чисел разной групировки. Ось y показывает время в миллисекундах, где меньше значение означает лучшую производительность.

Для групп размером 1000, существующий метод RangeSetBlaze::into_iter (красный) уже был в 30 раз быстрее, чем HashSet (оранжевый). Обратите внимание, что шкала логарифмическая. С использованием avx512f, новый алгоритм RangeSetBlaze::into_slice на основе SIMD (светло-синий) в 230 раз быстрее, чем HashSet. С использованием sse2 (темно-синий) скорость увеличивается в 220 раз. С использованием avx2 (желтый), она увеличивается в 180 раз. В сравнении с RangeSetBlaze::into_iter, avx512f RangeSetBlaze::into_slice на этом рейтинге в 7 раз быстрее.

Мы также должны рассмотреть худший случай, а именно, ввод данных без группировки. Я провел этот бенчмарк. Он показал, что существующий метод RangeSetBlaze::into_iter примерно в 2.2 раза медленнее, чем HashSet. Новый метод RangeSetBlaze::into_slice в 2.4 раза медленнее, чем HashSet.

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

С интеграцией SIMD-кода в наш проект, мы готовы его выпустить, верно? К сожалению, нет. Поскольку наш код зависит от Rust nightly, мы должны сделать его опциональным. Узнаем, как это сделать в следующем правиле.

Правило 9: Извлеките ваш лучший SIMD-алгоритм из проекта (на данный момент) с помощью опционального карго-функционала.

Наш прекрасный новый SIMD-код зависит от Rust nightly, который может меняться. Заставлять пользователей зависеть от Rust nightly было бы жестоко. (Кроме того, получать жалобы, когда что-то ломается, было бы неприятно.) Решением является скрытие SIMD-кода за карго-функционалом.

Функция, функция, функция – В контексте работы с SIMD и Rust слово “функция” используется тремя разными способами. Во-первых, “функции CPU/цели” – они описывают возможности процессора, включая поддержку SIMD-расширений. См. функцию “target-feature” и “is_x86_feature_detected!”. Во-вторых, “функциональные ворота ночной сборки” – Rust контролирует видимость новых языковых функций в ночной сборке Rust с помощью функциональных ворот. Например, #![feature(portable_simd)]. В-третьих, “функции грузовика” – они позволяют любому ящику Rust или библиотеке предлагать/ограничивать доступ к части их возможностей. Вы видите их в вашем Cargo.toml, когда, например, добавляете зависимость от itertools/use_std.

Вот шаги, которые предпринимает каркас range-set-blaze, чтобы сделать код, зависящий от ночной версии, опциональным:

• В Cargo.toml определите функцию грузовика, относящуюся к коду SIMD:

 [features]from_slice = [] 

• В начале файла lib.rs, сделайте функциональные ворота ночной версии portable_simd, зависимыми от функции грузовика from_slice:

 #![cfg_attr(feature = "from_slice", feature(portable_simd))] 

• Используйте атрибут условной компиляции, например, #[cfg(feature = “from_slice”)], чтобы включить код SIMD выборочно. Это включает тесты.

 /// Создает [`RangeSetBlaze`] из коллекции целых чисел. Он обычно работает намного/// быстрее, чем [`from_iter`][1]/[`collect`][1]./// На репрезентативном тесте, ускорение составило 6×.////// **Внимание: Требуется ночной компилятор. Кроме того, вы должны включить функцию `from_slice`/// в вашем `Cargo.toml`. Например, с помощью команды:**/// ```bash///  cargo add range-set-blaze --features "from_slice"/// ```////// **Внимание**: Компиляция с опцией `-C target-cpu=native` оптимизирует бинарный файл для вашей текущей архитектуры CPU,/// что может привести к проблемам совместимости на других машинах с разными архитектурами./// Это особенно важно при распространении бинарного файла или запуске его в различных средах./// [1]: struct.RangeSetBlaze.html#impl-FromIterator<T>-for-RangeSetBlaze<T>#[cfg(feature = "from_slice")]#[inline]pub fn from_slice(slice: impl AsRef<[T]>) -> Self {    T::from_slice(slice)} 

• Как показано в документации выше, добавьте предупреждения и осторожность в документацию.
• Используйте –features from_slice для проверки или тестирования вашего кода SIMD.

 cargo check --features from_slicecargo test --features from_slice 

• Используйте –all-features для запуска всех тестов, генерации всей документации и публикации всех функций грузовика:

 cargo test --all-features --doccargo doc --no-deps --all-features --opencargo publish --all-features --dry-run 

Заключение

Вот вам девять правил для добавления операций SIMD в ваш код Rust. Легкость этого процесса отражает отличный дизайн библиотеки core::simd. Всегда ли стоит использовать SIMD там, где это применимо? В конечном итоге, да, когда библиотека переходит из ночной версии Rust в стабильную. А пока используйте SIMD там, где его производительность важна, или делайте его использование опциональным.

Идеи по улучшению опыта использования SIMD в Rust? Качество core::simd уже высокое; главная потребность – стабилизировать его.

Спасибо, что присоединились ко мне в этом путешествии в программирование SIMD. Я надеюсь, что если у вас есть проблема, подходящая для SIMD, эти шаги помогут вам ускорить ее.

Пожалуйста, подпишитесь на Карла на VoAGI. Я пишу о научном программировании на Rust и Python, машинном обучении и статистике. Обычно я пишу около одной статьи в месяц.