batch.hpp

Go to the documentation of this file.

 
 
#ifndef MORTON_BATCH_HPP
#define MORTON_BATCH_HPP
 
#include <cstddef>
#include <cstdint>
#include <type_traits>
 
#include "morton/morton.hpp"
#include "morton/simd.hpp"
 
namespace morton {
namespace batch {
 
template <unsigned Dim, unsigned Bits>
void add(const typename Morton<Dim, Bits>::code_type* in,
         typename Morton<Dim, Bits>::code_type* out, std::size_t n, unsigned axis,
         typename Morton<Dim, Bits>::coord_type k) {
  using M = Morton<Dim, Bits>;
  using code = typename M::code_type;
  const code mask = M::axis_mask(axis);
  const code notM = code(~mask);
  const code keep = M::field_mask & notM;
  const code dk = M::deposit(k, axis);  // loop-invariant
#if MORTON_X86
  if constexpr (std::is_same_v<code, std::uint64_t>) {
    if (detail::cpu_has_avx512f()) {
      detail::avx512::add64(in, out, n, notM, dk, mask, keep);
      return;
    }
  }
#endif
  for (std::size_t i = 0; i < n; ++i) {
    code c = in[i];
    out[i] = code((c | notM) + dk) & mask;
    out[i] |= c & keep;
  }
}
 
template <unsigned Dim, unsigned Bits>
void sub(const typename Morton<Dim, Bits>::code_type* in,
         typename Morton<Dim, Bits>::code_type* out, std::size_t n, unsigned axis,
         typename Morton<Dim, Bits>::coord_type k) {
  using M = Morton<Dim, Bits>;
  using code = typename M::code_type;
  const code mask = M::axis_mask(axis);
  const code keep = M::field_mask & code(~mask);
  const code dk = M::deposit(k, axis);
#if MORTON_X86
  if constexpr (std::is_same_v<code, std::uint64_t>) {
    if (detail::cpu_has_avx512f()) {
      detail::avx512::sub64(in, out, n, dk, mask, keep);
      return;
    }
  }
#endif
  for (std::size_t i = 0; i < n; ++i) {
    code c = in[i];
    out[i] = (code((c & mask) - dk) & mask) | (c & keep);
  }
}
 
template <unsigned Dim, unsigned Bits>
void step(const typename Morton<Dim, Bits>::code_type* in,
          typename Morton<Dim, Bits>::code_type* out, std::size_t n, unsigned axis, int dir) {
  if (dir >= 0)
    add<Dim, Bits>(in, out, n, axis, 1);
  else
    sub<Dim, Bits>(in, out, n, axis, 1);
}
 
template <unsigned Bits>
void encode2(const typename Morton<2, Bits>::coord_type* x,
             const typename Morton<2, Bits>::coord_type* y,
             typename Morton<2, Bits>::code_type* out, std::size_t n) {
  using M = Morton<2, Bits>;
#if MORTON_X86
  if constexpr (Bits == 32) {  // coord_type == uint32, code_type == uint64
    if (detail::cpu_has_avx512f()) {
      detail::avx512::encode2(x, y, out, n);
      return;
    }
  }
#endif
  for (std::size_t i = 0; i < n; ++i)
    out[i] = M::encode(x[i], y[i]).code();
}
 
template <unsigned Bits>
void encode3(const typename Morton<3, Bits>::coord_type* x,
             const typename Morton<3, Bits>::coord_type* y,
             const typename Morton<3, Bits>::coord_type* z,
             typename Morton<3, Bits>::code_type* out, std::size_t n) {
  using M = Morton<3, Bits>;
#if MORTON_X86
  if constexpr (Bits == 21) {  // coord_type == uint32, code_type == uint64
    if (detail::cpu_has_avx512f()) {
      detail::avx512::encode3(x, y, z, out, n);
      return;
    }
  }
#endif
  for (std::size_t i = 0; i < n; ++i)
    out[i] = M::encode(x[i], y[i], z[i]).code();
}
 
template <unsigned Bits>
void decode2(const typename Morton<2, Bits>::code_type* in, typename Morton<2, Bits>::coord_type* x,
             typename Morton<2, Bits>::coord_type* y, std::size_t n) {
  using M = Morton<2, Bits>;
#if MORTON_X86
  if constexpr (Bits == 32) {
    if (detail::cpu_has_avx512f()) {
      detail::avx512::decode2(in, x, y, n);
      return;
    }
  }
#endif
  for (std::size_t i = 0; i < n; ++i) {
    auto a = M::from_code(in[i]).decode();
    x[i] = a[0];
    y[i] = a[1];
  }
}
 
template <unsigned Bits>
void decode3(const typename Morton<3, Bits>::code_type* in, typename Morton<3, Bits>::coord_type* x,
             typename Morton<3, Bits>::coord_type* y, typename Morton<3, Bits>::coord_type* z,
             std::size_t n) {
  using M = Morton<3, Bits>;
#if MORTON_X86
  if constexpr (Bits == 21) {
    if (detail::cpu_has_avx512f()) {
      detail::avx512::decode3(in, x, y, z, n);
      return;
    }
  }
#endif
  for (std::size_t i = 0; i < n; ++i) {
    auto a = M::from_code(in[i]).decode();
    x[i] = a[0];
    y[i] = a[1];
    z[i] = a[2];
  }
}
 
}  // namespace batch
}  // namespace morton
 
#endif  // MORTON_BATCH_HPP