AVX2, что является наиболее эффективным способом для упаковки влево на основе маски?

Если у вас есть входной массив и выходной массив, но вы хотите только написать те элементы, которые передают определенное условие, что было бы самым эффективным способом сделать это в AVX2?

Я видел в SSE, где это было сделано следующим образом: (От: https://deplinenoise.files.wordpress.com/2015/03/gdc2015_afredriksson_simd.pdf)

__m128i LeftPack_SSSE3(__m128 mask, __m128 val)
{
 // Move 4 sign bits of mask to 4-bit integer value.
 int mask = _mm_movemask_ps(mask);
 // Select shuffle control data
 __m128i shuf_ctrl = _mm_load_si128(&shufmasks[mask]);
 // Permute to move valid values to front of SIMD register
 __m128i packed = _mm_shuffle_epi8(_mm_castps_si128(val), shuf_ctrl);
 return packed;
}

Это кажется прекрасным для SSE, который имеет ширину 4 и, следовательно, ему нужен только 16-разрядный LUT, но для AVX, ширина которого 8, LUT становится довольно большим (256 записей, каждый 32 байта или 8 тыс.).

Я удивлен, что у AVX нет инструкции для упрощения этого процесса, например, в масках с упаковкой.

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

Кто-нибудь знает какие-либо трюки, чтобы сделать это с помощью AVX2? Или что является наиболее эффективным методом?

Вот иллюстрация проблемы левой упаковки из приведенного выше документа:

Left.Packing.Problem

Спасибо

Ответ 1

AVX2 + BMI2. Смотрите мой другой ответ для AVX512. (Обновление: pdep сохранен в 64-битных сборках.)

Мы можем использовать AVX2 vpermps (_mm256_permutevar8x32_ps) (или целочисленный эквивалент, vpermd), чтобы сделать vpermd через полосу.

Мы можем генерировать маски на лету, поскольку BMI2 pext (Parallel Bits Extract) предоставляет нам побитовую версию необходимой нам операции.

Помните, что pdep/pext очень медленно работают на процессорах AMD, например, задержка цикла 6 моп /18 и пропускная способность на Ryzen. Эта реализация будет ужасно работать на AMD. Для AMD вам лучше всего использовать 128-битные векторы, использующие LUT pshufb или vpermilps, или некоторые из предложенных в комментариях предложений по переменному смещению AVX2, если в качестве входной маски используется векторная маска (а не уже вычисленная битовая маска из памяти). AMD до Zen2 в любом случае имеет только 128-битные векторные исполнительные блоки, а 256-битные тасовки пересекают полосы медленно. Таким образом, 128-битные векторы очень привлекательны для нынешней AMD.


Для целочисленных векторов с 32-разрядными или более широкими элементами: 1) _mm256_movemask_ps(_mm256_castsi256_ps(compare_mask)).
Или 2) используйте _mm256_movemask_epi8 а затем измените первую константу PDEP с 0x0101010101010101 на 0x0F0F0F0F0F0F0F0F для разброса блоков из 4 смежных битов. Измените умножение на 0xFFU на extended_mask expanded_mask |= expanded_mask<<4; extended_mask expanded_mask |= expanded_mask<<4; или expanded_mask *= 0x11; (Не испытано). В любом случае, используйте маску тасования с VPERMD вместо VPERMPS.

Для 64-битных целых или double элементов все по-прежнему просто работает; Просто в маске сравнения всегда есть пары одинаковых 32-битных элементов, поэтому итоговый случайный порядок помещает обе половины каждого 64-битного элемента в нужное место. (Таким образом, вы все еще используете VPERMPS или VPERMD, потому что VPERMPD и VPERMQ доступны только с операндами непосредственного управления.)

Для 16-битных элементов вы можете адаптировать это с 128-битными векторами.


Алгоритм:

Начните с константы упакованных 3-битных индексов, где каждая позиция имеет свой собственный индекс. т.е. [ 7 6 5 4 3 2 1 0 ] где каждый элемент имеет ширину 3 бита. 0b111'110'101'...'010'001'000.

