morton.hpp

Go to the documentation of this file.

 
 
#ifndef MORTON_MORTON_HPP
#define MORTON_MORTON_HPP
 
#include <array>
#include <climits>
#include <cstddef>
#include <cstdint>
#include <type_traits>
 
#include "morton/wide_uint.hpp"
 
// ---- x86 ISA feature plumbing ---------------------------------------------
// MORTON_X86 : building for x86-64 with GCC/Clang for the host (not CUDA/HIP).
//   Gates the runtime CPU-feature queries and the AVX-512 batch kernels, which
//   use per-function `target` attributes so they need no global -mavx512f/-mbmi2
//   and never run unless the executing CPU actually supports the instructions.
#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) && \
    !defined(__CUDACC__) && !defined(__HIPCC__)
#define MORTON_X86 1
#else
#define MORTON_X86 0
#endif
 
// MORTON_ENABLE_RUNTIME_DISPATCH (opt-in, default 0): build a single binary
// *without* -mbmi2 that still uses PDEP/PEXT when the running CPU has BMI2 and
// the software fallback otherwise -- removing the build-time portability/speed
// trade-off (e.g. for distributable wheels). Left off by default so the plain
// non-BMI2 build stays strictly software (and pdep/pext-free). It self-disables
// when -mbmi2 is already on (the compile-time intrinsic path is used instead).
#ifndef MORTON_ENABLE_RUNTIME_DISPATCH
#define MORTON_ENABLE_RUNTIME_DISPATCH 0
#endif
#if MORTON_X86 && MORTON_ENABLE_RUNTIME_DISPATCH && !defined(__BMI2__)
#define MORTON_X86_RUNTIME_DISPATCH 1
#else
#define MORTON_X86_RUNTIME_DISPATCH 0
#endif
 
#if MORTON_X86 || defined(__BMI2__)
#include <immintrin.h>
#endif
 
// MORTON_HAS_INT128 and detail::uint128_t are defined by wide_uint.hpp (included above), which is
// the lowest header that needs the built-in; the alias carries the `__extension__` marker that
// keeps `-Wpedantic` consumers warning-free.
 
// Maximum supported code width (Dim * Bits). Codes <= 64 use a built-in
// integer, <= 128 use __uint128_t (where available), and wider codes use the
// software wide_uint<W> backend.
#ifndef MORTON_MAX_BITS
#define MORTON_MAX_BITS 256
#endif
 
// Mark functions callable from device code. Expands to nothing for an ordinary
// host C++ build, so the CPU library is completely unchanged.
//
// MORTON_HD (the host/device function marker) lives in its own header so that headers included
// above — wide_uint.hpp in particular — can annotate their operators device-callable too. It is
// already pulled in transitively via wide_uint.hpp; include it here as well for the
// direct/standalone case. See morton/hd.hpp for the Kokkos-vs-raw-CUDA-vs-host resolution.
#include "morton/hd.hpp"
 
