KaMPIng 0.1.2
Flexible and (near) zero-overhead C++ bindings for MPI
Loading...
Searching...
No Matches
flatten.hpp
1// This file is part of KaMPIng.
2//
3// Copyright 2024-2025 The KaMPIng Authors
4//
5// KaMPIng is free software : you can redistribute it and/or modify it under the terms of the GNU Lesser General Public
6// License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later
7// version. KaMPIng is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the
8// implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License
9// for more details.
10//
11// You should have received a copy of the GNU Lesser General Public License along with KaMPIng. If not, see
12// <https://www.gnu.org/licenses/>.
13
14#pragma once
15#include <numeric>
16#include <utility>
17#include <vector>
18
19#include <kamping/utils/traits.hpp>
20
22
23namespace kamping {
24namespace internal {
25
26/// @brief A wrapper around a functor \p F that makes it callable using the `call` method.
27template <typename F>
29 F f; ///< The functor to wrap.
30
31 /// @brief Calls the wrapped functor with the given arguments.
32 template <typename... Args>
33 auto call(Args&&... args) {
34 return f(std::forward<Args...>(args)...);
35 }
36};
37
38/// @brief A factory function for \ref CallableWrapper.
39template <typename F>
41 return CallableWrapper<F>{std::move(f)};
42}
43
44/// @brief Maps a container to is underlying nested container.
45template <typename T, typename Enable = void>
46struct FlatContainer {};
47
48/// @brief Maps a container to is underlying nested container.
49template <typename T>
50struct FlatContainer<T, std::enable_if_t<is_sparse_send_buffer_v<T>>> {
51 using type =
52 std::remove_const_t<std::tuple_element_t<1, typename T::value_type>>; ///< The type of the nested container.
53};
54
55/// @brief Maps a container to is underlying nested container.
56template <typename T>
57struct FlatContainer<T, std::enable_if_t<is_nested_send_buffer_v<T>>> {
58 using type = std::remove_const_t<typename T::value_type>; ///< The type of the nested container.
59};
60
61} // namespace internal
62
63/// @brief Flattens a container of containers or destination-container-pairs and returns the flattened buffer, send
64/// counts and send displacements as a tuple.
65///
66/// The container can be a range of pair-like types of destination and data (see \c is_sparse_send_buffer_v ) or
67/// a container of containers (see \c is_nested_send_buffer_v ).
68///
69/// @param nested_send_buf The nested container of send buffers.
70/// @param comm_size The size of the communicator, used as number of elements in the computed count buffers.
71/// @tparam CountContainer The type of the container to use for the send counts and send displacements.
72/// @tparam Container The type of the nested container.
73/// @tparam Enable SFINAE.
74///
75template <
76 template <typename...> typename CountContainer = std::vector,
77 typename Container,
78 typename Enable = std::enable_if_t<is_sparse_send_buffer_v<Container> || is_nested_send_buffer_v<Container>>>
79auto flatten(Container const& nested_send_buf, size_t comm_size)
80 -> std::tuple<typename internal::FlatContainer<Container>::type, CountContainer<int>, CountContainer<int>> {
83 for (auto const& [destination, message]: nested_send_buf) {
86 }
87 } else {
88 static_assert(is_nested_send_buffer_v<Container>);
89 size_t i = 0;
90 for (auto const& message: nested_send_buf) {
91 auto send_buf = kamping::send_buf(message);
92 send_counts[i] = asserting_cast<int>(send_buf.size());
93 i++;
94 }
95 }
96 CountContainer<int> send_displs(comm_size);
97 std::exclusive_scan(send_counts.begin(), send_counts.end(), send_displs.begin(), 0);
98 size_t total_send_count = asserting_cast<size_t>(send_displs.back() + send_counts.back());
99 typename internal::FlatContainer<Container>::type flat_send_buf;
100 flat_send_buf.resize(total_send_count);
101 if constexpr (is_sparse_send_buffer_v<Container>) {
102 for (auto const& [destination, message]: nested_send_buf) {
103 auto send_buf = kamping::send_buf(message).construct_buffer_or_rebind();
104 std::copy_n(
105 send_buf.data(),
106 send_buf.size(),
107 flat_send_buf.data() + send_displs[asserting_cast<size_t>(destination)]
108 );
109 }
110 } else {
111 static_assert(is_nested_send_buffer_v<Container>);
112 size_t i = 0;
113 for (auto const& message: nested_send_buf) {
114 auto send_buf = kamping::send_buf(message).construct_buffer_or_rebind();
115 std::copy_n(send_buf.data(), send_buf.size(), flat_send_buf.data() + send_displs[i]);
116 i++;
117 }
118 }
119 return std::tuple(std::move(flat_send_buf), std::move(send_counts), std::move(send_displs));
120}
121
122/// @brief Flattens a container of containers and and returns the flattened buffer, send counts and send displacement as
123/// a tuple.
124///
125/// The size of the computed count buffers is the size of the container.
126/// @param nested_send_buf The nested container of send buffers. Must satisfy \c is_nested_send_buffer_v.
127/// @tparam CountContainer The type of the container to use for the send counts and send displacements.
128/// @tparam Container The type of the nested container.
129/// @tparam Enable SFINAE.
130template <
131 template <typename...> typename CountContainer = std::vector,
132 typename Container,
133 typename Enable = std::enable_if_t<is_nested_send_buffer_v<Container>>>
134auto flatten(Container const& nested_send_buf)
135 -> std::tuple<typename internal::FlatContainer<Container>::type, CountContainer<int>, CountContainer<int>> {
136 return flatten<CountContainer>(nested_send_buf, nested_send_buf.size());
137}
138
139/// @brief Flattens a container of containers or destination-container-pairs and provides the flattened buffer, send
140/// counts and send displacements as parameters to be passed to an \c MPI call.
141///
142/// This returns a callable wrapper that can be called with a functor which accepts a parameter pack of these arguments.
143///
144/// Example:
145/// ```cpp
146/// Communicator comm;
147/// std::vector<std::vector<int>> nested_send_buf(comm.size()); // or std::unordered_map<int, std::vector<int>>
148/// auto [recv_buf, recv_counts, recv_displs] = with_flattened(nested_send_buf).call([&](auto... flattened) {
149/// return comm.alltoallv(std::move(flattened)..., recv_counts_out(), recv_displs_out());
150/// });
151/// ```
152///
153/// The container can be a range of pair-like types of destination and data (see \c is_sparse_send_buffer_v ) or
154/// a container of containers (see \c is_nested_send_buffer_v ).
155///
156/// @param nested_send_buf The nested container of send buffers.
157/// @param comm_size The size of the communicator, used as number of elements in the computed count buffers.
158/// @tparam CountContainer The type of the container to use for the send counts and send displacements.
159/// @tparam Container The type of the nested container.
160/// @tparam Enable SFINAE.
161///
162template <
163 template <typename...> typename CountContainer = std::vector,
164 typename Container,
165 typename Enable = std::enable_if_t<is_sparse_send_buffer_v<Container> || is_nested_send_buffer_v<Container>>>
166auto with_flattened(Container const& nested_send_buf, size_t comm_size) {
167 auto flat = flatten(nested_send_buf, comm_size);
168 return internal::make_callable_wrapper([flat_send_buf = std::move(std::get<0>(flat)),
169 send_counts = std::move(std::get<1>(flat)),
170 send_displs = std::move(std::get<2>(flat))](auto&& f) {
171 return std::apply(
172 std::forward<decltype(f)>(f),
173 std::tuple(
174 kamping::send_buf(flat_send_buf),
177 )
178 );
179 });
180}
181
182/// @brief Flattens a container of containers and provides the flattened buffer, send counts and send
183/// displacements as parameters to be passed to an \c MPI call.
184///
185/// This returns a callable wrapper that can be called with a functor which accepts a parameter pack of these arguments.
186/// The size of the computed count buffers is the size of the container.
187///
188/// Example:
189/// ```cpp
190/// Communicator comm;
191/// std::vector<std::vector<int>> nested_send_buf(comm.size());
192/// auto [recv_buf, recv_counts, recv_displs] = with_flattened(nested_send_buf).call([&](auto... flattened) {
193/// return comm.alltoallv(std::move(flattened)..., recv_counts_out(), recv_displs_out());
194/// });
195/// ```
196///
197/// @param nested_send_buf The nested container of send buffers. Must satisfy \c is_nested_send_buffer_v.
198/// @tparam CountContainer The type of the container to use for the send counts and send displacements.
199/// @tparam Container The type of the nested container.
200/// @tparam Enable SFINAE.
201///
202template <
203 template <typename...> typename CountContainer = std::vector,
204 typename Container,
205 typename Enable = std::enable_if_t<is_nested_send_buffer_v<Container>>>
206auto with_flattened(Container const& nested_send_buf) {
207 return with_flattened<CountContainer>(nested_send_buf, nested_send_buf.size());
208}
209} // namespace kamping
STL-compatible allocator for requesting memory using the builtin MPI allocator.
Definition allocator.hpp:32
auto send_displs(Container &&container)
Passes a container as send displacements to the underlying call, i.e. the container's storage must co...
Definition named_parameters.hpp:576
auto send_counts(Container &&container)
Passes a container as send counts to the underlying call, i.e. the container's storage must contain t...
Definition named_parameters.hpp:205
auto destination(int rank)
Passes rank as destination rank to the underlying call. This parameter is needed in point-to-point ex...
Definition named_parameters.hpp:1001
auto send_buf(internal::ignore_t< Data > ignore)
Generates a dummy send buf that wraps a nullptr.
Definition named_parameters.hpp:53
Factory methods for buffer wrappers.
auto make_callable_wrapper(F f)
A factory function for CallableWrapper.
Definition flatten.hpp:40
STL namespace.
A wrapper around a functor F that makes it callable using the call method.
Definition flatten.hpp:28
auto call(Args &&... args)
Calls the wrapped functor with the given arguments.
Definition flatten.hpp:33
F f
The functor to wrap.
Definition flatten.hpp:29
std::remove_const_t< typename T::value_type > type
The type of the nested container.
Definition flatten.hpp:58
std::remove_const_t< std::tuple_element_t< 1, typename T::value_type > > type
The type of the nested container.
Definition flatten.hpp:51
Maps a container to is underlying nested container.
Definition flatten.hpp:46