Предыдущая часть вызвала бурную дискуссию, в ходе которой выяснилось, что AVX/AVX2 на самом деле есть в десктопных CPU, нет только AVX512. Поэтому продолжаем знакомиться с SIMD, но уже с современной его частью — AVX. А так же разберём некоторые комментарии:

• медленнее ли _mm256_load_si256, чем прямое обращение к памяти?
• влияет ли на скорость использование AVX команд над SSE регистрами?
• действительно ли так плохо использовать _popcnt?

Немного про AVX

AVX/AVX2 — это более мощная версия SSE, которая расширяет большинство 128 битных SSE операций до 256 бит, плюс приносит ряд новых инструкций.

Из тонкостей реализации можно выделить то, что на уровне ассемблера AVX использует 3 аргумента, что позволяет не разрушать данные в первых двух. SSE сохраняет результат в одном из аргументов.

Так же нужно учитывать, что при прямой адресации данные должны быть выровнены по 32 байта, в SSE выравнивание по 16.

Дополненная версия бенчмарка

Изменения:

1. Количество элементов увеличено в 10 000 раз (до 10 240 000), чтобы гарантированно не вместиться в кэш процессора.
2. Выравнивание изменено с 16 байт на 32 для поддержки AVX.
3. Добавлены AVX реализации аналогичные SSE.

#include <benchmark/benchmark.h>
#include <x86intrin.h>
#include <cstring>

#define ARR_SIZE 10240000
#define VAL 50

static int16_t *getRandArr() {
    auto res = new int16_t[ARR_SIZE];
    for (int i = 0; i < ARR_SIZE; ++i) {
        res[i] = static_cast<int16_t>(rand() % (VAL * 2));
    }
    return res;
}
static auto arr = getRandArr();

static int16_t *getAllignedArr() {
    auto res = aligned_alloc(32, sizeof(int16_t) * ARR_SIZE);
    memcpy(res, arr, sizeof(int16_t) * ARR_SIZE);
    return static_cast<int16_t *>(res);
}
static auto allignedArr = getAllignedArr();

static void BM_Count(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;
        for (int i = 0; i < ARR_SIZE; ++i)
            if (arr[i] == VAL)
                ++cnt;
        benchmark::DoNotOptimize(cnt);
    }
}

BENCHMARK(BM_Count);