namespace morton {
 
namespace detail {
 
// Smallest unsigned type that holds at least NBits bits: a built-in where one
// exists, otherwise a wide_uint of the right number of 64-bit words.
template <unsigned NBits>
struct uint_for {
#if defined(MORTON_HAS_INT128)
  static constexpr unsigned builtin_max = 128;
#else
  static constexpr unsigned builtin_max = 64;
#endif
  using builtin = std::conditional_t<
      (NBits <= 8), std::uint8_t,
      std::conditional_t<
          (NBits <= 16), std::uint16_t,
          std::conditional_t<(NBits <= 32), std::uint32_t,
#if defined(MORTON_HAS_INT128)
                             std::conditional_t<(NBits <= 64), std::uint64_t, uint128_t>
#else
                             std::uint64_t
#endif
                             >>>;
  using type = std::conditional_t<(NBits <= builtin_max), builtin, wide_uint<(NBits + 63) / 64>>;
};
template <unsigned NBits>
using uint_for_t = typename uint_for<NBits>::type;
 
// True during constant evaluation. Lets a constexpr function avoid the
// (non-constexpr) BMI2 intrinsics when evaluated at compile time.
MORTON_HD constexpr bool is_consteval() noexcept {
#if defined(__cpp_if_consteval)
  if consteval {
    return true;
  } else {
    return false;
  }
#elif defined(__GNUC__) || defined(__clang__)
  return __builtin_is_constant_evaluated();
#else
  return false;
#endif
}
 
// 3^n, as a compile-time value (size of a Moore neighbourhood including self).
MORTON_HD constexpr unsigned pow3(unsigned n) {
  unsigned r = 1;
  for (unsigned i = 0; i < n; ++i)
    r *= 3;
  return r;
}
 
// Build the mask selecting the bits that belong to axis `d` in a code with
// `Dim` interleaved axes of `Bits` bits each: positions d, d+Dim, d+2*Dim, ...
template <unsigned Dim, unsigned Bits, typename Code>
MORTON_HD constexpr Code axis_mask(unsigned d) {
  Code m = 0;
  for (unsigned i = 0; i < Bits; ++i)
    m |= (Code(1) << (i * Dim + d));
  return m;
}
 
// Software deposit: spread the low `Bits` bits of x across positions
// d, d+Dim, ... (the portable fallback for PDEP). Works for any code width.
template <unsigned Dim, unsigned Bits, typename Code, typename Coord>
MORTON_HD constexpr Code spread_sw(Coord x, unsigned d) {
  Code r = 0;
  for (unsigned i = 0; i < Bits; ++i)
    r |= Code(Code(x >> i) & Code(1)) << (i * Dim + d);
  return r;
}
 
// Software extract: inverse of spread_sw (the portable fallback for PEXT).
template <unsigned Dim, unsigned Bits, typename Code, typename Coord>
MORTON_HD constexpr Coord compact_sw(Code c, unsigned d) {
  Coord r = 0;
  for (unsigned i = 0; i < Bits; ++i)
    r |= Coord(Coord(c >> (i * Dim + d)) & Coord(1)) << i;
  return r;
}
 
#if MORTON_X86
// Cached runtime CPU-feature queries. __builtin_cpu_supports returns a nonzero
// bitfield when the feature is present, so compare against 0. The first call
// triggers detection; results are memoised in function-local statics.
inline bool cpu_has_bmi2() {
  static const bool v = __builtin_cpu_supports("bmi2") != 0;
  return v;
}
inline bool cpu_has_avx2() {
  static const bool v = __builtin_cpu_supports("avx2") != 0;
  return v;
}
inline bool cpu_has_avx512f() {
  static const bool v = __builtin_cpu_supports("avx512f") != 0;
  return v;
}
#endif
 
#if MORTON_X86_RUNTIME_DISPATCH
// Compiled with BMI2 codegen for just these two functions (via the `target`
// attribute), so the rest of the binary stays portable; only ever called from
// the !is_consteval() runtime path guarded by cpu_has_bmi2().
__attribute__((target("bmi2"))) inline std::uint64_t pdep_u64_hw(std::uint64_t v, std::uint64_t m) {
  return _pdep_u64(v, m);
}
__attribute__((target("bmi2"))) inline std::uint64_t pext_u64_hw(std::uint64_t v, std::uint64_t m) {
  return _pext_u64(v, m);
}
#endif
 
}  // namespace detail
 
template <unsigned Dim, unsigned Bits>
class Morton {
  static_assert(Dim >= 1, "Dim must be >= 1");
  static_assert(Bits >= 1, "Bits must be >= 1");
  static_assert(Dim * Bits <= MORTON_MAX_BITS,
                "Dim * Bits exceeds MORTON_MAX_BITS (raise it if you need wider codes)");
 
 public:
  static constexpr unsigned dimensions = Dim;
  static constexpr unsigned bits_per_axis = Bits;
  static constexpr unsigned code_bits = Dim * Bits;
 
  using code_type = detail::uint_for_t<Dim * Bits>;
  using coord_type = detail::uint_for_t<Bits>;
 
