如何根据运行时未知的掩码长度,在__m128i寄存器中屏蔽其中一端?

3

我有一个看似简单的问题。将一个字符串加载到__m128i寄存器中(使用_mm_loadu_si128),然后找到字符串的长度(使用_mm_cmpistri)。现在,假设长度小于16,我想在字符串末尾的零后面只有零。实现这个目标的一种方法是将'len'个字节复制到另一个寄存器中,或者将原始寄存器与一个长度为8 * len的1s掩码进行AND运算。但是很难找到一个简单的方式来创建这样的掩码,使它仅取决于刚刚计算的长度。


pcmpeqb / pmovmskb / tzcnt 会给你一个位置,然后你可以使用它来索引一个滑动窗口到一个缓冲区中的 0xff, ..., 0 ... 或类似的内容,以获取一个 AND 掩码。例如:Vectorizing with unaligned buffers: using VMASKMOVPS: generating a mask from a misalignment count? Or not using that insn at all - Peter Cordes
@PeterCordes 谢谢您的回答,我还在努力理解。您是说“使用其他指令来代替cmpistri以查找0”吗?另外值得注意的是:编译-m32时,这些SSE4.2内置函数是否可用? - Jacek Ambroziak
我明白了,但是在m32中似乎没有_mm_crc32_u64可用。(为IvyBridge编译) - Jacek Ambroziak
是的,这与我链接的问题的答案的想法相同。如果您经常使用它,它将在缓存中热门。只要您安排数据,使其不跨越缓存行边界(例如alignas(32)),在Nehalem及更高版本的Intel CPU以及最近的AMD上,未对齐访问将不会有任何惩罚。如何在x86_64上准确基准测试未对齐访问速度 - Peter Cordes
2
@JacekAmbroziak 我认为这个是正确的(不使用pcmpistri),但实现了你所说的目标。如果您知道__m128i包含零,那么如果有一个技巧可以仅使用SIMD指令创建__m128i not_zero掩码,它可能也可以得到改进。 - Noah
显示剩余4条评论
2个回答

2
我会这样做。未经测试。
// Load 16 bytes and propagate the first zero towards the end of the register
inline __m128i loadNullTerminated( const char* pointer )
{
    // Load 16 bytes
    const __m128i chars = _mm_loadu_si128( ( const __m128i* )pointer );

    const __m128i zero = _mm_setzero_si128();
    // 0xFF for bytes that were '\0', 0 otherwise
    __m128i zeroBytes = _mm_cmpeq_epi8( chars, zero );

    // If you have long strings and expect most calls to not have any zeros, uncomment the line below.
    // You can return a flag to the caller, to know when to stop.
    // if( _mm_testz_si128( zeroBytes, zeroBytes ) ) return chars;

    // Propagate the first "0xFF" byte towards the end of the register.
    // Following 8 instructions are fast, 1 cycle latency/each.
    // Pretty sure _mm_movemask_epi8 / _BitScanForward / _mm_loadu_si128 is slightly slower even when the mask is in L1D
    zeroBytes = _mm_or_si128( zeroBytes, _mm_slli_si128( zeroBytes, 1 ) );
    zeroBytes = _mm_or_si128( zeroBytes, _mm_slli_si128( zeroBytes, 2 ) );
    zeroBytes = _mm_or_si128( zeroBytes, _mm_slli_si128( zeroBytes, 4 ) );
    zeroBytes = _mm_or_si128( zeroBytes, _mm_slli_si128( zeroBytes, 8 ) );
    // Now apply that mask
    return _mm_andnot_si128( zeroBytes, chars );
}

更新:这里还有另一种版本,使用了Noah有关那个int64 -1指令的想法。可能稍微快一些。反汇编。

