Search function using SIMD

Question

I wrote a search function, similar to std::find, that uses SIMD instructions. Since I am new to SIMD, I would appreciate comments on other SIMD instructions I have missed that would be useful for this use case, possible corner cases I have overlooked, and best practices in using SIMD and C++ in general.

#pragma once

#include <immintrin.h>

#include <cstddef>
#include <cstdint>
#include <iterator>
#include <memory>
#include <type_traits>

namespace simd {

using vec256i = __m256i;
using vec256f = __m256;
using vec256d = __m256d;

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) vec256i compare(const vec256i& lhs, const void* const rhs) {
  static_assert(WIDTH == 1 || WIDTH == 2 || WIDTH == 4 || WIDTH == 8);
  if constexpr (WIDTH == 1) {
    return _mm256_cmpeq_epi8(lhs, _mm256_loadu_epi8(rhs));
  } else if (WIDTH == 2) {
    return _mm256_cmpeq_epi16(lhs, _mm256_loadu_epi16(rhs));
  } else if (WIDTH == 4) {
    return _mm256_cmpeq_epi32(lhs, _mm256_load_epi32(rhs));
  } else if (WIDTH == 8) {
    return _mm256_cmpeq_epi64(lhs, _mm256_load_epi64(rhs));
  } else {
    return vec256i{0};
  }
}

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) std::uint32_t compress_mask(const vec256i& mask) {
  static_assert(WIDTH == 1 || WIDTH == 2 || WIDTH == 4 || WIDTH == 8);
  if constexpr (WIDTH == 1 || WIDTH == 2) {
    return _mm256_movemask_epi8(mask);
  } else if (WIDTH == 4) {
    return _mm256_movemask_ps(reinterpret_cast<vec256f>(mask));
  } else if (WIDTH == 8) {
    return _mm256_movemask_pd(reinterpret_cast<vec256d>(mask));
  } else {
    return 0;
  }
}

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) std::uint32_t ctz_ui128(const std::uint32_t& mask4,
                                                                            const std::uint32_t& mask3,
                                                                            const std::uint32_t& mask2,
                                                                            const std::uint32_t& mask1) {
  static_assert(WIDTH == 1 || WIDTH == 2);
  if (mask1 != 0) {
    return __builtin_ctz(mask1) / WIDTH;
  } else if (mask2 != 0) {
    return (__builtin_ctz(mask2) + 32) / WIDTH;
  } else if (mask3 != 0) {
    return (__builtin_ctz(mask3) + 64) / WIDTH;
  } else if (mask4 != 0) {
    return (__builtin_ctz(mask4) + 96) / WIDTH;
  } else {
    return 128;
  }
}

template <std::size_t WIDTH>
const void* find(const void* const current_void, const void* const value_void);

template <>
[[nodiscard]] inline __attribute__((always_inline)) const void* find<1>(const void* const current,
                                                                        const void* const value) {
  const auto* const current_i8{reinterpret_cast<const std::uint8_t*>(current)};
  const auto* const value_i8{reinterpret_cast<const std::uint8_t*>(value)};
  const auto value_vec{_mm256_set1_epi8(*value_i8)};

  const auto mask1{compare<1>(value_vec, current_i8)};
  const auto mask2{compare<1>(value_vec, current_i8 + 32)};
  const auto mask3{compare<1>(value_vec, current_i8 + 64)};
  const auto mask4{compare<1>(value_vec, current_i8 + 96)};
  const auto mask12{_mm256_or_si256(mask1, mask2)};
  const auto mask34{_mm256_or_si256(mask3, mask4)};
  const auto mask1234{_mm256_or_si256(mask12, mask34)};

  if (_mm256_testz_si256(mask1234, mask1234) == 0) {
    const auto first_occurrence{ctz_ui128<1>(compress_mask<1>(mask4), compress_mask<1>(mask3), compress_mask<1>(mask2),
                                             compress_mask<1>(mask1))};
    return reinterpret_cast<const void*>(current_i8 + first_occurrence);
  }

  return nullptr;
}