Используйте pext чтобы извлечь нужные нам индексы в непрерывную последовательность внизу целочисленного регистра. Например, если нам нужны индексы 0 и 2, наша контрольная маска для pext должна быть 0b000'...'111'000'111. pext будет захватывать группы индексов 010 и 000 которые совпадают с 1 битом в селекторе. Выбранные группы упаковываются в младшие биты вывода, поэтому на выходе будет 0b000'...'010'000. (т.е. [... 2 0 ])

См. 0b111000111 код, чтобы узнать, как сгенерировать ввод 0b111000111 для pext из маски входного вектора.

Теперь мы находимся в одной лодке со сжатым LUT: распаковываем до 8 упакованных индексов.

К тому времени, когда вы pdep все кусочки, будет всего три pext/pdep s. Я работал в обратном направлении от того, что я хотел, так что, вероятно, легче всего понять это и в этом направлении. (то есть начните с линии тасования и оттуда работайте задом наперед.)

Мы можем упростить распаковку, если будем работать с индексами по одному на байт вместо упакованных 3-битных групп. Поскольку у нас есть 8 индексов, это возможно только с 64-битным кодом.

Смотрите эту и 32-битную версию в Godbolt Compiler Explorer. Я использовал #ifdef поэтому он оптимально компилируется с -m64 или -m32. GCC тратит некоторые инструкции, но Clang делает действительно хороший код.

#include <stdint.h>
#include <immintrin.h>

// Uses 64bit pdep / pext to save a step in unpacking.
__m256 compress256(__m256 src, unsigned int mask /* from movmskps */)
{
  uint64_t expanded_mask = _pdep_u64(mask, 0x0101010101010101);  // unpack each bit to a byte
  expanded_mask *= 0xFF;    // mask |= mask<<1 | mask<<2 | ... | mask<<7;
  // ABC... -> AAAAAAAABBBBBBBBCCCCCCCC...: replicate each bit to fill its byte

  const uint64_t identity_indices = 0x0706050403020100;    // the identity shuffle for vpermps, packed to one index per byte
  uint64_t wanted_indices = _pext_u64(identity_indices, expanded_mask);

  __m128i bytevec = _mm_cvtsi64_si128(wanted_indices);
  __m256i shufmask = _mm256_cvtepu8_epi32(bytevec);

  return _mm256_permutevar8x32_ps(src, shufmask);
}

Это компилируется в код без загрузок из памяти, только с непосредственными константами. (Смотрите ссылку на Godbolt для этой и 32-битной версии).

    # clang 3.7.1 -std=gnu++14 -O3 -march=haswell
    mov     eax, edi                   # just to zero extend: goes away when inlining
    movabs  rcx, 72340172838076673     # The constants are hoisted after inlining into a loop
    pdep    rax, rax, rcx              # ABC       -> 0000000A0000000B....
    imul    rax, rax, 255              # 0000000A0000000B.. -> AAAAAAAABBBBBBBB..
    movabs  rcx, 506097522914230528
    pext    rax, rcx, rax
    vmovq   xmm1, rax
    vpmovzxbd       ymm1, xmm1         # 3c latency since this is lane-crossing
    vpermps ymm0, ymm1, ymm0
    ret

Таким образом, согласно числам Агнера Тумана, это 6 мопов (не считая констант, или расширяющегося нулями mov, который исчезает при встраивании). На Intel Haswell это задержка 16c (1 для vmovq, 3 для каждого pdep/imul/pext/vpmovzx/vpermps). Там нет параллелизма на уровне команд. В цикле, где это не является частью переносимой в цикле зависимости (как, например, тот, который я включил в ссылку Godbolt), узкое место, как мы надеемся, просто пропускная способность, сохраняя сразу несколько итераций этого в полете.

