Various performance optimization in 3D compression

include/SPECK3D_INT.h

@@ -3,39 +3,20 @@
 
 #include "SPECK_INT.h"
 
-#include <cstring>  // std::memcpy()
 #include <tuple>
 
 namespace sperr {
 
 class Set3D {
- private:
-  // The first 6 bytes of the morton offset in uint64_t. Because each set dimension is
-  // stored using 16-bit integers, these 48 bits are big enough too!
-  std::array<uint8_t, 6> m_morton = {0, 0, 0, 0, 0, 0};
-
  public:
-  //
-  // Publicly accessible public data members.
-  //
+  uint64_t morton_idx = 0;
   uint16_t start_x = 0;
   uint16_t start_y = 0;
   uint16_t start_z = 0;
   uint16_t length_x = 0;
   uint16_t length_y = 0;
   uint16_t length_z = 0;
 
- public:
-  //
-  // Member functions (intended to be inline)
-  //
-  auto get_morton() const -> uint64_t
-  {
-    auto tmp = uint64_t{0};
-    std::memcpy(&tmp, m_morton.data(), sizeof(m_morton));
-    return tmp;
-  }
-  void set_morton(uint64_t val) { std::memcpy(m_morton.data(), &val, sizeof(m_morton)); }
   void make_empty() { length_x = 0; }
   auto num_elem() const -> size_t { return (size_t{length_x} * length_y * length_z); }
 };

include/SPECK3D_INT_ENC.h

@@ -35,10 +35,20 @@ class SPECK3D_INT_ENC final : public SPECK3D_INT<T> {
   void m_process_P(size_t idx, size_t morton, size_t& counter, bool output) final;
   void m_process_P_lite(size_t idx) final;
   void m_additional_initialization() final;
+  void m_bitplane_init() final;
+  void m_refinement_extra() final;
 
   // Data structures and functions for morton data layout.
-  vecui_type m_morton_buf;
+  // `m_morton_buf` stores the MSB bit position of each coefficient (via m_msb_position()),
+  // rather than the full coefficient value. This shrinks the buffer from sizeof(T) to 1 byte
+  // per element, reducing cache pressure in the significance-testing hot path (m_process_S).
+  std::vector<int8_t> m_morton_buf;
+  // `m_morton_threshold` is the MSB position of `m_threshold`, updated each bitplane via
+  // m_bitplane_init(). Significance tests compare m_morton_buf entries against this value.
+  int8_t m_morton_threshold = -1;
   void m_deposit_set(Set3D);
+  // Returns the bit position of the most significant bit (0-based), or -1 for zero.
+  auto m_msb_position(uint_type v) const -> int8_t;
 };
 
 };  // namespace sperr

include/SPECK_INT.h

@@ -69,15 +69,17 @@ class SPECK_INT {
   virtual void m_clean_LIS() = 0;
   virtual void m_sorting_pass() = 0;
   virtual void m_initialize_lists() = 0;
+  virtual void m_bitplane_init() {}
+  virtual void m_refinement_extra() {}
   void m_refinement_pass_encode();
   void m_refinement_pass_decode();
 
   // Data members
-  uint8_t m_num_bitplanes = 0;
-  uint_type m_threshold = 0;
   uint64_t m_total_bits = 0;  // The number of bits of a complete SPECK stream.
   uint64_t m_avail_bits = 0;  // Decoding only. `m_avail_bits` <= `m_total_bits`
   size_t m_budget = std::numeric_limits<size_t>::max();
+  uint_type m_threshold = 0;
+  uint8_t m_num_bitplanes = 0;
 
   dims_type m_dims = {0, 0, 0};
   vecui_type m_coeff_buf;

include/sperr_helper.h

@@ -184,11 +184,6 @@ 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_pow2(const T* buf, size_t len, T threshold) -> bool;
-#endif
-
 };  // namespace sperr
 
 #endif

src/SPECK2D_INT_ENC.cpp

@@ -66,13 +66,8 @@ auto sperr::SPECK2D_INT_ENC<T>::m_decide_S_significance(const Set2D& set) const
 
   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;
 }
@@ -86,26 +81,16 @@ 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 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.
   //
   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.cpp

@@ -148,7 +148,7 @@ void sperr::SPECK3D_INT<T>::m_code_S(size_t idx1, size_t idx2)
 
     // Element (0, 0, 0)
     const auto id = set.start_z * m_dims[0] * m_dims[1] + set.start_y * m_dims[0] + set.start_x;
-    auto mort = set.get_morton();
+    auto mort = set.morton_idx;
     m_LIP_mask.wtrue(id);
     m_process_P(id, mort, sig_counter, need_decide);
 
