KokkosMesh.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 "KokkosTypes.h"

#ifdef MOOSE_KOKKOS_SCOPE
#include "KokkosUtils.h"
#endif

using ContiguousSubdomainID = SubdomainID;
using ContiguousBoundaryID = BoundaryID;
using ContiguousElementID = dof_id_type;
using ContiguousNodeID = dof_id_type;

class MooseMesh;

namespace Moose
{
namespace Kokkos
{

struct ElementInfo
{
  unsigned int type;
  ContiguousElementID id;
  ContiguousSubdomainID subdomain;
};

class Mesh
{
public:
  Mesh(const MooseMesh & mesh) : _mesh(mesh) {}
  const MooseMesh & getMesh() { return _mesh; }
  void update();

  ContiguousSubdomainID getContiguousSubdomainID(const SubdomainID subdomain) const;
  ContiguousBoundaryID getContiguousBoundaryID(const BoundaryID boundary) const;
  unsigned int getElementTypeID(const Elem * elem) const;
  ContiguousElementID getContiguousElementID(const Elem * elem) const;
  const auto & getElementTypeMap() const { return _maps->_elem_type_id_mapping; }
  const auto & getContiguousElementMap() const { return _maps->_local_elem_id_mapping; }
  const auto & getSubdomainContiguousElementIDs(const SubdomainID subdomain) const
  {
    return _maps->_subdomain_elem_ids.at(subdomain);
  }
  ContiguousNodeID getContiguousNodeID(const Node * node) const;
  const auto & getContiguousNodeMap() const { return _maps->_local_node_id_mapping; }
  const auto & getSubdomainContiguousNodeIDs(const SubdomainID subdomain) const
  {
    return _maps->_subdomain_node_ids.at(subdomain);
  }

#ifdef MOOSE_KOKKOS_SCOPE

  KOKKOS_FUNCTION const auto & getElementInfo(ContiguousElementID elem) const
  {
    return _elem_info[elem];
  }
  KOKKOS_FUNCTION ContiguousElementID getNeighbor(ContiguousElementID elem, unsigned int side) const
  {
    return _elem_neighbor(side, elem);
  }
  KOKKOS_FUNCTION unsigned int getNumSides(unsigned int elem_type) const
  {
    return _num_sides[elem_type];
  }
  KOKKOS_FUNCTION unsigned int getNumNodes(unsigned int elem_type) const
  {
    return _num_nodes[elem_type];
  }
  KOKKOS_FUNCTION unsigned int getNumNodes(unsigned int elem_type, unsigned int side) const
  {
    return _num_side_nodes[elem_type][side];
  }
  KOKKOS_FUNCTION ContiguousNodeID getContiguousNodeID(ContiguousElementID elem,
                                                       unsigned int node) const
  {
    return _nodes(node, elem);
  }
  KOKKOS_FUNCTION ContiguousNodeID getContiguousNodeID(ContiguousElementID elem,
                                                       unsigned int side,
                                                       unsigned int node) const
  {
    return _nodes_face(node, side, elem);
  }
  KOKKOS_FUNCTION Real3 getNodePoint(ContiguousNodeID node) const { return _points[node]; }
  KOKKOS_FUNCTION bool isBoundaryNode(ContiguousNodeID node, ContiguousBoundaryID boundary) const;
#endif

private:
  void initMap();
  void initElement();

  const MooseMesh & _mesh;
  struct MeshMap
  {
    std::unordered_map<SubdomainID, ContiguousSubdomainID> _subdomain_id_mapping;
    std::unordered_map<BoundaryID, ContiguousBoundaryID> _boundary_id_mapping;
    std::unordered_map<ElemType, unsigned int> _elem_type_id_mapping;
    std::unordered_map<const Elem *, ContiguousElementID> _local_elem_id_mapping;
    std::unordered_map<const Node *, ContiguousNodeID> _local_node_id_mapping;
    std::unordered_map<SubdomainID, std::unordered_set<ContiguousElementID>> _subdomain_elem_ids;
    std::unordered_map<SubdomainID, std::unordered_set<ContiguousNodeID>> _subdomain_node_ids;
  };
  std::shared_ptr<MeshMap> _maps;

  Array<ElementInfo> _elem_info;
  Array2D<ContiguousElementID> _elem_neighbor;
  Array<unsigned int> _num_sides;
  Array<unsigned int> _num_nodes;
  Array<Array<unsigned int>> _num_side_nodes;
  Array<Real3> _points;
  Array2D<ContiguousNodeID> _nodes;
  Array3D<ContiguousNodeID> _nodes_face;

  Array<Array<ContiguousNodeID>> _boundary_nodes;
};

#ifdef MOOSE_KOKKOS_SCOPE
KOKKOS_FUNCTION inline bool
Mesh::isBoundaryNode(ContiguousNodeID node, ContiguousBoundaryID boundary) const
{
  if (!_boundary_nodes[boundary].size())
    return false;

  auto begin = &_boundary_nodes[boundary].begin();
  auto end = &_boundary_nodes[boundary].end();
  auto target = Utils::find(node, begin, end);

  return target != end;
}
#endif

class MeshHolder
{
public:
  MeshHolder(const Mesh & mesh) : _mesh_host(mesh), _mesh_device(mesh) {}
  MeshHolder(const MeshHolder & holder)
    : _mesh_host(holder._mesh_host), _mesh_device(holder._mesh_host)
  {
  }

#ifdef MOOSE_KOKKOS_SCOPE

  KOKKOS_FUNCTION const Mesh & kokkosMesh() const
  {
    KOKKOS_IF_ON_HOST(return _mesh_host;)

    return _mesh_device;
  }
#endif

private:
  const Mesh & _mesh_host;
  const Mesh _mesh_device;
};

} // namespace Kokkos
} // namespace Moose