Skip to content

Commit b11eec3

Browse files
committed
Add for_each_extents example
1 parent b31a635 commit b11eec3

File tree

3 files changed

+117
-0
lines changed

3 files changed

+117
-0
lines changed

examples/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,3 +16,4 @@ add_subdirectory(dot_product)
1616
add_subdirectory(tiled_layout)
1717
add_subdirectory(restrict_accessor)
1818
add_subdirectory(aligned_accessor)
19+
add_subdirectory(for_each_extents)
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
mdspan_add_example(for_each_extents)
Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
#include <experimental/mdspan.hpp>
2+
#include <cassert>
3+
#include <iostream>
4+
#include <type_traits>
5+
6+
// "gcc trunk" on godbolt.org as of 2023/03/21
7+
// (> 12.2) does not define __cpp_lib_ranges_iota,
8+
// yet std::views::iota works just fine.
9+
#if defined(__cpp_lib_ranges_cartesian_product) // && defined(__cpp_lib_ranges_iota)
10+
# define MDSPAN_EXAMPLE_CAN_USE_STD_RANGES 1
11+
#endif
12+
13+
#if defined(MDSPAN_EXAMPLE_CAN_USE_STD_RANGES)
14+
15+
// GCC >= 13 ("gcc trunk" on godbolt.org as of 2023/03/21)
16+
// implements std::views::cartesian_product.
17+
// If you don't have it, you can use range-v3 instead.
18+
// Note that mixing std::views::iota with
19+
// ranges::views::cartesian_product doesn't work.
20+
// The range-v3 work-around looks like this.
21+
//
22+
// #include <range/v3/view/cartesian_product.hpp>
23+
// #include <range/v3/view/indices.hpp>
24+
// namespace ranges_views = ranges::views;
25+
26+
#include <ranges>
27+
namespace ranges_views = std::views;
28+
29+
namespace stdex = std::experimental;
30+
31+
auto print_args = [] <class ... Args> (Args&&... args) {
32+
((std::cout << std::forward<Args>(args) << '\n'), ...);
33+
};
34+
35+
template<size_t ... Lefts, size_t ... Rights>
36+
auto concat_index_sequence(std::index_sequence<Lefts...>,
37+
std::index_sequence<Rights...>)
38+
{
39+
return std::index_sequence<Lefts..., Rights...>{};
40+
}
41+
42+
auto reverse_index_sequence(std::index_sequence<> x)
43+
{
44+
return x;
45+
}
46+
47+
template<size_t First, size_t ... Rest>
48+
auto reverse_index_sequence(std::index_sequence<First, Rest...>)
49+
{
50+
return concat_index_sequence(
51+
reverse_index_sequence(std::index_sequence<Rest...>{}),
52+
std::index_sequence<First>{});
53+
}
54+
55+
template<size_t N>
56+
auto make_reverse_index_sequence()
57+
{
58+
return reverse_index_sequence(std::make_index_sequence<N>());
59+
}
60+
61+
template<class Callable, class IndexType,
62+
std::size_t ... Extents, std::size_t ... RankIndices>
63+
void for_each_in_extents_impl(Callable&& f,
64+
stdex::extents<IndexType, Extents...> e,
65+
std::index_sequence<RankIndices...> rank_sequence)
66+
{
67+
// In the layout_left case, caller passes in N-1, N-2, ..., 1, 0.
68+
// This reverses the order of the Cartesian product,
69+
// but also reverses the order of indices in each tuple.
70+
[&] <std::size_t ... Indices> (std::index_sequence<Indices...>) {
71+
auto v = std::views::cartesian_product(
72+
std::views::iota(IndexType(0), e.extent(Indices))...);
73+
for (const auto& tuple_of_indices : v) {
74+
// In the layout_left case, we undo the reversal of each tuple
75+
// by getting its elements in reverse order.
76+
[&] <std::size_t ... InnerIndices> (std::index_sequence<InnerIndices...>) {
77+
std::forward<Callable>(f)(std::get<InnerIndices>(tuple_of_indices)...);
78+
} (rank_sequence);
79+
}
80+
} (rank_sequence);
81+
}
82+
83+
template<class Callable, class IndexType, std::size_t ... Extents, class Layout>
84+
void for_each_in_extents(Callable&& f,
85+
stdex::extents<IndexType, Extents...> e,
86+
Layout)
87+
{
88+
using layout_type = std::remove_cvref_t<Layout>;
89+
if constexpr (std::is_same_v<layout_type, stdex::layout_left>) {
90+
for_each_in_extents_impl(std::forward<Callable>(f), e,
91+
make_reverse_index_sequence<e.rank()>());
92+
}
93+
else { // layout_right or any other layout
94+
for_each_in_extents_impl(std::forward<Callable>(f), e,
95+
std::make_index_sequence<e.rank()>());
96+
}
97+
}
98+
99+
#endif // defined(MDSPAN_EXAMPLE_CAN_USE_STD_RANGES)
100+
101+
int main() {
102+
103+
#if defined(MDSPAN_EXAMPLE_CAN_USE_STD_RANGES)
104+
stdex::extents<int, 2, 3> e;
105+
auto printer = [] (int i, int j) {
106+
std::cout << "(" << i << "," << j << ")\n";
107+
};
108+
std::cout << "layout_right:\n";
109+
for_each_in_extents(printer, e, stdex::layout_right{});
110+
std::cout << "\nlayout_left:\n";
111+
for_each_in_extents(printer, e, stdex::layout_left{});
112+
#endif // defined(MDSPAN_EXAMPLE_CAN_USE_STD_RANGES)
113+
114+
return 0;
115+
}

0 commit comments

Comments
 (0)