Line data Source code
1 0 : // Distributed under the MIT License.
2 : // See LICENSE.txt for details.
3 :
4 : #pragma once
5 :
6 : #include <cstddef>
7 : #include <type_traits>
8 : #include <utility>
9 :
10 : #include "DataStructures/DataVector.hpp"
11 : #include "DataStructures/Tags/MirrorView.hpp"
12 : #include "DataStructures/Tensor/AtIndex.hpp"
13 : #include "DataStructures/Variables.hpp"
14 : #include "Utilities/ErrorHandling/FloatingPointExceptions.hpp"
15 : #include "Utilities/Kokkos/KokkosCore.hpp"
16 : #include "Utilities/MakeSignalingNan.hpp"
17 : #include "Utilities/TMPL.hpp"
18 :
19 : #ifdef SPECTRE_KOKKOS
20 :
21 : /*!
22 : * \brief Variables specialization that is backed by a Kokkos::View, meaning
23 : * it can be allocated on device (GPU) memory.
24 : *
25 : * Given a list of tags that hold `Tensor<Kokkos::View<double*>>`, this class
26 : * allocates a larger `Kokkos::View<double*[num_components]>` to hold all the
27 : * tensor components in contiguous memory. The memory is either allocated on the
28 : * host or on a device, depending on the `Kokkos::View` (see Kokkos
29 : * documentation for details). Individual tensors can be accessed with reference
30 : * semantics, meaning the `Kokkos::View` returned by `get` are subviews that
31 : * point into the large memory allocation.
32 : *
33 : * This class doesn't support all the arithmetic operations that the `Variables`
34 : * class backed by a `DataVector` supports. This is because arithmetic
35 : * operations should be done on device using Kokkos parallel algorithms with
36 : * precise control over launching kernels on device.
37 : */
38 : template <typename... Tags>
39 : requires(sizeof...(Tags) > 0 and
40 : (Kokkos::is_view<typename Tags::type::type>::value and ...))
41 : class Variables<tmpl::list<Tags...>> {
42 : public:
43 0 : using tags_list = tmpl::list<Tags...>;
44 : static constexpr auto number_of_variables = sizeof...(Tags);
45 : static constexpr size_t number_of_independent_components =
46 : (... + Tags::type::size());
47 :
48 : // The 1D Kokkos::View that backs the tensors, e.g. Kokkos::View<double*>.
49 0 : using vector_type = typename tmpl::front<tags_list>::type::type;
50 : static_assert(vector_type::rank() == 1,
51 : "Variables can only be used with a 1D Kokkos::View. "
52 : "A second dimension over tensor components is added "
53 : "automatically.");
54 : // The value type of the Kokkos::View, e.g. double.
55 0 : using value_type = typename vector_type::value_type;
56 : // Store the data on device in a 2D Kokkos::View over grid points and tensor
57 : // components. The two common layouts are:
58 : // - SoA: tensor components x grid points
59 : // (typically better for derivatives because neighboring grid points
60 : // are contiguous for each tensor component)
61 : // - AoS: grid points x tensor components
62 : // (typically better for pointwise operations because all tensor components
63 : // for a grid point are contiguous)
64 : // We choose SoA ("LayoutLeft") for now, but this can be changed if needed and
65 : // should be measured for performance.
66 0 : using storage_type =
67 : Kokkos::View<value_type* [number_of_independent_components],
68 : Kokkos::LayoutLeft, typename vector_type::memory_space>;
69 : // The Kokkos memory space the data lives in, e.g. Kokkos::CudaSpace.
70 0 : using memory_space = typename vector_type::memory_space;
71 :
72 0 : Variables() = default;
73 0 : explicit Variables(const size_t number_of_grid_points) {
74 : initialize(number_of_grid_points);
75 : }
76 0 : explicit Variables(storage_type&& rhs) : storage_(std::move(rhs)) {
77 : set_reference_variable_data();
78 : }
79 0 : ~Variables() = default;
80 :
81 : /// Initialize the memory for the Variables. You can specify a `memory_space`
82 : /// to allocate the memory on a specific GPU. Does nothing if the number of
83 : /// grid points is unchanged. Fills the memory with NaN only if
84 : /// SPECTRE_NAN_INIT is defined.
85 1 : void initialize(const size_t number_of_grid_points,
86 : const memory_space& space = {}) {
87 : if (number_of_grid_points == this->number_of_grid_points()) {
88 : return;
89 : }
90 : storage_ = storage_type(
91 : Kokkos::view_alloc(Kokkos::WithoutInitializing, "Variables", space),
92 : number_of_grid_points);
93 : #if defined(SPECTRE_NAN_INIT)
94 : {
95 : // Kokkos compares the fill value to zero to decide whether it can
96 : // memset, which raises an FPE for a signaling NaN.
97 : const ScopedFpeState disable_fpes(false);
98 : Kokkos::deep_copy(storage_, make_signaling_NaN<value_type>());
99 : }
100 : #endif // SPECTRE_NAN_INIT
101 : set_reference_variable_data();
102 : }
103 :
104 0 : constexpr SPECTRE_ALWAYS_INLINE size_t size() const {
105 : return storage_.size();
106 : }
107 0 : constexpr SPECTRE_ALWAYS_INLINE size_t number_of_grid_points() const {
108 : return storage_.extent(0);
109 : }
110 :
111 0 : storage_type& view() { return storage_; }
112 0 : const storage_type& view() const { return storage_; }
113 :
114 : template <typename Tag, typename TagList>
115 1 : friend constexpr typename Tag::type& get( // NOLINT
116 : Variables<TagList>& v);
117 : template <typename Tag, typename TagList>
118 0 : friend constexpr const typename Tag::type& get( // NOLINT
119 : const Variables<TagList>& v);
120 :
121 : template <typename Space>
122 0 : auto create_mirror_view(const Space& space) const {
123 : return Variables<tmpl::list<::Tags::MirrorView<Tags, Space>...>>(
124 : Kokkos::create_mirror_view(space, storage_));
125 : }
126 :
127 : template <typename Space>
128 0 : auto create_mirror_view_and_copy(const Space& space) const {
129 : return Variables<tmpl::list<::Tags::MirrorView<Tags, Space>...>>(
130 : Kokkos::create_mirror_view_and_copy(space, storage_));
131 : }
132 :
133 : private:
134 0 : void set_reference_variable_data() {
135 : size_t variable_offset = 0;
136 : tmpl::for_each<tags_list>([this, &variable_offset](auto tag_v) {
137 : using Tag = tmpl::type_from<decltype(tag_v)>;
138 : auto& var = tuples::get<Tag>(reference_variable_data_);
139 : for (size_t i = 0; i < Tag::type::size(); ++i) {
140 : // The subview is contiguous if each tensor component is contiguous in
141 : // memory, which is the case for Kokkos::LayoutLeft (SoA). Otherwise,
142 : // the subview has a stride.
143 : var[i] = Kokkos::subview(storage_, Kokkos::ALL(), variable_offset);
144 : ++variable_offset;
145 : }
146 : });
147 : }
148 :
149 0 : storage_type storage_{};
150 : tuples::TaggedTuple<Tags...> reference_variable_data_;
151 : };
152 :
153 : /// Get each tensor at the given grid-point index.
154 : template <typename... VarTags, typename... Is>
155 1 : KOKKOS_FUNCTION auto make_at_index(
156 : const Variables<tmpl::list<VarTags...>>& vars, const Is&... i) {
157 : tuples::TaggedTuple<::Tags::AtIndex<VarTags>...> result{};
158 : (..., (get<::Tags::AtIndex<VarTags>>(result) =
159 : make_at_index(get<VarTags>(vars), i...)));
160 : return result;
161 : }
162 :
163 : /// @{
164 : /// Set each tensor at the given grid-point index.
165 : template <typename... VarTags, typename... Is>
166 1 : KOKKOS_FUNCTION void set_at_index(
167 : const gsl::not_null<Variables<tmpl::list<VarTags...>>*> vars,
168 : const tuples::TaggedTuple<::Tags::AtIndex<VarTags>...>& value,
169 : const Is&... i) {
170 : (..., (set_at_index(make_not_null(&get<VarTags>(*vars)),
171 : get<::Tags::AtIndex<VarTags>>(value), i...)));
172 : }
173 : template <typename... VarTags, typename... Is>
174 1 : KOKKOS_FUNCTION void set_at_index(
175 : const gsl::not_null<const Variables<tmpl::list<VarTags...>>*> vars,
176 : const tuples::TaggedTuple<::Tags::AtIndex<VarTags>...>& value,
177 : const Is&... i) {
178 : (..., (set_at_index(make_not_null(&get<VarTags>(*vars)),
179 : get<::Tags::AtIndex<VarTags>>(value), i...)));
180 : }
181 : /// @}
182 :
183 : /// Allocate a `Variables` of the same size as `vars_host` in device memory,
184 : /// without copying the data yet. You can specify a `memory_space` to allocate
185 : /// the memory on a specific GPU.
186 : template <typename... Tags>
187 1 : auto allocate_on_device(
188 : const Variables<tmpl::list<Tags...>>& vars_host,
189 : const typename Kokkos::DefaultExecutionSpace::memory_space& space = {}) {
190 : static_assert((std::is_same_v<typename Tags::type::type, DataVector> and ...),
191 : "Only a DataVector-backed Variables can be mirrored in device "
192 : "memory.");
193 : Variables<tmpl::list<::Tags::MirrorView<Tags>...>> vars_device{};
194 : if constexpr (sizeof...(Tags) > 0) {
195 : vars_device.initialize(vars_host.number_of_grid_points(), space);
196 : }
197 : return vars_device;
198 : }
199 :
200 : /// Copy the `vars_host` to an existing `Variables` in device memory. The device
201 : /// memory is only reallocated if the number of grid points changed.
202 : template <typename DeviceTags, typename... Tags>
203 1 : void copy_to_device(const gsl::not_null<Variables<DeviceTags>*> vars_device,
204 : const Variables<tmpl::list<Tags...>>& vars_host) {
205 : static_assert((std::is_same_v<typename Tags::type::type, DataVector> and ...),
206 : "Only a DataVector-backed Variables can be copied to device "
207 : "memory.");
208 : if constexpr (sizeof...(Tags) > 0) {
209 : using HostVars = Variables<tmpl::list<Tags...>>;
210 : using ValueType = typename HostVars::value_type;
211 : static constexpr size_t num_components =
212 : HostVars::number_of_independent_components;
213 : const size_t num_points = vars_host.number_of_grid_points();
214 : if (vars_device->number_of_grid_points() != num_points) {
215 : vars_device->initialize(num_points);
216 : }
217 : // View that wraps the host data (DataVector-backed Variables) in a
218 : // Kokkos::View without owning it so we can deep-copy the data. Relies on
219 : // the memory layout in the DataVector-backed Variables being the same as
220 : // the Kokkos view (LayoutLeft).
221 : using HostUnmanaged =
222 : Kokkos::View<const ValueType* [num_components], Kokkos::LayoutLeft,
223 : Kokkos::HostSpace, Kokkos::MemoryUnmanaged>;
224 : HostUnmanaged unmanaged_host_view(vars_host.data(), num_points);
225 : // Copy to device
226 : Kokkos::deep_copy(vars_device->view(), unmanaged_host_view);
227 : }
228 : }
229 :
230 : /// Copy the `vars_host` to device memory. You can specify a `memory_space` to
231 : /// allocate the memory on a specific GPU.
232 : template <typename... Tags>
233 1 : auto copy_to_device(
234 : const Variables<tmpl::list<Tags...>>& vars_host,
235 : const typename Kokkos::DefaultExecutionSpace::memory_space& space = {}) {
236 : auto vars_device = allocate_on_device(vars_host, space);
237 : copy_to_device(make_not_null(&vars_device), vars_host);
238 : return vars_device;
239 : }
240 :
241 : /// Copy a `Variables` from device to host memory.
242 : template <typename HostTags, typename DeviceTags>
243 1 : void copy_to_host(const gsl::not_null<Variables<HostTags>*> vars_host,
244 : const Variables<DeviceTags>& vars_device) {
245 : if constexpr (tmpl::size<DeviceTags>::value > 0) {
246 : using HostVars = Variables<HostTags>;
247 : static_assert(std::is_same_v<typename HostVars::vector_type, DataVector>,
248 : "Only a DataVector-backed Variables can be copied from "
249 : "device memory.");
250 : using ValueType = typename HostVars::value_type;
251 : static constexpr size_t num_components =
252 : HostVars::number_of_independent_components;
253 : const size_t num_points = vars_device.number_of_grid_points();
254 : // View that wraps the host data (DataVector-backed Variables) in a
255 : // Kokkos::View without owning it so we can deep-copy the data. Relies on
256 : // the memory layout in the DataVector-backed Variables being the same as
257 : // the Kokkos view (LayoutLeft).
258 : using HostUnmanaged =
259 : Kokkos::View<ValueType* [num_components], Kokkos::LayoutLeft,
260 : Kokkos::HostSpace, Kokkos::MemoryUnmanaged>;
261 : if (vars_host->number_of_grid_points() != num_points) {
262 : vars_host->initialize(num_points);
263 : }
264 : HostUnmanaged unmanaged_host_view(vars_host->data(), num_points);
265 : // Copy to host
266 : Kokkos::deep_copy(unmanaged_host_view, vars_device.view());
267 : }
268 : }
269 :
270 : #endif // SPECTRE_KOKKOS
|
