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