SSE/AVX寄存器的非零字节索引

Question

如果一个SSE/AVX寄存器的值是所有字节都是0或1，那么有什么办法可以有效地获得所有非零元素的索引吗？

例如，如果xmm值是 | r0 = 0 | r1 = 1 | r2 = 0 | r3 = 1 | r4 = 0 | r5 = 1 | r6 = 0 | ... | r14 = 0 | r15 = 1 | 结果应该是（1,3,5，...，15）。结果应放置在另一个_m128i变量或char [16]数组中。

如果有帮助，我们可以假设寄存器的值是所有字节都是0或某个恒定的非零值（不是必需的1）。

我很想知道是否有一个指令，或最好是C/C++内在。在任何SSE或AVX指令集中。

编辑1：

这是正确observed by @zx485是原来的问题还不够清楚。我正在寻找任何“连续”解决方案。

上面的例子0 1 0 1 0 1 0 1...应导致下列任一：

• 如果我们假设指数从1开始，然后0将是终止字节，其结果可能是

• 如果我们假设负字节是终止字节的结果可能是

001 003 005 007 009 011 013 015为0xFF 0xFF的0xFF的0xFF的0xFF的0xFF的0xFF的0xFF的

• 什么，这给出了作为连续字节，我们可以将其解释为原始值中的非零元素的索引

编辑2：

实际上，如@harold和@Peter Cordes建议在注释到原来的职位，可能的解决方案之一是先创建的掩模（例如与pmovmskb）并在那里检查非零指数。但那会导致循环。

Answer 1

如果你想要结果数组是“压缩的”，你的问题不清楚。我的意思是“压缩”，结果应该是连续的。因此，例如用于0 1 0 1 0 1 0 1...，有两种可能性：

非连续：

XMM0：000 001 000 003 000 005 000 007 000 009 000 011 000 013 000 015

连续：

XMM0：001 003 005 007 009 011 013 015 000 000 000 000 000 000 000

连续方法的一个问题是：您如何确定索引0或终止值？

我提供一个简单的解决方案，第一，非连续的方式，这应该是相当快：

.data 
    ddqZeroToFifteen    db 0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 
    ddqTestValue:     db 0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1 
.code 
    movdqa xmm0, xmmword ptr [ddqTestValue] 
    pxor xmm1, xmm1        ; zero XMM1 
    pcmpeqb xmm0, xmm1       ; set to -1 for all matching 
    pandn xmm0, xmmword ptr [ddqZeroToFifteen] ; invert and apply indices 

只是为了完整起见：第二，连续的方式，不是盖的在这个答案。

Answer 2

更新回答：新解决方案效率稍高。

您可以使用Bit Manipulation Instruction Set 2, 中的pext指令和其他几条SSE指令结合使用，而无需循环。

/* 
gcc -O3 -Wall -m64 -mavx2 -march=broadwell ind_nonz_avx.c 
*/ 

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