template <>
[[nodiscard]] inline __attribute__((always_inline)) const void* find<2>(const void* const current,
                                                                        const void* const value) {
  const auto* const current_i16{reinterpret_cast<const std::uint16_t*>(current)};
  const auto* const value_i16{reinterpret_cast<const std::uint16_t*>(value)};
  const auto value_vec{_mm256_set1_epi16(*value_i16)};

  const auto mask1{compare<2>(value_vec, current_i16)};
  const auto mask2{compare<2>(value_vec, current_i16 + 16)};
  const auto mask3{compare<2>(value_vec, current_i16 + 32)};
  const auto mask4{compare<2>(value_vec, current_i16 + 48)};
  const auto mask12{_mm256_or_si256(mask1, mask2)};
  const auto mask34{_mm256_or_si256(mask3, mask4)};
  const auto mask1234{_mm256_or_si256(mask12, mask34)};

  if (_mm256_testz_si256(mask1234, mask1234) == 0) {
    const auto first_occurrence{ctz_ui128<2>(compress_mask<2>(mask4), compress_mask<2>(mask3), compress_mask<2>(mask2),
                                             compress_mask<2>(mask1))};
    return reinterpret_cast<const void*>(current_i16 + first_occurrence);
  }

  return nullptr;
}

template <>
[[nodiscard]] inline __attribute__((always_inline)) const void* find<4>(const void* const current,
                                                                        const void* const value) {
  const auto* const current_i32{reinterpret_cast<const std::uint32_t*>(current)};
  const auto* const value_i32{reinterpret_cast<const std::uint32_t*>(value)};
  const auto value_vec{_mm256_set1_epi32(*value_i32)};

  const auto mask1{compare<4>(value_vec, current_i32)};
  const auto mask2{compare<4>(value_vec, current_i32 + 8)};
  const auto mask3{compare<4>(value_vec, current_i32 + 16)};
  const auto mask4{compare<4>(value_vec, current_i32 + 24)};
  const auto mask12{_mm256_or_si256(mask1, mask2)};
  const auto mask34{_mm256_or_si256(mask3, mask4)};
  const auto mask1234{_mm256_or_si256(mask12, mask34)};

  if (_mm256_testz_si256(mask1234, mask1234) == 0) {
    const auto compressed_mask{(compress_mask<4>(mask4) << 24) + (compress_mask<4>(mask3) << 16) +
                               (compress_mask<4>(mask2) << 8) + compress_mask<4>(mask1)};
    return reinterpret_cast<const void*>(current_i32 + __builtin_ctz(compressed_mask));
  }

  return nullptr;
}

template <>
[[nodiscard]] inline __attribute__((always_inline)) const void* find<8>(const void* const current,
                                                                        const void* const value) {
  const auto* const current_i64{reinterpret_cast<const std::uint64_t*>(current)};
  const auto* const value_i64{reinterpret_cast<const std::uint64_t*>(value)};
  const auto value_vec{_mm256_set1_epi64x(*value_i64)};

  const auto mask1{compare<8>(value_vec, current_i64)};
  const auto mask2{compare<8>(value_vec, current_i64 + 4)};
  const auto mask3{compare<8>(value_vec, current_i64 + 8)};
  const auto mask4{compare<8>(value_vec, current_i64 + 12)};
  const auto mask12{_mm256_or_si256(mask1, mask2)};
  const auto mask34{_mm256_or_si256(mask3, mask4)};
  const auto mask1234{_mm256_or_si256(mask12, mask34)};

  if (_mm256_testz_si256(mask1234, mask1234) == 0) {
    const auto compressed_mask{(compress_mask<8>(mask4) << 12) + (compress_mask<8>(mask3) << 8) +
                               (compress_mask<8>(mask2) << 4) + compress_mask<8>(mask1)};
    return reinterpret_cast<const void*>(current_i64 + __builtin_ctz(compressed_mask));
  }

  return nullptr;
}

template <std::contiguous_iterator ITERATOR_T>
[[nodiscard]] ITERATOR_T find(ITERATOR_T begin_it, ITERATOR_T end_it,
                              const typename std::iterator_traits<ITERATOR_T>::value_type& value) {
  using value_t = typename std::iterator_traits<ITERATOR_T>::value_type;
  static_assert(std::is_scalar_v<value_t>);

  constexpr auto WIDTH{sizeof(value_t)};
  const value_t* const begin{std::to_address(begin_it)};
  const value_t* const end{std::to_address(end_it)};
  const value_t* current{std::to_address(begin_it)};

  // Scalar comparison until the pointer is aligned to 32 bytes.
  while (current != end && reinterpret_cast<std::uintptr_t>(current) % 32 != 0) {
    if (*current == value) {
      return begin_it + (current - begin);
    }
    ++current;
  }

  // SIMD comparison.
  while (current + (128 / WIDTH) <= end) {
    const auto* const current_void{reinterpret_cast<const void*>(current)};
    const auto* const value_void{reinterpret_cast<const void*>(&value)};
    if (const auto* const found{find<WIDTH>(current_void, value_void)}; found != nullptr) {
      return begin_it + (reinterpret_cast<const value_t*>(found) - begin);
    }
    current += 128 / WIDTH;
  }

  // Scalar tail comparison.
  while (current != end) {
    if (*current == value) {
      return begin_it + (current - begin);
    }
    ++current;
  }

  return end_it;
}

}  // namespace simd

