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::Kokkos
{

struct ElementInfo
{
  unsigned int type = libMesh::invalid_uint;
  ContiguousElementID id = libMesh::DofObject::invalid_id;
  ContiguousSubdomainID subdomain = std::numeric_limits<ContiguousSubdomainID>::max();
};

class Mesh
{
public:
  Mesh(MooseMesh & mesh) : _mesh(mesh) {}
  const MooseMesh & getMesh() { return _mesh; }
  [[nodiscard]] bool initialized() const { return _initialized; }
  void update();

  auto getNumSubdomains() const { return _maps->subdomain_id_mapping.size(); }
  auto getNumLocalElementTypes() const { return _maps->elem_type_id_mapping.size(); }
  const auto & getElementTypeMap() const { return _maps->elem_type_id_mapping; }
  auto getNumLocalElements() const { return _num_local_elems; }
  auto getNumSubdomainLocalElements(const SubdomainID subdomain) const
  {
    auto range = libmesh_map_find(_maps->subdomain_elem_id_ranges, subdomain);
    return range.second - range.first;
  }
  auto getNumLocalNodes() const { return _num_local_nodes; }
  const auto & getLocalNodes() const { return _maps->local_nodes; }
  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;
  auto getSubdomainContiguousElementIDRange(const SubdomainID subdomain) const
  {
    const auto & range = libmesh_map_find(_maps->subdomain_elem_id_ranges, subdomain);
    return libMesh::make_range(range.first, range.second);
  }
  ContiguousNodeID getContiguousNodeID(const Node * node) const;
  const auto & getSubdomainContiguousNodeIDs(const SubdomainID subdomain) const
  {
    return libmesh_map_find(_maps->subdomain_node_ids, subdomain);
  }
  const auto & getBoundaryContiguousNodeIDs(const BoundaryID boundary) const
  {
    return libmesh_map_find(_maps->boundary_node_ids, boundary);
  }
#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 dof_id_type getExtraElementID(ContiguousElementID elem, unsigned int index) const
  {
    KOKKOS_ASSERT(index < _extra_elem_ids.n(1));

    return _extra_elem_ids(elem, index);
  }
  KOKKOS_FUNCTION dof_id_type getStartingContiguousElementID(ContiguousSubdomainID subdomain) const
  {
    return _starting_elem_id[subdomain];
  }
  KOKKOS_FUNCTION ContiguousNodeID getContiguousNodeID(ContiguousElementID elem,
                                                       unsigned int node) const
  {
    return _nodes(node, elem);
  }
  KOKKOS_FUNCTION ContiguousNodeID getContiguousNodeID(ElementInfo info,
                                                       unsigned int side,
                                                       unsigned int node) const
  {
    return _nodes(_local_side_node[info.type](node, side), info.id);
  }
  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();

  MooseMesh & _mesh;
  bool _initialized = false;

  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::vector<Node *> local_nodes;
    std::unordered_map<const Node *, ContiguousNodeID> semi_local_node_id_mapping;
    std::unordered_map<SubdomainID, std::pair<ContiguousElementID, ContiguousElementID>>
        subdomain_elem_id_ranges;
    std::unordered_map<SubdomainID, std::vector<ContiguousNodeID>> subdomain_node_ids;
    std::unordered_map<BoundaryID, std::set<ContiguousNodeID>> boundary_node_ids;
  };
  std::shared_ptr<MeshMap> _maps;

  unsigned int _elem_id_integer = libMesh::invalid_uint;
  unsigned int _node_id_integer = libMesh::invalid_uint;
  dof_id_type _num_local_elems = 0;
  dof_id_type _num_local_nodes = 0;
  Array<ElementInfo> _elem_info;
  Array2D<ContiguousElementID> _elem_neighbor;
  Array<ContiguousElementID> _starting_elem_id;
  Array2D<dof_id_type> _extra_elem_ids;
  Array<unsigned int> _num_sides;
  Array<unsigned int> _num_nodes;
  Array<Array<unsigned int>> _num_side_nodes;
  Array<Array2D<unsigned int>> _local_side_node;
  Array<Real3> _points;
  Array2D<ContiguousNodeID> _nodes;
  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(
        if (!_mesh_host.initialized()) mooseError(
            "kokkosMesh() was called too early. Kokkos mesh is available after problem "
            "initialization. Override initialSetup() if you need to setup your object data "
            "using the Kokkos mesh.");

        return _mesh_host;)

    return _mesh_device;
  }
#endif

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

} // namespace Moose::Kokkos