Simpler AVX2 code, and other clean ups

include/Bitmask.h

@@ -40,11 +40,8 @@ class Bitmask {
 
   // Functions to perform bulk tests.
   //
-  // Two versions of the `has_true()` function. Both versions return -1 in case of no true found.
-  //   1) Position == false: it returns 1 indicating finding a true.
-  //   2) Position == true:  it returns the offset relative to `start` of the first true.
-  template <bool Position>
-  auto has_true(size_t start, size_t len) const -> int64_t;
+  auto has_true(size_t start, size_t len) const -> bool;
+  auto find_true(size_t start, size_t len) const -> int64_t;
   auto count_true() const -> size_t;  // How many 1's in this mask?
 
   // Functions for write

include/sperr_helper.h

@@ -184,8 +184,10 @@ auto chunk_volume(dims_type vol_dim, dims_type chunk_dim) -> std::vector<std::ar
 template <typename T>
 auto calc_mean_var(const T*, size_t len, size_t omp_nthreads = 0) -> std::array<T, 2>;
 
+#ifdef __AVX2__
 template <typename T>
-auto any_ge(const T* buf, size_t len, T threshold) -> bool;
+auto any_ge_pow2(const T* buf, size_t len, T threshold) -> bool;
+#endif
 
 };  // namespace sperr
 

src/Bitmask.cpp

@@ -41,46 +41,71 @@ auto sperr::Bitmask::rbit(size_t idx) const -> bool
   return word;
 }
 
-template <bool Position>
-auto sperr::Bitmask::has_true(size_t start, size_t len) const -> int64_t
+auto sperr::Bitmask::has_true(size_t start, size_t len) const -> bool
 {
   auto long_idx = start >> 6;
   auto processed_bits = int64_t{0};
   auto word = m_buf[long_idx];
-  auto answer = uint64_t{0};
 
   // Collect the remaining bits from the start long.
   auto begin_idx = start & 63;
   auto nbits = std::min(size_t{64}, begin_idx + len);
   for (auto i = begin_idx; i < nbits; i++) {
-    answer |= word & (uint64_t{1} << i);
-    if constexpr (Position) {
-      if (answer != 0)
-        return processed_bits;
-    }
+    if (word & (uint64_t{1} << i))
+      return true;
     processed_bits++;
   }
-  if constexpr (!Position) {
-    if (answer != 0)
-      return 1;
+
+  // Examine the subsequent full longs.
+  while (processed_bits + 64 <= len) {
+    word = m_buf[++long_idx];
+    if (word) {
+      return true;
+    }
+    processed_bits += 64;
+  }
+
+  // Examine the remaining bits
+  if (processed_bits < len) {
+    nbits = len - processed_bits;
+    assert(nbits < 64);
+    word = m_buf[++long_idx];
+    for (int64_t i = 0; i < nbits; i++) {
+      if (word & (uint64_t{1} << i))
+        return true;
+    }
+  }
+
+  return false;
+}
+
+auto sperr::Bitmask::find_true(size_t start, size_t len) const -> int64_t
+{
+  auto long_idx = start >> 6;
+  auto processed_bits = int64_t{0};
+  auto word = m_buf[long_idx];
+
+  // Collect the remaining bits from the start long.
+  auto begin_idx = start & 63;
+  auto nbits = std::min(size_t{64}, begin_idx + len);
+  for (auto i = begin_idx; i < nbits; i++) {
+    if (word & (uint64_t{1} << i))
+      return processed_bits;
+    processed_bits++;
   }
 
   // Examine the subsequent full longs.
   while (processed_bits + 64 <= len) {
     word = m_buf[++long_idx];
     if (word) {
-      if constexpr (Position) {
 #if __cplusplus >= 202002L
-        int64_t i = std::countr_zero(word);
-        return processed_bits + i;
+      int64_t i = std::countr_zero(word);
+      return processed_bits + i;
 #else
-        for (int64_t i = 0; i < 64; i++)
-          if (word & (uint64_t{1} << i))
-            return processed_bits + i;
+      for (int64_t i = 0; i < 64; i++)
+        if (word & (uint64_t{1} << i))
+          return processed_bits + i;
 #endif
-      }
-      else
-        return 1;
     }
     processed_bits += 64;
   }
@@ -90,24 +115,14 @@ auto sperr::Bitmask::has_true(size_t start, size_t len) const -> int64_t
     nbits = len - processed_bits;
     assert(nbits < 64);
     word = m_buf[++long_idx];
-    answer = 0;
     for (int64_t i = 0; i < nbits; i++) {
-      answer |= word & (uint64_t{1} << i);
-      if constexpr (Position) {
-        if (answer != 0)
-          return processed_bits + i;
-      }
-    }
-    if constexpr (!Position) {
-      if (answer != 0)
-        return 1;
+      if (word & (uint64_t{1} << i))
+        return processed_bits + i;
     }
   }
 
   return -1;
 }