Это может управлять пропускной способностью один на 3 цикла, узкое место на порту 1 для pdep/pext/imul. Конечно, с нагрузками/хранилищами и накладными расходами цикла (включая сравнение, movmsk и popcnt) общая пропускная способность uop может легко стать проблемой. (Например, цикл фильтра в моей ссылке на Godbolt равен 14 моп с clang, с -fno-unroll-loops чтобы его было легче читать. Он может выдержать одну итерацию на 4c, не отставая от -fno-unroll-loops интерфейса, если нам повезет, но я думаю, что clang не смог учесть ложную зависимость popcnt от своего вывода, поэтому он будет узким местом на 3/5 задержки латентности функции compress256.)

gcc умножает на 0xFF с несколькими инструкциями, используя сдвиг влево на 8 и sub. Для этого требуются дополнительные инструкции mov, но конечным результатом является умножение с задержкой 2. (Haswell обрабатывает mov на этапе переименования регистра с нулевой задержкой.)


Поскольку все оборудование, поддерживающее AVX2, также поддерживает BMI2, вероятно, нет смысла предоставлять версию для AVX2 без BMI2.

Если вам нужно сделать это в очень длинном цикле, LUT, вероятно, того стоит, если начальные ошибки кэширования амортизируются в течение достаточного количества итераций с меньшими накладными расходами на простую распаковку записи LUT. Вам все еще нужно использовать movmskps, поэтому вы можете открыть маску и использовать ее в качестве индекса LUT, но вы сохраните pdep/imul/pexp.

Вы можете распаковать записи LUT с той же самой целочисленной последовательностью, которую я использовал, но @Froglegs set1()/vpsrlvd/vpand, вероятно, лучше, когда запись LUT начинается в памяти и не нуждается в целочисленных регистрах. (32-битная широковещательная загрузка не требует ALU-моп на процессорах Intel). Однако переменное смещение составляет 3 мопа на Haswell (но только 1 на Skylake).

Ответ 2

Если вы ориентируетесь на AMD Zen, этот метод может оказаться предпочтительным из-за очень медленного pdepand pext на ризене (18 циклов каждый).

Я придумал этот метод, который использует сжатую LUT, которая составляет 768 (заполнение +1) байтов вместо 8k. Это требует широковещательной передачи одного скалярного значения, которое затем сдвигается на разную величину в каждой полосе, а затем маскируется на младшие 3 бита, что обеспечивает 0-7 LUT.

Вот внутренняя версия, а также код для построения LUT.

//Generate Move mask via: _mm256_movemask_ps(_mm256_castsi256_ps(mask)); etc
__m256i MoveMaskToIndices(u32 moveMask) {
    u8 *adr = g_pack_left_table_u8x3 + moveMask * 3;
    __m256i indices = _mm256_set1_epi32(*reinterpret_cast<u32*>(adr));//lower 24 bits has our LUT

   // __m256i m = _mm256_sllv_epi32(indices, _mm256_setr_epi32(29, 26, 23, 20, 17, 14, 11, 8));

    //now shift it right to get 3 bits at bottom
    //__m256i shufmask = _mm256_srli_epi32(m, 29);

    //Simplified version suggested by wim
    //shift each lane so desired 3 bits are a bottom
    //There is leftover data in the lane, but _mm256_permutevar8x32_ps  only examines the first 3 bits so this is ok
    __m256i shufmask = _mm256_srlv_epi32 (indices, _mm256_setr_epi32(0, 3, 6, 9, 12, 15, 18, 21));
    return shufmask;
}

u32 get_nth_bits(int a) {
    u32 out = 0;
    int c = 0;
    for (int i = 0; i < 8; ++i) {
        auto set = (a >> i) & 1;
        if (set) {
            out |= (i << (c * 3));
            c++;
        }
    }
    return out;
}
u8 g_pack_left_table_u8x3[256 * 3 + 1];

void BuildPackMask() {
    for (int i = 0; i < 256; ++i) {
        *reinterpret_cast<u32*>(&g_pack_left_table_u8x3[i * 3]) = get_nth_bits(i);
    }
}

Вот сборка, сгенерированная MSVC:

  lea ecx, DWORD PTR [rcx+rcx*2]
  lea rax, OFFSET FLAT:unsigned char * g_pack_left_table_u8x3 ; g_pack_left_table_u8x3
  vpbroadcastd ymm0, DWORD PTR [rcx+rax]
  vpsrlvd ymm0, ymm0, YMMWORD PTR [email protected]000000

