From 761b07dffb9c31864c4eeab35565a7d324d449d3 Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Tue, 12 May 2026 11:12:36 -0400 Subject: [PATCH 1/6] add btas::zb::RangeNd: zero-based, packed range for tiny inner tiles Tensor-of-Tensor applications use inner tiles whose Range is structurally degenerate (lobound is always 0, upbound equals extent, strides are derivable, offset is 0, always contiguous) and small (10-100 elements, rank <= 6). The general-purpose RangeNd pays ~304 B for capabilities the inner case never uses, dominating per-tile memory. btas::zb::RangeNd stores only the extents (std::array) plus a 1 B rank, fitting in 14 B (align 2). Strides are synthesized on demand via ordinal_view returned by value; lobound() materializes zeros from a static buffer. Models TWG.BoxRange via range_traits, is_index, and boxrange_iteration_order specializations, and provides MADNESS archive load/store that serializes rank + extents only (ordinal is derivable). Unit tests in unittest/zb_range_test.cc cover sizeof/alignof contracts, concept membership, construction, ordinal mapping, iteration, equality, swap, and non-default template parameters (6 cases, 60 assertions). --- btas/zb/range.h | 402 ++++++++++++++++++++++++++++++++++++++ unittest/CMakeLists.txt | 1 + unittest/zb_range_test.cc | 161 +++++++++++++++ 3 files changed, 564 insertions(+) create mode 100644 btas/zb/range.h create mode 100644 unittest/zb_range_test.cc diff --git a/btas/zb/range.h b/btas/zb/range.h new file mode 100644 index 00000000..25e7c17c --- /dev/null +++ b/btas/zb/range.h @@ -0,0 +1,402 @@ +/* + * zb/range.h + * + * Slimmed-down version of RangeNd for zero-based indexing apps. State is just (extent[], rank); lobound is + * structurally zero, upbound = extent, strides are derived row-major on demand. + * + * sizeof(zb::RangeNd<6, int16_t, int32_t>) == 14 (vs ~304 for btas::Range). + */ + +#ifndef BTAS_ZB_RANGE_H_ +#define BTAS_ZB_RANGE_H_ + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace btas { +namespace zb { + +/// Packed zero-based index/extent vector. Stores at most \c MaxRank entries of +/// type \c Int plus a 1-byte size; size <= MaxRank invariant is enforced. +/// +/// Models the TWG.Index concept (\c btas::is_index) — exposes nested +/// \c value_type plus member \c begin / \c end / \c operator[] / \c size . +template +class index { + public: + using value_type = Int; + using size_type = std::size_t; + using reference = Int&; + using const_reference = const Int&; + using pointer = Int*; + using const_pointer = const Int*; + using iterator = Int*; + using const_iterator = const Int*; + + static constexpr size_type max_size_v = MaxRank; + + constexpr index() noexcept : data_{}, size_(0) {} + + /// fixed-size construction with default-initialized elements + constexpr explicit index(size_type n) : data_{}, size_(check_size(n)) {} + + /// fixed-size construction filled with \p v + constexpr index(size_type n, Int v) : data_{}, size_(check_size(n)) { + for (size_type i = 0; i < size_; ++i) data_[i] = v; + } + + template + constexpr index(std::initializer_list il) : data_{}, size_(check_size(il.size())) { + size_type i = 0; + for (auto v : il) data_[i++] = static_cast(v); + } + + /// from any iterable container (size must be <= MaxRank) + template >::value && + !std::is_same_v, index>>> + constexpr index(const C& c) : data_{}, size_(0) { + using std::begin; + using std::end; + auto first = begin(c); + auto last = end(c); + size_type n = 0; + for (auto it = first; it != last; ++it) ++n; + size_ = check_size(n); + size_type i = 0; + for (auto it = first; it != last; ++it) data_[i++] = static_cast(*it); + } + + constexpr size_type size() const noexcept { return size_; } + constexpr bool empty() const noexcept { return size_ == 0; } + static constexpr size_type max_size() noexcept { return MaxRank; } + + constexpr reference operator[](size_type i) noexcept { return data_[i]; } + constexpr const_reference operator[](size_type i) const noexcept { return data_[i]; } + + constexpr reference at(size_type i) { assert(i < size_); return data_[i]; } + constexpr const_reference at(size_type i) const { assert(i < size_); return data_[i]; } + + constexpr iterator begin() noexcept { return data_.data(); } + constexpr iterator end() noexcept { return data_.data() + size_; } + constexpr const_iterator begin() const noexcept { return data_.data(); } + constexpr const_iterator end() const noexcept { return data_.data() + size_; } + constexpr const_iterator cbegin() const noexcept { return data_.data(); } + constexpr const_iterator cend() const noexcept { return data_.data() + size_; } + + constexpr pointer data() noexcept { return data_.data(); } + constexpr const_pointer data() const noexcept { return data_.data(); } + + /// resize to \p n (elements past the old size are value-initialized) + constexpr void resize(size_type n) { + auto new_size = check_size(n); + for (size_type i = size_; i < new_size; ++i) data_[i] = Int{}; + size_ = new_size; + } + + friend constexpr bool operator==(const index& a, const index& b) noexcept { + if (a.size_ != b.size_) return false; + for (std::uint8_t i = 0; i < a.size_; ++i) + if (a.data_[i] != b.data_[i]) return false; + return true; + } + friend constexpr bool operator!=(const index& a, const index& b) noexcept { + return !(a == b); + } + + private: + static constexpr std::uint8_t check_size(size_type n) { + assert(n <= MaxRank); + return static_cast(n); + } + + std::array data_; + std::uint8_t size_; +}; + +/// Lightweight value type returned by \c RangeNd::ordinal() . Synthesizes +/// strides from extent at construction; offset is always 0 and the range is +/// always contiguous, by construction. +template +class ordinal_view { + public: + using value_type = Ord; + using stride_type = std::array; + + ordinal_view() noexcept : stride_{}, rank_(0) {} + + template + ordinal_view(const Extents& ext, std::size_t rank) : stride_{}, rank_(rank) { + using std::cbegin; + auto it = cbegin(ext); + Ord vol{1}; + for (std::ptrdiff_t i = static_cast(rank) - 1; i >= 0; --i) { + stride_[i] = vol; + vol *= static_cast(*(it + i)); + } + } + + std::size_t rank() const noexcept { return rank_; } + const stride_type& stride() const noexcept { return stride_; } + const Ord* stride_data() const noexcept { return stride_.data(); } + constexpr Ord offset() const noexcept { return Ord{0}; } + constexpr bool contiguous() const noexcept { return true; } + + template + std::enable_if_t::value, Ord> + operator()(const Index& idx) const { + assert(static_cast(idx.size()) == rank_); + using std::cbegin; + Ord o{0}; + auto it = cbegin(idx); + for (std::size_t i = 0; i < rank_; ++i) + o += static_cast(*(it + i)) * stride_[i]; + return o; + } + + private: + stride_type stride_; + std::size_t rank_; +}; + +/// Zero-based row-major N-dim range optimized for applications with zero-based indexing. +/// +/// \tparam MaxRank static cap on rank (default 6) +/// \tparam Ext per-dim extent integer type (default int16_t) +/// \tparam Ord ordinal integer type (default int32_t) +template +class RangeNd { + public: + static_assert(MaxRank > 0 && MaxRank < 256, "MaxRank must lie in (0, 256)"); + static_assert(std::is_integral_v, "Ext must be an integer type"); + static_assert(std::is_integral_v && std::is_signed_v, + "Ord must be a signed integer type"); + + static constexpr ::blas::Layout order = ::blas::Layout::RowMajor; + static constexpr std::size_t max_rank = MaxRank; + + using extent_type = index; + using index_type = extent_type; + using index1_type = Ext; + using index_element_type = Ext; + using extent_element_type = Ext; + using ordinal_type = Ord; + using size_type = std::size_t; + + using value_type = index_type; + using reference = index_type&; + using const_reference = const index_type&; + + using iterator = btas::RangeIterator; + using const_iterator = iterator; + friend class btas::RangeIterator; + + /// Default constructor: rank-0 range. + constexpr RangeNd() noexcept = default; + + /// Construct from an extent container (rank inferred from size). + template >::value && + !std::is_same_v, RangeNd>>> + RangeNd(const C& ext) : extent_(ext) {} + + /// Construct from an initializer list of extents. + template >> + RangeNd(std::initializer_list il) : extent_(il) {} + + /// Construct from a pack of integer extents (>=2 to avoid clashing with + /// the container-taking constructor; pass a one-element \c {e0} for rank-1). + template && std::is_integral_v && + (std::is_integral_v && ...)>> + RangeNd(I0 e0, I1 e1, Is... es) + : extent_({static_cast(e0), static_cast(e1), + static_cast(es)...}) {} + + // + // Rank, extent, lo/up bounds + // + + std::size_t rank() const noexcept { return extent_.size(); } + + /// Volume; matches \c btas::BaseRangeNd::area() which returns 0 for rank 0. + size_type area() const noexcept { + if (extent_.size() == 0) return 0; + size_type v = 1; + for (std::size_t i = 0; i < extent_.size(); ++i) + v *= static_cast(extent_[i]); + return v; + } + size_type volume() const noexcept { return area(); } + + const extent_type& extent() const noexcept { return extent_; } + Ext extent(std::size_t i) const noexcept { return extent_[i]; } + const Ext* extent_data() const noexcept { return extent_.data(); } + + /// Lower bound: always zeros. Returns by value (small, fixed-size object). + index_type lobound() const { return index_type(rank(), Ext{0}); } + Ext lobound(std::size_t) const noexcept { return Ext{0}; } + const Ext* lobound_data() const noexcept { return zero_buffer(); } + + /// Upper bound: equal to extent (zero-based range). + const extent_type& upbound() const noexcept { return extent_; } + Ext upbound(std::size_t i) const noexcept { return extent_[i]; } + const Ext* upbound_data() const noexcept { return extent_.data(); } + + // + // Ordinal mapping (synthesized row-major; nothing stored) + // + + ordinal_view ordinal() const { + return ordinal_view(extent_, rank()); + } + + template + std::enable_if_t::value, Ord> ordinal(const I& idx) const { + assert(static_cast(idx.size()) == rank()); + using std::cbegin; + auto it = cbegin(idx); + const auto r = rank(); + Ord o{0}; + Ord vol{1}; + for (std::ptrdiff_t i = static_cast(r) - 1; i >= 0; --i) { + o += static_cast(*(it + i)) * vol; + vol *= static_cast(extent_[i]); + } + return o; + } + + // + // Iteration + // + + const_iterator begin() const { return const_iterator(lobound(), this); } + const_iterator end() const { return const_iterator(extent_, this); } + const_iterator cbegin() const { return begin(); } + const_iterator cend() const { return end(); } + + /// Row-major in-place increment used by \c RangeIterator . After the final + /// valid index, idx == upbound() (== extent_), which matches \c end() . + void increment(index_type& idx) const { + const auto r = rank(); + if (r == 0) return; + for (std::ptrdiff_t d = static_cast(r) - 1; d >= 0; --d) { + ++idx[d]; + if (idx[d] < extent_[d]) return; + idx[d] = Ext{0}; + } + for (std::size_t d = 0; d < r; ++d) idx[d] = extent_[d]; + } + + // + // Equality, swap + // + + friend bool operator==(const RangeNd& a, const RangeNd& b) noexcept { + return a.extent_ == b.extent_; + } + friend bool operator!=(const RangeNd& a, const RangeNd& b) noexcept { + return !(a == b); + } + + void swap(RangeNd& other) noexcept { + using std::swap; + swap(extent_, other.extent_); + } + + private: + /// Static MaxRank-sized zero buffer; \c lobound_data() returns a pointer + /// into it. Callers iterate only the first \c rank() bytes. + static const Ext* zero_buffer() noexcept { + static const std::array z{}; + return z.data(); + } + + extent_type extent_{}; +}; + +template +inline void swap(RangeNd& a, + RangeNd& b) noexcept { + a.swap(b); +} + +} // namespace zb + +// +// Trait specializations placing zb::RangeNd into the BTAS Range concept. +// + +template +struct range_traits> { + static constexpr ::blas::Layout order = ::blas::Layout::RowMajor; + using index_type = typename zb::RangeNd::index_type; + using ordinal_type = Ord; + static constexpr bool is_general_layout = false; +}; + +template +class boxrange_iteration_order> { + public: + enum { + row_major = boxrange_iteration_order::row_major, + other = boxrange_iteration_order::other, + column_major = boxrange_iteration_order::column_major + }; + static constexpr int value = row_major; +}; + +} // namespace btas + +// +// MADNESS archive load/store specializations. Includes only the rank + extents; +// the ordinal is fully derivable from extent. Caller must have included +// before instantiating these. +// +namespace madness { +namespace archive { + +template +struct ArchiveLoadImpl> { + static inline void load(const Archive& ar, + btas::zb::RangeNd& r) { + std::uint8_t rank{}; + ar& rank; + typename btas::zb::RangeNd::extent_type ext( + static_cast(rank)); + for (std::uint8_t i = 0; i < rank; ++i) ar& ext[i]; + r = btas::zb::RangeNd(ext); + } +}; + +template +struct ArchiveStoreImpl> { + static inline void store(const Archive& ar, + const btas::zb::RangeNd& r) { + const std::uint8_t rank = static_cast(r.rank()); + ar& rank; + for (std::uint8_t i = 0; i < rank; ++i) ar& r.extent(i); + } +}; + +} // namespace archive +} // namespace madness + +#endif // BTAS_ZB_RANGE_H_ diff --git a/unittest/CMakeLists.txt b/unittest/CMakeLists.txt index a3241119..2f9753a9 100644 --- a/unittest/CMakeLists.txt +++ b/unittest/CMakeLists.txt @@ -10,6 +10,7 @@ set(btas_test_src_files tensor_lapack_test.cc tensor_test.cc tensorview_test.cc + zb_range_test.cc ztensor_cp_test.cc test.cc ) diff --git a/unittest/zb_range_test.cc b/unittest/zb_range_test.cc new file mode 100644 index 00000000..346db03d --- /dev/null +++ b/unittest/zb_range_test.cc @@ -0,0 +1,161 @@ +#include "test.h" + +#include +#include +#include + +#include "btas/zb/range.h" + +using btas::zb::RangeNd; +// NB: do NOT `using btas::zb::index` at file scope — macOS POSIX +// declares ::index(const char*, int) and the names collide. + +// Sizeof contract: this is the whole point of zb::RangeNd. Default +// MaxRank=6 / Ext=int16_t / Ord=int32_t must fit in 14 bytes (12 B extents + +// 1 B size, alignment 2). +static_assert(sizeof(RangeNd<>) == 14, "zb::RangeNd<> must be 14 bytes"); +static_assert(alignof(RangeNd<>) == 2, "zb::RangeNd<> must align to 2"); + +// Index concept membership: TWG.Index is what btas::Tensor and the +// expression layer check via SFINAE. +static_assert(btas::is_index::index_type>::value, + "zb::index must model btas::is_index"); +static_assert(btas::is_boxrange>::value, + "zb::RangeNd must model btas::is_boxrange"); +static_assert(btas::boxrange_iteration_order>::value == + btas::boxrange_iteration_order::row_major, + "zb::RangeNd must be row-major"); + +TEST_CASE("zb::index basics") { + using Idx = btas::zb::index<6, std::int16_t>; + + SECTION("default") { + Idx a; + CHECK(a.size() == 0); + CHECK(a.empty()); + CHECK(a == Idx{}); + } + + SECTION("initializer list") { + Idx a{2, 3, 4}; + CHECK(a.size() == 3); + CHECK(a[0] == 2); + CHECK(a[1] == 3); + CHECK(a[2] == 4); + } + + SECTION("from container") { + std::vector v{5, 6}; + Idx a(v); + CHECK(a.size() == 2); + CHECK(a[0] == 5); + CHECK(a[1] == 6); + } + + SECTION("equality") { + CHECK(Idx{1, 2} == Idx{1, 2}); + CHECK(Idx{1, 2} != Idx{1, 2, 3}); + CHECK(Idx{1, 2} != Idx{2, 1}); + } +} + +TEST_CASE("zb::RangeNd construction and accessors") { + SECTION("default is rank-0, area()==0") { + RangeNd<> r; + CHECK(r.rank() == 0); + CHECK(r.area() == 0); // matches btas::BaseRangeNd::area() convention + } + + SECTION("from variadic extents") { + RangeNd<> r(2, 3, 4); + CHECK(r.rank() == 3); + CHECK(r.area() == 24); + CHECK(r.extent(0) == 2); + CHECK(r.extent(1) == 3); + CHECK(r.extent(2) == 4); + } + + SECTION("from initializer list") { + RangeNd<> r{5, 7}; + CHECK(r.rank() == 2); + CHECK(r.area() == 35); + } + + SECTION("lobound is zeros, upbound is extent") { + RangeNd<> r(2, 3, 4); + auto lo = r.lobound(); + CHECK(lo.size() == 3); + CHECK(lo[0] == 0); + CHECK(lo[1] == 0); + CHECK(lo[2] == 0); + CHECK(r.upbound() == r.extent()); + CHECK(r.upbound_data() == r.extent_data()); + } + + SECTION("lobound_data points to MaxRank zeros") { + RangeNd<> r(2, 3); + auto* p = r.lobound_data(); + for (std::size_t i = 0; i < RangeNd<>::max_rank; ++i) CHECK(p[i] == 0); + } +} + +TEST_CASE("zb::RangeNd ordinal mapping is row-major contiguous") { + RangeNd<> r(2, 3, 4); + + SECTION("formula") { + // row-major: ord(i,j,k) = i*(3*4) + j*4 + k + using idx_t = RangeNd<>::index_type; + CHECK(r.ordinal(idx_t{0, 0, 0}) == 0); + CHECK(r.ordinal(idx_t{0, 0, 1}) == 1); + CHECK(r.ordinal(idx_t{0, 1, 0}) == 4); + CHECK(r.ordinal(idx_t{1, 0, 0}) == 12); + CHECK(r.ordinal(idx_t{1, 2, 3}) == 23); + } + + SECTION("ordinal_view exposes strides and is contiguous") { + auto ov = r.ordinal(); + CHECK(ov.rank() == 3); + CHECK(ov.contiguous()); + CHECK(ov.offset() == 0); + CHECK(ov.stride()[0] == 12); + CHECK(ov.stride()[1] == 4); + CHECK(ov.stride()[2] == 1); + using idx_t = RangeNd<>::index_type; + CHECK(ov(idx_t{1, 2, 3}) == 23); + } +} + +TEST_CASE("zb::RangeNd iteration covers volume in row-major order") { + RangeNd<> r(2, 3); + std::size_t count = 0; + std::int32_t expected_ordinal = 0; + for (auto it = r.begin(); it != r.end(); ++it, ++count, ++expected_ordinal) { + CHECK(r.ordinal(*it) == expected_ordinal); + } + CHECK(count == r.area()); +} + +TEST_CASE("zb::RangeNd equality and swap") { + RangeNd<> a(2, 3, 4); + RangeNd<> b(2, 3, 4); + RangeNd<> c(2, 3, 5); + + CHECK(a == b); + CHECK(a != c); + + using std::swap; + swap(a, c); + CHECK(a == RangeNd<>(2, 3, 5)); + CHECK(c == RangeNd<>(2, 3, 4)); +} + +TEST_CASE("zb::RangeNd with non-default template parameters") { + using R4 = RangeNd<4, std::int32_t, std::int64_t>; + static_assert(sizeof(R4) == 20, "expected packed size for MaxRank=4, int32"); + static_assert(R4::max_rank == 4); + R4 r(10, 20, 30); + CHECK(r.rank() == 3); + CHECK(r.area() == 6000); + CHECK(r.ordinal(typename R4::index_type{1, 2, 3}) == + 1 * 20 * 30 + 2 * 30 + 3); +} From 8fd73bfa3528fbbd76cf5f402a8953248361307d Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Tue, 12 May 2026 12:28:21 -0400 Subject: [PATCH 2/6] zb::RangeNd: add (lobound, upbound) ctor; use BTAS_ASSERT Two follow-ups discovered while wiring btas::Tensor with zb::RangeNd as the inner tile in TiledArray ToT tests: * btas::Tensor's (range, storage) ctor instantiates range_type(range.lobound(), range.upbound()) even when the runtime branch wouldn't take it, so zb::RangeNd needs that constructor to satisfy template instantiation. The ctor asserts that lobound is all zeros (zero-based ranges) and takes upbound as the extent. * Replace plain assert() with BTAS_ASSERT per BTAS convention, and include for it. All 60 zb::RangeNd unit-test assertions still pass. --- btas/zb/range.h | 29 +++++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/btas/zb/range.h b/btas/zb/range.h index 25e7c17c..906f71a2 100644 --- a/btas/zb/range.h +++ b/btas/zb/range.h @@ -12,6 +12,7 @@ #include +#include #include #include #include @@ -19,7 +20,6 @@ #include #include -#include #include #include #include @@ -87,8 +87,8 @@ class index { constexpr reference operator[](size_type i) noexcept { return data_[i]; } constexpr const_reference operator[](size_type i) const noexcept { return data_[i]; } - constexpr reference at(size_type i) { assert(i < size_); return data_[i]; } - constexpr const_reference at(size_type i) const { assert(i < size_); return data_[i]; } + constexpr reference at(size_type i) { BTAS_ASSERT(i < size_); return data_[i]; } + constexpr const_reference at(size_type i) const { BTAS_ASSERT(i < size_); return data_[i]; } constexpr iterator begin() noexcept { return data_.data(); } constexpr iterator end() noexcept { return data_.data() + size_; } @@ -119,7 +119,7 @@ class index { private: static constexpr std::uint8_t check_size(size_type n) { - assert(n <= MaxRank); + BTAS_ASSERT(n <= MaxRank); return static_cast(n); } @@ -158,7 +158,7 @@ class ordinal_view { template std::enable_if_t::value, Ord> operator()(const Index& idx) const { - assert(static_cast(idx.size()) == rank_); + BTAS_ASSERT(static_cast(idx.size()) == rank_); using std::cbegin; Ord o{0}; auto it = cbegin(idx); @@ -230,6 +230,23 @@ class RangeNd { : extent_({static_cast(e0), static_cast(e1), static_cast(es)...}) {} + /// Construct from lobound/upbound pair. \c lobound must be all zeros + /// (zero-based ranges); \c upbound becomes the extent. Useful because + /// downstream code such as @c btas::Tensor 's @c (range, storage) ctor + /// instantiates @c range_type(lobound, upbound) even when the runtime + /// branch would not take that path. + template >::value && + is_index>::value>> + RangeNd(const Lo& lobound, const Up& upbound) : extent_(upbound) { + (void)lobound; + using std::cbegin; + using std::cend; + BTAS_ASSERT(std::all_of(cbegin(lobound), cend(lobound), + [](auto v) { return v == 0; }) && + "btas::zb::RangeNd: lobound must be all zeros"); + } + // // Rank, extent, lo/up bounds // @@ -270,7 +287,7 @@ class RangeNd { template std::enable_if_t::value, Ord> ordinal(const I& idx) const { - assert(static_cast(idx.size()) == rank()); + BTAS_ASSERT(static_cast(idx.size()) == rank()); using std::cbegin; auto it = cbegin(idx); const auto r = rank(); From 7e64fbad97c76f316f313f4c8ed3fca5445da15f Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Tue, 12 May 2026 16:25:36 -0400 Subject: [PATCH 3/6] add Tensor (Range, generator) ctor Mirrors TA::Tensor's range+lambda ctor: builds each element from `gen` called on the element's multi-index. Useful when a btas::Tensor is the inner tile in a TiledArray ToT and generic code (e.g. MPQC's jacobi_update) wants tile-type-agnostic per-index construction. SFINAE keeps the new ctor distinct from existing (Range, value), (Range, iterator), and (Range, Storage) overloads via is_invocable_r_v. --- btas/tensor.h | 19 +++++++++++++++++++ unittest/tensor_test.cc | 19 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/btas/tensor.h b/btas/tensor.h index 3195f2fe..5e089077 100644 --- a/btas/tensor.h +++ b/btas/tensor.h @@ -144,6 +144,25 @@ namespace btas { } } + /// construct from \c range object, fill each element from \c gen called on + /// the element's multi-index. \c gen must be callable with the range's + /// iteration value (its multi-index) and return a value convertible to + /// \c value_type. + template ::value && + std::is_invocable_r_v< + value_type, F, + decltype(*std::begin(std::declval()))>>> + Tensor(const Range& range, F&& gen) + : range_(range.lobound(), range.upbound()) { + array_adaptor::resize(storage_, range_.area()); + auto out_it = begin(); + for (auto&& idx : range_) { + *out_it++ = gen(idx); + } + } + /// construct from \c range and \c storage template Tensor(const Range& range, const Storage& storage, diff --git a/unittest/tensor_test.cc b/unittest/tensor_test.cc index d6b6e46a..aac58fb3 100644 --- a/unittest/tensor_test.cc +++ b/unittest/tensor_test.cc @@ -134,6 +134,25 @@ TEST_CASE("Tensor Constructors") { // range + vector of values CHECK_NOTHROW(DTensor(r1, T1.data())); + + // range + generator lambda: gen(multi-index) → value + Range r2(3, 4); + DTensor T2(r2, [](auto const& idx) -> double { + return 10.0 * idx[0] + idx[1]; + }); + CHECK(T2.rank() == 2); + CHECK(T2.extent(0) == 3); + CHECK(T2.extent(1) == 4); + CHECK(T2.size() == 3 * 4); + for (auto i = 0u; i != 3; ++i) + for (auto j = 0u; j != 4; ++j) + CHECK(T2(i, j) == Approx(10.0 * i + j)); + + // generator returning an int (convertible to double) should also work + Tensor Ti2(r2, [](auto const& idx) { return idx[0] + idx[1]; }); + CHECK(Ti2.extent(0) == 3); + CHECK(Ti2.extent(1) == 4); + CHECK(Ti2(2, 3) == 5); } SECTION("Fixed Rank Tensor") { From 20842a7df95fa912e2f9e94a114466735c5aaa12 Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Tue, 12 May 2026 22:31:44 -0400 Subject: [PATCH 4/6] zb::RangeNd: expose stride(); permute uses if constexpr MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit zb::RangeNd intentionally stores no strides — its 14 B packed representation is the reason it exists, and contiguous row-major strides are derivable from extent alone. But BTAS's generic permute(X, p, Y), TA's tile ops, and other consumers expect r.stride() to be callable. Add a stride() member that synthesizes row-major strides on demand and returns them by value — nothing is stored, the packed footprint is preserved. Callers that need a pointer (e.g. btas::permute forwarding into btas::Range(lo, up, stride)) bind the temporary as a const& and copy. Also turn the if/else in btas::permute(X, p, Y) into if constexpr so the strided-fallback branch is discarded for ranges whose is_general_layout is true — keeps permute compilable against ranges that genuinely lack stride storage (none today, but a defensible cleanup). --- btas/generic/permute.h | 2 +- btas/zb/range.h | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/btas/generic/permute.h b/btas/generic/permute.h index f0afbc03..45bd8278 100644 --- a/btas/generic/permute.h +++ b/btas/generic/permute.h @@ -36,7 +36,7 @@ namespace btas { ++itrY; } }; - if (r_is_permutable) + if constexpr (r_is_permutable) do_perm(X, Y, permute(r, p)); else { do_perm(X, Y, permute(btas::Range(r.lobound(), r.upbound(), r.stride()), p)); diff --git a/btas/zb/range.h b/btas/zb/range.h index 906f71a2..a9b11fff 100644 --- a/btas/zb/range.h +++ b/btas/zb/range.h @@ -285,6 +285,16 @@ class RangeNd { return ordinal_view(extent_, rank()); } + /// Row-major strides synthesized on demand. Returned by value (not by + /// reference) so nothing is stored in the range itself — preserves the + /// packed footprint. Callers that need a pointer (e.g. the BTAS generic + /// permute, which forwards r.stride() into a btas::Range ctor) bind the + /// temporary to a const& and copy from it. + using stride_type = typename ordinal_view::stride_type; + stride_type stride() const noexcept { + return ordinal().stride(); + } + template std::enable_if_t::value, Ord> ordinal(const I& idx) const { BTAS_ASSERT(static_cast(idx.size()) == rank()); From 245e49f117981d6124e0f1aa0d1ae72f1c16318b Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Tue, 12 May 2026 23:04:42 -0400 Subject: [PATCH 5/6] zb::RangeNd: reorder template params; add column-major support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bring zb::RangeNd's template signature in line with btas::RangeNd: template <::blas::Layout _Order, typename Ext, typename Ord, std::size_t MaxRank> (layout first, MaxRank last so the common zb::RangeNd<> form stays a single token). The ordinal_view, ordinal(idx), and increment() paths became layout-aware via if constexpr — row-major iterates last dim fastest, column-major first dim fastest. range_traits and boxrange_iteration_order specializations now reflect the runtime layout. Also address review feedback / harden the zero-based path: - Add missing and headers in zb/range.h, and in tensor.h (no longer rely on transitive includes). - static_assert(MaxRank < 256) on btas::zb::index, mirroring the same check on RangeNd, since the size_ field is uint8_t. - BTAS_ASSERT rank <= MaxRank in the MADNESS archive load to fail deterministically on malformed archives instead of silently truncating into the uint8_t size_. - Add a column-major test case to zb_range_test.cc and update the non-default-params test to use the new template ordering. --- btas/tensor.h | 1 + btas/zb/range.h | 146 ++++++++++++++++++++++++++------------ unittest/zb_range_test.cc | 25 ++++++- 3 files changed, 125 insertions(+), 47 deletions(-) diff --git a/btas/tensor.h b/btas/tensor.h index 5e089077..357906cd 100644 --- a/btas/tensor.h +++ b/btas/tensor.h @@ -14,6 +14,7 @@ #include #include #include +#include #include namespace btas { diff --git a/btas/zb/range.h b/btas/zb/range.h index a9b11fff..1610b1cc 100644 --- a/btas/zb/range.h +++ b/btas/zb/range.h @@ -2,9 +2,9 @@ * zb/range.h * * Slimmed-down version of RangeNd for zero-based indexing apps. State is just (extent[], rank); lobound is - * structurally zero, upbound = extent, strides are derived row-major on demand. + * structurally zero, upbound = extent, strides are derived from extent (per the requested layout) on demand. * - * sizeof(zb::RangeNd<6, int16_t, int32_t>) == 14 (vs ~304 for btas::Range). + * sizeof(zb::RangeNd<>) == 14 (vs ~304 for btas::Range). */ #ifndef BTAS_ZB_RANGE_H_ @@ -19,12 +19,14 @@ #include #include +#include #include #include #include #include #include #include +#include namespace btas { namespace zb { @@ -37,6 +39,10 @@ namespace zb { template class index { public: + // size_ is stored as a uint8_t so MaxRank must fit + static_assert(MaxRank > 0 && MaxRank < 256, + "btas::zb::index: MaxRank must lie in (0, 256)"); + using value_type = Int; using size_type = std::size_t; using reference = Int&; @@ -128,13 +134,14 @@ class index { }; /// Lightweight value type returned by \c RangeNd::ordinal() . Synthesizes -/// strides from extent at construction; offset is always 0 and the range is -/// always contiguous, by construction. -template +/// strides from extent at construction (row- or column-major per \c _Order ); +/// offset is always 0 and the range is always contiguous, by construction. +template <::blas::Layout _Order, std::size_t MaxRank, typename Ord> class ordinal_view { public: using value_type = Ord; using stride_type = std::array; + static constexpr ::blas::Layout order = _Order; ordinal_view() noexcept : stride_{}, rank_(0) {} @@ -143,9 +150,19 @@ class ordinal_view { using std::cbegin; auto it = cbegin(ext); Ord vol{1}; - for (std::ptrdiff_t i = static_cast(rank) - 1; i >= 0; --i) { - stride_[i] = vol; - vol *= static_cast(*(it + i)); + if constexpr (_Order == ::blas::Layout::RowMajor) { + // last dim is fastest: stride[N-1]=1, stride[i]=stride[i+1]*extent[i+1] + for (std::ptrdiff_t i = static_cast(rank) - 1; i >= 0; + --i) { + stride_[i] = vol; + vol *= static_cast(*(it + i)); + } + } else { + // first dim is fastest: stride[0]=1, stride[i]=stride[i-1]*extent[i-1] + for (std::size_t i = 0; i < rank; ++i) { + stride_[i] = vol; + vol *= static_cast(*(it + i)); + } } } @@ -172,14 +189,20 @@ class ordinal_view { std::size_t rank_; }; -/// Zero-based row-major N-dim range optimized for applications with zero-based indexing. +/// Zero-based N-dim range optimized for applications with zero-based indexing. /// -/// \tparam MaxRank static cap on rank (default 6) +/// Template parameter order matches \c btas::RangeNd (layout first, index/ +/// ordinal types next), with the static rank cap moved to the end so the +/// common default form \c zb::RangeNd<> stays a single token. +/// +/// \tparam _Order data layout (default RowMajor) /// \tparam Ext per-dim extent integer type (default int16_t) /// \tparam Ord ordinal integer type (default int32_t) -template + typename Ord = std::int32_t, + std::size_t MaxRank = 6> class RangeNd { public: static_assert(MaxRank > 0 && MaxRank < 256, "MaxRank must lie in (0, 256)"); @@ -187,7 +210,7 @@ class RangeNd { static_assert(std::is_integral_v && std::is_signed_v, "Ord must be a signed integer type"); - static constexpr ::blas::Layout order = ::blas::Layout::RowMajor; + static constexpr ::blas::Layout order = _Order; static constexpr std::size_t max_rank = MaxRank; using extent_type = index; @@ -278,19 +301,19 @@ class RangeNd { const Ext* upbound_data() const noexcept { return extent_.data(); } // - // Ordinal mapping (synthesized row-major; nothing stored) + // Ordinal mapping (synthesized from extent per _Order; nothing stored) // - ordinal_view ordinal() const { - return ordinal_view(extent_, rank()); + ordinal_view<_Order, MaxRank, Ord> ordinal() const { + return ordinal_view<_Order, MaxRank, Ord>(extent_, rank()); } - /// Row-major strides synthesized on demand. Returned by value (not by + /// Strides synthesized on demand per \c _Order . Returned by value (not by /// reference) so nothing is stored in the range itself — preserves the /// packed footprint. Callers that need a pointer (e.g. the BTAS generic /// permute, which forwards r.stride() into a btas::Range ctor) bind the /// temporary to a const& and copy from it. - using stride_type = typename ordinal_view::stride_type; + using stride_type = typename ordinal_view<_Order, MaxRank, Ord>::stride_type; stride_type stride() const noexcept { return ordinal().stride(); } @@ -303,9 +326,16 @@ class RangeNd { const auto r = rank(); Ord o{0}; Ord vol{1}; - for (std::ptrdiff_t i = static_cast(r) - 1; i >= 0; --i) { - o += static_cast(*(it + i)) * vol; - vol *= static_cast(extent_[i]); + if constexpr (_Order == ::blas::Layout::RowMajor) { + for (std::ptrdiff_t i = static_cast(r) - 1; i >= 0; --i) { + o += static_cast(*(it + i)) * vol; + vol *= static_cast(extent_[i]); + } + } else { + for (std::size_t i = 0; i < r; ++i) { + o += static_cast(*(it + i)) * vol; + vol *= static_cast(extent_[i]); + } } return o; } @@ -319,15 +349,26 @@ class RangeNd { const_iterator cbegin() const { return begin(); } const_iterator cend() const { return end(); } - /// Row-major in-place increment used by \c RangeIterator . After the final - /// valid index, idx == upbound() (== extent_), which matches \c end() . + /// In-place increment used by \c RangeIterator , layout-aware per \c _Order . + /// After the final valid index, idx == upbound() (== extent_), which matches + /// \c end() . void increment(index_type& idx) const { const auto r = rank(); if (r == 0) return; - for (std::ptrdiff_t d = static_cast(r) - 1; d >= 0; --d) { - ++idx[d]; - if (idx[d] < extent_[d]) return; - idx[d] = Ext{0}; + if constexpr (_Order == ::blas::Layout::RowMajor) { + // last dim varies fastest + for (std::ptrdiff_t d = static_cast(r) - 1; d >= 0; --d) { + ++idx[d]; + if (idx[d] < extent_[d]) return; + idx[d] = Ext{0}; + } + } else { + // first dim varies fastest + for (std::size_t d = 0; d < r; ++d) { + ++idx[d]; + if (idx[d] < extent_[d]) return; + idx[d] = Ext{0}; + } } for (std::size_t d = 0; d < r; ++d) idx[d] = extent_[d]; } @@ -359,9 +400,10 @@ class RangeNd { extent_type extent_{}; }; -template -inline void swap(RangeNd& a, - RangeNd& b) noexcept { +template <::blas::Layout _Order, typename Ext, typename Ord, + std::size_t MaxRank> +inline void swap(RangeNd<_Order, Ext, Ord, MaxRank>& a, + RangeNd<_Order, Ext, Ord, MaxRank>& b) noexcept { a.swap(b); } @@ -371,23 +413,27 @@ inline void swap(RangeNd& a, // Trait specializations placing zb::RangeNd into the BTAS Range concept. // -template -struct range_traits> { - static constexpr ::blas::Layout order = ::blas::Layout::RowMajor; - using index_type = typename zb::RangeNd::index_type; +template <::blas::Layout _Order, typename Ext, typename Ord, + std::size_t MaxRank> +struct range_traits> { + static constexpr ::blas::Layout order = _Order; + using index_type = + typename zb::RangeNd<_Order, Ext, Ord, MaxRank>::index_type; using ordinal_type = Ord; static constexpr bool is_general_layout = false; }; -template -class boxrange_iteration_order> { +template <::blas::Layout _Order, typename Ext, typename Ord, + std::size_t MaxRank> +class boxrange_iteration_order> { public: enum { row_major = boxrange_iteration_order::row_major, other = boxrange_iteration_order::other, column_major = boxrange_iteration_order::column_major }; - static constexpr int value = row_major; + static constexpr int value = + (_Order == ::blas::Layout::RowMajor) ? row_major : column_major; }; } // namespace btas @@ -400,23 +446,31 @@ class boxrange_iteration_order> { namespace madness { namespace archive { -template -struct ArchiveLoadImpl> { +template +struct ArchiveLoadImpl> { static inline void load(const Archive& ar, - btas::zb::RangeNd& r) { + btas::zb::RangeNd<_Order, Ext, Ord, MaxRank>& r) { std::uint8_t rank{}; ar& rank; - typename btas::zb::RangeNd::extent_type ext( + // Guard against malformed archives: rank must fit in MaxRank, + // otherwise the underlying btas::zb::index would truncate the size_ + // field and corrupt subsequent reads. + BTAS_ASSERT(static_cast(rank) <= MaxRank && + "btas::zb::RangeNd archive load: rank exceeds MaxRank"); + typename btas::zb::RangeNd<_Order, Ext, Ord, MaxRank>::extent_type ext( static_cast(rank)); for (std::uint8_t i = 0; i < rank; ++i) ar& ext[i]; - r = btas::zb::RangeNd(ext); + r = btas::zb::RangeNd<_Order, Ext, Ord, MaxRank>(ext); } }; -template -struct ArchiveStoreImpl> { - static inline void store(const Archive& ar, - const btas::zb::RangeNd& r) { +template +struct ArchiveStoreImpl> { + static inline void store( + const Archive& ar, + const btas::zb::RangeNd<_Order, Ext, Ord, MaxRank>& r) { const std::uint8_t rank = static_cast(r.rank()); ar& rank; for (std::uint8_t i = 0; i < rank; ++i) ar& r.extent(i); diff --git a/unittest/zb_range_test.cc b/unittest/zb_range_test.cc index 346db03d..39cf682e 100644 --- a/unittest/zb_range_test.cc +++ b/unittest/zb_range_test.cc @@ -150,7 +150,7 @@ TEST_CASE("zb::RangeNd equality and swap") { } TEST_CASE("zb::RangeNd with non-default template parameters") { - using R4 = RangeNd<4, std::int32_t, std::int64_t>; + using R4 = RangeNd<::blas::Layout::RowMajor, std::int32_t, std::int64_t, 4>; static_assert(sizeof(R4) == 20, "expected packed size for MaxRank=4, int32"); static_assert(R4::max_rank == 4); R4 r(10, 20, 30); @@ -159,3 +159,26 @@ TEST_CASE("zb::RangeNd with non-default template parameters") { CHECK(r.ordinal(typename R4::index_type{1, 2, 3}) == 1 * 20 * 30 + 2 * 30 + 3); } + +TEST_CASE("zb::RangeNd column-major layout") { + using RC = RangeNd<::blas::Layout::ColMajor>; + RC r(10, 20, 30); + CHECK(r.rank() == 3); + CHECK(r.area() == 6000); + // col-major: ordinal = i0 + i1*ext0 + i2*ext0*ext1 + CHECK(r.ordinal(typename RC::index_type{1, 2, 3}) == + 1 + 2 * 10 + 3 * 10 * 20); + // strides[0]=1, strides[1]=ext0=10, strides[2]=ext0*ext1=200 + auto s = r.stride(); + CHECK(s[0] == 1); + CHECK(s[1] == 10); + CHECK(s[2] == 200); + // Iteration covers volume in column-major order. + std::size_t count = 0; + typename RC::index_type prev; + for (auto&& idx : r) { + (void)idx; + ++count; + } + CHECK(count == r.area()); +} From bfb525f3c57a81d8dd8b517700bca6cc46a5649b Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Tue, 12 May 2026 23:30:11 -0400 Subject: [PATCH 6/6] zb::RangeNd, tensor: address Copilot review comments - tensor.h (range, generator) ctor: dispatch via std::invoke, and constrain the generator against range_type's iteration value (not the input Range's), so the SFINAE-accepted set matches what compiles in the body. - zb::RangeNd (lobound, upbound) ctor: assert rank-equality before checking the all-zero lobound condition, so a size mismatch fails loudly instead of silently extending extent_ from upbound alone. - zb::RangeNd::zero_buffer comment: "first rank() bytes" -> "elements". - zb_range_test.cc: drop the unused `prev` index in the column-major test and instead validate that traversal is in column-major order by checking r.ordinal(idx) == iteration count for each visited index. --- btas/tensor.h | 4 ++-- btas/zb/range.h | 10 ++++++---- unittest/zb_range_test.cc | 6 +++--- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/btas/tensor.h b/btas/tensor.h index 357906cd..75e55c87 100644 --- a/btas/tensor.h +++ b/btas/tensor.h @@ -154,13 +154,13 @@ namespace btas { btas::is_boxrange::value && std::is_invocable_r_v< value_type, F, - decltype(*std::begin(std::declval()))>>> + decltype(*std::begin(std::declval()))>>> Tensor(const Range& range, F&& gen) : range_(range.lobound(), range.upbound()) { array_adaptor::resize(storage_, range_.area()); auto out_it = begin(); for (auto&& idx : range_) { - *out_it++ = gen(idx); + *out_it++ = std::invoke(std::forward(gen), idx); } } diff --git a/btas/zb/range.h b/btas/zb/range.h index 1610b1cc..9a9017a6 100644 --- a/btas/zb/range.h +++ b/btas/zb/range.h @@ -262,12 +262,14 @@ class RangeNd { typename = std::enable_if_t>::value && is_index>::value>> RangeNd(const Lo& lobound, const Up& upbound) : extent_(upbound) { - (void)lobound; using std::cbegin; using std::cend; + using std::size; + BTAS_ASSERT(size(lobound) == size(upbound) && + "btas::zb::RangeNd: lobound and upbound must have equal rank"); BTAS_ASSERT(std::all_of(cbegin(lobound), cend(lobound), - [](auto v) { return v == 0; }) && - "btas::zb::RangeNd: lobound must be all zeros"); + [](auto v) { return v == 0; }) && + "btas::zb::RangeNd: lobound must be all zeros"); } // @@ -391,7 +393,7 @@ class RangeNd { private: /// Static MaxRank-sized zero buffer; \c lobound_data() returns a pointer - /// into it. Callers iterate only the first \c rank() bytes. + /// into it. Callers iterate only the first \c rank() elements. static const Ext* zero_buffer() noexcept { static const std::array z{}; return z.data(); diff --git a/unittest/zb_range_test.cc b/unittest/zb_range_test.cc index 39cf682e..ac0f261b 100644 --- a/unittest/zb_range_test.cc +++ b/unittest/zb_range_test.cc @@ -173,11 +173,11 @@ TEST_CASE("zb::RangeNd column-major layout") { CHECK(s[0] == 1); CHECK(s[1] == 10); CHECK(s[2] == 200); - // Iteration covers volume in column-major order. + // Iteration covers volume — and the ordinal of each visited index must + // equal the iteration count, which validates column-major traversal order. std::size_t count = 0; - typename RC::index_type prev; for (auto&& idx : r) { - (void)idx; + CHECK(r.ordinal(idx) == static_cast(count)); ++count; } CHECK(count == r.area());