__m128i loadNullTerminated_v2( const char* pointer )
{
    // Load 16 bytes
    const __m128i chars = _mm_loadu_si128( ( const __m128i* )pointer );

    const __m128i zero = _mm_setzero_si128();
    // 0xFF for bytes that were '\0', 0 otherwise
    const __m128i zeroBytes = _mm_cmpeq_epi8( chars, zero );

    // If you have long strings and expect most calls to not have any zeros, uncomment the line below.
    // You can return a flag to the caller, to know when to stop.
    // if( _mm_testz_si128( eq_zero, eq_zero ) ) return chars;

    // Using the fact that v-1 == v+(-1), and -1 has all bits set
    const __m128i ones = _mm_cmpeq_epi8( zero, zero );
    __m128i mask = _mm_add_epi64( zeroBytes, ones );
    // This instruction makes a mask filled with lowest valid bytes in each 64-bit lane
    mask = _mm_andnot_si128( zeroBytes, mask );

    // Now need to propagate across 64-bit lanes

    // ULLONG_MAX if there were no zeros in the corresponding 8-byte long pieces of the string
    __m128i crossLaneMask = _mm_cmpeq_epi64( zeroBytes, zero );
    // Move the lower 64-bit lanes of noZeroes64 into higher position
    crossLaneMask = _mm_unpacklo_epi64( mask, crossLaneMask );
    // Update the mask.
    // Lower 8 bytes will not change because _mm_unpacklo_epi64 copied that part from the mask.
    // However, upper lane may become zeroed out.
    // Happens when _mm_cmpeq_epi64 detected at least 1 '\0' in any of the first 8 characters.
    mask = _mm_and_si128( mask, crossLaneMask );

    // Apply that mask
    return _mm_and_si128( mask, chars );
}

8个pslldq/por依赖链的周期可能比pmovmskb / bsf / load低1或2个周期的延迟。但是这些操作需要8个uops,如果您没有AVX进行非破坏性复制和移位,则需要另外4个movdqa uops。(您可以使用_mm_shuffle_epi32来进行4字节和8字节粒度的复制和移位,以避免这种情况,因为将低4或8字节保持不变而不是移入零以提供OR是可以的) - Peter Cordes
1
也许我们可以在这里做一些更聪明的事情:(SRC-1) XOR (SRC)(如blsmsk)可以设置qword中最低位以下的所有位,如果我们使用psubq(或带有-1的paddq)。 如果我们可以从两个8字节半扩展到16字节向量,我们只需与其进行AND运算即可。 - Peter Cordes
1
@PeterCordes 的确,这是一个权衡。根据我的经验,不访问内存而是计算更多指令的代码总体上往往会稍微快一些。特别是一旦集成到更大的系统中:微基准测试对 L1D 缓存没有任何其他用途。 - Soonts
1
对于一些在较大操作中仅使用一次的东西,避免内存可能是最好的选择。对于每个内部循环仅使用一次(且不在延迟关键路径上)的东西,L1d缺失可以在多次迭代中分摊。权衡的明智选择取决于用例。SIMD代码通常需要一些向量常量,而32字节表等效于2个常量(除非加载潜在地在关键路径上)。顺便说一句,这使得@Noah的解决方案并没有更好。您可以通过pcmpeqd / psrlq按63便宜地生成1,1,并从中获得0,1,但编译器将constprop&load。 - Peter Cordes
1
@Noah,你最近发布的解决方案存在一个bug。它会分别屏蔽__m128i的低位和高位。4a,61,63,65,6b,00,00,00,63,65,6b,20,00,00,00,00 - Jacek Ambroziak
显示剩余10条评论

0
static const __m128i ZERO =
    _MM_SETR_EPI32(0u, 0u, 0u, 0u);

static const __m128i INDEXES =
    _MM_SETR_EPI8(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);

static const __m128i ONES = _MM_SETR_EPI32(0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF);

_Alignas(32) static unsigned char MASK_SOURCE[32] =
    {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};

static __m128i mask_string1(__m128i input, uint32_t *const plen) {
  const __m128i zeros = _mm_cmpeq_epi8(input, ZERO);
  if (_mm_testz_si128(zeros, zeros)) {
    *plen = 16;
    return input;
  } else {
    const uint32_t length = _tzcnt_u32(_mm_movemask_epi8(zeros));
    *plen = length;
    return
        length < 15 ?
        _mm_and_si128(input, _mm_loadu_si128((__m128i_u *) (MASK_SOURCE + (16 - length)))) :
        input;
  }
}

