ndarray.h

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <array>
#include <atomic>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <initializer_list>
#include <memory>
#include <new>
#include <span>
#include <sstream>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "clustering/ndarray_range.h"
 
namespace clustering {
namespace detail {

struct BorrowedTag {};

inline std::atomic<std::uint64_t> &alignedAllocCallCount() noexcept {
  static std::atomic<std::uint64_t> counter{0};
  return counter;
}

template <class T, std::size_t Align> class AlignedAllocator : public std::allocator<T> {
public:
  using value_type = T;
  using size_type = std::size_t;
  using pointer = T *;
  using const_pointer = const T *;
 
  // is_always_equal + POCMA/POCCA/POCS = true_type lets std::vector pointer-steal on move and
  // perform a fresh allocation on copy-assign without per-instance allocator state checks.
  using is_always_equal = std::true_type;
  using propagate_on_container_move_assignment = std::true_type;
  using propagate_on_container_copy_assignment = std::true_type;
  using propagate_on_container_swap = std::true_type;
 
  template <typename U> struct rebind {
    using other = AlignedAllocator<U, Align>;
  };
 
  AlignedAllocator() noexcept = default;
 
  template <typename U> AlignedAllocator(const AlignedAllocator<U, Align> & /*unused*/) noexcept {}
 
  pointer allocate(size_type n) {
    // std::aligned_alloc(Align, 0) is implementation-defined; bypass it explicitly.
    if (n == 0) {
      return nullptr;
    }
    const size_type bytes = n * sizeof(T);
    const size_type aligned_bytes = (bytes + Align - 1) / Align * Align;
    alignedAllocCallCount().fetch_add(1, std::memory_order_relaxed);
    if (auto *ptr = static_cast<pointer>(std::aligned_alloc(Align, aligned_bytes))) {
      return ptr;
    }
    throw std::bad_alloc();
  }
 
  void deallocate(pointer ptr, size_type /*n*/) noexcept { std::free(ptr); }
};
 
template <class T, class U, std::size_t Align>
bool operator==(const AlignedAllocator<T, Align> &, const AlignedAllocator<U, Align> &) noexcept {
  return true;
}
 
template <class T, class U, std::size_t Align>
bool operator!=(const AlignedAllocator<T, Align> &, const AlignedAllocator<U, Align> &) noexcept {
  return false;
}
 
} // namespace detail

enum class NDArrayStorage : std::uint8_t { Owned, Borrowed };

enum class Layout : std::uint8_t { Contig, MaybeStrided };

template <class T, std::size_t N, Layout L = Layout::Contig> class NDArray {
  static_assert(N >= 1, "NDArray rank must be >= 1");
  // Widened from {float, double} so integer widths (signed/unsigned) are permitted as label /
  // index / count storage. bool stays excluded: std::vector<bool> is a specialization without
  // contiguous T-addressable storage, which would silently break data(), alignedData, and the
  // AlignedAllocator cache-line alignment invariant. Distance / reduction / GEMM primitives
  // carry their own float/double gates so integer NDArrays cannot reach numeric math without a
  // compile error.
  static_assert(std::is_arithmetic_v<T> && !std::is_same_v<T, bool>,
                "NDArray element type must be arithmetic and not bool");
 
  // All NDArray template instantiations share friendship so view-producing verbs can construct
  // a result with a different rank or layout via the private BorrowedTag constructor.
  template <class U, std::size_t M, Layout LL> friend class NDArray;
 
public:
  class BaseAccessor {
  public:
    const T *data() const { return this->m_ndarray->data() + this->m_index; }
 
  protected:
    NDArray<T, N, Layout::Contig> *m_ndarray; 
    std::size_t m_index;                      
    std::size_t m_dim;                        

    BaseAccessor(NDArray<T, N, Layout::Contig> *ndarray, std::size_t index, std::size_t dim)
        : m_ndarray(ndarray), m_index(index), m_dim(dim) {}
  };

  class ConstAccessor : public BaseAccessor {
  public:
    ConstAccessor(const NDArray<T, N, Layout::Contig> &ndarray, std::size_t index, std::size_t dim)
        : BaseAccessor(const_cast<NDArray<T, N, Layout::Contig> *>(&ndarray), index, dim) {}

    ConstAccessor(const ConstAccessor &other) = default;

    ConstAccessor operator[](std::size_t index) const noexcept {
      assert(this->m_dim < N && index < this->m_ndarray->dim(this->m_dim + 1));
      // Contig invariant lets the chain collapse to m_index * shape[dim+1] + index at every step,
      // matching the baseline hot-loop asm clang can vectorise into an 8x-unrolled aligned load.
      const size_t new_index = (this->m_index * this->m_ndarray->dim(this->m_dim + 1)) + index;
      return ConstAccessor(*this->m_ndarray, new_index, this->m_dim + 1);
    }

    operator T() const noexcept { return this->m_ndarray->flatIndex(this->m_index); }

    [[nodiscard]] size_t index() const noexcept { return this->m_index; }
  };

  class Accessor : public ConstAccessor {
  public:
    Accessor(NDArray<T, N, Layout::Contig> &ndarray, std::size_t index, std::size_t dim)
        : ConstAccessor(ndarray, index, dim) {}

    Accessor operator[](std::size_t index) noexcept {
      assert(this->m_dim < N && index < this->m_ndarray->dim(this->m_dim + 1));
      const size_t new_index = (this->m_index * this->m_ndarray->dim(this->m_dim + 1)) + index;
      return Accessor(*this->m_ndarray, new_index, this->m_dim + 1);
    }

    Accessor &operator=(T value) noexcept {
      assert(this->m_ndarray->m_mutable && "write to read-only borrow");
      this->m_ndarray->flatIndex(this->m_index) = value;
      return *this;
    }
  };

  template <Layout L2 = L>
    requires(L2 == Layout::Contig)
  NDArray(std::initializer_list<std::size_t> dims)
      : m_data(nullptr), m_base(nullptr), m_offset(0), m_storage(NDArrayStorage::Owned),
        m_mutable(true) {
    assert(dims.size() == N);
    std::size_t i = 0;
    std::size_t size = 1;
    for (auto d : dims) {
      m_shape[i++] = d;
      size *= d;
    }
    m_strides = computeContiguousStrides(m_shape);
    m_vec.resize(size);
    m_data = m_vec.data();
    m_base = m_data;
  }

  template <Layout L2 = L>
    requires(L2 == Layout::Contig)
  explicit NDArray(std::array<std::size_t, N> shape)
      : m_data(nullptr), m_base(nullptr), m_shape(shape), m_offset(0),
        m_storage(NDArrayStorage::Owned), m_mutable(true) {
    std::size_t size = 1;
    for (std::size_t k = 0; k < N; ++k) {
      size *= m_shape[k];
    }
    m_strides = computeContiguousStrides(m_shape);
    m_vec.resize(size);
    m_data = m_vec.data();
    m_base = m_data;
  }
 
  // Storage-aware special members: Owned arrays re-seat m_data against this->m_vec (the move
  // stole or the copy just populated it), while Borrowed arrays carry an empty m_vec and must
  // preserve the external pointer from the source.
  NDArray(const NDArray &other)
      : m_vec(other.m_vec), m_shape(other.m_shape), m_strides(other.m_strides),
        m_offset(other.m_offset), m_storage(other.m_storage), m_mutable(other.m_mutable) {
    m_data = (m_storage == NDArrayStorage::Owned) ? m_vec.data() : other.m_data;
    m_base = (m_storage == NDArrayStorage::Owned) ? m_vec.data() : other.m_base;
  }

  NDArray(NDArray &&other) noexcept
      : m_vec(std::move(other.m_vec)), m_shape(other.m_shape), m_strides(other.m_strides),
        m_offset(other.m_offset), m_storage(other.m_storage), m_mutable(other.m_mutable) {
    m_data = (m_storage == NDArrayStorage::Owned) ? m_vec.data() : other.m_data;
    m_base = (m_storage == NDArrayStorage::Owned) ? m_vec.data() : other.m_base;
    other.m_data = nullptr;
    other.m_base = nullptr;
  }

  NDArray &operator=(const NDArray &other) {
    if (this == &other) {
      return *this;
    }
    m_vec = other.m_vec;
    m_shape = other.m_shape;
    m_strides = other.m_strides;
    m_offset = other.m_offset;
    m_storage = other.m_storage;
    m_mutable = other.m_mutable;
    m_data = (m_storage == NDArrayStorage::Owned) ? m_vec.data() : other.m_data;
    m_base = (m_storage == NDArrayStorage::Owned) ? m_vec.data() : other.m_base;
    return *this;
  }

  NDArray &operator=(NDArray &&other) noexcept {
    if (this == &other) {
      return *this;
    }
    m_vec = std::move(other.m_vec);
    m_shape = other.m_shape;
    m_strides = other.m_strides;
    m_offset = other.m_offset;
    m_storage = other.m_storage;
    m_mutable = other.m_mutable;
    m_data = (m_storage == NDArrayStorage::Owned) ? m_vec.data() : other.m_data;
    m_base = (m_storage == NDArrayStorage::Owned) ? m_vec.data() : other.m_base;
    other.m_data = nullptr;
    other.m_base = nullptr;
    return *this;
  }
 
private:
  NDArray(clustering::detail::BorrowedTag, T *data, T *base, std::array<std::size_t, N> shape,
          std::array<std::ptrdiff_t, N> strides, std::ptrdiff_t offset, bool isMutable) noexcept
      : m_data(data), m_base(base), m_vec(), m_shape(shape), m_strides(strides), m_offset(offset),
        m_storage(NDArrayStorage::Borrowed), m_mutable(isMutable) {
    if constexpr (L == Layout::Contig) {
      assert(offset == 0 && strides == computeContiguousStrides(shape) &&
             "Contig NDArray requires contiguous strides and zero offset");
    }
  }
 
public:
  template <Layout L2 = L>
    requires(L2 == Layout::Contig)
  Accessor operator[](std::size_t index) noexcept {
    assert(index < m_shape[0]);
    return Accessor(*this, index, 0);
  }

  template <Layout L2 = L>
    requires(L2 == Layout::Contig)
  ConstAccessor operator[](std::size_t index) const noexcept {
    assert(index < m_shape[0]);
    return ConstAccessor(*this, index, 0);
  }

  template <class... Ix> T &operator()(Ix... ix) noexcept {
    static_assert(sizeof...(Ix) == N, "operator() requires exactly N indices");
    assert(m_mutable && "write to read-only borrow");
    return m_data[computeElementOffset(std::index_sequence_for<Ix...>{}, ix...)];
  }

  template <class... Ix> const T &operator()(Ix... ix) const noexcept {
    static_assert(sizeof...(Ix) == N, "operator() requires exactly N indices");
    return m_data[computeElementOffset(std::index_sequence_for<Ix...>{}, ix...)];
  }

  T &flatIndex(std::size_t index) noexcept {
    assert(m_mutable && "write to read-only borrow");
    return m_data[index];
  }

  const T &flatIndex(std::size_t index) const noexcept { return m_data[index]; }

  size_t dim(std::size_t index) const noexcept { return m_shape[index]; }

  std::ptrdiff_t strideAt(std::size_t index) const noexcept { return m_strides[index]; }

  [[nodiscard]] bool isContiguous() const noexcept {
    if constexpr (L == Layout::Contig) {
      return true;
    } else {
      return m_offset == 0 && m_strides == computeContiguousStrides(m_shape);
    }
  }

  [[nodiscard]] bool isMutable() const noexcept { return m_mutable; }

  [[nodiscard]] bool isOwned() const noexcept { return m_storage == NDArrayStorage::Owned; }

  const T *data() const noexcept { return m_data; }

  T *data() noexcept {
    assert(m_mutable && "write to read-only borrow");
    return m_data;
  }

  [[nodiscard]] T *baseData() const noexcept { return m_base; }

  template <std::size_t A> bool isAligned() const noexcept {
    return (reinterpret_cast<std::uintptr_t>(m_data) % A) == 0;
  }

  template <std::size_t A> T *alignedData() noexcept {
    assert(isAligned<A>() && "alignedData<A>() requires A-byte aligned data");
    return static_cast<T *>(__builtin_assume_aligned(m_data, A));
  }

  template <std::size_t A> const T *alignedData() const noexcept {
    assert(isAligned<A>() && "alignedData<A>() requires A-byte aligned data");
    return static_cast<const T *>(__builtin_assume_aligned(m_data, A));
  }

  template <Layout L2 = L>
    requires(L2 == Layout::Contig)
  static NDArray borrow(T *ptr, std::array<std::size_t, N> shape) noexcept {
    return NDArray(clustering::detail::BorrowedTag{}, ptr, ptr, shape,
                   computeContiguousStrides(shape), 0, true);
  }

  template <Layout L2 = L>
    requires(L2 == Layout::Contig)
  static NDArray borrow(const T *ptr, std::array<std::size_t, N> shape) noexcept {
    auto *mutPtr = const_cast<T *>(ptr);
    return NDArray(clustering::detail::BorrowedTag{}, mutPtr, mutPtr, shape,
                   computeContiguousStrides(shape), 0, false);
  }

  template <Layout L2 = L>
    requires(L2 == Layout::MaybeStrided)
  static NDArray borrow(T *ptr, std::array<std::size_t, N> shape,
                        std::array<std::ptrdiff_t, N> strides) noexcept {
    return NDArray(clustering::detail::BorrowedTag{}, ptr, ptr, shape, strides, 0, true);
  }

  template <Layout L2 = L>
    requires(L2 == Layout::MaybeStrided)
  static NDArray borrow(const T *ptr, std::array<std::size_t, N> shape,
                        std::array<std::ptrdiff_t, N> strides) noexcept {
    auto *mutPtr = const_cast<T *>(ptr);
    return NDArray(clustering::detail::BorrowedTag{}, mutPtr, mutPtr, shape, strides, 0, false);
  }

  template <std::size_t M = N>
    requires(M == 1 && L == Layout::Contig)
  static NDArray borrow1D(T *ptr, std::size_t n) noexcept {
    return borrow(ptr, std::array<std::size_t, 1>{n});
  }

  template <std::size_t M = N>
    requires(M == 1 && L == Layout::Contig)
  static NDArray borrow1D(const T *ptr, std::size_t n) noexcept {
    return borrow(ptr, std::array<std::size_t, 1>{n});
  }

  template <Layout L2 = L>
    requires(L2 == Layout::MaybeStrided)
  static NDArray borrowBytes(T *ptr, std::array<std::size_t, N> shape,
                             std::array<std::ptrdiff_t, N> stridesInBytes,
                             bool isMutable) noexcept {
    std::array<std::ptrdiff_t, N> element_strides{};
    for (std::size_t k = 0; k < N; ++k) {
      assert(stridesInBytes[k] % static_cast<std::ptrdiff_t>(sizeof(T)) == 0 &&
             "borrowBytes requires byte strides divisible by sizeof(T)");
      element_strides[k] = stridesInBytes[k] / static_cast<std::ptrdiff_t>(sizeof(T));
    }
    return NDArray(clustering::detail::BorrowedTag{}, ptr, ptr, shape, element_strides, 0,
                   isMutable);
  }

  template <std::size_t M = N>
    requires(M == 1 && L == Layout::Contig)
  static NDArray fromSpan(std::span<T> s) noexcept {
    return borrow(s.data(), std::array<std::size_t, 1>{s.size()});
  }

  template <std::size_t M = N>
    requires(M == 1 && L == Layout::Contig)
  static NDArray fromSpan(std::span<const T> s) noexcept {
    return borrow(s.data(), std::array<std::size_t, 1>{s.size()});
  }

  template <std::size_t M = N>
    requires(M == 2)
  NDArray<T, 2, Layout::MaybeStrided> t() noexcept {
    return NDArray<T, 2, Layout::MaybeStrided>(
        clustering::detail::BorrowedTag{}, m_data, m_base,
        std::array<std::size_t, 2>{m_shape[1], m_shape[0]},
        std::array<std::ptrdiff_t, 2>{m_strides[1], m_strides[0]}, m_offset, m_mutable);
  }

  template <std::size_t M = N>
    requires(M == 2)
  NDArray<T, 2, Layout::MaybeStrided> t() const noexcept {
    return NDArray<T, 2, Layout::MaybeStrided>(
        clustering::detail::BorrowedTag{}, const_cast<T *>(m_data), const_cast<T *>(m_base),
        std::array<std::size_t, 2>{m_shape[1], m_shape[0]},
        std::array<std::ptrdiff_t, 2>{m_strides[1], m_strides[0]}, m_offset, false);
  }

  template <std::size_t M = N>
    requires(M > 1)
  NDArray<T, N - 1, L> row(std::size_t i) noexcept {
    assert(i < m_shape[0]);
    std::array<std::size_t, N - 1> new_shape{};
    std::array<std::ptrdiff_t, N - 1> new_strides{};
    for (std::size_t k = 0; k + 1 < N; ++k) {
      new_shape[k] = m_shape[k + 1];
      new_strides[k] = m_strides[k + 1];
    }
    return NDArray<T, N - 1, L>(clustering::detail::BorrowedTag{},
                                m_data + m_offset + (static_cast<std::ptrdiff_t>(i) * m_strides[0]),
                                m_base, new_shape, new_strides, 0, m_mutable);
  }

  template <std::size_t M = N>
    requires(M > 1)
  NDArray<T, N - 1, L> row(std::size_t i) const noexcept {
    assert(i < m_shape[0]);
    std::array<std::size_t, N - 1> new_shape{};
    std::array<std::ptrdiff_t, N - 1> new_strides{};
    for (std::size_t k = 0; k + 1 < N; ++k) {
      new_shape[k] = m_shape[k + 1];
      new_strides[k] = m_strides[k + 1];
    }
    return NDArray<T, N - 1, L>(clustering::detail::BorrowedTag{},
                                const_cast<T *>(m_data) + m_offset +
                                    (static_cast<std::ptrdiff_t>(i) * m_strides[0]),
                                const_cast<T *>(m_base), new_shape, new_strides, 0, false);
  }

  template <std::size_t M = N>
    requires(M == 2)
  NDArray<T, 1, Layout::MaybeStrided> col(std::size_t j) noexcept {
    assert(j < m_shape[1]);
    return NDArray<T, 1, Layout::MaybeStrided>(
        clustering::detail::BorrowedTag{},
        m_data + m_offset + (static_cast<std::ptrdiff_t>(j) * m_strides[1]), m_base,
        std::array<std::size_t, 1>{m_shape[0]}, std::array<std::ptrdiff_t, 1>{m_strides[0]}, 0,
        m_mutable);
  }

  template <std::size_t M = N>
    requires(M == 2)
  NDArray<T, 1, Layout::MaybeStrided> col(std::size_t j) const noexcept {
    assert(j < m_shape[1]);
    return NDArray<T, 1, Layout::MaybeStrided>(
        clustering::detail::BorrowedTag{},
        const_cast<T *>(m_data) + m_offset + (static_cast<std::ptrdiff_t>(j) * m_strides[1]),
        const_cast<T *>(m_base), std::array<std::size_t, 1>{m_shape[0]},
        std::array<std::ptrdiff_t, 1>{m_strides[0]}, 0, false);
  }

  NDArray<T, N, Layout::MaybeStrided> slice(std::size_t axis, std::size_t begin,
                                            std::size_t end) noexcept {
    assert(axis < N && begin <= end && end <= m_shape[axis]);
    std::array<std::size_t, N> new_shape = m_shape;
    new_shape[axis] = end - begin;
    return NDArray<T, N, Layout::MaybeStrided>(
        clustering::detail::BorrowedTag{},
        m_data + m_offset + (static_cast<std::ptrdiff_t>(begin) * m_strides[axis]), m_base,
        new_shape, m_strides, 0, m_mutable);
  }

  NDArray<T, N, Layout::MaybeStrided> slice(std::size_t axis, std::size_t begin,
                                            std::size_t end) const noexcept {
    assert(axis < N && begin <= end && end <= m_shape[axis]);
    std::array<std::size_t, N> new_shape = m_shape;
    new_shape[axis] = end - begin;
    return NDArray<T, N, Layout::MaybeStrided>(
        clustering::detail::BorrowedTag{},
        const_cast<T *>(m_data) + m_offset + (static_cast<std::ptrdiff_t>(begin) * m_strides[axis]),
        const_cast<T *>(m_base), new_shape, m_strides, 0, false);
  }

  NDArray<T, N, Layout::MaybeStrided> slice(const std::array<Range, N> &ranges) noexcept {
    std::array<std::size_t, N> new_shape{};
    std::array<std::ptrdiff_t, N> new_strides{};
    std::ptrdiff_t advance = 0;
    for (std::size_t k = 0; k < N; ++k) {
      const std::size_t end = std::min(ranges[k].end, m_shape[k]);
      const std::size_t begin = ranges[k].begin;
      const std::ptrdiff_t step = ranges[k].step;
      assert(begin <= end && step > 0);
      new_shape[k] = step == 1 ? (end - begin)
                               : (end - begin + static_cast<std::size_t>(step) - 1) /
                                     static_cast<std::size_t>(step);
      new_strides[k] = m_strides[k] * step;
      advance += static_cast<std::ptrdiff_t>(begin) * m_strides[k];
    }
    return NDArray<T, N, Layout::MaybeStrided>(clustering::detail::BorrowedTag{},
                                               m_data + m_offset + advance, m_base, new_shape,
                                               new_strides, 0, m_mutable);
  }

  NDArray<T, N, Layout::MaybeStrided> slice(const std::array<Range, N> &ranges) const noexcept {
    std::array<std::size_t, N> new_shape{};
    std::array<std::ptrdiff_t, N> new_strides{};
    std::ptrdiff_t advance = 0;
    for (std::size_t k = 0; k < N; ++k) {
      const std::size_t end = std::min(ranges[k].end, m_shape[k]);
      const std::size_t begin = ranges[k].begin;
      const std::ptrdiff_t step = ranges[k].step;
      assert(begin <= end && step > 0);
      new_shape[k] = step == 1 ? (end - begin)
                               : (end - begin + static_cast<std::size_t>(step) - 1) /
                                     static_cast<std::size_t>(step);
      new_strides[k] = m_strides[k] * step;
      advance += static_cast<std::ptrdiff_t>(begin) * m_strides[k];
    }
    return NDArray<T, N, Layout::MaybeStrided>(
        clustering::detail::BorrowedTag{}, const_cast<T *>(m_data) + m_offset + advance,
        const_cast<T *>(m_base), new_shape, new_strides, 0, false);
  }

  NDArray<T, N, Layout::MaybeStrided> permute(const std::array<std::size_t, N> &perm) noexcept {
    std::array<std::size_t, N> new_shape{};
    std::array<std::ptrdiff_t, N> new_strides{};
    for (std::size_t k = 0; k < N; ++k) {
      assert(perm[k] < N);
      new_shape[k] = m_shape[perm[k]];
      new_strides[k] = m_strides[perm[k]];
    }
    return NDArray<T, N, Layout::MaybeStrided>(clustering::detail::BorrowedTag{}, m_data, m_base,
                                               new_shape, new_strides, m_offset, m_mutable);
  }

  NDArray<T, N, Layout::MaybeStrided>
  permute(const std::array<std::size_t, N> &perm) const noexcept {
    std::array<std::size_t, N> new_shape{};
    std::array<std::ptrdiff_t, N> new_strides{};
    for (std::size_t k = 0; k < N; ++k) {
      assert(perm[k] < N);
      new_shape[k] = m_shape[perm[k]];
      new_strides[k] = m_strides[perm[k]];
    }
    return NDArray<T, N, Layout::MaybeStrided>(clustering::detail::BorrowedTag{},
                                               const_cast<T *>(m_data), const_cast<T *>(m_base),
                                               new_shape, new_strides, m_offset, false);
  }

  template <std::size_t M>
  NDArray<T, M, Layout::Contig> view(std::array<std::size_t, M> shape) noexcept {
    assert(isContiguous() && "view<M> requires a contiguous source");
    assert(productOfShape(shape) == numel() && "view<M> must preserve element count");
    return NDArray<T, M, Layout::Contig>(
        clustering::detail::BorrowedTag{}, m_data, m_base, shape,
        NDArray<T, M, Layout::Contig>::computeContiguousStrides(shape), 0, m_mutable);
  }

  template <std::size_t M>
  NDArray<T, M, Layout::Contig> view(std::array<std::size_t, M> shape) const noexcept {
    assert(isContiguous() && "view<M> requires a contiguous source");
    assert(productOfShape(shape) == numel() && "view<M> must preserve element count");
    return NDArray<T, M, Layout::Contig>(
        clustering::detail::BorrowedTag{}, const_cast<T *>(m_data), const_cast<T *>(m_base), shape,
        NDArray<T, M, Layout::Contig>::computeContiguousStrides(shape), 0, false);
  }

  template <std::size_t M> NDArray<T, M, Layout::Contig> reshape(std::array<std::size_t, M> shape) {
    assert(productOfShape(shape) == numel() && "reshape<M> must preserve element count");
    if (isContiguous()) {
      return NDArray<T, M, Layout::Contig>(
          clustering::detail::BorrowedTag{}, m_data, m_base, shape,
          NDArray<T, M, Layout::Contig>::computeContiguousStrides(shape), 0, m_mutable);
    }
    NDArray<T, M, Layout::Contig> result(shape);
    copyToContiguous(result.data());
    return result;
  }

  template <std::size_t M>
  NDArray<T, M, Layout::Contig> reshape(std::array<std::size_t, M> shape) const {
    assert(productOfShape(shape) == numel() && "reshape<M> must preserve element count");
    if (isContiguous()) {
      return NDArray<T, M, Layout::Contig>(
          clustering::detail::BorrowedTag{}, const_cast<T *>(m_data), const_cast<T *>(m_base),
          shape, NDArray<T, M, Layout::Contig>::computeContiguousStrides(shape), 0, false);
    }
    NDArray<T, M, Layout::Contig> result(shape);
    copyToContiguous(result.data());
    return result;
  }

  NDArray<T, N, Layout::Contig> contiguous() {
    if (isContiguous()) {
      return NDArray<T, N, Layout::Contig>(
          clustering::detail::BorrowedTag{}, m_data, m_base, m_shape,
          NDArray<T, N, Layout::Contig>::computeContiguousStrides(m_shape), 0, m_mutable);
    }
    NDArray<T, N, Layout::Contig> result(m_shape);
    copyToContiguous(result.data());
    return result;
  }

  NDArray<T, N, Layout::Contig> contiguous() const {
    if (isContiguous()) {
      return NDArray<T, N, Layout::Contig>(
          clustering::detail::BorrowedTag{}, const_cast<T *>(m_data), const_cast<T *>(m_base),
          m_shape, NDArray<T, N, Layout::Contig>::computeContiguousStrides(m_shape), 0, false);
    }
    NDArray<T, N, Layout::Contig> result(m_shape);
    copyToContiguous(result.data());
    return result;
  }

  NDArray<T, N, Layout::Contig> clone() const {
    NDArray<T, N, Layout::Contig> result(m_shape);
    copyToContiguous(result.data());
    return result;
  }

  std::string debugDump() const {
    std::stringstream ss;
    ss << "NDarray<" << typeid(T).name() << ", " << N << ">(";
    for (auto d : m_shape) {
      ss << d << ", ";
    }
    ss << ")\n";
    ss << "data: [";
    const std::size_t total = numel();
    if (total > 0) {
      std::array<std::size_t, N> idx{};
      for (std::size_t flat = 0; flat < total; ++flat) {
        std::ptrdiff_t off = m_offset;
        for (std::size_t k = 0; k < N; ++k) {
          off += static_cast<std::ptrdiff_t>(idx[k]) * m_strides[k];
        }
        ss << m_data[off] << ", ";
        for (std::size_t k = N; k-- > 0;) {
          if (++idx[k] < m_shape[k]) {
            break;
          }
          idx[k] = 0;
        }
      }
    }
    ss << "]\n";
    ss << "size: " << total << "\n";
    return ss.str();
  }
 
  // Equality between NDArrays has three plausible semantics (element-wise, storage-identity,
  // deep-value); none is obviously correct, so the operator is deleted to force callers to pick
  // an explicit intent (@c math::arrayEqual, @c sameStorage, or an explicit shape-and-element
  // comparison).
  friend bool operator==(const NDArray &, const NDArray &) = delete;
  friend bool operator!=(const NDArray &, const NDArray &) = delete;
 
private:
  static std::array<std::ptrdiff_t, N>
  computeContiguousStrides(const std::array<std::size_t, N> &shape) {
    std::array<std::ptrdiff_t, N> s{};
    s[N - 1] = 1;
    for (std::size_t k = N - 1; k > 0; --k) {
      s[k - 1] = s[k] * static_cast<std::ptrdiff_t>(shape[k]);
    }
    return s;
  }
 
  template <std::size_t M>
  static std::size_t productOfShape(const std::array<std::size_t, M> &shape) noexcept {
    std::size_t size = 1;
    for (std::size_t k = 0; k < M; ++k) {
      size *= shape[k];
    }
    return size;
  }
 
  std::size_t numel() const noexcept { return productOfShape(m_shape); }
 
  // Walk this array in row-major order and write elements densely to @p dst. Fast-path
  // @c memcpy when already contiguous, else advance a multi-index cursor and index through
  // @c m_offset + sum_k idx[k] * m_strides[k]. Shared by @c reshape, @c contiguous, @c clone.
  void copyToContiguous(T *dst) const noexcept {
    const std::size_t total = numel();
    if (total == 0) {
      return;
    }
    if (isContiguous()) {
      std::memcpy(dst, m_data + m_offset, total * sizeof(T));
      return;
    }
    std::array<std::size_t, N> idx{};
    for (std::size_t flat = 0; flat < total; ++flat) {
      std::ptrdiff_t off = m_offset;
      for (std::size_t k = 0; k < N; ++k) {
        off += static_cast<std::ptrdiff_t>(idx[k]) * m_strides[k];
      }
      dst[flat] = m_data[off];
      for (std::size_t k = N; k-- > 0;) {
        if (++idx[k] < m_shape[k]) {
          break;
        }
        idx[k] = 0;
      }
    }
  }
 
  template <std::size_t... Ks, class... Ix>
  std::size_t computeElementOffset(std::index_sequence<Ks...>, Ix... ix) const noexcept {
    std::ptrdiff_t off = m_offset;
    ((off += static_cast<std::ptrdiff_t>(ix) * m_strides[Ks]), ...);
    return static_cast<std::size_t>(off);
  }
 
  T *m_data;
  T *m_base;
  std::vector<T, clustering::detail::AlignedAllocator<T, 64>> m_vec;
  std::array<std::size_t, N> m_shape;
  std::array<std::ptrdiff_t, N> m_strides;
  std::ptrdiff_t m_offset;
  NDArrayStorage m_storage;
  bool m_mutable;
};

template <class T, std::size_t NA, Layout LA, std::size_t NB, Layout LB>
bool sameStorage(const NDArray<T, NA, LA> &a, const NDArray<T, NB, LB> &b) noexcept {
  return a.baseData() == b.baseData();
}
 
} // namespace clustering