【代码分享】使用 avx2 + 查表法，优化凯撒加密

作者:张富春(ahfuzhang)，转载时请注明作者和引用链接，谢谢！

• cnblogs博客
• zhihu
• Github
• 公众号:一本正经的瞎扯
  
  

---

接上一篇：【代码分享】使用 avx512 + 查表法，优化凯撒加密

好不容易捣鼓出来了 avx512 指令集的查表法代码，可是部署的时候发现服务器不支持 avx512 指令集。

终于，avx2 版本的查表法终于写出来了。上代码：

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <immintrin.h>
#include <avx2intrin.h>
#include <inttypes.h>
#include <assert.h>

typedef uint32_t Table[26][256];

void initTable(Table *table){
    for (int i=0; i<26; i++){
        for (int j=0; j<256; j++){
            if (j>='a' && j<='z'){
                (*table)[i][j] = (uint8_t)((j-'a'+i)%26 + 'a');
            } else if (j>='A' && j<='Z'){
                (*table)[i][j] = (uint8_t)((j-'A'+i)%26 + 'A');
            } else {
                (*table)[i][j] = j;
            }
        }
    }
}

void caesarEncodeAVX2(void* out1, void* in1, uint64_t len, uint64_t rot, void* table1){
	uint8_t* out = (uint8_t*)out1;
	uint8_t* in = (uint8_t*)in1;
	Table* table = (Table*)table1;  // 这些代码为了解决  cgo 中函数调用的警告
	//
	rot = rot % 26;
    uint32_t* line = (uint32_t*)((*table)[rot]);
    #define batchSize  16
    uint64_t tailLen = len & 0x0f;
    uint8_t* end = in + len - tailLen;
    uint8_t* start = in;
    const __m256i index_mask = _mm256_set_epi32(
        7, 6, 3, 2,
        5,4, 1, 0
    );
    for (; start<end;  out += batchSize, start += batchSize){
        __m128i src = _mm_loadu_si128(start);  // 加载 16 个字符，但只处理前面 8 个字符
        __m256i srcI32 = _mm256_cvtepu8_epi32(src);  // 前  8  个字节，转换为 32 字节
        __m256i foundHead = _mm256_i32gather_epi32(line, srcI32, 4);  // 8  字节查表
        // 处理后面 8 字节
        src = _mm_srli_si128(src, 8);  // 移动 8 字节，也就是 64 位
        srcI32 = _mm256_cvtepu8_epi32(src);  //  后  8  个字节，转换为 32 字节
        __m256i foundTail = _mm256_i32gather_epi32(line, srcI32, 4);  // 8  字节查表
        // 进行压缩
        __m256i found16 = _mm256_packus_epi32(foundHead, foundTail);  // 16 个  16 位的值
        /*
              _mm256_packus_epi32 这个指令非常的恶心：aaaa bbbb 合并成了  aa bb aa bb
              于是还要用下面这个指令把顺序换过来
        */
        found16 =  _mm256_permutevar8x32_epi32 (found16, index_mask);  // 换过来了，正常了
        found16 =  _mm256_packus_epi16(found16, _mm256_setzero_si256());
        found16 =  _mm256_permutevar8x32_epi32 (found16, index_mask);  // 换过来了，正常了
        //存储
        __m128i result = _mm256_castsi256_si128(found16);
        _mm_storeu_si128 (out, result);
    }
    //
    end = in + len;
    for (; start<end; start++, out++){
        *out = (uint8_t)line[*start];
    }
}

// 逐字节处理的版本
void caesarEncode(uint8_t* out, uint8_t* in, int len, int rot, Table* t){
    rot = rot % 26;
    uint8_t* end = in + len;
    uint32_t* line = (uint32_t*)((*t)[rot]);
    for (;in<end; in++, out++){
        *out = (uint8_t)line[*in];
    }
}

编译：

gcc -o caesar caesar.c -g -w -mavx -mavx2 -O3

测试了一下各个版本的性能表现：

逐字节处理版本：total:201429 us, avg: 100.715 ns/op
AVX2 版本：       total:51059 us, avg: 25.529 ns/op
AVX512  版本：   total:49756 us, avg: 24.878 ns/op

avx512 版本只比 avx2 快了一点点，鸡肋！