Huawei | Использование инструкций по ускорению вычислений SIMD на языке Rust
Автор: Ли Юань
1. Введение в SIMD
SIMD расшифровывается как Single Instruction Multiple Data, который представляет собой поток нескольких данных с одной инструкцией, который представляет собой технологию оптимизации вычислительной производительности, основанную на определенном наборе инструкций ЦП. Как следует из названия, это относится к вычислению нескольких данных одновременно во время выполнения инструкции ЦП. Он может повысить производительность в несколько раз или даже в десятки раз в сценариях вычислений с интенсивным использованием данных, таких как научные вычисления и мультимедийные приложения.
2. Знакомство с официальной библиотекой ускорения SIMD языка Rust
Язык Rust — это язык программирования, который может выбирать различные механизмы компиляции.Конечно, большинство проектов Rust в отрасли в настоящее время используют llvm, который компилятор выбирает по умолчанию, в качестве механизма компиляции. Стоит отметить, что сам llvm интегрировал различные наборы инструкций, включая SIMD большинства основных архитектур ЦП. Это обеспечивает естественное удобство языка Rust для использования SIMD, потому что Rust может напрямую использовать интерфейс функции SIMD, предоставляемый llvm в виде статической компоновки в компиляторе и даже пользовательского кода, вместо того, чтобы быть написанным самими разработчиками, такими как Go, C и другие языки ассемблерный код.
Язык Rust предоставляет две библиотеки ускорения simd в официальной группе проекта github:stdarchа такжеstdsimd. Вот адрес их репозитория на github. stdarch предоставляет свой собственный выделенный набор инструкций ускорения simd для каждой архитектуры ЦП по модульному принципу, например, AVX, AVX512, SSE, SSE2 и другие наборы инструкций для архитектуры x86; наборы инструкций NEON, SVE для платформы ARM/Aarch64; и Набор инструкций simd RISCV, WASM и других архитектур и т. д., пользователи должны понимать архитектуру ЦП, которую они используют, и функции набора инструкций simd, предоставляемые этой архитектурой. И stdsimd предоставляет большое количество интерфейсов абстрактных функций simd, общих для различных платформ, таких как сложение векторов, вычитание, умножение и деление, смещение, преобразование типов и т. д. Читателям не обязательно знать об используемой ими аппаратной архитектуре и наборе инструкций, это относительно удобнее в использовании, но будут некоторые ограничения на использование функций. Два проекта имеют разные отправные точки для функционального проектирования, и сопровождающие соответствующие проекты также различны. Их конструкция и использование будут подробно описаны ниже.
3. Использование мультиархитектурной библиотеки общего ускорения stdsimd
stdsimd предоставляет общий интерфейс ускорения simd для каждой платформы, и его реализация зависит от встроенного в платформу интерфейса, предоставляемого компилятором Rust. Набор интерфейсов также является инкапсуляцией наборов инструкций каждой платформы, предоставляемых llvm, поэтому отношения между ними должны быть следующими:
stdsimd— Пакет →Компилятор ржавчины— Пакет →llvm
Проект stdsimd еще не интегрирован в стандартную библиотеку Rust, потому что его функции интегрированы не полностью.Читатели могут использовать его, клонируя исходный код в свои проекты или используя версию stdsimd от сообщества.packed_simd(В файле Cargo.toml добавитьpacked_simd = { version = "0.3.4", package = "packed_simd_2" }). Использование ниже также описано на основе версии сообщества.
Проект Packed_simd предоставляет серию векторных типов данных Simd, то есть вектор, состоящий из элементов NT, и предоставляет им простые для понимания псевдонимы типов, такие как тип f32x4, который представляет Simd . Функции ускорения SIMD, предоставляемые Packed_simd, также реализованы на основе этого векторного типа данных.
Packed_simd предоставляет следующие типы данных SIMD (element_width представляет размер и количество данных, например 32x4, 64x8):
-
i{element_width}: целочисленный тип со знаком -
u{element_width}: беззнаковый целочисленный тип -
f{element_width}: плавающий тип -
m{element_width}: логический тип -
*{const,mut} T: изменяемый или неизменяемый указатель типа SIMD
По умолчанию операции над векторными структурами являются «вертикальными», т.е. они применяются к каждому векторному каналу независимо от других векторов, как в этом примере:
let a = i32x4::new(1, 2, 3, 4);
let b = i32x4::new(5, 6, 7, 8);
assert_eq!(a + b, i32x4::new(6, 8, 10, 12));
В примере объявляются два вектора i32x4 и вычисляется их сумма с помощью перегрузки оператора сложения. С другой стороны, конечно, предусмотрены и «горизонтальные» операции, такие как следующий пример:
assert_eq!(a.wrapping_sum(), 10);
В общем, «вертикальные» операции всегда самые быстрые, тогда как «горизонтальные» операции относительно медленные. Тем не менее, при вычислении суммы массива самым быстрым способом является использование нескольких «вертикальных» операций плюс одна «горизонтальная» операция, например:
fn reduce(x: &[i32]) -> i32 {
assert!(x.len() % 4 == 0);
let mut sum = i32x4::splat(0); // [0, 0, 0, 0]
for i in (0..x.len()).step_by(4) {
sum += i32x4::from_slice_unaligned(&x[i..]);
}
sum.wrapping_sum()
}
let x = [0, 1, 2, 3, 4, 5, 6, 7];
assert_eq!(reduce(&x), 28);
Вот еще несколько распространенных вариантов использования:
// 生成元素全为0的i32x4向量:
let a = i32x4::splat(0);
// 由数组中的前4个元素生成i32x4向量:
let mut arr = [0, 0, 0, 1, 2, 3, 4, 5];
let b = i32x4::from_slice_unaligned(&arr);
// 读取向量中的元素:
assert_eq!(b.extract(3), 1);
// 替换向量中对应位置的元素:
let a = a.replace(3, 1);
assert_eq!(a, b);
// 将向量写入数组中:
let a = a.replace(2, 1);
a.write_to_slice_unaligned(&mut arr[4..]);
assert_eq!(arr, [0, 0, 0, 1, 0, 0, 1, 1]);
Кроме того, packed_simd также предоставляет условные операции над векторами.Например, следующий код указывает, что операция +1 соответствующего элемента в векторе выполняется в зависимости от того, является ли элемент в m истинным:
let a = i32x4::new(1, 1, 2, 2);
// 将a中的前两个元素进行+1操作.
let m = m16x4::new(true, true, false, false);
let a = m.select(a + 1, a);
assert_eq!(a, i32x4::splat(2));
Отсюда могут быть получены более гибкие методы использования, такие как формирование нового вектора из большего значения каждой позиции в двух векторах.
let a = i32x4::new(1, 1, 3, 3);
let b = i32x4::new(2, 2, 0, 0);
// ge: 大于等于计算，生成bool元素类型的向量
let m = a.ge(i32x4::splat(2));
if m.any() {
// 根据m中的结果选择a或b中的元素
let d = m.select(a, b);
assert_eq!(d, i32x4::new(2, 2, 3, 3));
}
Выше приведено базовое использование stdsimd (packed_simd) Этот проект позволяет разработчикам легко пользоваться эффектом ускорения SIMD через структуру данных типа SIMD. Но у проекта есть и определенные недостатки, например, пользователь должен вручную выбирать длину вектора. Поскольку большинство архитектур ЦП обеспечивают как минимум 128-битный набор инструкций SIMD, всегда разумно выбирать длину вектора 128 бит. Но когда ЦП предоставляет более продвинутый набор инструкций SIMD (например, AVX512), выбор более длинного набора инструкций даст лучшие результаты. Поэтому, когда у разработчика есть определенная архитектура ЦП и запас знаний, связанных с SIMD, при использовании это будет иметь эффект мультипликатора.
4. Использование специальной библиотеки ускорения инструкций stdarch
stdarch интегрирован в стандартную библиотеку языка Rust и может передаваться в кодеuse std::archзаявление для использования. Здесь следует отметить, что только две архитектуры, x86_64 и x86, выпустили стабильные версии, поэтому другие архитектуры, такие как arm, aarch64 и т. д., должны переключить компилятор Rust на ночную версию (введите команду rustup default nightly в командной строке ) перед компиляцией и использованием. Поэтому далее в качестве примера для ознакомления в основном используется стабильная версия x86_64 (x86).
stdarch инкапсулирует множество наборов инструкций SIMD, предоставляемых llvm, в статической ссылке и предоставляет наборы инструкций SIMD для различных основных архитектур в модулях, как показано ниже. Интерфейсы функций SIMD, доступные для каждой архитектуры, можно просмотреть, щелкнув соответствующую ссылку.
По сравнению со stdsimd, stdarch предъявляет более высокие требования к знанию разработчиками архитектуры ЦП. Поскольку stdarch предоставляет тысячи SIMD-инструкций с различными функциями для каждой основной архитектуры ЦП, разработчикам приходится вручную определять, какая инструкция наиболее необходима.
Например следующий пример:
#[cfg(
all(
any(target_arch = "x86", target_arch = "x86_64"),
target_feature = "avx2"
)
)]
fn foo() {
#[cfg(target_arch = "x86")]
use std::arch::x86::_mm256_add_epi64;
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::_mm256_add_epi64;
unsafe {
_mm256_add_epi64(...);
}
}
Этот код сначала использует функцию обнаружения функций ЦП, изначально предоставляемую языком Rust, то есть макрос атрибута target_arch, чтобы определить, является ли среда разработки x86_64 или x86, а затем использует макрос атрибута target_feature, чтобы определить, является ли набор инструкций avx2 доступный. При выполнении вышеуказанных условий будет скомпилирована следующая функция foo. Внутри функции foo соответствующая инструкция simd выбирается в зависимости от архитектуры процессора x86_64 или x86.
Или разработчики могут использовать операторы обнаружения динамических функций.is_x86_feature_detected!,Следующим образом:
fn foo() {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
{
if is_x86_feature_detected!("avx2") {
return unsafe { foo_avx2() };
}
}
// return without using AVX2
}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
unsafe fn foo_avx2() {
#[cfg(target_arch = "x86")]
use std::arch::x86::_mm256_add_epi64;
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::_mm256_add_epi64;
_mm256_add_epi64(...);
}
В самом stdarch существует много подобного кода условной компиляции. Поэтому соответствующий модуль набора инструкций доступен только тогда, когда он соответствует потребностям среды. Например, его можно использовать под архитектуру x86_64.use std::arch::x86_64утверждение, но не может использоватьuse std::arch::x86_64илиuse std::arch::armутверждение.
Ниже представлено конкретное использование stdarch на конкретном примере, то есть реализация simd функции шестнадцатеричного кодирования. В этом примере в основном используется набор инструкций SSE4.1 для x86 и x86_64.
Конкретная реализация кода выглядит следующим образом, а различные инструкции и способы использования SIMD можно найти в комментариях или в связанной документации соответствующего модуля (x86 или x86_64) выше.
fn main() {
let mut dst = [0; 32];
hex_encode(b"\x01\x02\x03", &mut dst);
assert_eq!(&dst[..6], b"010203");
let mut src = [0; 16];
for i in 0..16 {
src[i] = (i + 1) as u8;
}
hex_encode(&src, &mut dst);
assert_eq!(&dst, b"0102030405060708090a0b0c0d0e0f10");
}
pub fn hex_encode(src: &[u8], dst: &mut [u8]) {
let len = src.len().checked_mul(2).unwrap();
assert!(dst.len() >= len);
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
{
if is_x86_feature_detected!("sse4.1") {
return unsafe { hex_encode_sse41(src, dst) };
}
}
hex_encode_fallback(src, dst)
}
#[target_feature(enable = "sse4.1")]
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
unsafe fn hex_encode_sse41(mut src: &[u8], dst: &mut [u8]) {
#[cfg(target_arch = "x86")]
use std::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;
// 生成包含16个int8类型的向量，并将全部值设为字符'0'的ascii编号
let ascii_zero = _mm_set1_epi8(b'0' as i8);
// 生成包含16个int8类型的向量，并将全部值设为整数9
let nines = _mm_set1_epi8(9);
// 生成包含16个int8类型的向量，并将全部值设为字符'a'的ascii编号减去10
let ascii_a = _mm_set1_epi8((b'a' - 9 - 1) as i8);
// 生成包含16个int8类型的向量，并将全部值设为二进制数00001111
let and4bits = _mm_set1_epi8(0xf);
let mut i = 0_isize;
while src.len() >= 16 {
// 从指针中读取128位整数，组成一个128位的向量（可以转化为int8x16、int32x4等形式的向量）
let invec = _mm_loadu_si128(src.as_ptr() as *const _);
// 将该128位向量类型转化为int8x16类型的向量，并将其中每个元素和二进制数00001111进行与操作
let masked1 = _mm_and_si128(invec, and4bits);
// 将该128位向量类型转化为int8x16类型的向量，再将每个元素逻辑右移4位，随后将其中每个元素和二进制数00001111进行与操作
let masked2 = _mm_and_si128(_mm_srli_epi64(invec, 4), and4bits);
// 向量对应元素比较大小，获取向量中所有大于9的元素的位置
let cmpmask1 = _mm_cmpgt_epi8(masked1, nines);
let cmpmask2 = _mm_cmpgt_epi8(masked2, nines);
// _mm_blendv_epi8表示生成一个新的向量，该向量中的元素是根据cmpmask1中对应位置是否为true选择ascii_zero或者ascii_a中的元素
// _mm_add_epi8则表示向量对应位置元素相加，结果表示最终生成的十六进制编码的ascii编号
let masked1 = _mm_add_epi8(
masked1,
_mm_blendv_epi8(ascii_zero, ascii_a, cmpmask1),
);
let masked2 = _mm_add_epi8(
masked2,
_mm_blendv_epi8(ascii_zero, ascii_a, cmpmask2),
);
// 生成一个新的向量，其中偶数位置元素（从0开始）来自于masked2，奇数位置元素来自于masked1
// 该向量共有256位，所以将前128位放入res1中，后128位放入res2中
let res1 = _mm_unpacklo_epi8(masked2, masked1);
let res2 = _mm_unpackhi_epi8(masked2, masked1);
// 将结果向量写入目标指针中
_mm_storeu_si128(dst.as_mut_ptr().offset(i * 2) as *mut _, res1);
_mm_storeu_si128(
dst.as_mut_ptr().offset(i * 2 + 16) as *mut _,
res2,
);
src = &src[16..];
i += 16;
}
let i = i as usize;
hex_encode_fallback(src, &mut dst[i * 2..]);
}
fn hex_encode_fallback(src: &[u8], dst: &mut [u8]) {
fn hex(byte: u8) -> u8 {
static TABLE: &[u8] = b"0123456789abcdef";
TABLE[byte as usize]
}
for (byte, slots) in src.iter().zip(dst.chunks_mut(2)) {
slots[0] = hex((*byte >> 4) & 0xf);
slots[1] = hex(*byte & 0xf);
}
}
Здесь использование инструкций ускорения SIMD в stdarch просто представлено в этом конкретном примере. Видно, что использование выделенных инструкций требует от разработчиков гораздо более высокого опыта SIMD, чем stdsimd, но предоставляемые функции и применимые сценарии также будут более полными.
Выше приведено простое введение в официальную библиотеку ускорения SIMD на языке Rust, и я надеюсь, что оно вдохновит и поможет читателям в их обучении и развитии.