-template auto sperr::Bitmask::has_true<true>(size_t, size_t) const -> int64_t;
-template auto sperr::Bitmask::has_true<false>(size_t, size_t) const -> int64_t;
 
 auto sperr::Bitmask::count_true() const -> size_t
 {

src/SPECK2D_INT_ENC.cpp

@@ -64,21 +64,15 @@ auto sperr::SPECK2D_INT_ENC<T>::m_decide_S_significance(const Set2D& set) const
 {
   assert(!set.is_empty());
 
-  // Only use SIMD implementation with 16 or more elements.
-  if (set.length_x < 16) {
-    for (auto y = set.start_y; y < (set.start_y + set.length_y); y++) {
-      auto first = m_coeff_buf.data() + y * m_dims[0] + set.start_x;
-      if (std::any_of(first, first + set.length_x,
-                      [thld = m_threshold](auto v) { return v >= thld; }))
-        return true;
-    }
-  }
-  else {
-    for (auto y = set.start_y; y < (set.start_y + set.length_y); y++) {
-      auto first = m_coeff_buf.data() + y * m_dims[0] + set.start_x;
-      if (sperr::any_ge(first, set.length_x, m_threshold))
-        return true;
-    }
+  for (auto y = set.start_y; y < (set.start_y + set.length_y); y++) {
+    auto first = m_coeff_buf.data() + y * m_dims[0] + set.start_x;
+#ifdef __AVX2__
+    if (sperr::any_ge_pow2(first, set.length_x, m_threshold))
+      return true;
+#else
+    if (std::any_of(first, first + set.length_x, [th = m_threshold](auto v) { return v >= th; }))
+      return true;
+#endif
   }
   return false;
 }
@@ -91,27 +85,27 @@ auto sperr::SPECK2D_INT_ENC<T>::m_decide_I_significance() const -> bool
   //
   assert(m_I.length_x == m_dims[0]);
   auto first = m_coeff_buf.data() + size_t{m_I.start_y} * size_t{m_I.length_x};
-  auto last = m_coeff_buf.data() + m_coeff_buf.size();
-  if (sperr::any_ge(first, last - first, m_threshold))
+  auto len = m_coeff_buf.size() - size_t{m_I.start_y} * size_t{m_I.length_x};
+#ifdef __AVX2__
+  if (sperr::any_ge_pow2(first, len, m_threshold))
+    return true;
+#else
+  if (std::any_of(first, first + len, [thld = m_threshold](auto v) { return v >= thld; }))
     return true;
+#endif
 
   // Second, test the rectangle that's directly to the right of the missing top-left corner.
-  // Only use SIMD implementation with 16 or more elements.
   //
-  auto len = m_dims[0] - m_I.start_x;
-  if (len < 16) {
-    for (auto y = 0u; y < m_I.start_y; y++) {
-      first = m_coeff_buf.data() + y * m_dims[0] + m_I.start_x;
-      if (std::any_of(first, first + len, [thld = m_threshold](auto v) { return v >= thld; }))
-        return true;
-    }
-  }
-  else {
-    for (auto y = 0u; y < m_I.start_y; y++) {
-      first = m_coeff_buf.data() + y * m_dims[0] + m_I.start_x;
-      if (sperr::any_ge(first, len, m_threshold))
-        return true;
-    }
+  len = m_dims[0] - m_I.start_x;
+  for (auto y = 0u; y < m_I.start_y; y++) {
+    first = m_coeff_buf.data() + y * m_dims[0] + m_I.start_x;
+#ifdef __AVX2__
+    if (sperr::any_ge_pow2(first, len, m_threshold))
+      return true;
+#else
+    if (std::any_of(first, first + len, [thld = m_threshold](auto v) { return v >= thld; }))
+      return true;
+#endif
   }
   return false;
 }

src/SPECK3D_INT_ENC.cpp

@@ -167,11 +167,12 @@ void sperr::SPECK3D_INT_ENC<T>::m_process_S(size_t idx1, size_t idx2, size_t& co
   // If need to output, it means the current set has unknown significance.
   if (output) {
     auto first = m_morton_buf.data() + set.get_morton();
-    if (set.num_elem() < 16)
-      is_sig = std::any_of(first, first + set.num_elem(),
-                           [thld = m_threshold](auto v) { return v >= thld; });
-    else
-      is_sig = sperr::any_ge(first, set.num_elem(), m_threshold);
+#ifdef __AVX2__
+    is_sig = sperr::any_ge_pow2(first, set.num_elem(), m_threshold);
+#else
+    is_sig = std::any_of(first, first + set.num_elem(),
+                         [thld = m_threshold](auto v) { return v >= thld; });
+#endif
     m_bit_buffer.wbit(is_sig);
   }
 

src/sperr_helper.cpp

@@ -642,87 +642,34 @@ auto sperr::calc_mean_var(const T* arr, size_t len, size_t omp_nthreads) -> std:
 template auto sperr::calc_mean_var(const float*, size_t, size_t) -> std::array<float, 2>;
 template auto sperr::calc_mean_var(const double*, size_t, size_t) -> std::array<double, 2>;
 
+#ifdef __AVX2__
 template <typename T>
-auto sperr::any_ge(const T* buf, size_t len, T thld) -> bool
+auto sperr::any_ge_pow2(const T* buf, size_t len, T thld) -> bool
 {
-#ifdef __AVX2__
-  if constexpr (sizeof(T) == 8) {  // uint64_t
-    // We use a trick here: for unsigned integers, A_unsigned > B_unsigned is equivalent to
-    // (A_unsigned XOR sign_bit)_signed > (B_unsigned XOR sign_bit)_signed.
-    //
-    const size_t simd_width = 4;
-    auto sign_flip_mask = _mm256_set1_epi64x(0x8000000000000000ULL);
-    auto thld_flipped = _mm256_set1_epi64x(thld ^ 0x8000000000000000ULL);
-
-    size_t i = 0;
-    for (; i + simd_width <= len; i += simd_width) {
-      auto data_vec = _mm256_lddqu_si256(reinterpret_cast<const __m256i*>(buf + i));
-      auto data_flipped = _mm256_xor_si256(data_vec, sign_flip_mask);
-      auto cmp_mask = _mm256_cmpgt_epi64(thld_flipped, data_flipped);  // threshold > data
-      int all_true = _mm256_movemask_epi8(cmp_mask);
-
-      if (all_true != 0xFFFFFFFF)
-        return true;
-    }
-
-    return std::any_of(buf + i, buf + len, [thld](auto v) { return v >= thld; });
-  }
-  else if constexpr (sizeof(T) == 4) {  // uint32_t
-    const size_t simd_width = 8;
-    auto thld_vec = _mm256_set1_epi32(thld);
-
-    size_t i = 0;
-    for (; i + simd_width <= len; i += simd_width) {
-      auto data_vec = _mm256_lddqu_si256(reinterpret_cast<const __m256i*>(buf + i));
-      auto min_vec = _mm256_min_epu32(data_vec, thld_vec);
-      auto cmp_vec = _mm256_cmpeq_epi32(min_vec, thld_vec);
-      int all_zeros = _mm256_testz_si256(cmp_vec, cmp_vec);
-
-      if (!all_zeros)
-        return true;
-    }
-
-    return std::any_of(buf + i, buf + len, [thld](auto v) { return v >= thld; });
-  }
-  else if constexpr (sizeof(T) == 2) {  // uint16_t
-    const size_t simd_width = 16;
-    auto thld_vec = _mm256_set1_epi16(thld);
-
-    size_t i = 0;
-    for (; i + simd_width <= len; i += simd_width) {
-      auto data_vec = _mm256_lddqu_si256(reinterpret_cast<const __m256i*>(buf + i));
-      auto min_vec = _mm256_min_epu16(data_vec, thld_vec);
-      auto cmp_vec = _mm256_cmpeq_epi16(min_vec, thld_vec);
-      int all_zeros = _mm256_testz_si256(cmp_vec, cmp_vec);
-
-      if (!all_zeros)
-        return true;
-    }
-
-    return std::any_of(buf + i, buf + len, [thld](auto v) { return v >= thld; });
-  }
-  else {  // uint8_t
-    const size_t simd_width = 32;
-    auto thld_vec = _mm256_set1_epi8(thld);
-
-    size_t i = 0;
-    for (; i + simd_width <= len; i += simd_width) {
-      auto data_vec = _mm256_lddqu_si256(reinterpret_cast<const __m256i*>(buf + i));
-      auto min_vec = _mm256_min_epu8(data_vec, thld_vec);
-      auto cmp_vec = _mm256_cmpeq_epi8(min_vec, thld_vec);
-      int all_zeros = _mm256_testz_si256(cmp_vec, cmp_vec);
-
-      if (!all_zeros)
-        return true;
-    }
+  assert((thld > 0) && (thld & (thld - 1)) == 0);
+
+  const size_t simd_width = 32 / sizeof(T);
+  T mask_val = ~(thld - 1);
+  __m256i mask_vec;
+  if constexpr (sizeof(T) == 8)
+    mask_vec = _mm256_set1_epi64x(mask_val);
+  else if constexpr (sizeof(T) == 4)
+    mask_vec = _mm256_set1_epi32(mask_val);
+  else if constexpr (sizeof(T) == 2)
+    mask_vec = _mm256_set1_epi16(mask_val);
+  else
+    mask_vec = _mm256_set1_epi8(mask_val);
 
-    return std::any_of(buf + i, buf + len, [thld](auto v) { return v >= thld; });
+  size_t i = 0;
+  for (; i + simd_width <= len; i += simd_width) {
+    auto data_vec = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(buf + i));
+    if (!_mm256_testz_si256(data_vec, mask_vec))
+      return true;
   }
-#else
-  return std::any_of(buf, buf + len, [thld](auto v) { return v >= thld; });
-#endif
+  return std::any_of(buf + i, buf + len, [thld](auto v) { return v >= thld; });
 }
-template auto sperr::any_ge(const uint8_t*, size_t, uint8_t) -> bool;
-template auto sperr::any_ge(const uint16_t*, size_t, uint16_t) -> bool;
-template auto sperr::any_ge(const uint32_t*, size_t, uint32_t) -> bool;
-template auto sperr::any_ge(const uint64_t*, size_t, uint64_t) -> bool;
+template auto sperr::any_ge_pow2(const uint8_t*, size_t, uint8_t) -> bool;
+template auto sperr::any_ge_pow2(const uint16_t*, size_t, uint16_t) -> bool;
+template auto sperr::any_ge_pow2(const uint32_t*, size_t, uint32_t) -> bool;
+template auto sperr::any_ge_pow2(const uint64_t*, size_t, uint64_t) -> bool;
+#endif

test_scripts/CMakeLists.txt

@@ -1,6 +1,12 @@
 add_executable(        sperr_helper sperr_helper_unit_test.cpp )
 target_link_libraries( sperr_helper PUBLIC SPERR GTest::gtest_main )
 
+if(ENABLE_AVX2)
+  target_compile_options(sperr_helper PRIVATE
+                        "$<$<CXX_COMPILER_ID:MSVC>:/arch:AVX2>"
+                        "$<$<OR:$<CXX_COMPILER_ID:GNU>,$<CXX_COMPILER_ID:Clang>>:-mavx2;-mfma>")
+endif()
+
 add_executable(        bitstream bitstream_unit_test.cpp )
 target_link_libraries( bitstream PUBLIC SPERR GTest::gtest_main )
 

test_scripts/bitstream_unit_test.cpp

@@ -374,11 +374,11 @@ TEST(Bitmask, has_true)
 
       // Loop over all range length
       for (size_t len = 0; len < mask_size - start; len++) {
-        auto ans1 = mask.has_true<false>(start, len);
-        auto ans2 = -1l;
+        auto ans1 = mask.has_true(start, len);
+        auto ans2 = false;
         for (size_t i = start; i < start + len; i++)
           if (mask.rbit(i)) {
-            ans2 = 1;
+            ans2 = true;
             break;
           }
         EXPECT_EQ(ans1, ans2);
@@ -402,7 +402,7 @@ TEST(Bitmask, has_true_position)
 
       // Loop over all range length
       for (size_t len = 0; len < mask_size - start; len++) {
-        auto ans1 = mask.has_true<true>(start, len);
+        auto ans1 = mask.find_true(start, len);
         auto ans2 = -1l;
         for (size_t i = start; i < start + len; i++)
           if (mask.rbit(i)) {