Ответ 3

Смотрите мой другой ответ для AVX2 + BMI2 без LUT.

Поскольку вы упоминаете о проблеме масштабируемости для AVX512: не беспокойтесь, есть инструкция AVX512F именно для этого:

VCOMPRESSPS - Храните разреженные упакованные значения с плавающей запятой одинарной точности в плотной памяти. (Существуют также версии для двойных и 32- или 64-битных целочисленных элементов (vpcompressq), но не для байта или слова (16-бит)). Это как BMI2 pdep/pext, но для векторных элементов вместо битов в целочисленном регистре.

Пункт назначения может быть векторным регистром или операндом памяти, в то время как источником является вектор и регистр маски. С регистром dest он может объединять или обнулять старшие биты. С помощью элемента dest памяти "в область памяти назначения записывается только непрерывный вектор".

Чтобы выяснить, как далеко продвинется указатель на следующий вектор, попкорн маску.

Допустим, вы хотите отфильтровать все, кроме значений> = 0 из массива:

#include <stdint.h>
#include <immintrin.h>
size_t filter_non_negative(float *__restrict__ dst, const float *__restrict__ src, size_t len) {
    const float *endp = src+len;
    float *dst_start = dst;
    do {
        __m512      sv  = _mm512_loadu_ps(src);
        __mmask16 keep = _mm512_cmp_ps_mask(sv, _mm512_setzero_ps(), _CMP_GE_OQ);  // true for src >= 0.0, false for unordered and src < 0.0
        _mm512_mask_compressstoreu_ps(dst, keep, sv);   // clang is missing this intrinsic, which can't be emulated with a separate store

        src += 16;
        dst += _mm_popcnt_u64(keep);   // popcnt_u64 instead of u32 helps gcc avoid a wasted movsx, but is potentially slower on some CPUs
    } while (src < endp);
    return dst - dst_start;
}

