Line data Source code
1 : //* This file is part of the MOOSE framework
2 : //* https://www.mooseframework.org
3 : //*
4 : //* All rights reserved, see COPYRIGHT for full restrictions
5 : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
6 : //*
7 : //* Licensed under LGPL 2.1, please see LICENSE for details
8 : //* https://www.gnu.org/licenses/lgpl-2.1.html
9 :
10 : #pragma once
11 :
12 : #include "KokkosArray.h"
13 :
14 : #include <memory>
15 : #include <map>
16 :
17 : namespace Moose
18 : {
19 : namespace Kokkos
20 : {
21 :
22 : #ifdef MOOSE_KOKKOS_SCOPE
23 : constexpr uint32_t FNV_PRIME = 0x01000193; // 16777619
24 : constexpr uint32_t FNV_OFFSET_BASIS = 0x811C9DC5; // 2166136261
25 :
26 : template <typename T>
27 : KOKKOS_FUNCTION uint32_t
28 89648 : fnv1aHash(const T & key, uint32_t hash)
29 : {
30 89648 : auto bytes = reinterpret_cast<const uint8_t *>(&key);
31 :
32 448240 : for (size_t i = 0; i < sizeof(T); ++i)
33 : {
34 358592 : hash ^= bytes[i];
35 358592 : hash *= FNV_PRIME;
36 : }
37 :
38 89648 : return hash;
39 : }
40 :
41 : template <typename T>
42 : KOKKOS_FUNCTION uint32_t
43 89648 : fnv1aHash(const T & key)
44 : {
45 89648 : return fnv1aHash(key, FNV_OFFSET_BASIS);
46 : }
47 :
48 : template <typename T1, typename T2>
49 : struct Pair;
50 :
51 : template <typename T1, typename T2>
52 : KOKKOS_FUNCTION uint32_t
53 : fnv1aHash(const Pair<T1, T2> & key)
54 : {
55 : return fnv1aHash(key.second, fnv1aHash(key.first));
56 : }
57 : #endif
58 :
59 : template <typename T1, typename T2>
60 : class Map;
61 :
62 : template <typename T1, typename T2>
63 : void dataStore(std::ostream & stream, Map<T1, T2> & map, void * context);
64 : template <typename T1, typename T2>
65 : void dataLoad(std::istream & stream, Map<T1, T2> & map, void * context);
66 :
67 : /**
68 : * The Kokkos wrapper class for standard map.
69 : * The map can only be populated on host.
70 : * Make sure to call copyToDevice() or copyToDeviceNested() after populating the map on host.
71 : */
72 : template <typename T1, typename T2>
73 : class Map
74 : {
75 : public:
76 : /**
77 : * Get the beginning writeable iterator of the host map
78 : * @returns The beginning iterator
79 : */
80 : auto begin() { return get().begin(); }
81 : /**
82 : * Get the beginning const iterator of the host map
83 : * @returns The beginning iterator
84 : */
85 : auto begin() const { return get().cbegin(); }
86 : /**
87 : * Get the end writeable iterator of the host map
88 : * @returns The end iterator
89 : */
90 : auto end() { return get().end(); }
91 : /**
92 : * Get the end const iterator of the host map
93 : * @returns The end iterator
94 : */
95 : auto end() const { return get().cend(); }
96 : /**
97 : * Get the underlying writeable host map
98 : * @returns The writeable host map
99 : */
100 48 : auto & get()
101 : {
102 48 : if (!_map_host)
103 12 : _map_host = std::make_shared<std::map<T1, T2>>();
104 :
105 48 : return *_map_host;
106 : }
107 : /**
108 : * Get the underlying const host map
109 : * @returns The const host map
110 : */
111 96 : const auto & get() const
112 : {
113 96 : if (!_map_host)
114 0 : _map_host = std::make_shared<std::map<T1, T2>>();
115 :
116 96 : return *_map_host;
117 : }
118 : /**
119 : * Call host map's operator[]
120 : * @param key The key
121 : * @returns The writeable reference of the value
122 : */
123 24 : T2 & operator[](const T1 & key) { return get()[key]; }
124 :
125 : #ifdef MOOSE_KOKKOS_SCOPE
126 : /**
127 : * Copy host map to device
128 : */
129 : void copyToDevice();
130 : /**
131 : * Copy host map to device, perform nested copy for Kokkos arrays
132 : */
133 : void copyToDeviceNested();
134 : /**
135 : * Swap with another Kokkos map
136 : * @param map The Kokkos map to be swapped
137 : */
138 : void swap(Map<T1, T2> & map);
139 :
140 : /**
141 : * Get the size of map
142 : * @returns The size of map
143 : */
144 179296 : KOKKOS_FUNCTION dof_id_type size() const
145 : {
146 179296 : KOKKOS_IF_ON_HOST(return get().size();)
147 :
148 179200 : return _keys.size();
149 : }
150 : /**
151 : * Find the index of a key
152 : * @param key The key
153 : * @returns The index of the key, invalid_id if the key does not exist
154 : */
155 : KOKKOS_FUNCTION dof_id_type find(const T1 & key) const;
156 : /**
157 : * Get whether a key exists
158 : * @param key The key
159 : * @returns Whether the key exists
160 : */
161 44800 : KOKKOS_FUNCTION bool exists(const T1 & key) const { return find(key) != invalid_id; }
162 : /**
163 : * Get the key of an index
164 : * @param idx The index returned by find()
165 : * @returns The const reference of the key
166 : */
167 : KOKKOS_FUNCTION const T1 & key(dof_id_type idx) const
168 : {
169 : KOKKOS_ASSERT(idx != invalid_id);
170 :
171 : return _keys[idx];
172 : }
173 : /**
174 : * Get the value of an index
175 : * @param idx The index returned by find()
176 : * @returns The const reference of the value
177 : */
178 : KOKKOS_FUNCTION const T2 & value(dof_id_type idx) const
179 : {
180 : KOKKOS_ASSERT(idx != invalid_id);
181 :
182 : return _values[idx];
183 : }
184 : /**
185 : * Get the value corresponding to a key
186 : * @param key The key
187 : * @returns The const reference of the value
188 : */
189 : ///@{
190 44800 : KOKKOS_FUNCTION const T2 & operator[](const T1 & key) const
191 : {
192 44800 : KOKKOS_IF_ON_HOST(return get().at(key);)
193 :
194 44800 : auto idx = find(key);
195 :
196 : KOKKOS_ASSERT(idx != invalid_id);
197 :
198 44800 : return _values[idx];
199 : }
200 : // Due to a stupid NVCC compiler bug, one cannot do var[i][j] for a variable whose type is
201 : // Array<Map<...>> (the second operator[] seems to be interpreted as if it is the operator of
202 : // Array), while var[i](j) works. Until we figure out what is going on, one can use the following
203 : // operator as a workaround.
204 : KOKKOS_FUNCTION const T2 & operator()(const T1 & key) const { return operator[](key); }
205 : ///@}
206 : #endif
207 :
208 : static const dof_id_type invalid_id = libMesh::DofObject::invalid_id;
209 :
210 : private:
211 : #ifdef MOOSE_KOKKOS_SCOPE
212 : /**
213 : * Internal method to copy host map to device
214 : */
215 : void copy();
216 : #endif
217 :
218 : /**
219 : * Standard map on host
220 : * Stored as a shared pointer to avoid deep copy
221 : */
222 : mutable std::shared_ptr<std::map<T1, T2>> _map_host;
223 : /**
224 : * Keys on device
225 : */
226 : Array<T1> _keys;
227 : /**
228 : * Values on device
229 : */
230 : Array<T2> _values;
231 : /**
232 : * Beginning offset into device arrays of each bucket
233 : */
234 : Array<dof_id_type> _offset;
235 :
236 : friend void dataStore<T1, T2>(std::ostream &, Map<T1, T2> &, void *);
237 : friend void dataLoad<T1, T2>(std::istream &, Map<T1, T2> &, void *);
238 : };
239 :
240 : #ifdef MOOSE_KOKKOS_SCOPE
241 : template <typename T1, typename T2>
242 : void
243 12 : Map<T1, T2>::copy()
244 : {
245 12 : _keys.create(size());
246 12 : _values.create(size());
247 12 : _offset.create(size() + 1);
248 12 : _offset = 0;
249 :
250 36 : for (const auto & [key, value] : get())
251 : {
252 24 : auto bucket = fnv1aHash(key) % size();
253 :
254 24 : _offset[bucket]++;
255 : }
256 :
257 12 : std::exclusive_scan(_offset.begin(), _offset.end(), _offset.begin(), 0);
258 :
259 12 : _offset.copyToDevice();
260 :
261 12 : std::vector<dof_id_type> idx(size(), 0);
262 :
263 36 : for (const auto & [key, value] : get())
264 : {
265 24 : auto bucket = fnv1aHash(key) % size();
266 :
267 24 : _keys[_offset[bucket] + idx[bucket]] = key;
268 24 : _values[_offset[bucket] + idx[bucket]] = value;
269 24 : idx[bucket]++;
270 : }
271 12 : }
272 :
273 : template <typename T1, typename T2>
274 : void
275 12 : Map<T1, T2>::copyToDevice()
276 : {
277 12 : copy();
278 :
279 12 : _keys.copyToDevice();
280 12 : _values.copyToDevice();
281 12 : }
282 :
283 : template <typename T1, typename T2>
284 : void
285 : Map<T1, T2>::copyToDeviceNested()
286 : {
287 : copy();
288 :
289 : _keys.copyToDeviceNested();
290 : _values.copyToDeviceNested();
291 : }
292 :
293 : template <typename T1, typename T2>
294 : void
295 : Map<T1, T2>::swap(Map<T1, T2> & map)
296 : {
297 : get().swap(map.get());
298 : _keys.swap(map._keys);
299 : _values.swap(map._values);
300 : _offset.swap(map._offset);
301 : }
302 :
303 : template <typename T1, typename T2>
304 : KOKKOS_FUNCTION dof_id_type
305 89600 : Map<T1, T2>::find(const T1 & key) const
306 : {
307 89600 : if (!size())
308 0 : return invalid_id;
309 :
310 89600 : auto bucket = fnv1aHash(key) % size();
311 89600 : auto begin = _offset[bucket];
312 89600 : auto end = _offset[bucket + 1];
313 :
314 89600 : for (dof_id_type i = begin; i < end; ++i)
315 89600 : if (_keys[i] == key)
316 89600 : return i;
317 :
318 0 : return invalid_id;
319 : }
320 :
321 : template <typename T1, typename T2>
322 : void
323 : dataStore(std::ostream & stream, Map<T1, T2> & map, void * context)
324 : {
325 : using ::dataStore;
326 :
327 : dataStore(stream, map.get(), context);
328 : dataStore(stream, map._keys, context);
329 : dataStore(stream, map._values, context);
330 : dataStore(stream, map._offset, context);
331 : }
332 :
333 : template <typename T1, typename T2>
334 : void
335 : dataLoad(std::istream & stream, Map<T1, T2> & map, void * context)
336 : {
337 : using ::dataLoad;
338 :
339 : dataLoad(stream, map.get(), context);
340 : dataLoad(stream, map._keys, context);
341 : dataLoad(stream, map._values, context);
342 : dataLoad(stream, map._offset, context);
343 : }
344 : #endif
345 :
346 : } // namespace Kokkos
347 : } // namespace Moose
|
