| Taskflow: A General-purpose Task-parallel Programming System

Loading...
Searching...
No Matches
iterator.hpp
1#pragma once
2
3#include <array>
4#include <cassert>
5#include <concepts>
6#include <cstddef>
7#include <tuple>
8#include <type_traits>
9
10 namespace tf {
11
28 template <std::integral T>
29 constexpr bool is_index_range_invalid(T beg, T end, T step) {
30 return ((step == T{0} && beg != end) ||
31 (beg < end && step <= T{0}) || // positive range
32 (beg > end && step >= T{0})); // negative range
33}
34
70 template <std::integral T>
71 constexpr size_t distance(T beg, T end, T step) {
72 if constexpr (std::is_unsigned_v<T>) {
73 return end > beg
74 ? static_cast<size_t>((end - beg + step - T{1}) / step)
75 : size_t{0};
76 } else {
77 return static_cast<size_t>(
78 std::max(T{0}, (end - beg + step + (step > T{0} ? T{-1} : T{1})) / step)
79 );
80 }
81}
82
83// ============================================================================
84// IndexRanges<T, N>
85//
86// A single class template representing an N-dimensional index range, the
87// Cartesian product of N independent 1D ranges (begin, end, step). Each
88// dimension is stored as a std::tuple<T, T, T>, so dim(d) returns a tuple
89// that can be read or mutated directly, including via structured bindings:
90//
91// auto& [beg, end, step] = ranges.dim(0);
92//
93// Members that only make sense for a single dimension (begin(), end(),
94// step_size(), reset(), unravel()) are gated with `requires (N == 1)`.
95// Members that only make sense for more than one dimension (ceil(), floor(),
96// upper_slice(), lower_slice()) are
97// gated with `requires (N > 1)`. This keeps
98// everything in one class body — instead of a primary template plus a
99// partial specialization — which Doxygen can parse and cross-reference
100// cleanly.
101//
102// tf::IndexRange<T> is an alias for tf::IndexRanges<T, 1>, the common 1D
103// case.
104//
105// Iteration order for N > 1 is row-major (the last dimension is innermost /
106// fastest), matching the natural loop nesting:
107//
108// for i in dim[0]: // outermost
109// for j in dim[1]:
110// ...
111// for k in dim[N-1]: // innermost
112//
113// Flat index 0 corresponds to (beg[0], beg[1], ..., beg[N-1]).
114// ============================================================================
115
187 template <std::integral T, size_t N = 1>
188 class IndexRanges {
189
190 public:
191
195 using index_type = T;
196
200 static constexpr size_t rank = N;
201
202 // --------------------------------------------------------------------------
203 // Construction
204 // --------------------------------------------------------------------------
205
213 IndexRanges() = default;
214
226 explicit IndexRanges(T beg, T end, T step_size) requires (N == 1)
227 : _dims{ std::tuple<T, T, T>{beg, end, step_size} } {}
228
251 template <typename... Ranges>
252 requires (sizeof...(Ranges) == N) &&
253 (std::same_as<std::decay_t<Ranges>, IndexRanges<T, 1>> && ...)
254 explicit IndexRanges(Ranges&&... ranges)
255 : _dims{ ranges.dim(0)... } {}
256
273 explicit IndexRanges(const std::array<std::tuple<T, T, T>, N>& dims) : _dims{dims} {}
274
275 // --------------------------------------------------------------------------
276 // Dimension access
277 // --------------------------------------------------------------------------
278
297 const std::tuple<T, T, T>& dim(size_t d) const { return _dims[d]; }
298
330 std::tuple<T, T, T>& dim(size_t d) { return _dims[d]; }
331
332 // --------------------------------------------------------------------------
333 // 1D convenience accessors (only available when N == 1)
334 // --------------------------------------------------------------------------
335
346 T begin() const requires (N == 1) { return std::get<0>(_dims[0]); }
347
358 T end() const requires (N == 1) { return std::get<1>(_dims[0]); }
359
370 T step_size() const requires (N == 1) { return std::get<2>(_dims[0]); }
371
396 IndexRanges& reset(T beg, T end, T step_size) requires (N == 1);
397
410 IndexRanges& begin(T new_begin) requires (N == 1);
411
424 IndexRanges& end(T new_end) requires (N == 1);
425
438 IndexRanges& step_size(T new_step_size) requires (N == 1);
439
465 IndexRanges unravel(size_t part_beg, size_t part_end) const requires (N == 1);
466
467 // --------------------------------------------------------------------------
468 // Size queries (available for any N)
469 // --------------------------------------------------------------------------
470
491 size_t size(size_t d) const { return std::apply(distance<T>, _dims[d]); }
492
521 size_t size() const;
522
523 // --------------------------------------------------------------------------
524 // N-dimensional algorithms (only available when N > 1)
525 // --------------------------------------------------------------------------
526
561 size_t ceil(size_t chunk_size) const requires (N > 1);
562
591 size_t floor(size_t chunk_size) const requires (N > 1);
592
619 IndexRanges upper_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1);
620
640 IndexRanges lower_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1);
641
655 IndexRanges empty_box() const requires (N > 1);
656
657private:
658
659 std::array<std::tuple<T, T, T>, N> _dims;
660};
661
662// ============================================================================
663// Out-of-class definitions — 1D convenience accessors and size queries
664// ============================================================================
665
666template <std::integral T, size_t N>
667 IndexRanges<T, N>&
668 IndexRanges<T, N>::reset(T beg, T end, T step_size) requires (N == 1) {
669 _dims[0] = {beg, end, step_size};
670 return *this;
671}
672
673 template <std::integral T, size_t N>
674 IndexRanges<T, N>&
675 IndexRanges<T, N>::begin(T new_begin) requires (N == 1) {
676 std::get<0>(_dims[0]) = new_begin;
677 return *this;
678}
679
680 template <std::integral T, size_t N>
681 IndexRanges<T, N>&
682 IndexRanges<T, N>::end(T new_end) requires (N == 1) {
683 std::get<1>(_dims[0]) = new_end;
684 return *this;
685}
686
687 template <std::integral T, size_t N>
688 IndexRanges<T, N>&
689 IndexRanges<T, N>::step_size(T new_step_size) requires (N == 1) {
690 std::get<2>(_dims[0]) = new_step_size;
691 return *this;
692}
693
694 template <std::integral T, size_t N>
695 IndexRanges<T, N>
696 IndexRanges<T, N>::unravel(size_t part_beg, size_t part_end) const requires (N == 1) {
697 auto [beg, end, step] = _dims[0];
698 return IndexRanges(
699 static_cast<T>(part_beg) * step + beg,
700 static_cast<T>(part_end) * step + beg,
701 step
702 );
703}
704
705 template <std::integral T, size_t N>
706 size_t IndexRanges<T, N>::size() const {
707 if constexpr (N == 1) {
708 return size(0);
709 } else {
710 // Compile-time-unrolled recursion: D is a template parameter, so each
711 // depth is a distinct instantiation and the recursion fully unrolls
712 // (no runtime loop). `self` is passed explicitly because C++20 lambdas
713 // cannot otherwise name themselves for recursion.
714 auto compute = [this]<size_t D>(auto& self, size_t total) -> size_t {
715 if constexpr (D == N) {
716 return total;
717 } else {
718 size_t s = this->size(D);
719 if (s == 0) return D == 0 ? 0 : total; // outermost zero -> 0, inner zero -> outer product
720 return self.template operator()<D + 1>(self, total * s);
721 }
722 };
723 return compute.template operator()<0>(compute, 1);
724 }
725}
726
727// ============================================================================
728// Out-of-class definitions — N-dimensional algorithms (N > 1)
729// ============================================================================
730
731 template <std::integral T, size_t N>
732 IndexRanges<T, N>
733 IndexRanges<T, N>::empty_box() const requires (N > 1) {
734 std::array<std::tuple<T, T, T>, N> box_dims = _dims;
735 box_dims[0] = std::tuple<T, T, T>{
736 std::get<0>(_dims[0]), std::get<0>(_dims[0]), std::get<2>(_dims[0])
737 };
738 return IndexRanges(box_dims);
739}
740
741 template <std::integral T, size_t N>
742 size_t IndexRanges<T, N>::ceil(size_t chunk_size) const requires (N > 1) {
743 // Pass 1: cache all dim sizes (one call each) and find d_active.
744 size_t dim_sizes[N];
745 size_t d_active = N;
746 for (size_t d = 0; d < N; ++d) {
747 dim_sizes[d] = size(d);
748 if (dim_sizes[d] == 0) { d_active = d; break; }
749 }
750 if (d_active == 0) return 0;
751
752 // Pass 2: walk suffix products backward within [0, d_active) only.
753 // Return the first suffix product >= chunk_size, capped at size().
754 size_t inner_volume = 1;
755 if (inner_volume >= chunk_size) return inner_volume;
756 for (size_t d = d_active; d-- > 0; ) {
757 inner_volume *= dim_sizes[d];
758 if (inner_volume >= chunk_size) return inner_volume;
759 }
760 return inner_volume; // == size() of active dims, capped
761}
762
763 template <std::integral T, size_t N>
764 size_t IndexRanges<T, N>::floor(size_t chunk_size) const requires (N > 1) {
765 // Pass 1: cache all dim sizes (one call each) and find d_active.
766 size_t dim_sizes[N];
767 size_t d_active = N;
768 for (size_t d = 0; d < N; ++d) {
769 dim_sizes[d] = size(d);
770 if (dim_sizes[d] == 0) { d_active = d; break; }
771 }
772 if (d_active == 0) return 0;
773
774 // Pass 2: walk suffix products backward within [0, d_active) only.
775 // Return the largest suffix product <= chunk_size.
776 size_t inner_volume = 1;
777 size_t last_fit = 1;
778 for (size_t d = d_active; d-- > 0; ) {
779 size_t next = inner_volume * dim_sizes[d];
780 if (next > chunk_size) break;
781 inner_volume = next;
782 last_fit = inner_volume;
783 }
784 return last_fit;
785}
786
787 template <std::integral T, size_t N>
788 IndexRanges<T, N>
789 IndexRanges<T, N>::upper_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1) {
790
791 if (chunk_size == 0) {
792 return empty_box();
793 }
794
795 // Pass 1 (forward): cache dim sizes and find d_active — the first zero-size
796 // dimension. Only [0, d_active) participates in the flat iteration space;
797 // dims [d_active, N) are copied as full extent into the returned box.
798 size_t dim_sizes[N];
799 size_t d_active = N;
800 for (size_t d = 0; d < N; ++d) {
801 dim_sizes[d] = size(d);
802 if (dim_sizes[d] == 0) { d_active = d; break; }
803 }
804
805 if (d_active == 0) return *this;
806
807 // Pass 2a (backward): decode flat_beg into per-dimension coords.
808 size_t coords[N] = {};
809 size_t temp = flat_beg;
810 for (size_t d = d_active; d-- > 0; ) {
811 coords[d] = temp % dim_sizes[d];
812 temp /= dim_sizes[d];
813 }
814
815 // Pass 2b (backward, fused with grow_dim search): find grow_dim.
816 // Ceil variant: commit grow_dim BEFORE the budget check (round up).
817 size_t grow_dim = d_active - 1;
818 size_t inner_volume = 1;
819 size_t active_inner_vol = 1;
820
821 for (size_t d = d_active; d-- > 0; ) {
822 // trailing-zeros rule: inner coord non-zero -> can't expand further out
823 if (d + 1 < d_active && coords[d + 1] != 0) break;
824
825 // ceil: commit BEFORE budget check
826 grow_dim = d;
827 active_inner_vol = inner_volume;
828
829 if (inner_volume >= chunk_size) break;
830
831 inner_volume *= dim_sizes[d];
832 }
833
834 // Steps along grow_dim: ceil division, at least 1 for forward progress.
835 size_t steps_left = dim_sizes[grow_dim] - coords[grow_dim];
836 size_t steps_needed = (std::max)(size_t{1}, chunk_size / active_inner_vol);
837 size_t steps_to_take = (std::min)(steps_left, steps_needed);
838
839 // Pass 3: construct the sub-box in three clean segments:
840 // [0, grow_dim) — locked dims (one element each)
841 // [grow_dim] — the grow dimension
842 // [grow_dim+1, N) — full inner/inactive extent
843 std::array<std::tuple<T, T, T>, N> box_dims;
844
845 for (size_t d = 0; d < grow_dim; ++d) {
846 const T beg = std::get<0>(_dims[d]);
847 const T step = std::get<2>(_dims[d]);
848 box_dims[d] = std::tuple<T, T, T>{
849 beg + static_cast<T>(coords[d]) * step,
850 beg + static_cast<T>(coords[d] + 1) * step,
851 step
852 };
853 }
854
855 {
856 const T beg = std::get<0>(_dims[grow_dim]);
857 const T step = std::get<2>(_dims[grow_dim]);
858 box_dims[grow_dim] = std::tuple<T, T, T>{
859 beg + static_cast<T>(coords[grow_dim]) * step,
860 beg + static_cast<T>(coords[grow_dim] + steps_to_take) * step,
861 step
862 };
863 }
864
865 for (size_t d = grow_dim + 1; d < N; ++d) {
866 box_dims[d] = _dims[d]; // full inner or inactive extent
867 }
868
869 return IndexRanges<T, N>(box_dims);
870}
871
872 template <std::integral T, size_t N>
873 IndexRanges<T, N>
874 IndexRanges<T, N>::lower_slice(size_t flat_beg, size_t chunk_size) const requires (N > 1) {
875
876 if (chunk_size == 0) {
877 return empty_box();
878 }
879
880 // Pass 1 (forward): cache dim sizes and find d_active.
881 size_t dim_sizes[N];
882 size_t d_active = N;
883 for (size_t d = 0; d < N; ++d) {
884 dim_sizes[d] = size(d);
885 if (dim_sizes[d] == 0) { d_active = d; break; }
886 }
887
888 if (d_active == 0) return *this;
889
890 // Pass 2a (backward): decode flat_beg into per-dimension coords.
891 size_t coords[N] = {};
892 size_t temp = flat_beg;
893 for (size_t d = d_active; d-- > 0; ) {
894 coords[d] = temp % dim_sizes[d];
895 temp /= dim_sizes[d];
896 }
897
898 // Pass 2b (backward, fused with grow_dim search): find grow_dim.
899 // Floor variant: commit grow_dim AFTER the budget check (round down).
900 size_t grow_dim = d_active - 1;
901 size_t inner_volume = 1;
902 size_t active_inner_vol = 1;
903
904 for (size_t d = d_active; d-- > 0; ) {
905 // trailing-zeros rule
906 if (d + 1 < d_active && coords[d + 1] != 0) break;
907
908 // floor: budget check BEFORE committing
909 size_t next_vol = inner_volume * dim_sizes[d];
910 if (next_vol > chunk_size && inner_volume > 1) break;
911
912 grow_dim = d;
913 active_inner_vol = inner_volume;
914 inner_volume = next_vol;
915 }
916
917 // Steps along grow_dim: floor division, at least 1 for forward progress.
918 size_t steps_left = dim_sizes[grow_dim] - coords[grow_dim];
919 size_t steps_needed = (std::max)(size_t{1}, chunk_size / active_inner_vol);
920 size_t steps_to_take = (std::min)(steps_left, steps_needed);
921
922 // Pass 3: construct the sub-box in three clean segments.
923 std::array<std::tuple<T, T, T>, N> box_dims;
924
925 for (size_t d = 0; d < grow_dim; ++d) {
926 const T beg = std::get<0>(_dims[d]);
927 const T step = std::get<2>(_dims[d]);
928 box_dims[d] = std::tuple<T, T, T>{
929 beg + static_cast<T>(coords[d]) * step,
930 beg + static_cast<T>(coords[d] + 1) * step,
931 step
932 };
933 }
934
935 {
936 const T beg = std::get<0>(_dims[grow_dim]);
937 const T step = std::get<2>(_dims[grow_dim]);
938 box_dims[grow_dim] = std::tuple<T, T, T>{
939 beg + static_cast<T>(coords[grow_dim]) * step,
940 beg + static_cast<T>(coords[grow_dim] + steps_to_take) * step,
941 step
942 };
943 }
944
945 for (size_t d = grow_dim + 1; d < N; ++d) {
946 box_dims[d] = _dims[d]; // full inner or inactive extent
947 }
948
949 return IndexRanges<T, N>(box_dims);
950}
951
952// ----------------------------------------------------------------------------
953// IndexRange<T> — alias for the common 1D case
954// ----------------------------------------------------------------------------
955
970 template <std::integral T>
971 using IndexRange = IndexRanges<T, 1>;
972
973// ==========================================
974// traits
975// ==========================================
976
981 template <typename>
982 constexpr bool is_index_ranges_v = false;
983
992 template <typename T, size_t N>
993 constexpr bool is_index_ranges_v<IndexRanges<T, N>> = true;
994
1003 template <typename R>
1004 concept IndexRangesLike = is_index_ranges_v<std::decay_t<std::unwrap_ref_decay_t<R>>>;
1005
1016 template <typename R>
1017 concept IndexRanges1DLike = IndexRangesLike<R> &&
1018 (std::decay_t<std::unwrap_ref_decay_t<R>>::rank == 1);
1019
1031 template <typename R>
1032 concept IndexRangesMDLike = IndexRangesLike<R> &&
1033 (std::decay_t<std::unwrap_ref_decay_t<R>>::rank > 1);
1034
1035} // end of namespace tf -----------------------------------------------------
tf::IndexRanges
class to create an N-dimensional index range of integral indices
Definition iterator.hpp:188
tf::IndexRanges::reset
IndexRanges & reset(T beg, T end, T step_size)
updates the range with a new starting index, ending index, and step size (only available when N == 1)
Definition iterator.hpp:668
tf::IndexRanges::upper_slice
IndexRanges upper_slice(size_t flat_beg, size_t chunk_size) const
returns the smallest perfectly-aligned hyperbox reachable from flat_beg, rounding up to the next hype...
Definition iterator.hpp:789
tf::IndexRanges< T, 1 >::end
T end() const
Definition iterator.hpp:358
tf::IndexRanges< T, 1 >::dim
const std::tuple< T, T, T > & dim(size_t d) const
Definition iterator.hpp:297
tf::IndexRanges::end
IndexRanges & end(T new_end)
updates the ending index of the range (only available when N == 1)
Definition iterator.hpp:682
tf::IndexRanges::ceil
size_t ceil(size_t chunk_size) const
returns the smallest hyperplane-aligned chunk size that is >= chunk_size, capped at size() when chunk...
Definition iterator.hpp:742
tf::IndexRanges::dim
std::tuple< T, T, T > & dim(size_t d)
returns the (begin, end, step) tuple for dimension d (mutable)
Definition iterator.hpp:330
tf::IndexRanges::unravel
IndexRanges unravel(size_t part_beg, size_t part_end) const
maps a contiguous index partition back to the corresponding subrange (only available when N == 1)
Definition iterator.hpp:696
tf::IndexRanges::step_size
IndexRanges & step_size(T new_step_size)
updates the step size of the range (only available when N == 1)
Definition iterator.hpp:689
tf::IndexRanges< T, 1 >::rank
static constexpr size_t rank
Definition iterator.hpp:200
tf::IndexRanges::IndexRanges
IndexRanges()=default
constructs an index range without initialization
tf::IndexRanges::size
size_t size(size_t d) const
returns the number of iterations along dimension d
Definition iterator.hpp:491
tf::IndexRanges::IndexRanges
IndexRanges(Ranges &&... ranges)
constructs an N-D index range from N 1D ranges
Definition iterator.hpp:254
tf::IndexRanges::lower_slice
IndexRanges lower_slice(size_t flat_beg, size_t chunk_size) const
returns the largest perfectly-aligned hyperbox reachable from flat_beg whose size does not exceed chu...
Definition iterator.hpp:874
tf::IndexRanges::size
size_t size() const
returns the number of active flat iterations
Definition iterator.hpp:706
tf::IndexRanges::index_type
T index_type
alias for the index type
Definition iterator.hpp:195
tf::IndexRanges::IndexRanges
IndexRanges(const std::array< std::tuple< T, T, T >, N > &dims)
constructs an index range from an array of (begin, end, step) tuples
Definition iterator.hpp:273
tf::IndexRanges::floor
size_t floor(size_t chunk_size) const
returns the largest hyperplane-aligned chunk size that is <= chunk_size (only available when N > 1)
Definition iterator.hpp:764
tf::IndexRanges::empty_box
IndexRanges empty_box() const
returns a box whose size() is 0 (only available when N > 1)
Definition iterator.hpp:733
tf::IndexRanges::IndexRanges
IndexRanges(T beg, T end, T step_size)
constructs a 1D index range (only available when N == 1)
Definition iterator.hpp:226
tf::IndexRanges::begin
IndexRanges & begin(T new_begin)
updates the starting index of the range (only available when N == 1)
Definition iterator.hpp:675
tf::IndexRanges::begin
T begin() const
queries the starting index of the range (only available when N == 1)
Definition iterator.hpp:346
tf::IndexRanges< T, 1 >::step_size
T step_size() const
Definition iterator.hpp:370
tf::IndexRanges1DLike
concept to check if a type is a tf::IndexRanges<T, 1> (i.e., tf::IndexRange<T>)
Definition iterator.hpp:1017
tf::IndexRangesLike
concept to check if a type is a tf::IndexRanges, regardless of dimensionality
Definition iterator.hpp:1004
tf::IndexRangesMDLike
concept to check if a type is a tf::IndexRanges<T, N> with rank > 1
Definition iterator.hpp:1032
tf
taskflow namespace
Definition small_vector.hpp:20
tf::IndexRange
IndexRanges< T, 1 > IndexRange
alias for the common 1D case of tf::IndexRanges
Definition iterator.hpp:971
tf::is_index_ranges_v
constexpr bool is_index_ranges_v
base type trait to detect if a type is a tf::IndexRanges
Definition iterator.hpp:982
tf::distance
constexpr size_t distance(T beg, T end, T step)
calculates the number of iterations in the given index range
Definition iterator.hpp:71
tf::is_index_range_invalid
constexpr bool is_index_range_invalid(T beg, T end, T step)
checks if the given index range is invalid
Definition iterator.hpp:29