基于 BitMask 在数组中设置值的固有特性

Intrinsic to set value in array based on a BitMask

是否有一个内在函数可以在输入数组中的所有位置设置单个值,其中相应位置在提供的 BitMask 中具有 1 位?

10101010 是位掩码

值为 121

它会将位置 0、2、4、6 设置为值 121

使用 AVX512,是的。屏蔽存储是 AVX512 中的第一个 class 操作。

使用位掩码作为向量存储到数组的 AVX512 掩码,使用 _mm512_mask_storeu_epi8 (void* mem_addr, __mmask64 k, __m512i a) vmovdqu8。 (AVX512BW。对于 AVX512F,您只能使用 32 或 64 位元素大小。)

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

void set_value_in_selected_elements(char *array, uint64_t bitmask, uint8_t value) {
    __m512i broadcastv = _mm512_set1_epi8(value);
    // integer types are implicitly convertible to/from __mmask types
    // the compiler emits the KMOV instruction for you.
    _mm512_mask_storeu_epi8 (array, bitmask, broadcastv);
}

这会将 (with gcc7.3 -O3 -march=skylake-avx512) 编译为:

    vpbroadcastb    zmm0, edx
    kmovq   k1, rsi
    vmovdqu8        ZMMWORD PTR [rdi]{k1}, zmm0
    vzeroupper
    ret

如果要在位图为零的元素中写入零,请使用零掩码移动从掩码创建常量并存储它,或者使用 AVX512BW 或 DQ 创建 0 / -1 向量__m512i _mm512_movm_epi8(__mmask64 )。其他元件尺寸可供选择。但是使用掩码存储可以在数组大小不是向量宽度的倍数时安全地使用它,因为未修改的元素不会被读取/重写或任何其他内容;他们真的没有受到影响。 (不过,如果任何未触及的元素在真实存储中出现故障,CPU 可以采用较慢的微代码支持。)


没有 AVX512,您仍然要求 "an intrinsic"(单数)。

pdep,您可以使用它来将位图扩展为字节图。有关使用 _pdep_u64(mask, 0x0101010101010101);mask 中的每个位解压缩为一个字节的示例,请参见 。这在 uint64_t 中为您提供了 8 个字节。在 C 中,如果你在它和一个数组之间使用 union,那么它会给你一个包含 0 / 1 个元素的数组。 (但是当然索引数组将需要编译器发出移位指令,如果它没有首先将它溢出到某个地方。你可能只想 memcpyuint64_t 到一个永久数组中。)

但在更一般的情况下(更大的位图),或者当您想基于位掩码混合新值时甚至有 8 个元素,您应该使用多个内在函数来实现 pmovmskb 的逆函数,并用它来混合。 (请参阅下面没有 pdep 的部分


一般来说,如果您的数组适合 64 位(例如 8 元素字符数组),您可以使用 pdep。或者,如果它是一个 4 位半字节数组,那么您可以使用 16 位掩码而不是 8 位掩码。

否则没有单一指令,因此没有内在指令。对于较大的位图,您可以将其处理为 8 位块并将 8 字节块存储到数组中。


如果您的数组元素宽度超过 8 位(并且您没有 AVX512),您可能仍应使用 pdep 将位扩展为字节,但是然后使用 [v]pmovzx 从字节扩展到双字或向量中的任何内容。例如

// only the low 8 bits of the input matter
__m256i bits_to_dwords(unsigned bitmap) {
    uint64_t mask_bytes = _pdep_u64(bitmap, 0x0101010101010101);  // expand bits to bytes
    __m128i byte_vec = _mm_cvtsi64x_si128(mask_bytes);
    return _mm256_cvtepu8_epi32(byte_vec);
}

如果你想保留元素不变而不是将它们设置为零,位掩码有零,或者使用以前的内容而不是分配/存储。

这在C / C++中表达起来相当不方便(与asm相比)。要将 uint64_t 中的 8 个字节复制到 char 数组中,您可以(并且应该)只使用 memcpy(以避免由于指针别名或未对齐 uint64_t* 而导致的任何未定义行为)。这将使用现代编译器编译为单个 8 字节存储。

但是要对它们进行或运算,您要么必须在 uint64_t 的字节上编写一个循环,要么将您的 char 数组转换为 uint64_t*。这 通常 工作正常,因为 char* 可以为任何东西起别名,所以稍后读取 char 数组没有任何严格别名的 UB。但是未对齐的 uint64_t* 可以 即使在 x86 上也会导致问题,如果编译器假定它在自动矢量化时 对齐的话。


赋值不是 0 / 1

使用乘以 0xFF 将 0/1 字节的掩码转换为 0 / -1 掩码,然后将其与 uint64_t 相乘,将您的值广播到所有字节职位。

如果你想保留元素不变而不是将它们设置为零或 value=121,你应该使用 SSE2 / SSE4 或 AVX2,即使你的数组​​有字节元素。加载旧内容,vpblendvbset1(121),使用字节掩码作为控制向量。

vpblendvb 仅使用每个字节的高位,因此您的 pdep 常量可以 0x8080808080808080 将输入位分散到 high 每个字节的位,而不是低位。 (所以你不需要乘以 0xFF 来得到一个 AND 掩码)。

如果您的元素是双字或更大,您可以使用 _mm256_maskstore_epi32。 (在将掩码从字节扩展到双字时,使用 pmovsx 而不是 zx 来复制符号位)。这可以是性能胜过可变混合 + 始终读取/重写。 Is it possible to use SIMD instruction for replace?.


没有pdep

pdep 在 Ryzen 上非常慢,即使在 Intel 上也可能不是最佳选择。

另一种方法是将位掩码转换为矢量掩码: is there an inverse instruction to the movemask instruction in intel avx2?
How to perform the inverse of _mm256_movemask_epi8 (VPMOVMSKB)?.

即将您的位图广播到向量的每个位置(或将其洗牌以便位图的正确位在相应的字节中),并使用 SIMD AND 屏蔽该字节的适当位。然后对 AND 掩码使用 pcmpeqb/w/d 来查找设置了位的元素。

如果您不想在位图为零的位置存储零,您可能想要加载/混合/存储。

使用比较遮罩来混合您的 value,例如使用 _mm_blendv_epi8 或 256 位 AVX2 版本。您可以处理 16 位块中的位图,生成 16 字节向量,只需 pshufb 将其字节发送到正确的元素。

多个线程在同一个数组上同时执行此操作是不安全的,即使它们的位图不相交,除非您使用屏蔽存储。