__m128i nonz_index(__m128i x){ 
    /* Set some constants that will (hopefully) be hoisted out of a loop after inlining. */ 
    uint64_t indx_const = 0xFEDCBA;      /* 16 4-bit integers, all possible indices from 0 o 15               */ 
    __m128i cntr   = _mm_set_epi8(64,60,56,52,48,44,40,36,32,28,24,20,16,12,8,4); 
    __m128i pshufbcnst = _mm_set_epi8(0x80,0x80,0x80,0x80,0x80,0x80,0x80,0x80, 0x0E,0x0C,0x0A,0x08,0x06,0x04,0x02,0x00); 
    __m128i cnst0F  = _mm_set1_epi8(0x0F); 

    __m128i msk   = _mm_cmpeq_epi8(x,_mm_setzero_si128()); /* Generate 16x8 bit mask.                      */ 
      msk   = _mm_srli_epi64(msk,4);     /* Pack 16x8 bit mask to 16x4 bit mask.                   */ 
      msk   = _mm_shuffle_epi8(msk,pshufbcnst);   /* Pack 16x8 bit mask to 16x4 bit mask, continued.                */ 
    uint64_t msk64  = ~ _mm_cvtsi128_si64x(msk);     /* Move to general purpose register and invert 16x4 bit mask.              */ 

                     /* Compute the termination byte nonzmsk separately.                */ 
    int64_t nnz64  = _mm_popcnt_u64(msk64);     /* Count the nonzero bits in msk64.                    */ 
    __m128i nnz   = _mm_set1_epi8(nnz64);      /* May generate vmovd + vpbroadcastb if AVX2 is enabled.               */ 
    __m128i nonzmsk  = _mm_cmpgt_epi8(cntr,nnz);     /* nonzmsk is a mask of the form 0xFF, 0xFF, ..., 0xFF, 0, 0, ...,0 to mark the output positions without an index */ 

    uint64_t indx64  = _pext_u64(indx_const,msk64);    /* parallel bits extract. pext shuffles indx_const such that indx64 contains the nnz64 4-bit indices that we want.*/ 
    __m128i indx   = _mm_cvtsi64x_si128(indx64);    /* Use a few integer instructions to unpack 4-bit integers to 8-bit integers.          */ 
    __m128i indx_024  = indx;          /* Even indices.                         */ 
    __m128i indx_135  = _mm_srli_epi64(indx,4);     /* Odd indices.                         */ 
      indx   = _mm_unpacklo_epi8(indx_024,indx_135);  /* Merge odd and even indices.                     */ 
      indx   = _mm_and_si128(indx,cnst0F);    /* Mask out the high bits 4,5,6,7 of every byte.                 */ 

      return _mm_or_si128(indx,nonzmsk);      /* Merge indx with nonzmsk .                      */ 
} 


int main(){ 
    int i; 
    char w[16],xa[16]; 
    __m128i x; 

    /* Example with bytes 15, 12, 7, 5, 4, 3, 2, 1, 0 set. */ 
    x = _mm_set_epi8(1,0,0,1, 0,0,0,0, 1,0,1,1, 1,1,1,1); 

    /* Other examples. */ 
    /* 
    x = _mm_set_epi8(1,1,1,1, 1,1,1,1, 1,1,1,1, 1,1,1,1); 
    x = _mm_set_epi8(0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0); 
    x = _mm_set_epi8(1,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0); 
    x = _mm_set_epi8(0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,1); 
    */ 
    __m128i indices = nonz_index(x); 
    _mm_storeu_si128((__m128i *)w,indices); 
    _mm_storeu_si128((__m128i *)xa,x); 

    printf("counter 15..0 ");for (i=15;i>-1;i--) printf(" %2d ",i);  printf("\n\n"); 
    printf("example xmm: ");for (i=15;i>-1;i--) printf(" %2d ",xa[i]); printf("\n"); 
    printf("result in dec ");for (i=15;i>-1;i--) printf(" %2hhd ",w[i]); printf("\n"); 
    printf("result in hex ");for (i=15;i>-1;i--) printf(" %2hhX ",w[i]); printf("\n"); 

    return 0; 
} 

大约需要五条指令才能在不需要的位置得到0xFF（终止字节）。 请注意，函数nonz_index（返回索引并且仅返回终止字节的位置，实际上不插入终止字节）可能会便宜得多，并且可能适合在特定应用程序中使用。 第一个终止字节的位置是nnz64>>2。

结果是：

$ ./a.out 
counter 15..0 15 14 13 12 11 10 9 8 7 6 5 4 3 2 1 0 

example xmm: 1 0 0 1 0 0 0 0 1 0 1 1 1 1 1 1 
result in dec -1 -1 -1 -1 -1 -1 -1 15 12 7 5 4 3 2 1 0 
result in hex FF FF FF FF FF FF FF F C 7 5 4 3 2 1 0 

的pext指令支持英特尔的Haswell处理器或更新的版本。