@@ -200,7 +200,7 @@ void sperr::SPECK3D_INT<T>::m_code_S(size_t idx1, size_t idx2)
       if (it->num_elem() == 1) {
         auto idx = it->start_z * m_dims[0] * m_dims[1] + it->start_y * m_dims[0] + it->start_x;
         m_LIP_mask.wtrue(idx);
-        m_process_P(idx, it->get_morton(), sig_counter, need_decide);
+        m_process_P(idx, it->morton_idx, sig_counter, need_decide);
       }
       else {
         m_LIS[next_lev].emplace_back(*it);
@@ -230,14 +230,14 @@ auto sperr::SPECK3D_INT<T>::m_partition_S_XYZ(Set3D set, uint16_t lev) const
 
   auto subsets = std::tuple<std::array<Set3D, 8>, uint16_t>();
   std::get<1>(subsets) = lev;
-  auto morton_offset = set.get_morton();
+  auto morton_offset = set.morton_idx;
 
   //
   // The actual figuring out where it starts/ends part...
   //
   // subset (0, 0, 0)
   auto& sub0 = std::get<0>(subsets)[0];
-  sub0.set_morton(morton_offset);
+  sub0.morton_idx = morton_offset;
   sub0.start_x = set.start_x;
   sub0.start_y = set.start_y;
   sub0.start_z = set.start_z;
@@ -248,7 +248,7 @@ auto sperr::SPECK3D_INT<T>::m_partition_S_XYZ(Set3D set, uint16_t lev) const
   // subset (1, 0, 0)
   auto& sub1 = std::get<0>(subsets)[1];
   morton_offset += sub0.num_elem();
-  sub1.set_morton(morton_offset);
+  sub1.morton_idx = morton_offset;
   sub1.start_x = set.start_x + split_x[0];
   sub1.start_y = set.start_y;
   sub1.start_z = set.start_z;
@@ -259,7 +259,7 @@ auto sperr::SPECK3D_INT<T>::m_partition_S_XYZ(Set3D set, uint16_t lev) const
   // subset (0, 1, 0)
   auto& sub2 = std::get<0>(subsets)[2];
   morton_offset += sub1.num_elem();
-  sub2.set_morton(morton_offset);
+  sub2.morton_idx = morton_offset;
   sub2.start_x = set.start_x;
   sub2.start_y = set.start_y + split_y[0];
   sub2.start_z = set.start_z;
@@ -270,7 +270,7 @@ auto sperr::SPECK3D_INT<T>::m_partition_S_XYZ(Set3D set, uint16_t lev) const
   // subset (1, 1, 0)
   auto& sub3 = std::get<0>(subsets)[3];
   morton_offset += sub2.num_elem();
-  sub3.set_morton(morton_offset);
+  sub3.morton_idx = morton_offset;
   sub3.start_x = set.start_x + split_x[0];
   sub3.start_y = set.start_y + split_y[0];
   sub3.start_z = set.start_z;
@@ -281,7 +281,7 @@ auto sperr::SPECK3D_INT<T>::m_partition_S_XYZ(Set3D set, uint16_t lev) const
   // subset (0, 0, 1)
   auto& sub4 = std::get<0>(subsets)[4];
   morton_offset += sub3.num_elem();
-  sub4.set_morton(morton_offset);
+  sub4.morton_idx = morton_offset;
   sub4.start_x = set.start_x;
   sub4.start_y = set.start_y;
   sub4.start_z = set.start_z + split_z[0];
@@ -292,7 +292,7 @@ auto sperr::SPECK3D_INT<T>::m_partition_S_XYZ(Set3D set, uint16_t lev) const
   // subset (1, 0, 1)
   auto& sub5 = std::get<0>(subsets)[5];
   morton_offset += sub4.num_elem();
-  sub5.set_morton(morton_offset);
+  sub5.morton_idx = morton_offset;
   sub5.start_x = set.start_x + split_x[0];
   sub5.start_y = set.start_y;
   sub5.start_z = set.start_z + split_z[0];
@@ -303,7 +303,7 @@ auto sperr::SPECK3D_INT<T>::m_partition_S_XYZ(Set3D set, uint16_t lev) const
   // subset (0, 1, 1)
   auto& sub6 = std::get<0>(subsets)[6];
   morton_offset += sub5.num_elem();
-  sub6.set_morton(morton_offset);
+  sub6.morton_idx = morton_offset;
   sub6.start_x = set.start_x;
   sub6.start_y = set.start_y + split_y[0];
   sub6.start_z = set.start_z + split_z[0];
@@ -314,7 +314,7 @@ auto sperr::SPECK3D_INT<T>::m_partition_S_XYZ(Set3D set, uint16_t lev) const
   // subset (1, 1, 1)
   auto& sub7 = std::get<0>(subsets)[7];
   morton_offset += sub6.num_elem();
-  sub7.set_morton(morton_offset);
+  sub7.morton_idx = morton_offset;
   sub7.start_x = set.start_x + split_x[0];
   sub7.start_y = set.start_y + split_y[0];
   sub7.start_z = set.start_z + split_z[0];
@@ -330,6 +330,7 @@ auto sperr::SPECK3D_INT<T>::m_partition_S_XY(Set3D set, uint16_t lev) const
     -> std::tuple<std::array<Set3D, 4>, uint16_t>
 {
   // This partition scheme is only used during initialization; no need to calculate morton offset.
+  // The correct morton offset will be calculated during initialization.
 
   const auto split_x = std::array<int, 2>{set.length_x - set.length_x / 2, set.length_x / 2};
   const auto split_y = std::array<int, 2>{set.length_y - set.length_y / 2, set.length_y / 2};
@@ -392,6 +393,7 @@ auto sperr::SPECK3D_INT<T>::m_partition_S_Z(Set3D set, uint16_t lev) const
     -> std::tuple<std::array<Set3D, 2>, uint16_t>
 {
   // This partition scheme is only used during initialization; no need to calculate morton offset.
+  // The correct morton offset will be calculated during initialization.
 
   const auto split_z = std::array<int, 2>{set.length_z - set.length_z / 2, set.length_z / 2};
   if (split_z[1] != 0)