avx2 使用vpshufb指令做字符分类

最近看到simdjson的论文，其中使用vpshufb指令做的字符匹配和分类，感觉这个方法很有收获，想分享以下。

先说simdjson中为啥用这个指令，它需要从字符数组中提取出6种控制字符(‘:’, \’, ‘:’, ‘"’, ‘{’, ‘}’),以及空格换行等4种无实际意义的字符(‘\r’,‘\n’,‘\t’,‘ ’)。

vpshufb是汇编指令，在C++中可以用__m256i _mm256_shuffle_epi8 (__m256i a, __m256i b)。

__m256i _mm256_shuffle_epi8 (__m256i a, __m256i b)

算法过程:
FOR j := 0 to 15
    i := j*8
    IF b[i+7] == 1
        dst[i+7:i] := 0
    ELSE
        index[3:0] := b[i+3:i]
        dst[i+7:i] := a[index*8+7:index*8]
    FI
    IF b[128+i+7] == 1
        dst[128+i+7:128+i] := 0
    ELSE
        index[3:0] := b[128+i+3:128+i]
        dst[128+i+7:128+i] := a[128+index*8+7:128+index*8]
    FI
ENDFOR
dst[MAX:256] := 0

简单解释一下, __m256i可以当做有256个bit位，假设i=8,那么b[i+3:i]就是b[11:8]，也就是第8到11bit,这四个bit组成一个无符号整数。四个比特最大是1111，也就是十六进制的f。

_mm256_shuffle_epi8可以当作是一种查表方式，但是只能用每个字节的低四位去查。从b中获取每个字节的低四位。其值只能是0-15中的一种，把这个值当作索引，然后从a中找这个索引对应的值。
总之，可以认为_mm256_shuffle_epi8是一种映射关系，将一个数映射到另一个数。

下边是_mm256_shuffle_epi8的两个使用例子。

假如有一个字节数组，存的都是ascii，你需要从中找到小写字母。
a-z对应的十六进制表示是：61-6f ， 70-7a (这里没写成61-7a是因为后边处理需要)

将61-6f分成一组，是因为它们高四位都是6，70-7a高四位都是7。因为高四位一样的数可以被分到一个组里边。

算法思路是将一个字节的高四位，低四位分别使用_mm256_shuffle_epi8做映射，两个映射后的结果做与运算，最终结果不为0，就代表是需要的字符。

关键就是如何构造这个映射表了。
以字母a和b举例：

a (0110 0001)和b(0110 0010),他们高四位相同，所以他们高四位的映射结果也是一样的。

对于高四位:

我们假设0110通过_mm256_shuffle_epi8映射后是，0000 0001

对于低四位:

那么我们把a和b的低四位映射后，只需要最低为含有1(xxxx xxx1),比如0000 0001或0000 0011等,那么这两个数的与运算的结果一定是0000 0001.

如果是匹配a,b,p,q四种字符

a (0110 0001)，b(0110 0010)，p(0111 0000),q(0111 0001)

对于高四位:

假设0110通过_mm256_shuffle_epi8映射后是是0000 0001
假设0111通过_mm256_shuffle_epi8映射后是是0000 0010

对于低四位:

把a和b的低四位映射后，只需要最低位含有1(xxxx xxx1),p和q映射后需要第二位是1(xxxx xx1x),但是你会发现，a和q的低四位是一样的,也就是他们映射后也是相同的。所以映射结果是:
0000-->0010,0001-->0011,0010-->0001。

下边不推到了，只是分享例子
对字母a-z A-Z
高4bit：

0110 和 0100 ，映射到  0001
0111 和 0101 ，映射到  0001
其他映射到0

低4bit:

0000映射到 0010
0001到1010 映射到 0011
1011到1111 映射到 0001

其实对字母的提取，完全可以用 'a'>= 字符 && 字符<='z'这样的方式来处理，真正用_mm256_shuffle_epi8的话，还是要对一些更复杂的分类做处理.我是用来对HTML中一些结构提取时:
我需要获取三种类型的字符
第一种: a-z A-Z ! ? / （因为<+字母可能是一个html标签的开始，</ <?可能是注释 </可能是闭合标签）
第二种: = " ' > (字符串用"或'包裹 标签属性用=号链接 >可能是标签闭合)
第三种: \n \f \r \t \0 空格 (这些是没意义的字符，需要被跳过)

直接上代码:

void strExtractTestAvx2(char* ms ,unsigned int length) {

    //低四位的转换
    __m256i vpshufbCharMaskn = _mm256_setr_epi64x(0x1303030303130bc2L, 0x0d21a18101838303L, 0x1303030303130bc2L, 0x0d21a18101838303L);
    //高四位的转换
    __m256i vpshufbCharMaskh = _mm256_setr_epi64x(0x0201020124580080L, 0, 0x0201020124580080L, 0);
    //字符类型判断
    __m256i alphaMask = _mm256_set1_epi8(0b00001111);
    __m256i constructMask = _mm256_set1_epi8(0b00110000);
    __m256i emptyMask = _mm256_set1_epi8(0b11000000);
    //全0
    __m256i zero = _mm256_set1_epi64x(0);
    //每个字节高四位置0
    __m256i vpshufbMask = _mm256_set1_epi64x(0x0f0f0f0f0f0f0f0fL);

    const unsigned int length1 = length - 31;

    for (int i = 0; i < length1; i += 32)
    {

        __m256i ymm1 = _mm256_lddqu_si256((__m256i*) & ms[i]);

        __m256i ymm2 = _mm256_srli_epi16(ymm1, 4);
        ymm2 = _mm256_and_si256(ymm2, vpshufbMask);
        __m256i ymm3 = _mm256_shuffle_epi8(vpshufbCharMaskn, ymm1);
        ymm2 = _mm256_shuffle_epi8(vpshufbCharMaskh, ymm2);
        ymm3 = _mm256_and_si256(ymm3, ymm2);
        
        __m256i  alpha = _mm256_and_si256(ymm3, alphaMask);
        __m256i construct = _mm256_and_si256(ymm3, constructMask);
        __m256i  empty = _mm256_and_si256(ymm3, emptyMask);
        // <
        __m256i less = _mm256_cmpeq_epi8(ymm1, lessSignMask);

        //字母 ? ! /
        alpha = _mm256_cmpgt_epi8(alpha, zero);
        
        //= " ' >
        construct = _mm256_cmpgt_epi8(construct, zero);

        empty = _mm256_srli_epi64(empty, 1);
        //6种空格
        empty = _mm256_cmpgt_epi8(empty, zero);

    }
}