Is there a way to mask one end of a __m128i register based on mask length that is not known at compile time?

Question

I have a seemingly simple problem. Load a string into an __m128i register (with _mm_loadu_si128), then find the string's length (with _mm_cmpistri). Now, assuming the length is under 16, I would like to have only zeros after the first, string ending, zero. One way to achieve that would be to copy just 'len' bytes to another register, or to AND the original register with a mask of 1s of length 8 * len. But it is not easy to find the simple way to create such mask that would depend on just computed length.

Answer 1

I would do it like this. Untested.

// 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 );
}

Update: here’s another version, uses Noah’s idea about that int64 -1 instruction. Might be slightly faster. Disassembly.

__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 );
}

Answer 2

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;
  }
}