c++ - 紧凑的 AVX2 寄存器，因此选定的整数根据掩码是连续的

This question already has answers here:

AVX2 what is the most efficient way to pack left based on a mask?

(5 个回答)


4年前关闭。




在问题 Optimizing Array Compaction 中，最佳答案指出:



Haswell (AVX2) 可以做到这一点吗？或者它是否需要 AVX512 的一种口味？

我有一个包含 int32s 的 AVX2 vector ，以及一个比较结果的相应 vector 。我想以某种方式对其进行洗牌，以便在掩码中设置相应 msb 的元素(比较 true)在 vector 的低端是连续的。

我能看到的最好的方法是使用 _mm256_movemask_ps/vmovmskps(没有 *d 变体？)获取位掩码，然后在 256 AVX2 vector 查找表中使用它来获取跨车道的随机掩码 _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");
}

关于c++ - 紧凑的 AVX2 寄存器，因此选定的整数根据掩码是连续的，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/25074197/