https://mooseframework.inl.gov
KokkosMap.h
Go to the documentation of this file.
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 fnv1aHash(const T & key, uint32_t hash)
29 {
30  auto bytes = reinterpret_cast<const uint8_t *>(&key);
31 
32  for (size_t i = 0; i < sizeof(T); ++i)
33  {
34  hash ^= bytes[i];
35  hash *= FNV_PRIME;
36  }
37 
38  return hash;
39 }
40 
41 template <typename T>
42 KOKKOS_FUNCTION uint32_t
43 fnv1aHash(const T & key)
44 {
45  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 
72 template <typename T1, typename T2>
73 class Map
74 {
75 public:
80  auto begin() { return get().begin(); }
85  auto begin() const { return get().cbegin(); }
90  auto end() { return get().end(); }
95  auto end() const { return get().cend(); }
100  auto & get()
101  {
102  if (!_map_host)
103  _map_host = std::make_shared<std::map<T1, T2>>();
104 
105  return *_map_host;
106  }
111  const auto & get() const
112  {
113  if (!_map_host)
114  _map_host = std::make_shared<std::map<T1, T2>>();
115 
116  return *_map_host;
117  }
123  T2 & operator[](const T1 & key) { return get()[key]; }
124 
125 #ifdef MOOSE_KOKKOS_SCOPE
126 
129  void copyToDevice();
133  void copyToDeviceNested();
138  void swap(Map<T1, T2> & map);
139 
144  KOKKOS_FUNCTION dof_id_type size() const
145  {
146  KOKKOS_IF_ON_HOST(return get().size();)
147 
148  return _keys.size();
149  }
155  KOKKOS_FUNCTION dof_id_type find(const T1 & key) const;
161  KOKKOS_FUNCTION bool exists(const T1 & key) const { return find(key) != invalid_id; }
167  KOKKOS_FUNCTION const T1 & key(dof_id_type idx) const
168  {
169  KOKKOS_ASSERT(idx != invalid_id);
170 
171  return _keys[idx];
172  }
178  KOKKOS_FUNCTION const T2 & value(dof_id_type idx) const
179  {
180  KOKKOS_ASSERT(idx != invalid_id);
181 
182  return _values[idx];
183  }
189  KOKKOS_FUNCTION const T2 & operator[](const T1 & key) const
191  {
192  KOKKOS_IF_ON_HOST(return get().at(key);)
193 
194  auto idx = find(key);
195 
196  KOKKOS_ASSERT(idx != invalid_id);
197 
198  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); }
206 #endif
207 
209 
210 private:
211 #ifdef MOOSE_KOKKOS_SCOPE
212 
215  void copy();
216 #endif
217 
222  mutable std::shared_ptr<std::map<T1, T2>> _map_host;
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
244 {
245  _keys.create(size());
246  _values.create(size());
247  _offset.create(size() + 1);
248  _offset = 0;
249 
250  for (const auto & [key, value] : get())
251  {
252  auto bucket = fnv1aHash(key) % size();
253 
254  _offset[bucket]++;
255  }
256 
257  std::exclusive_scan(_offset.begin(), _offset.end(), _offset.begin(), 0);
258 
259  _offset.copyToDevice();
260 
261  std::vector<dof_id_type> idx(size(), 0);
262 
263  for (const auto & [key, value] : get())
264  {
265  auto bucket = fnv1aHash(key) % size();
266 
267  _keys[_offset[bucket] + idx[bucket]] = key;
268  _values[_offset[bucket] + idx[bucket]] = value;
269  idx[bucket]++;
270  }
271 }
272 
273 template <typename T1, typename T2>
274 void
276 {
277  copy();
278 
279  _keys.copyToDevice();
280  _values.copyToDevice();
281 }
282 
283 template <typename T1, typename T2>
284 void
286 {
287  copy();
288 
289  _keys.copyToDeviceNested();
290  _values.copyToDeviceNested();
291 }
292 
293 template <typename T1, typename T2>
294 void
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 Map<T1, T2>::find(const T1 & key) const
306 {
307  if (!size())
308  return invalid_id;
309 
310  auto bucket = fnv1aHash(key) % size();
311  auto begin = _offset[bucket];
312  auto end = _offset[bucket + 1];
313 
314  for (dof_id_type i = begin; i < end; ++i)
315  if (_keys[i] == key)
316  return i;
317 
318  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 {
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 {
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
T2 & operator[](const T1 &key)
Call host map&#39;s operator[].
Definition: KokkosMap.h:123
constexpr uint32_t FNV_PRIME
Definition: KokkosMap.h:23
auto & get()
Get the underlying writeable host map.
Definition: KokkosMap.h:100
void copyToDevice()
Copy host map to device.
Definition: KokkosMap.h:275
Array< T1 > _keys
Keys on device.
Definition: KokkosMap.h:226
void swap(std::vector< T > &data, const std::size_t idx0, const std::size_t idx1, const libMesh::Parallel::Communicator &comm)
Swap function for serial or distributed vector of data.
Definition: Shuffle.h:494
std::shared_ptr< std::map< T1, T2 > > _map_host
Standard map on host Stored as a shared pointer to avoid deep copy.
Definition: KokkosMap.h:222
static const dof_id_type invalid_id
Definition: KokkosMap.h:208
void copyToDeviceNested()
Copy host map to device, perform nested copy for Kokkos arrays.
Definition: KokkosMap.h:285
Array< T2 > _values
Values on device.
Definition: KokkosMap.h:230
auto begin() const
Get the beginning const iterator of the host map.
Definition: KokkosMap.h:85
KOKKOS_FUNCTION const T2 & operator()(const T1 &key) const
Definition: KokkosMap.h:204
The Kokkos wrapper class for standard map.
Definition: KokkosMap.h:60
KOKKOS_FUNCTION uint32_t fnv1aHash(const T &key, uint32_t hash)
Definition: KokkosMap.h:28
constexpr uint32_t FNV_OFFSET_BASIS
Definition: KokkosMap.h:24
auto end() const
Get the end const iterator of the host map.
Definition: KokkosMap.h:95
Array< dof_id_type > _offset
Beginning offset into device arrays of each bucket.
Definition: KokkosMap.h:234
void dataLoad(std::istream &stream, Map< T1, T2 > &map, void *context)
Definition: KokkosMap.h:335
KOKKOS_FUNCTION bool exists(const T1 &key) const
Get whether a key exists.
Definition: KokkosMap.h:161
void copy()
Internal method to copy host map to device.
Definition: KokkosMap.h:243
Real value(unsigned n, unsigned alpha, unsigned beta, Real x)
static const dof_id_type invalid_id
KOKKOS_FUNCTION const T2 & value(dof_id_type idx) const
Get the value of an index.
Definition: KokkosMap.h:178
auto end()
Get the end writeable iterator of the host map.
Definition: KokkosMap.h:90
void dataStore(std::ostream &stream, Array< T, dimension > &array, void *context)
Definition: KokkosArray.h:887
KOKKOS_FUNCTION dof_id_type find(const T1 &key) const
Find the index of a key.
Definition: KokkosMap.h:305
void dataLoad(std::istream &stream, Array< T, dimension > &array, void *context)
Definition: KokkosArray.h:931
KOKKOS_FUNCTION const T1 & key(dof_id_type idx) const
Get the key of an index.
Definition: KokkosMap.h:167
void swap(Map< T1, T2 > &map)
Swap with another Kokkos map.
Definition: KokkosMap.h:295
auto begin()
Get the beginning writeable iterator of the host map.
Definition: KokkosMap.h:80
void dataStore(std::ostream &stream, Map< T1, T2 > &map, void *context)
Definition: KokkosMap.h:323
KOKKOS_FUNCTION dof_id_type size() const
Get the size of map.
Definition: KokkosMap.h:144
MOOSE now contains C++17 code, so give a reasonable error message stating what the user can do to add...
unsigned int idx(const ElemType type, const unsigned int nx, const unsigned int i, const unsigned int j)
uint8_t dof_id_type