 private:
  static constexpr unsigned code_type_bits = sizeof(code_type) * CHAR_BIT;
  static constexpr unsigned coord_type_bits = sizeof(coord_type) * CHAR_BIT;
 
 public:
  static constexpr code_type field_mask = (code_bits >= code_type_bits)
                                              ? code_type(~code_type(0))
                                              : code_type((code_type(1) << code_bits) - 1);
 
  static constexpr coord_type coord_max = (Bits >= coord_type_bits)
                                              ? coord_type(~coord_type(0))
                                              : coord_type((coord_type(1) << Bits) - 1);
 
  static constexpr unsigned octants = unsigned(1) << Dim;  // children per node
 
  // ---- construction -----------------------------------------------------
 
  MORTON_HD constexpr Morton() noexcept : code_(0) {}
 
  MORTON_HD static constexpr Morton from_code(code_type raw) noexcept {
    Morton m;
    m.code_ = raw & field_mask;
    return m;
  }
 
  template <typename... Cs, typename = std::enable_if_t<sizeof...(Cs) == Dim>>
  MORTON_HD static constexpr Morton encode(Cs... coords) {
    coord_type tmp[Dim] = {static_cast<coord_type>(coords)...};
    Morton m;
    for (unsigned d = 0; d < Dim; ++d)
      m.code_ |= deposit(tmp[d], d);
    return m;
  }
 
  MORTON_HD static constexpr Morton encode(const std::array<coord_type, Dim>& c) {
    Morton m;
    for (unsigned d = 0; d < Dim; ++d)
      m.code_ |= deposit(c[d], d);
    return m;
  }
 
  // ---- access -----------------------------------------------------------
 
  MORTON_HD constexpr code_type code() const noexcept { return code_; }
 
  MORTON_HD constexpr coord_type get(unsigned d) const { return extract(code_, d); }
 
  MORTON_HD constexpr std::array<coord_type, Dim> decode() const {
    std::array<coord_type, Dim> out{};
    for (unsigned d = 0; d < Dim; ++d)
      out[d] = extract(code_, d);
    return out;
  }
 
  MORTON_HD constexpr void set(unsigned d, coord_type value) {
    code_ = (code_ & keep_mask(d)) | deposit(value, d);
  }
 
  // ---- the headline: arithmetic in Morton space -------------------------
  //
  // Each of these touches only the bits of one axis and leaves the other
  // axes' interleaved bits exactly where they are -- no decode/encode.
  // All wrap modulo 2^Bits on the affected axis.
 
  MORTON_HD constexpr void add(unsigned d, coord_type k) {
    const code_type M = axis_mask(d);
    code_type s = (code_ | ~M) + deposit(k, d);
    code_ = (s & M) | (code_ & M_complement(d));
  }
 
  MORTON_HD constexpr void sub(unsigned d, coord_type k) {
    const code_type M = axis_mask(d);
    code_type s = (code_ & M) - deposit(k, d);
    code_ = (s & M) | (code_ & M_complement(d));
  }
 
  MORTON_HD constexpr void inc(unsigned d) {
    const code_type M = axis_mask(d);
    code_type s = (code_ | ~M) + lsb(d);
    code_ = (s & M) | (code_ & M_complement(d));
  }
 
  MORTON_HD constexpr void dec(unsigned d) {
    const code_type M = axis_mask(d);
    code_type s = (code_ & M) - lsb(d);
    code_ = (s & M) | (code_ & M_complement(d));
  }
 
  MORTON_HD constexpr Morton neighbor(unsigned d, int dir) const {
    Morton m = *this;
    if (dir >= 0)
      m.inc(d);
    else
      m.dec(d);
    return m;
  }
 
  // ---- saturating / checked arithmetic (do not wrap) --------------------
 
  MORTON_HD constexpr void add_sat(unsigned d, coord_type k) {
    coord_type cur = get(d);
    coord_type room = coord_type(coord_max - cur);
    set(d, (k > room) ? coord_max : coord_type(cur + k));
  }
 
