基于AVX512实现64位布尔向量与512位字节向量的点积问题

问题描述 #

我有两个向量：

• vec1：包含64个布尔值，以uint64_t类型存储
• vec2：包含64个字节（共512位），以uint8_t*数组形式存储

目标是计算二者的点积：对vec1中位为1的对应vec2元素求和。希望用AVX512指令集高效实现，替代逐位逐字节的循环，但卡在了几个关键步骤上。

完整解决方案 #

步骤1：将64位布尔向量扩展为512位掩码向量 #

要把uint64_t类型的vec1转换成512位向量（原vec1某一位为1，对应位置的字节全为0xFF；为0则全为0x00），可以直接使用AVX512的位扩展指令：

// 将uint64_t掩码扩展为512位字节掩码
__m512i expanded_mask = _mm512_broadcastmb_epi8(_mm_set1_epi64x(vec1));

_mm512_broadcastmb_epi8指令会自动把输入的64位掩码每一位，扩展为对应的8位字节，完全符合需求的转换效果。

步骤2：字节级AND运算并转换为16位整数 #

先加载vec2到512位寄存器，再和扩展后的掩码做字节级AND，保留vec1中位为1对应的vec2元素；接着用_mm512_maddubs_epi16把相邻两个字节的结果合并为16位整数：

// 加载vec2到512位寄存器
__m512i vec2_m512 = _mm512_loadu_si512((const __m512i*)vec2);
// 保留vec1中位为1对应的vec2字节
__m512i masked_vec2 = _mm512_and_si512(vec2_m512, expanded_mask);
// 将相邻字节值累加为16位整数（第二个参数设为1，仅做求和）
__m512i result16 = _mm512_maddubs_epi16(masked_vec2, _mm512_set1_epi8(1));

步骤3：将16位整数累加为最终结果 #

先把16位整数成对相加为32位整数，再通过AVX512的水平求和指令完成最终累加：

// 将16位整数转换为32位整数（成对相加）
__m512i result32 = _mm512_madd_epi16(result16, _mm512_set1_epi16(1));
// 对512位中的32位整数做水平求和，得到128位寄存器中的4个32位整数
__m128i sum128 = _mm512_reduce_add_epi32(result32);
// 把4个32位整数相加得到最终结果
uint32_t sum = _mm_cvtsi128_si32(sum128);
sum += _mm_extract_epi32(sum128, 1);
sum += _mm_extract_epi32(sum128, 2);
sum += _mm_extract_epi32(sum128, 3);

完整可运行代码 #

#include <immintrin.h>
#include <cstdint>

uint32_t dot_product(uint64_t vec1, const uint8_t* vec2) {
    // 步骤1：扩展64位掩码到512位字节掩码
    __m512i expanded_mask = _mm512_broadcastmb_epi8(_mm_set1_epi64x(vec1));
    
    // 步骤2：加载vec2并做掩码过滤+字节转16位求和
    __m512i vec2_m512 = _mm512_loadu_si512((const __m512i*)vec2);
    __m512i masked_vec2 = _mm512_and_si512(vec2_m512, expanded_mask);
    __m512i result16 = _mm512_maddubs_epi16(masked_vec2, _mm512_set1_epi8(1));
    
    // 步骤3：累加得到最终结果
    __m512i result32 = _mm512_madd_epi16(result16, _mm512_set1_epi16(1));
    __m128i sum128 = _mm512_reduce_add_epi32(result32);
    
    uint32_t sum = _mm_cvtsi128_si32(sum128);
    sum += _mm_extract_epi32(sum128, 1);
    sum += _mm_extract_epi32(sum128, 2);
    sum += _mm_extract_epi32(sum128, 3);
    
    return sum;
}

内容的提问来源于stack exchange，提问作者Sacha