Низкая производительность SIMD - нет встраивания

Рассмотрим следующие примеры вычисления суммы массива i32:

Пример 1: простой цикл for

pub fn vec_sum_for_loop_i32(src: &[i32]) -> i32 {
    let mut sum = 0;
    for c in src {
        sum += *c;
    }

    sum
}

Пример 2: Явная сумма SIMD:

use std::arch::x86_64::*;
// #[inline]
pub fn vec_sum_simd_direct_loop(src: &[i32]) -> i32 {
    #[cfg(debug_assertions)]
    assert!(src.as_ptr() as u64 % 64 == 0);
    #[cfg(debug_assertions)]
    assert!(src.len() % (std::mem::size_of::<__m256i>() / std::mem::size_of::<i32>()) == 0);

    let p_src = src.as_ptr();
    let batch_size = std::mem::size_of::<__m256i>() / std::mem::size_of::<i32>();

    #[cfg(debug_assertions)]
    assert!(src.len() % batch_size == 0);

    let result: i32;
    unsafe {
        let mut offset: isize = 0;
        let total: isize = src.len() as isize;
        let mut curr_sum = _mm256_setzero_si256();

        while offset < total {
            let curr = _mm256_load_epi32(p_src.offset(offset));
            curr_sum = _mm256_add_epi32(curr_sum, curr);
            offset += 8;
        }

        // this can be reduced with hadd.
        let a0 = _mm256_extract_epi32::<0>(curr_sum);
        let a1 = _mm256_extract_epi32::<1>(curr_sum);
        let a2 = _mm256_extract_epi32::<2>(curr_sum);
        let a3 = _mm256_extract_epi32::<3>(curr_sum);
        let a4 = _mm256_extract_epi32::<4>(curr_sum);
        let a5 = _mm256_extract_epi32::<5>(curr_sum);
        let a6 = _mm256_extract_epi32::<6>(curr_sum);
        let a7 = _mm256_extract_epi32::<7>(curr_sum);

        result = a0 + a1 + a2 + a3 + a4 + a5 + a6 + a7;
    }

    result
}

Когда я попытался протестировать код, первый пример получил ~ 23 ГБ / с (что близко к теоретическому максимуму для моей скорости ОЗУ). Второй пример получил 8GB/s.

При взгляде на сборку с грузовым asm первый пример превращается в развернутые SIMD-оптимизированные циклы:

.LBB11_7:
 sum += *c;
 movdqu  xmm2, xmmword, ptr, [rcx, +, 4*rax]
 paddd   xmm2, xmm0
 movdqu  xmm0, xmmword, ptr, [rcx, +, 4*rax, +, 16]
 paddd   xmm0, xmm1
 movdqu  xmm1, xmmword, ptr, [rcx, +, 4*rax, +, 32]
 movdqu  xmm3, xmmword, ptr, [rcx, +, 4*rax, +, 48]
 movdqu  xmm4, xmmword, ptr, [rcx, +, 4*rax, +, 64]
 paddd   xmm4, xmm1
 paddd   xmm4, xmm2
 movdqu  xmm2, xmmword, ptr, [rcx, +, 4*rax, +, 80]
 paddd   xmm2, xmm3
 paddd   xmm2, xmm0
 movdqu  xmm0, xmmword, ptr, [rcx, +, 4*rax, +, 96]
 paddd   xmm0, xmm4
 movdqu  xmm1, xmmword, ptr, [rcx, +, 4*rax, +, 112]
 paddd   xmm1, xmm2
 add     rax, 32
 add     r11, -4
 jne     .LBB11_7
.LBB11_8:
 test    r10, r10
 je      .LBB11_11
 lea     r11, [rcx, +, 4*rax]
 add     r11, 16
 shl     r10, 5
 xor     eax, eax

Во втором примере нет разворачивания цикла и даже встроенного кода в _mm256_add_epi32:

...
movaps  xmmword, ptr, [rbp, +, 320], xmm7
 movaps  xmmword, ptr, [rbp, +, 304], xmm6
 and     rsp, -32
 mov     r12, rdx
 mov     rdi, rcx
 lea     rcx, [rsp, +, 32]
 let mut curr_sum = _mm256_setzero_si256();
 call    core::core_arch::x86::avx::_mm256_setzero_si256
 movaps  xmm6, xmmword, ptr, [rsp, +, 32]
 movaps  xmm7, xmmword, ptr, [rsp, +, 48]
 while offset < total {
 test    r12, r12
 jle     .LBB13_3
 xor     esi, esi
 lea     rbx, [rsp, +, 384]
 lea     r14, [rsp, +, 64]
 lea     r15, [rsp, +, 96]
.LBB13_2:
 let curr = _mm256_load_epi32(p_src.offset(offset));
 mov     rcx, rbx
 mov     rdx, rdi
 call    core::core_arch::x86::avx512f::_mm256_load_epi32
 curr_sum = _mm256_add_epi32(curr_sum, curr);
 movaps  xmmword, ptr, [rsp, +, 112], xmm7
 movaps  xmmword, ptr, [rsp, +, 96], xmm6
 mov     rcx, r14
 mov     rdx, r15
 mov     r8, rbx
 call    core::core_arch::x86::avx2::_mm256_add_epi32
 movaps  xmm6, xmmword, ptr, [rsp, +, 64]
 movaps  xmm7, xmmword, ptr, [rsp, +, 80]
 offset += 8;
 add     rsi, 8
 while offset < total {
 add     rdi, 32
 cmp     rsi, r12
...

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

примечание: // this can be reduced with hadd. - да, если вы оптимизируете размер кода в байтах или количество инструкций, а не скорость. В противном случае см. Самый быстрый способ сделать горизонтальную векторную сумму SSE (или другое сокращение). Если вам нужен простой исходный код, сохраните его в массиве tmp и переберите его.

Peter Cordes 09.04.2022 10:53

Вы скомпилировали с включенной оптимизацией и эквивалентом gcc/clang -march=haswell или -mavx2? (Или -mavx512vl вместо _mm256_load_epi32)? Если нет, это может быть тем, что побеждает встраивание. Я попробовал godbolt.org/z/ffrr5P87z после того, как понял, что ваш минимальный воспроизводимый пример отсутствует #![feature(stdsimd)] вверху, и да, он плохо компилируется, но это без параметров арки. Ассемблер использует только инструкции SSE2 даже при вызовах функций AVX2, а не vmovaps.

Peter Cordes 09.04.2022 10:56

Извините за отсутствующую функцию (stdsimd) в репо. Я скомпилировал код с rustflags = "-C target-cpu=native" в config.toml. Я также попробовал использовать явный avx2 через RUSTFLAGS = "-C target-feature=+avx2" `

Klark 09.04.2022 11:11

Это не имеет ничего общего со stdsimd. Вам нужно сделать репортер статическим или динамическим обнаружением функции процессора. -C target-cpu=native должен отлично работать для включения avx2 везде. См.: doc.rust-lang.org/std/arch/index.html#cpu-feature-detection

BurntSushi5 09.04.2022 14:08

Н.Б. На большинстве архитектур [v]paddd имеет задержку 1 цикл, но пропускную способность 1/3 или 1/2, поэтому вам следует использовать как минимум два отдельных аккумулятора, если вы в любом случае не ограничены оперативной памятью или не имеете других операций, которые выполняются параллельно. (Ваша проблема, по-видимому, больше в том, что SIMD-вызовы, конечно, не встроены)

chtz 09.04.2022 14:08

@chtz: источник памяти vpaddd имеет пропускную способность только 2/такт, ограниченную пропускной способностью загрузочного порта (uops.info). За исключением 3/час на озере Олдер, я не знал, что они добавили еще один загрузочный порт! Но да, 2 или 3 помогут для данных, горячих в кеше L1d. Или, более того, вы не сталкиваетесь с узким местом в задержке и пропускной способности интерфейса. LLVM обычно разворачивает такие крошечные циклы на 4 при автовекторизации, что является хорошим выбором. Но OTOH, для небольших данных, которые помещаются в L1d, ваш цикл не будет выполнять количество итераций огромный, поэтому накладные расходы при запуске становятся проблемой, если только они не всегда составляют около 16 КБ.

Peter Cordes 09.04.2022 14:13
Почему Python в конце концов умрет
Почему Python в конце концов умрет
Последние 20 лет были действительно хорошими для Python. Он прошел путь от "просто языка сценариев" до основного языка, используемого для написания...
3
6
69
1
Перейти к ответу Данный вопрос помечен как решенный

Ответы 1

Ответ принят как подходящий

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

У меня отлично работает с -O -C target-cpu=skylake-avx512 (https://godbolt.org/z/csY5or43T), поэтому он может встроить даже использованную вами загрузку AVX512VL, _mm256_load_epi321, а затем оптимизировать ее в операнд источника памяти для vpaddd ymm0, ymm0, ymmword ptr [rdi + 4*rax] (AVX2) внутри узкого цикла.

В GCC/clang в данном случае вместо работающего, но медленного ассесмента выдает ошибку вроде "сбой встраивания при вызове always_inline foobar". (См. это для деталей). Это то, с чем Rust, вероятно, должен разобраться, прежде чем это будет готово для прайм-тайма, либо быть как MSVC и фактически встроить инструкцию в функцию, используя встроенную функцию, либо отказаться от компиляции, как GCC/clang.

Сноска 1: См. Как эмулировать _mm256_loadu_epi32 с помощью gcc или clang?, если вы не хотели использовать AVX512.

С -O -C target-cpu=skylake (только AVX2) он встраивает все остальное, включая vpaddd ymm, но по-прежнему вызывает функцию, которая копирует 32 байта из памяти в память с помощью AVX vmovaps. Он требует AVX512VL для встраивания встроенного кода, но позже в процессе оптимизации он понимает, что без маскирования это просто 256-битная загрузка, которую он должен делать без раздутой инструкции AVX-512. Это немного глупо, что Intel даже предоставила версию _mm256_mask[z]_loadu_epi32 без маскировки, для которой требуется AVX-512. Или глупо, что gcc/clang/rustc считают это встроенным AVX512.

Упс, да я не увидел, что инструкция требует AVX512F. Процессор, который я использовал для тестирования, имеет только AVX2. Комментарий из связанной ветки был полезен - Just use _mm256_loadu_si256 like a normal person. :). Удаление _mm256_load_epi32 решило проблему.

Klark 09.04.2022 19:14

Вы можете включить его для каждой функции с помощью #[target_feature(enable = "avx2")] (к сожалению, он пока не поддерживает AVX-512). rust.godbolt.org/z/TTfa3fErx

Chayim Friedman 10.04.2022 01:08

Другие вопросы по теме