Как реализовать операции вставки и извлечения 16- и 32-битных целых чисел с помощью AVX-512?

Связанный:

• Как мне написать QuadWord из AVX512, регистр zmm26, в регистр rax? - большая часть этого также применима к извлечению 32- и 16-битных элементов.
• Переместите int64_t в максимальное quadword вектора AVX2 __m256i показывает вставку, и снова большая часть этого должна работать для 32 и, вероятно, 16-битных элементов. (Хотя vpblendw повторяет управление смешением для обеих полос, в отличие от vpblendd). И это не использует преимущества AVX512, например, трансляция с маской слияния.
• Загрузка xmm из регистров GP показывает, как AVX512 может использовать широковещательные передачи с маской слияния. Но я не стал писать встроенные функции для ассемблера.

В AVX есть инструкции для вставки и извлечения 16- и 32-битных целых чисел в векторы __m256i:

Нет, внутренние _mm256_insert_epi16 и epi32 "фальшивые"; их нужно эмулировать с помощью нескольких инструкций, точно так же _mm_set_epi32(a,b,c,d) не является неотъемлемой частью какой-либо отдельной инструкции.

IDK, почему Intel решила предоставить их для версий AVX1 / 2, но не для версий AVX512; возможно, они думали, что это заставит людей писать неэффективный код, если они предполагали, что это стоит всего лишь одного перетасовки.

vpinsrd ymm_dst, ymm_src, r/m32, imm8 (или ZMM), к сожалению, не существует, только xmm. (https://www.felixcloutier.com/x86/pinsrb:pinsrd:pinsrq). Версию XMM нельзя использовать на __m256i, потому что она обнуляет старшие 128 бит. См. Использование регистров ymm в качестве места хранения, подобного памяти (Вы можете вставить в младшие 128 бит YMM, используя устаревшую кодировку SSE pinsrd xmm, r/m32, imm, но это катастрофически медленно для Haswell из-за того, как там работают штрафы за переход SSE / AVX. Но нормально для Skylake или Ryzen. Тем не менее, компиляторы будут никогда не испускать это.)

_mm256_insert_epi32 может компилироваться с AVX2 для широковещательной загрузки и vpblendd для вставки двойного слова из памяти. Или, что еще хуже, с целым числом, которое было в регистре, компилятор может vmovd его в регистр xmm, транслировать его в YMM, а затем смешивать. (Как я показал, как делать вручную в Переместите int64_t в старшие четверные слова вектора AVX2 __m256i)

«Подходящая» реализация зависит от окружающего кода.

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

Однако для одной вставки AVX512F на самом деле имеет несколько хороших возможностей: он имеет тасование с двумя входами, например vpermt2d, которое вы можете использовать для вставки элемента снизу одного x / y / zmm в любую позицию в другой вектор (принимая все остальные элементы назначения из этого другого вектора в качестве источника).

Но наиболее полезной здесь является скрытая трансляция: uops.info подтверждает что VPBROADCASTW zmm0{k1}, eax - это однократная инструкция с задержкой в ​​3 цикла от вектора к вектору (для слияния) и от маски к вектору. И ‹= 5 циклов задержки от eax до результата слияния. Единственная проблема заключается в настройке маски, но, надеюсь, ее можно вывести из цикла для инвариантной позиции вставки.

#include <immintrin.h>
#include <stdint.h>
__m512i _mm512_insert32(__m512i target, uint32_t x, const int pos)
{
    return _mm512_mask_set1_epi32(target, 1UL<<pos, x);
}

компилирует на Godbolt к этому asm:

# gcc8.3 -O3 -march=skylake-avx512
_mm512_insert32(long long __vector(8), unsigned int, int):
        mov     eax, 1
        shlx    eax, eax, esi
        kmovw   k1, eax                    # mask = 1<<pos
        vpbroadcastd    zmm0{k1}, edi
        ret

(gcc9 тратит лишние инструкции, копируя ESI без причины).

С константой времени компиляции pos вы получите код вида mov eax,2 / kmovw k1, eax; маскированная трансляция, вероятно, по-прежнему лучший выбор.

Это работает для 8, 16, 32 или 64-битных элементов. 8 и 16, конечно, требуют AVX512BW для узкой трансляции vpbroadcastb/w, а 32 и 64 требуется только AVX512F.

Извлекать:

Просто перетасуйте нужный элемент в конец __m512i, где вы можете использовать _mm_cvtsi128_si32. (После _mm512_castsi512_si128). Полезное перемешивание - это valignd для сдвига или поворота по элементам двойного слова, позволяя вам эффективно переместить любой элемент в нижнюю часть вектора без необходимости в векторном элементе управления. https://www.felixcloutier.com/x86/valignd:valignq

Чтобы завершить ответ Питера, вот реализации 16- и 32-битных методов вставки / извлечения:

#if defined(__GNUC__)

int _mm512_cvtsi512_si32(__m512i a)
{
    __v16si b = (__v16si) a;
    return b[0];
}

#endif

__m512i _mm512_insert_epi16(__m512i target, const std::int16_t x, const int index)
{
    return _mm512_mask_set1_epi16(target, 1UL << index, x);
}
static inline __m512i _mm512_insert_epi32(__m512i target, const std::int32_t x, const int index)
{
    return _mm512_mask_set1_epi32(target, 1UL << index, x);
}

template <int index>
int _mm512_extract_epi32(__m512i target)
{
    return _mm512_cvtsi512_si32(_mm512_alignr_epi32(target, target, index));
}
template <int index>
int  _mm512_extract_epi16(__m512i target)
{
    return (_mm512_extract_epi32<index / 2>(target) >> (index % 2 ? 16 : 0)) & 0xFFFF;
}

См. пример.