Example usage:

#include <cassert>
#include <algorithm>
#include <numeric>
#include <vector>

#include "find.hpp"

int main() {
  std::vector<int> elements(1'000'000, 0);
  std::iota(elements.begin(), elements.end(), 0);

  const auto simd_it{simd::find(elements.cbegin(), elements.cend(), 1337)};
  const auto std_it{std::find(elements.cbegin(), elements.cend(), 1337)};
  assert(simd_it == std_it);
}

Answer 1

Sneaky AVX512 dependence

Note that _mm256_loadu_epi8 and such ("typed integer loads" or whatever you want to call them) are AVX512. It looks like this code is aimed at AVX2 though, and the corresponding AVX2 load is _mm256_loadu_si256 regardless of the size of the element. _mm256_loadu_epi8 and friends probably exist (I think) because there are masked versions of them, the element size matters for the masked versions but not the "basic" version that just loads 32 bytes.

Passing vectors by const-reference

That's unnecessary. Think of vectors as "slightly bigger integers", they fit in registers, you can copy them for almost free. Passing them by const-reference is fine as long as the compiler optimizes away the indirection, but it's not really a good thing, more of an "not necessarily a bad thing, when the compiler cooperates".

Passing uint32_t by const-reference is even more unusual, and once again the best you can hope for is that the compiler doesn't do what you told it to do, which is not a great position to be in.

reinterpret_cast<vec256f>

Does that work? The usual way to express vector reinterpretation is with the "cast" family of intrinsics such as _mm256_castsi256_ps.

__builtin_ctz

__builtin_ctz is fine but you may like to know that as of C++20 there is a <bit> header that defines std::countr_zero .. or, looking at your user name, perhaps that's not your preference ;)

Head and tail handling

This is a common problem when working with SIMD, and there is nothing inherently wrong with using some scalar code at the start and end, but it's not the only option.

You can also consider one of these tricks:

• Use an unaligned load at the start. Once upon a time unaligned loads used to be really quite bad, worth avoiding at significant cost, but not today, and there would be only a couple of them. You're not stuck with unaligned loads throughout, because it's fine if the first aligned(*) load (in the main loop) partially overlaps the unaligned load. There is some wasted work there, but also the potential to still be faster than a scalar loop.
  By the way, be careful with this: this technique does not handle an input array that's shorter than a vector.
• Similarly, use an unaligned load at the end, also partially overlapping, with the last aligned load.
• You can use aligned loads, but then ignore/discard the data that was loaded from before the start of the input and after the end of the input. Be careful with zeroing out the invalid parts when searching for a zero.

---

*: by "aligned load" I mean that the address is aligned. Back in the old days, it used to be that the unaligned load instruction was always slow, even if the address was actually aligned. That hasn't been the case for over a decade. A load can be considered aligned if the address is aligned, the type of instruction doesn't really matter. Some compilers refuse to emit the aligned instruction even when use its intrinsic, opting to always emit the unaligned instruction.

Answer 2

Below is the updated code related to @harold's answer to my question.

Specifically, I changed the following things:

1. I replaced AVX512 instructions like _mm256_loadu_epi8 with the AVX2 instruction _mm256_load_si256.
2. I removed references to vectors and integers and used pass-by-value instead.
3. When I convert vectors from one type to another, I use the appropriate intrinsic, for example _mm256_castsi256_ps.
4. Since I use C++20, I replaced __builtin_ctz() with std::countr_zero().
5. I removed scalar comparisons at the beginning and end and replaced them with a single unaligned SIMD comparison.
6. Finally, I refactored the code to minimize code duplication when adding the unaligned SIMD comparison.

#pragma once

#include <immintrin.h>

#include <bit>
#include <cstddef>
#include <cstdint>
#include <iterator>
#include <memory>
#include <type_traits>