static void BM_SSE_COUNT_SET_EPI(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto sseVal = _mm_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 8) {
            cnt += _popcnt32(
                    _mm_movemask_epi8(
                            _mm_cmpeq_epi16(
                                    sseVal,
                                    _mm_set_epi16(arr[i + 7], arr[i + 6], arr[i + 5], arr[i + 4],
                                                  arr[i + 3], arr[i + 2], arr[i + 1], arr[i])
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_SSE_COUNT_SET_EPI);

static void BM_SSE_COUNT_LOADU(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto sseVal = _mm_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 8) {
            cnt += _popcnt32(
                    _mm_movemask_epi8(
                            _mm_cmpeq_epi16(
                                    sseVal,
                                    _mm_loadu_si128((__m128i *) &arr[i])
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_SSE_COUNT_LOADU);

static void BM_SSE_COUNT_DIRECT(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto sseVal = _mm_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 8) {
            cnt += _popcnt32(
                    _mm_movemask_epi8(
                            _mm_cmpeq_epi16(
                                    sseVal,
                                    *(__m128i *) &allignedArr[i]
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_SSE_COUNT_DIRECT);

#ifdef __AVX2__

static void BM_AVX_COUNT_LOADU(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto avxVal = _mm256_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 16) {
            cnt += _popcnt32(
                    _mm256_movemask_epi8(
                            _mm256_cmpeq_epi16(
                                    avxVal,
                                    _mm256_loadu_si256((__m256i *) &arr[i])
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_AVX_COUNT_LOADU);

static void BM_AVX_COUNT_LOAD(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto avxVal = _mm256_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 16) {
            cnt += _popcnt32(
                    _mm256_movemask_epi8(
                            _mm256_cmpeq_epi16(avxVal,
                                               _mm256_load_si256((__m256i *) &allignedArr[i])
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_AVX_COUNT_LOAD);

static void BM_AVX_COUNT_DIRECT(benchmark::State &state) {
    for (auto _ : state) {
        int64_t cnt = 0;

        auto avxVal = _mm256_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 16) {
            cnt += _popcnt32(
                    _mm256_movemask_epi8(
                            _mm256_cmpeq_epi16(
                                    avxVal,
                                    *(__m256i *) &allignedArr[i]
                            )
                    )
            );
        }
        benchmark::DoNotOptimize(cnt >> 1);
    }
}

BENCHMARK(BM_AVX_COUNT_DIRECT);

#endif

BENCHMARK_MAIN();

Новые результаты выглядят так (-O0):

---------------------------------------------------------------------
Benchmark                              Time           CPU Iterations
---------------------------------------------------------------------
BM_Count                        17226622 ns   17062958 ns         41
BM_SSE_COUNT_SET_EPI             8901343 ns    8814845 ns         79
BM_SSE_COUNT_LOADU               3664778 ns    3664766 ns        185
BM_SSE_COUNT_DIRECT              3468436 ns    3468423 ns        202
BM_AVX_COUNT_LOADU               2090817 ns    2090796 ns        343
BM_AVX_COUNT_LOAD                1904424 ns    1904419 ns        364
BM_AVX_COUNT_DIRECT              1814875 ns    1814854 ns        385

Итого суммарное ускорение в 9+ раз, AVX ожидаемо быстрей SSE почти в 2 раза.

Медленнее ли _mm256_load_si256, чем прямое обращение к памяти?

Однозначного ответа нет. С -O0 медленнее прямого обращения, но быстрее _mm256_loadu_si256:

---------------------------------------------------------------------
Benchmark                              Time           CPU Iterations
---------------------------------------------------------------------
BM_AVX_COUNT_LOADU               2090817 ns    2090796 ns        343
BM_AVX_COUNT_LOAD                1904424 ns    1904419 ns        364
BM_AVX_COUNT_DIRECT              1814875 ns    1814854 ns        385

С -O3 быстрее, чем прямое обращение к памяти, но всё ещё ожидаемо медленней _mm256_loadu_si256.

---------------------------------------------------------------------
Benchmark                              Time           CPU Iterations
---------------------------------------------------------------------
BM_AVX_COUNT_LOADU                992319 ns     992368 ns        701
BM_AVX_COUNT_LOAD                 956120 ns     956166 ns        712
BM_AVX_COUNT_DIRECT              1027624 ns    1027674 ns        730

В продакшн коде всё-таки лучше использовать _mm256_load_si256 вместо прямого обращения, этот вариант компилятор умеет лучше оптимизировать.

Влияет ли на скорость использование AVX команд над SSE регистрами?

Короткий ответ — нет. Для эксперимента я собрал и запустил бенчмарк с -mavx2 и с -msse4.2.

-mavx2

_popcnt32(_mm_movemask_epi8(_mm_cmpeq_epi16(...))) превращается в

vpcmpeqw %xmm1,%xmm0,%xmm0
vpmovmskb %xmm0,%edx
popcnt %edx,%edx

Результаты:

------------------------------------------------------------
Benchmark                     Time           CPU Iterations
------------------------------------------------------------
BM_SSE_COUNT_SET_EPI    9376699 ns    9376767 ns         75
BM_SSE_COUNT_LOADU      4425510 ns    4425560 ns        159
BM_SSE_COUNT_DIRECT     3938604 ns    3938648 ns        177

-msse4.2

_popcnt32(_mm_movemask_epi8(_mm_cmpeq_epi16(...))) превращается в

pcmpeqw %xmm1,%xmm0
pmovmskb %xmm0,%edx
popcnt %edx,%edx

Результаты:

------------------------------------------------------------
Benchmark                     Time           CPU Iterations
------------------------------------------------------------
BM_SSE_COUNT_SET_EPI    9309352 ns    9309375 ns         76
BM_SSE_COUNT_LOADU      4382183 ns    4382195 ns        159
BM_SSE_COUNT_DIRECT     3944579 ns    3944590 ns        176

bonus

AVX команды _popcnt32(_mm256_movemask_epi8(_mm256_cmpeq_epi16(...)))превращаются в

vpcmpeqw %ymm1,%ymm0,%ymm0
vpmovmskb %ymm0,%edx
popcnt %edx,%edx

Действительно ли так плохо использовать _popcnt?

В одном из комментариев Antervis написал:

А еще, ты несколько недоработал алгоритм. Зачем делать через movemask + popcnt? Для массивов не более 2^18 элементов можно сначала собирать поэлементную сумму:
auto cmp = _mm_cmpeq_epi16(sseVal, sseArr);
cmp = _mm_and_si128(cmp, _mm_set1_epi16(1));
sum = _mm_add_epi16(sum, cmp);

а потом, в конце цикла, сделать одно горизонтальное сложение (не забывая про переполнение).

Я сделал бенчмарк

static void BM_AVX_COUNT_DIRECT_WO_POPCNT(benchmark::State &state) {
    auto avxVal1 = _mm256_set1_epi16(1);
    for (auto _ : state) {
        auto sum = _mm256_set1_epi16(0);

        auto avxVal = _mm256_set1_epi16(VAL);
        for (int i = 0; i < ARR_SIZE; i += 16) {
            sum = _mm256_add_epi16(
                    sum,
                    _mm256_and_si256(
                            avxVal1,
                            _mm256_cmpeq_epi16(
                                    avxVal,
                                    *(__m256i *) &allignedArr[i])
                    )
            );
        }

        auto arrSum = (uint16_t *) &sum;
        size_t cnt = 0;
        for (int j = 0; j < 16; ++j)
            cnt += arrSum[j];

        benchmark::DoNotOptimize(cnt >> 1);
    }
}

и он оказался медленней c -O0:

---------------------------------------------------------------------
Benchmark                              Time           CPU Iterations
---------------------------------------------------------------------
BM_AVX_COUNT_DIRECT              1814821 ns    1814785 ns        392
BM_AVX_COUNT_DIRECT_WO_POPCNT    2386289 ns    2386227 ns        287

и немного быстрее с -O3:

---------------------------------------------------------------------
Benchmark                              Time           CPU Iterations
---------------------------------------------------------------------
BM_AVX_COUNT_DIRECT               960941 ns     960924 ns        722
BM_AVX_COUNT_DIRECT_WO_POPCNT     948611 ns     948596 ns        732