KokkosMap.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#pragma once

#include "KokkosArray.h"

#include <memory>
#include <map>

namespace Moose::Kokkos
{

#ifdef MOOSE_KOKKOS_SCOPE
constexpr uint32_t FNV_PRIME = 0x01000193;        // 16777619
constexpr uint32_t FNV_OFFSET_BASIS = 0x811C9DC5; // 2166136261

template <typename T>
KOKKOS_FUNCTION uint32_t
fnv1aHash(const T & key, uint32_t hash)
{
  auto bytes = reinterpret_cast<const uint8_t *>(&key);

  for (size_t i = 0; i < sizeof(T); ++i)
  {
    hash ^= bytes[i];
    hash *= FNV_PRIME;
  }

  return hash;
}

template <typename T>
KOKKOS_FUNCTION uint32_t
fnv1aHash(const T & key)
{
  return fnv1aHash(key, FNV_OFFSET_BASIS);
}

template <typename T1, typename T2>
struct Pair;

template <typename T1, typename T2>
KOKKOS_FUNCTION uint32_t
fnv1aHash(const Pair<T1, T2> & key)
{
  return fnv1aHash(key.second, fnv1aHash(key.first));
}
#endif

template <typename T1, typename T2, template <typename...> class MapType>
class MapBase;

template <typename T1, typename T2, template <typename...> class MapType>
void dataStore(std::ostream & stream, MapBase<T1, T2, MapType> & map, void * context);
template <typename T1, typename T2, template <typename...> class MapType>
void dataLoad(std::istream & stream, MapBase<T1, T2, MapType> & map, void * context);

template <typename T1, typename T2, template <typename...> class MapType>
class MapBase
{
public:
  auto begin() { return get().begin(); }
  auto begin() const { return get().cbegin(); }
  auto end() { return get().end(); }
  auto end() const { return get().cend(); }
  auto & get()
  {
    if (!_map_host)
      _map_host = std::make_shared<MapType<T1, T2>>();

    return *_map_host;
  }
  const auto & get() const
  {
    if (!_map_host)
      _map_host = std::make_shared<MapType<T1, T2>>();

    return *_map_host;
  }
  void clear();
  T2 & operator[](const T1 & key) { return get()[key]; }

#ifdef MOOSE_KOKKOS_SCOPE

  void copyToDevice();
  void copyToDeviceNested();
  void swap(MapBase<T1, T2, MapType> & map);

  KOKKOS_FUNCTION dof_id_type size() const
  {
    KOKKOS_IF_ON_HOST(return get().size();)

    return _keys.size();
  }
  KOKKOS_FUNCTION dof_id_type find(const T1 & key) const;
  KOKKOS_FUNCTION bool exists(const T1 & key) const { return find(key) != invalid_id; }
  KOKKOS_FUNCTION const T1 & key(dof_id_type idx) const
  {
    KOKKOS_ASSERT(idx != invalid_id);

    return _keys[idx];
  }
  KOKKOS_FUNCTION const T2 & value(dof_id_type idx) const
  {
    KOKKOS_ASSERT(idx != invalid_id);

    return _values[idx];
  }
  KOKKOS_FUNCTION const T2 & operator[](const T1 & key) const
  {
    KOKKOS_IF_ON_HOST(return get().at(key);)

    auto idx = find(key);

    KOKKOS_ASSERT(idx != invalid_id);

    return _values[idx];
  }
  // Due to a stupid NVCC compiler bug, one cannot do var[i][j] for a variable whose type is
  // Array<Map<...>> (the second operator[] seems to be interpreted as if it is the operator of
  // Array), while var[i](j) works. Until we figure out what is going on, one can use the following
  // operator as a workaround.
  KOKKOS_FUNCTION const T2 & operator()(const T1 & key) const { return operator[](key); }
#endif

  static const dof_id_type invalid_id = libMesh::DofObject::invalid_id;

private:
#ifdef MOOSE_KOKKOS_SCOPE

  void copy();
#endif

  mutable std::shared_ptr<MapType<T1, T2>> _map_host;
  Array<T1> _keys;
  Array<T2> _values;
  Array<dof_id_type> _offset;

  friend void dataStore<T1, T2, MapType>(std::ostream &, MapBase<T1, T2, MapType> &, void *);
  friend void dataLoad<T1, T2, MapType>(std::istream &, MapBase<T1, T2, MapType> &, void *);
};

template <typename T1, typename T2, template <typename...> class MapType>
void
MapBase<T1, T2, MapType>::clear()
{
  get().clear();

  _keys.destroy();
  _values.destroy();
  _offset.destroy();
}

#ifdef MOOSE_KOKKOS_SCOPE
template <typename T1, typename T2, template <typename...> class MapType>
void
MapBase<T1, T2, MapType>::copy()
{
  _keys.create(size());
  _values.create(size());
  _offset.create(size() + 1);
  _offset = 0;

  for (const auto & [key, value] : get())
  {
    auto bucket = fnv1aHash(key) % size();

    _offset[bucket]++;
  }

  std::exclusive_scan(_offset.begin(), _offset.end(), _offset.begin(), 0);

  _offset.copyToDevice();

  std::vector<dof_id_type> idx(size(), 0);

  for (const auto & [key, value] : get())
  {
    auto bucket = fnv1aHash(key) % size();

    _keys[_offset[bucket] + idx[bucket]] = key;
    _values[_offset[bucket] + idx[bucket]] = value;
    idx[bucket]++;
  }
}

template <typename T1, typename T2, template <typename...> class MapType>
void
MapBase<T1, T2, MapType>::copyToDevice()
{
  copy();

  _keys.copyToDevice();
  _values.copyToDevice();
}

template <typename T1, typename T2, template <typename...> class MapType>
void
MapBase<T1, T2, MapType>::copyToDeviceNested()
{
  copy();

  _keys.copyToDeviceNested();
  _values.copyToDeviceNested();
}

template <typename T1, typename T2, template <typename...> class MapType>
void
MapBase<T1, T2, MapType>::swap(MapBase<T1, T2, MapType> & map)
{
  get().swap(map.get());
  _keys.swap(map._keys);
  _values.swap(map._values);
  _offset.swap(map._offset);
}

template <typename T1, typename T2, template <typename...> class MapType>
KOKKOS_FUNCTION dof_id_type
MapBase<T1, T2, MapType>::find(const T1 & key) const
{
  if (!size())
    return invalid_id;

  auto bucket = fnv1aHash(key) % size();
  auto begin = _offset[bucket];
  auto end = _offset[bucket + 1];

  for (dof_id_type i = begin; i < end; ++i)
    if (_keys[i] == key)
      return i;

  return invalid_id;
}

template <typename T1, typename T2, template <typename...> class MapType>
void
dataStore(std::ostream & stream, MapBase<T1, T2, MapType> & map, void * context)
{
  using ::dataStore;

  dataStore(stream, map.get(), context);
  dataStore(stream, map._keys, context);
  dataStore(stream, map._values, context);
  dataStore(stream, map._offset, context);
}

template <typename T1, typename T2, template <typename...> class MapType>
void
dataLoad(std::istream & stream, MapBase<T1, T2, MapType> & map, void * context)
{
  using ::dataLoad;

  dataLoad(stream, map.get(), context);
  dataLoad(stream, map._keys, context);
  dataLoad(stream, map._values, context);
  dataLoad(stream, map._offset, context);
}
#endif

template <typename T1, typename T2>
using Map = MapBase<T1, T2, std::map>;

template <typename T1, typename T2>
using UnorderedMap = MapBase<T1, T2, std::unordered_map>;

} // namespace Moose::Kokkos