namespace simd {

namespace details {

using vec256i = __m256i;
using vec256f = __m256;
using vec256d = __m256d;

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) vec256i compare(const vec256i lhs, const vec256i rhs) {
  static_assert(WIDTH == 1 || WIDTH == 2 || WIDTH == 4 || WIDTH == 8);
  if constexpr (WIDTH == 1) {
    return _mm256_cmpeq_epi8(lhs, rhs);
  } else if (WIDTH == 2) {
    return _mm256_cmpeq_epi16(lhs, rhs);
  } else if (WIDTH == 4) {
    return _mm256_cmpeq_epi32(lhs, rhs);
  } else if (WIDTH == 8) {
    return _mm256_cmpeq_epi64(lhs, rhs);
  }
}

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) vec256i set_vector(const void* const value) {
  static_assert(WIDTH == 1 || WIDTH == 2 || WIDTH == 4 || WIDTH == 8);
  if constexpr (WIDTH == 1) {
    return _mm256_set1_epi8(*reinterpret_cast<const std::uint8_t*>(value));
  } else if (WIDTH == 2) {
    return _mm256_set1_epi16(*reinterpret_cast<const std::uint16_t*>(value));
  } else if (WIDTH == 4) {
    return _mm256_set1_epi32(*reinterpret_cast<const std::uint32_t*>(value));
  } else if (WIDTH == 8) {
    return _mm256_set1_epi64x(*reinterpret_cast<const std::uint64_t*>(value));
  }
}

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) std::uint32_t compress_mask(const vec256i mask) {
  static_assert(WIDTH == 1 || WIDTH == 2 || WIDTH == 4 || WIDTH == 8);
  if constexpr (WIDTH == 1 || WIDTH == 2) {
    return _mm256_movemask_epi8(mask);
  } else if (WIDTH == 4) {
    return _mm256_movemask_ps(_mm256_castsi256_ps(mask));
  } else if (WIDTH == 8) {
    return _mm256_movemask_pd(_mm256_castsi256_pd(mask));
  }
}

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) std::uint32_t countr_zero_ui128(const std::uint32_t mask0,
                                                                                    const std::uint32_t mask1,
                                                                                    const std::uint32_t mask2,
                                                                                    const std::uint32_t mask3) {
  static_assert(WIDTH == 1 || WIDTH == 2);
  if (mask0 != 0) {
    return std::countr_zero(mask0) / WIDTH;
  } else if (mask1 != 0) {
    return (std::countr_zero(mask1) + 32) / WIDTH;
  } else if (mask2 != 0) {
    return (std::countr_zero(mask2) + 64) / WIDTH;
  } else if (mask3 != 0) {
    return (std::countr_zero(mask3) + 96) / WIDTH;
  } else {
    return 128;
  }
}

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) std::uint32_t get_offset_first_occurrence(const vec256i mask0,
                                                                                              const vec256i mask1,
                                                                                              const vec256i mask2,
                                                                                              const vec256i mask3) {
  static_assert(WIDTH == 1 || WIDTH == 2 || WIDTH == 4 || WIDTH == 8);
  const auto compressed_mask0{compress_mask<WIDTH>(mask0)};
  const auto compressed_mask1{compress_mask<WIDTH>(mask1)};
  const auto compressed_mask2{compress_mask<WIDTH>(mask2)};
  const auto compressed_mask3{compress_mask<WIDTH>(mask3)};
  if constexpr (WIDTH == 1 || WIDTH == 2) {
    return countr_zero_ui128<WIDTH>(compressed_mask0, compressed_mask1, compressed_mask2, compressed_mask3);
  } else if (WIDTH == 4 || WIDTH == 8) {
    return std::countr_zero((compressed_mask3 << (96 / WIDTH)) + (compressed_mask2 << (64 / WIDTH)) +
                            (compressed_mask1 << (32 / WIDTH)) + compressed_mask0);
  }
}

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) const void* advance(const void* const begin,
                                                                        const std::uint32_t offset) {
  static_assert(WIDTH == 1 || WIDTH == 2 || WIDTH == 4 || WIDTH == 8);
  if constexpr (WIDTH == 1) {
    return reinterpret_cast<const std::uint8_t*>(begin) + offset;
  } else if (WIDTH == 2) {
    return reinterpret_cast<const std::uint16_t*>(begin) + offset;
  } else if (WIDTH == 4) {
    return reinterpret_cast<const std::uint32_t*>(begin) + offset;
  } else if (WIDTH == 8) {
    return reinterpret_cast<const std::uint64_t*>(begin) + offset;
  }
}

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) const void* find_aligned(const void* const begin,
                                                                             const void* const value) {
  static_assert(WIDTH == 1 || WIDTH == 2 || WIDTH == 4 || WIDTH == 8);
  const auto* const begin_vector{reinterpret_cast<const vec256i*>(begin)};
  const auto value_vector{set_vector<WIDTH>(value)};

  const auto mask0{compare<WIDTH>(value_vector, _mm256_load_si256(begin_vector))};
  const auto mask1{compare<WIDTH>(value_vector, _mm256_load_si256(begin_vector + 1))};
  const auto mask2{compare<WIDTH>(value_vector, _mm256_load_si256(begin_vector + 2))};
  const auto mask3{compare<WIDTH>(value_vector, _mm256_load_si256(begin_vector + 3))};

  const auto mask01{_mm256_or_si256(mask0, mask1)};
  const auto mask23{_mm256_or_si256(mask2, mask3)};
  const auto mask0123{_mm256_or_si256(mask01, mask23)};

  if (_mm256_testz_si256(mask0123, mask0123) == 0) {
    const auto offset{get_offset_first_occurrence<WIDTH>(mask0, mask1, mask2, mask3)};
    return advance<WIDTH>(begin, offset);
  }

  return nullptr;
}

