Компактный регистр AVX2, чтобы выбранные целые числа были смежными в соответствии с маской
Вопрос
В вопросе Оптимизация сжатия массива, верхний ответ гласит:
Регистры SSE/AVX с новейшими наборами инструкций обеспечивают лучший подход.Мы можем использовать результат PMOVMSKB напрямую, преобразуя его в регистр управления для чего-то вроде PSHUFB.
Возможно ли это с Haswell (AVX2)?Или для этого требуется одна из разновидностей AVX512?
У меня есть вектор AVX2, содержащий int32, и соответствующий вектор результата сравнения.Я хочу как-то перетасовать его, чтобы элементы с соответствующим старшим битом, установленным в маске (сравните true), были смежными в нижнем конце вектора.
Лучшее, что я вижу, это получить битовую маску с помощью _mm256_movemask_ps/vmovmskps (нет варианта *d?), а затем использовать ее в векторной таблице поиска 256 AVX2, чтобы получить маску перемешивания для перекрестной полосы _mm256_permutevar8x32_epi32/vpermd.
Решение
Первое, что нужно сделать, это найти быструю скалярную функцию.Вот версия, которая не использует ветку.
inline int compact(int *x, int *y, const int n) {
int cnt = 0;
for(int i=0; i<n; i++) {
int cut = x[i]!=0;
y[cnt] = cut*x[i];
cnt += cut;
}
return cnt;
}
Лучший результат с SIMD, вероятно, зависит от распределения нулей.Если он редкий или густой.Следующий код должен хорошо работать для разреженного или плотного распределения.Например, длинные серии нулей и ненулей.Если распределение будет более равномерным, я не знаю, принесет ли этот код какую-либо пользу.Но в любом случае это даст правильный результат.
Вот версия AVX2, которую я тестировал.
int compact_AVX2(int *x, int *y, int n) {
int i =0, cnt = 0;
for(i=0; i<n-8; i+=8) {
__m256i x4 = _mm256_loadu_si256((__m256i*)&x[i]);
__m256i cmp = _mm256_cmpeq_epi32(x4, _mm256_setzero_si256());
int mask = _mm256_movemask_epi8(cmp);
if(mask == -1) continue; //all zeros
if(mask) {
cnt += compact(&x[i],&y[cnt], 8);
}
else {
_mm256_storeu_si256((__m256i*)&y[cnt], x4);
cnt +=8;
}
}
cnt += compact(&x[i], &y[cnt], n-i); // cleanup for n not a multiple of 8
return cnt;
}
Вот версия SSE2, которую я тестировал.
int compact_SSE2(int *x, int *y, int n) {
int i =0, cnt = 0;
for(i=0; i<n-4; i+=4) {
__m128i x4 = _mm_loadu_si128((__m128i*)&x[i]);
__m128i cmp = _mm_cmpeq_epi32(x4, _mm_setzero_si128());
int mask = _mm_movemask_epi8(cmp);
if(mask == 0xffff) continue; //all zeroes
if(mask) {
cnt += compact(&x[i],&y[cnt], 4);
}
else {
_mm_storeu_si128((__m128i*)&y[cnt], x4);
cnt +=4;
}
}
cnt += compact(&x[i], &y[cnt], n-i); // cleanup for n not a multiple of 4
return cnt;
}
Вот полный тест
#include <stdio.h>
#include <stdlib.h>
#if defined (__GNUC__) && ! defined (__INTEL_COMPILER)
#include <x86intrin.h>
#else
#include <immintrin.h>
#endif
#define N 50
inline int compact(int *x, int *y, const int n) {
int cnt = 0;
for(int i=0; i<n; i++) {
int cut = x[i]!=0;
y[cnt] = cut*x[i];
cnt += cut;
}
return cnt;
}
int compact_SSE2(int *x, int *y, int n) {
int i =0, cnt = 0;
for(i=0; i<n-4; i+=4) {
__m128i x4 = _mm_loadu_si128((__m128i*)&x[i]);
__m128i cmp = _mm_cmpeq_epi32(x4, _mm_setzero_si128());
int mask = _mm_movemask_epi8(cmp);
if(mask == 0xffff) continue; //all zeroes
if(mask) {
cnt += compact(&x[i],&y[cnt], 4);
}
else {
_mm_storeu_si128((__m128i*)&y[cnt], x4);
cnt +=4;
}
}
cnt += compact(&x[i], &y[cnt], n-i); // cleanup for n not a multiple of 4
return cnt;
}
int compact_AVX2(int *x, int *y, int n) {
int i =0, cnt = 0;
for(i=0; i<n-8; i+=8) {
__m256i x4 = _mm256_loadu_si256((__m256i*)&x[i]);
__m256i cmp = _mm256_cmpeq_epi32(x4, _mm256_setzero_si256());
int mask = _mm256_movemask_epi8(cmp);
if(mask == -1) continue; //all zeros
if(mask) {
cnt += compact(&x[i],&y[cnt], 8);
}
else {
_mm256_storeu_si256((__m256i*)&y[cnt], x4);
cnt +=8;
}
}
cnt += compact(&x[i], &y[cnt], n-i); // cleanup for n not a multiple of 8
return cnt;
}
int main() {
int x[N], y[N];
for(int i=0; i<N; i++) x[i] = rand()%10;
//int cnt = compact_SSE2(x,y,N);
int cnt = compact_AVX2(x,y,N);
for(int i=0; i<N; i++) printf("%d ", x[i]); printf("\n");
for(int i=0; i<cnt; i++) printf("%d ", y[i]); printf("\n");
}