You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

如何使用掩码打包__m128i元素(AVX2 intrinsics实现)

AVX2 实现方案

核心实现逻辑:

  • 直接用浮点比较指令生成判断掩码,避免减法操作带来的浮点异常、NaN处理错误问题
  • 提取比较结果得到4位位掩码,每个bit对应一个待打包整数是否满足rij < sij条件
  • 4位掩码仅存在16种组合,通过预定义的重排控制字查表,单指令完成整数打包,高位取值无要求的场景下性能最优。
完整实现代码
#include <immintrin.h>
#include <stdint.h>

// 16种掩码对应的32位元素重排控制字,高位填充值无特殊要求
const uint8_t shuffle_ctrl_tbl[16] = {
    0x00, // 0b0000 无匹配元素
    0x00, // 0b0001 仅j0匹配
    0x01, // 0b0010 仅j1匹配
    0x04, // 0b0011 j0、j1匹配
    0x02, // 0b0100 仅j2匹配
    0x08, // 0b0101 j0、j2匹配
    0x09, // 0b0110 j1、j2匹配 → 对应示例场景
    0x24, // 0b0111 j0、j1、j2匹配
    0x03, // 0b1000 仅j3匹配
    0x0c, // 0b1001 j0、j3匹配
    0x0d, // 0b1010 j1、j3匹配
    0x34, // 0b1011 j0、j1、j3匹配
    0x0e, // 0b1100 j2、j3匹配
    0x38, // 0b1101 j0、j2、j3匹配
    0x39, // 0b1110 j1、j2、j3匹配
    0xe4  // 0b1111 全部元素匹配
};

__m128i pack_matched_j()
{
    int j0 = 190;
    int j1 = 191;
    int j2 = 192;
    int j3 = 193;
    __m128i jv = _mm_set_epi32(j3, j2, j1, j0);
    __m256d rij = _mm256_set_pd(2.8, 1.8, 2.1, 3.4);
    __m256d sij = _mm256_set1_pd(2.5);

    // 生成小于比较掩码
    __m256d cmp_res = _mm256_cmp_pd(rij, sij, _CMP_LT_OQ);
    // 提取4位位掩码,bit N对应jN是否满足条件
    int mask = _mm256_movemask_pd(cmp_res);
    // 查表重排得到打包结果
    __m128i res = _mm_shuffle_epi32(jv, shuffle_ctrl_tbl[mask]);

    // 若需要高位强制为0,可打开下面这行代码
    // res = _mm_and_si128(res, _mm_set_epi32(0, 0, 0xFFFFFFFF, 0xFFFFFFFF));
    return res;
}
结果说明

示例场景下比较得到mask=0b0110=6,查表得到控制字0x09,重排后__m128i按高32位到低32位的顺序为[190, 190, 192, 191],低两位完全符合预期,高两位为不影响逻辑的填充值,满足[X,X,192,191]的要求。
该实现总指令数极少,16字节的查找表会被编译器优化为常量池访问,是AVX2平台下该需求的最优实现。

内容的提问来源于stack exchange,提问作者user18186610

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.08.30 18:57:25