template <std::size_t WIDTH>
[[nodiscard]] inline __attribute__((always_inline)) const void* find_unaligned(const void* const begin,
                                                                               const void* const value) {
  static_assert(WIDTH == 1 || WIDTH == 2 || WIDTH == 4 || WIDTH == 8);
  const auto* const begin_vector{reinterpret_cast<const vec256i*>(begin)};
  const auto value_vector{set_vector<WIDTH>(value)};

  const auto mask{compare<WIDTH>(value_vector, _mm256_loadu_si256(begin_vector))};
  if (_mm256_testz_si256(mask, mask) == 0) {
    const auto offset{get_offset_first_occurrence<WIDTH>(mask, mask, mask, mask)};
    return advance<WIDTH>(begin, offset);
  }

  return nullptr;
}

}  // namespace details

template <std::contiguous_iterator ITERATOR>
[[nodiscard]] ITERATOR find(ITERATOR begin_it, const ITERATOR end_it,
                            const typename std::iterator_traits<ITERATOR>::value_type value) {
  using value_t = typename std::iterator_traits<ITERATOR>::value_type;
  static_assert(std::is_scalar_v<value_t>);

  constexpr auto WIDTH{sizeof(value_t)};
  const auto* const value_void{reinterpret_cast<const void*>(&value)};
  const auto* const begin{std::to_address(begin_it)};
  const auto* const end{std::to_address(end_it)};
  const auto* current{std::to_address(begin_it)};

  if (std::distance(begin_it, end_it) * WIDTH < 256) {
    for (; begin_it != end_it; ++begin_it) {
      if (*begin_it == value) {
        return begin_it;
      }
    }
    return end_it;
  }

  if (auto address{reinterpret_cast<std::uintptr_t>(current)}; address % 32 != 0) {
    const auto* const current_void{reinterpret_cast<const void*>(current)};
    if (const auto* const found_value{details::find_unaligned<WIDTH>(current_void, value_void)};
        found_value != nullptr) {
      return begin_it + (reinterpret_cast<const value_t*>(found_value) - begin);
    }

    address += 32 - (address % 32);
    current = reinterpret_cast<const value_t*>(address);
  }

  while (current + (128 / WIDTH) <= end) {
    const auto* const current_void{reinterpret_cast<const void*>(current)};
    if (const auto* const found_value{details::find_aligned<WIDTH>(current_void, value_void)}; found_value != nullptr) {
      return begin_it + (reinterpret_cast<const value_t*>(found_value) - begin);
    }
    current += 128 / WIDTH;
  }

  if (current != end) {
    current = end - (32 / WIDTH);
    const auto* current_void{reinterpret_cast<const void*>(current)};
    if (const auto* const found_value{details::find_unaligned<WIDTH>(current_void, value_void)};
        found_value != nullptr) {
      return begin_it + (reinterpret_cast<const value_t*>(found_value) - begin);
    }
  }

  return end_it;
}

}  // namespace simd
```