  MORTON_HD constexpr void sub_sat(unsigned d, coord_type k) {
    coord_type cur = get(d);
    set(d, (k > cur) ? coord_type(0) : coord_type(cur - k));
  }
 
  MORTON_HD constexpr bool try_add(unsigned d, coord_type k) {
    coord_type cur = get(d);
    if (k > coord_type(coord_max - cur))
      return false;
    set(d, coord_type(cur + k));
    return true;
  }
 
  MORTON_HD constexpr bool try_sub(unsigned d, coord_type k) {
    coord_type cur = get(d);
    if (k > cur)
      return false;
    set(d, coord_type(cur - k));
    return true;
  }
 
  // ---- neighbour sets ----------------------------------------------------
 
  MORTON_HD constexpr std::array<Morton, 2 * Dim> face_neighbors() const {
    std::array<Morton, 2 * Dim> out{};
    for (unsigned d = 0; d < Dim; ++d) {
      out[2 * d] = neighbor(d, -1);
      out[2 * d + 1] = neighbor(d, +1);
    }
    return out;
  }
 
  MORTON_HD constexpr std::array<Morton, detail::pow3(Dim) - 1> all_neighbors() const {
    std::array<Morton, detail::pow3(Dim) - 1> out{};
    unsigned n = 0;
    for (unsigned i = 0; i < detail::pow3(Dim); ++i) {
      int off[Dim];
      unsigned t = i;
      bool zero = true;
      for (unsigned d = 0; d < Dim; ++d) {
        off[d] = int(t % 3) - 1;
        if (off[d] != 0)
          zero = false;
        t /= 3;
      }
      if (zero)
        continue;  // skip self
      Morton m = *this;
      for (unsigned d = 0; d < Dim; ++d) {
        if (off[d] > 0)
          m.inc(d);
        else if (off[d] < 0)
          m.dec(d);
      }
      out[n++] = m;
    }
    return out;
  }
 
  // ---- octree / hierarchy navigation ------------------------------------
  //
  // A cell at `level` covers a 2^level block per axis and has its low
  // level*Dim code bits zero. These helpers express that hierarchy directly.
 
  MORTON_HD constexpr Morton ancestor(unsigned level) const {
    if (level * Dim >= code_bits)
      return Morton{};
    code_type lowmask = (code_type(1) << (level * Dim)) - 1;
    return from_code(code_ & ~lowmask);
  }
 
  MORTON_HD constexpr unsigned child_index(unsigned level) const {
    return unsigned((code_ >> (level * Dim)) & (octants - 1));
  }
 
  MORTON_HD constexpr Morton child(unsigned level, unsigned octant) const {
    unsigned shift = (level - 1) * Dim;
    code_type cleared = code_ & ~(code_type(octants - 1) << shift);
    return from_code(cleared | (code_type(octant) << shift));
  }
 
  // ---- Z-order successor / predecessor ----------------------------------
 
  MORTON_HD constexpr Morton& operator++() noexcept {
    code_ = (code_ + 1) & field_mask;
    return *this;
  }
  MORTON_HD constexpr Morton operator++(int) noexcept {
    Morton t = *this;
    ++*this;
    return t;
  }
  MORTON_HD constexpr Morton& operator--() noexcept {
    code_ = (code_ - 1) & field_mask;
    return *this;
  }
  MORTON_HD constexpr Morton operator--(int) noexcept {
    Morton t = *this;
    --*this;
    return t;
  }
 
  // ---- comparison (Z-order) ---------------------------------------------
 
  friend MORTON_HD constexpr bool operator==(Morton a, Morton b) noexcept {
    return a.code_ == b.code_;
  }
  friend MORTON_HD constexpr bool operator!=(Morton a, Morton b) noexcept {
    return a.code_ != b.code_;
  }
  friend MORTON_HD constexpr bool operator<(Morton a, Morton b) noexcept {
    return a.code_ < b.code_;
  }
  friend MORTON_HD constexpr bool operator<=(Morton a, Morton b) noexcept {
    return a.code_ <= b.code_;
  }
  friend MORTON_HD constexpr bool operator>(Morton a, Morton b) noexcept {
    return a.code_ > b.code_;
  }
  friend MORTON_HD constexpr bool operator>=(Morton a, Morton b) noexcept {
    return a.code_ >= b.code_;
  }
 