Это компилируется (с gcc4.9 или новее) в ( %23include+ //%23include+ size_t filter_non_negative(float+*__restrict__ dst, const float+*__restrict__ src, size_t len) {%0A++++const float+*endp+= src+len;%0A++++float+*dst_start+= dst;%0A++++do {%0A++++ __m512 sv = _mm512_loadu_ps(src);%0A++++ __mmask16+keep = _mm512_cmp_ps_mask(sv,+_mm512_setzero_ps(),+_CMP_GE_OQ)%3B++//true for src+>= 0.0, false for unordered and src+< 0.0%0A++++ _mm512_mask_compressstoreu_ps(dst,+keep, sv)%3B++//clang is missing this intrinsic, which can!'t be emulated with a separate store %0A++++ src++= 16;%0A++++ dst += _mm_popcnt_u64(keep)%3B++//popcnt_u64 instead of u32 helps gcc avoid a wasted movsx, but is potentially slower on some CPUs%0A++++} while (src+< endp);%0A++++return dst - dst_start; } ')),filterAsm:(commentOnly:!t,directives:!t,intel:!t,labels:!t),version:3 rel="nofollow noreferrer">Godbolt Compiler Explorer):

 # Output from gcc6.1, with -O3 -march=haswell -mavx512f.  Same with other gcc versions
    lea     rcx, [rsi+rdx*4]             # endp
    mov     rax, rdi
    vpxord  zmm1, zmm1, zmm1             # vpxor  xmm1, xmm1,xmm1 would save a byte, using VEX instead of EVEX
.L2:
    vmovups zmm0, ZMMWORD PTR [rsi]
    add     rsi, 64
    vcmpps  k1, zmm0, zmm1, 29           # AVX512 compares have mask regs as a destination
    kmovw   edx, k1                      # There are some insns to add/or/and mask regs, but not popcnt
    movzx   edx, dx                      # gcc is dumb and doesn't know that kmovw already zero-extends to fill the destination.
    vcompressps     ZMMWORD PTR [rax]{k1}, zmm0
    popcnt  rdx, rdx
    ## movsx   rdx, edx         # with _popcnt_u32, gcc is dumb.  No casting can get gcc to do anything but sign-extend.  You'd expect (unsigned) would mov to zero-extend, but no.
    lea     rax, [rax+rdx*4]             # dst += ...
    cmp     rcx, rsi
    ja      .L2

    sub     rax, rdi
    sar     rax, 2                       # address math -> element count
    ret

Производительность: 256-битные векторы могут быть быстрее на Skylake-X/Cascade Lake

Теоретически цикл, который загружает растровое изображение и фильтрует один массив в другой, должен работать с 1 вектором на 3 такта в SKX/CSLX, независимо от ширины вектора, узким местом на порту 5. (kmovb/w/d/q k1, eax работает на p5, а vcompressps в память - это 2p5 + хранилище, согласно IACA и тестированию http://uops.info/).

@ZachB в комментариях сообщает, что на практике цикл с использованием ZMM _mm512_mask_compressstoreu_ps немного медленнее, чем _mm256_mask_compressstoreu_ps на реальном оборудовании CSLX. (Я не уверен, был ли это микробенчмарк, который позволил бы 256-битной версии выйти из "512-битного векторного режима" и повысить тактовую частоту, или был ли окружающий 512-битный код.)

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

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

Было бы интересно vcompressps в регистре с отдельными полными векторами, перекрывающими друг друга. Это, вероятно, тот же мопс, но в магазине может возникнуть микроплавкость, когда это отдельная инструкция. И если есть какая-то разница между магазинами в масках и перекрывающимися магазинами, это выявит это.


Другая идея, обсуждаемая в комментариях ниже, заключалась в использовании vpermt2ps для создания полных векторов для выровненных хранилищ. Это было бы трудно сделать без ветвления, и ветвление, когда мы заполняем вектор, вероятно, будет неверно предсказано, если только у битовой маски нет довольно регулярного паттерна или больших серий all-0 и all-1.

Реализация без ответвлений с переносимой в цикле цепочкой зависимостей из 4 или 6 циклов по построенному вектору может быть возможной, с vpermt2ps и смесью или чем-то, что может заменить его, когда он "заполнен". С выровненным вектором сохраняйте каждую итерацию, но перемещая выходной указатель, только если вектор заполнен.

Это, вероятно, медленнее, чем vcompressps с невыровненными хранилищами на текущих процессорах Intel.

Ответ 4

В случае, если кто-то заинтересован в этом, это решение для SSE2, которое использует инструкцию LUT вместо данных LUT, а также таблицу перехода. С AVX это потребует 256 случаев.

Каждый раз, когда вы вызываете LeftPack_SSE2 ниже, он использует по существу три команды: jmp, shufps, jmp. Пять из шестнадцати случаев не нуждаются в изменении вектора.

static inline __m128 LeftPack_SSE2(__m128 val, int mask)  {
  switch(mask) {
  case  0:
  case  1: return val;
  case  2: return _mm_shuffle_ps(val,val,0x01);
  case  3: return val;
  case  4: return _mm_shuffle_ps(val,val,0x02);
  case  5: return _mm_shuffle_ps(val,val,0x08);
  case  6: return _mm_shuffle_ps(val,val,0x09);
  case  7: return val;
  case  8: return _mm_shuffle_ps(val,val,0x03);
  case  9: return _mm_shuffle_ps(val,val,0x0c);
  case 10: return _mm_shuffle_ps(val,val,0x0d);
  case 11: return _mm_shuffle_ps(val,val,0x34);
  case 12: return _mm_shuffle_ps(val,val,0x0e);
  case 13: return _mm_shuffle_ps(val,val,0x38);
  case 14: return _mm_shuffle_ps(val,val,0x39);
  case 15: return val;
  }
}

__m128 foo(__m128 val, __m128 maskv) {
  int mask = _mm_movemask_ps(maskv);
  return LeftPack_SSE2(val, mask);
}