static __m128i mask_string2(__m128i input, uint32_t *const plen) {
  __m128i zeros = _mm_cmpeq_epi8(input, ZERO);
  if (_mm_testz_si128(zeros, zeros)) {
    *plen = 16;
    return input;
  } else {
    const uint32_t length = _tzcnt_u32(_mm_movemask_epi8(zeros));
    *plen = length;
    if (length < 15) {
      zeros = _mm_or_si128(zeros, _mm_slli_si128(zeros, 1));
      zeros = _mm_or_si128(zeros, _mm_slli_si128(zeros, 2));
      zeros = _mm_or_si128(zeros, _mm_slli_si128(zeros, 4));
      zeros = _mm_or_si128(zeros, _mm_slli_si128(zeros, 8));
      // Now apply that mask
      return _mm_andnot_si128(zeros, input);
    } else {
      return input;
    }
  }
}

static __m128i mask_string3(__m128i input, uint32_t *const plen) {
  const __m128i zeros = _mm_cmpeq_epi8(input, ZERO);
  if (_mm_testz_si128(zeros, zeros)) {
    *plen = 16;
    return input;
  } else {
    const uint32_t length = _tzcnt_u32(_mm_movemask_epi8(zeros));
    *plen = length;
    return
        length < 15 ?
        _mm_andnot_si128(_mm_cmpgt_epi8(INDEXES, _mm_set1_epi8(length)), input) :
        input;
  }
}

__m128i set_zeros_3(__m128i v, uint32_t *plen) {
  // cmp zeros
  __m128i eq_zero = _mm_cmpeq_epi8(ZERO, v);
  if (_mm_testz_si128(eq_zero, eq_zero)) {
    *plen = 16;
    return v;
  } else {
    *plen = _tzcnt_u32(_mm_movemask_epi8(eq_zero));
#ifdef COND
    if (_mm_testz_si128(eq_zero, eq_zero)) {
          return;
      }
#endif
    __m128i eq_zero64 = _mm_cmpeq_epi64(eq_zero, ZERO);
    __m128i mask64_1 = _mm_unpacklo_epi64(ONES, eq_zero64);
    // add(-1) / sub(1)
    __m128i partial_mask = _mm_add_epi64(eq_zero, ONES);

#if defined __AVX512F__ && defined __AVX512VL__
    __m128i result =
          _mm_ternarylogic_epi64(partial_mask, mask64_1, v, (1 << 7));
#else
    __m128i mask = _mm_and_si128(mask64_1, partial_mask);
    __m128i result = _mm_and_si128(mask, v);
#endif
    return result;
  }
}

我已经发布了4种不同的掩码字符串方法的代码,这些方法被加载到16字节向量中。我们关注以下两点:1)字符串长度,2)清理后的字符串(终止0之后没有随机字节)。当输入字符串超过15个字符时,向量中将没有0。在这种情况下,我们需要循环遍历字符串,直到遇到零为止。但对于长度为[1, 15]的字符串,我们将学习其长度并清理该字符串。这些方法来自Peter Cordes、Soonts和Noah。 - Jacek Ambroziak
1
你最好将_mm_movemask_epi8的输出分配给一个变量,并在该变量上测试0。_mm_testz_si128的延迟为3个uops,与_mm_movemask_epi8+testl; jz相同。_mm_movemask_epi8位于port0上,因此它占用了一个执行单元,否则该单元将用于关键路径依赖链。 - Noah
我已经这样做了,但是在最终的时间上没有观察到任何差异。所有的方法都执行“相同”的操作,这是一个有点令人沮丧的结果。所有的方法都很好用,选择纯粹是“美学”问题。我觉得这很奇怪。 - Jacek Ambroziak
如果大多数时间没有0,则您可能每次都会触发if语句。我看到的最好结果(显著优于其他方式)是使用 此代码(基本上适用于所有测试数据,除了在没有0的情况下需要帮助)。 - Noah
@Noah 实际上,在我 4.82 亿字符串文件中,约有 77% 的字符串长度小于 16 个字符。这非常好,因为这么多字符串可以只用第一个“批处理”中的前 16 个字符来处理。在当前速度测试中,我甚至没有执行任何循环,只是处理第一个批次。我想知道我们是否可以从这个竞赛没有获胜者的事实中学到一些重要的东西。如果您想运行速度测试,请随意(我可以发布在 GitHub 上)。现在不确定这是否值得任何人花费时间。 - Jacek Ambroziak
这是一个有趣的问题,至少可以让你尝试解决。请在此处发布链接。 - Noah

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接