  // ---- low-level deposit/extract ----------------------------------------
 
  MORTON_HD static constexpr code_type deposit(coord_type x, unsigned d) {
    if constexpr (code_bits <= 64) {
      // PDEP is a host x86 instruction; never emit it in CUDA device code
      // (where __CUDA_ARCH__ is defined), even if the host was -mbmi2.
#if defined(__BMI2__) && !defined(__CUDA_ARCH__)
      if (!detail::is_consteval())
        return code_type(_pdep_u64(std::uint64_t(x), std::uint64_t(axis_mask(d))));
#elif MORTON_X86_RUNTIME_DISPATCH
      // Single-binary path: use PDEP when the running CPU has BMI2.
      if (!detail::is_consteval() && detail::cpu_has_bmi2())
        return code_type(detail::pdep_u64_hw(std::uint64_t(x), std::uint64_t(axis_mask(d))));
#endif
    }
    return detail::spread_sw<Dim, Bits, code_type, coord_type>(x, d);
  }
 
  MORTON_HD static constexpr coord_type extract(code_type c, unsigned d) {
    if constexpr (code_bits <= 64) {
#if defined(__BMI2__) && !defined(__CUDA_ARCH__)
      if (!detail::is_consteval())
        return coord_type(_pext_u64(std::uint64_t(c), std::uint64_t(axis_mask(d))));
#elif MORTON_X86_RUNTIME_DISPATCH
      if (!detail::is_consteval() && detail::cpu_has_bmi2())
        return coord_type(detail::pext_u64_hw(std::uint64_t(c), std::uint64_t(axis_mask(d))));
#endif
    }
    return detail::compact_sw<Dim, Bits, code_type, coord_type>(c, d);
  }
 
  MORTON_HD static constexpr code_type axis_mask(unsigned d) {
#if defined(__CUDA_ARCH__)
    // On device, compute the mask (avoids referencing the host static
    // array from device code); it is a short compile-time-sized loop.
    return detail::axis_mask<Dim, Bits, code_type>(d);
#else
    return axis_masks_[d];
#endif
  }
 
 private:
  MORTON_HD static constexpr code_type M_complement(unsigned d) {
    return field_mask & ~axis_mask(d);
  }
  MORTON_HD static constexpr code_type keep_mask(unsigned d) { return M_complement(d); }
  MORTON_HD static constexpr code_type lsb(unsigned d) { return code_type(1) << d; }
 
  MORTON_HD static constexpr std::array<code_type, Dim> make_masks() {
    std::array<code_type, Dim> a{};
    for (unsigned d = 0; d < Dim; ++d)
      a[d] = detail::axis_mask<Dim, Bits, code_type>(d);
    return a;
  }
  static constexpr std::array<code_type, Dim> axis_masks_ = make_masks();
 
  code_type code_;
};
 
// Convenient aliases for the common cases.
using Morton2D32 = Morton<2, 32>;  // 2D, 32 bits/axis  -> 64-bit code
using Morton2D16 = Morton<2, 16>;  // 2D, 16 bits/axis  -> 32-bit code
using Morton3D21 = Morton<3, 21>;  // 3D, 21 bits/axis  -> 63-bit code
using Morton3D16 = Morton<3, 16>;  // 3D, 16 bits/axis  -> 48-bit code
#if defined(MORTON_HAS_INT128)
using Morton3D32 = Morton<3, 32>;  // 3D, 32 bits/axis  -> 96-bit code
using Morton2D64 = Morton<2, 64>;  // 2D, 64 bits/axis  -> 128-bit code
#endif
 
}  // namespace morton
 
#endif  // MORTON_MORTON_HPP